Aller au contenu principal
Dernière mise à jour le

Migrer depuis Slurm

info

Aperçu

Cette fonctionnalité est en Aperçu public.

Ce guide montre comment migrer des workloads d'entraînement distribué de Slurm vers AI Runtime. Il explique les concepts fondamentaux, montre comment surveiller et contrôler les exécutions, et détaille la traduction d'un script batch Slurm.

Dans AI Runtime, vous définissez le type de GPU, le nombre total de GPU, l’environnement et la commande à exécuter sur chaque nœud dans une configuration de charge de travail YAML. Lorsque vous soumettez cette configuration avec databricks air run --file train.yaml, AI Runtime gère la planification, le provisionnement et le nettoyage des ressources pour l'exécution.

Avant de vous start, installez la CLI Databricks, configurez l’authentification et suivez le guide de démarrage rapide.

Concepts Slurm et AI Runtime​

Ressources GPU​

Dans Slurm, un job reçoit une attribution de nœuds et de GPU provenant d'une partition. Avec AI Runtime, vous demandez un type de GPU et un nombre total de GPU pour chaque exécution. Par default, le service provisionne ce compute à la demande.

Définissez compute.num_accelerators sur le nombre total de GPU . Le paramètre compute.accelerator_type détermine le nombre de GPU par nœud . Par exemple, GPU_8xH100 fournit 8 GPU H100 par nœud, de sorte qu’en demandant 16 GPU, vous obtenez 2 nœuds. Le total doit être un multiple du nombre de GPU par nœud.

Processus d’entraînement​

Dans Slurm, sbatch soumet un job qui demande des nœuds et des GPU, et srun lance des tâches sur les nœuds alloués. Pour l'entraînement PyTorch distribué, un schéma courant consiste à utiliser srun pour start un torchrun lanceur par nœud. Chaque lanceur lance un processus d'entraînement par GPU.

AI Runtime exécute le même command une fois sur chaque nœud. Utilisez torchrun dans cette commande pour lancer les processus d'entraînement. AI Runtime fournit le rang de chaque nœud ainsi que les informations de connexion dont les nœuds ont besoin pour se coordonner. torchrun attribue à chaque processus d'entraînement son RANK, son LOCAL_RANK et son WORLD_SIZE. Avec 2 nœuds et 8 GPU par nœud, cela vous donne 2 lanceurs et 16 processus d'entraînement.

Stockage et nouvelles tentatives​

Utilisez les volumes Unity Catalog (/Volumes/<catalog>/<schema>/<volume>/...) pour les datasets partagés et les points de contrôle qui doivent persister après une exécution. Traitez le disque local de chaque nœud comme un espace de travail temporaire.

timeout_minutes limite chaque tentative, et max_retries contrôle le nombre de nouvelles tentatives pour une charge de travail ayant échoué. Chaque tentative start la commande again. Pour continuer l’entraînement à partir d'un point de contrôle, votre code d'entraînement doit charger l'état sauvegardé. Consultez la section Améliorer les performances et la résilience de l'entraînement sur AI Runtime pour en savoir plus sur les points de contrôle et les schémas de récupération.

Surveiller et contrôler les exécutions​

Slurm

AI Runtime

squeue

databricks air list runs

sacct, scontrol show job <id>

databricks air get run <run-id>

tail -f slurm-<id>.out

databricks air logs <run-id> --node <n>

scancel <id>

databricks air cancel <run-id>

Slurm

AI Runtime

squeue

databricks air list runs

sacct, scontrol show job <id>

databricks air get run <run-id>

tail -f slurm-<id>.out

databricks air logs <run-id> --node <n>

scancel <id>

databricks air cancel <run-id>

Chaque workload dispose d'une exécution MLflow avec des logs et des métriques système collectées automatiquement. Utilisez MLflow dans votre code d'entraînement pour consigner les parameters, les métriques d'entraînement et les artefacts. Consultez Suivre les exécutions avec MLflow et la page d'exécution des jobs.

Exemple : traduire un script sbatch​

Un script de lancement Slurm pour 2 nœuds dotés de 8 GPU chacun (16 GPU au total) :

Bash
#!/bin/bash
#SBATCH --job-name=llama-sft
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --time=02:00:00

srun bash -c '
  torchrun \
    --nnodes="$SLURM_NNODES" \
    --node_rank="$SLURM_NODEID" \
    --nproc_per_node=8 \
    --master_addr="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)" \
    --master_port=29500 \
    train.py
'

L’équivalent AI Runtime train.yaml:

YAML
experiment_name: llama-sft

environment:
  version: '4'
  dependencies:
    - transformers>=4.45
    - datasets>=3.0

# 16 GPUs across 2 nodes (GPU_8xH100 = 8 H100 per node).
compute:
  num_accelerators: 16
  accelerator_type: GPU_8xH100

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  cd "$CODE_SOURCE_PATH"
  # AI Runtime sets these rendezvous variables on each node.
  torchrun \
    --nnodes="$NUM_NODES" \
    --node_rank="$NODE_RANK" \
    --nproc_per_node="${LOCAL_WORLD_SIZE:-8}" \
    --master_addr="$MASTER_ADDR" \
    --master_port="$MASTER_PORT" \
    train.py

timeout_minutes: 120
max_retries: 1

Envoyez et suivez-le :

Bash
databricks air run --file train.yaml --watch

Pour obtenir une version exécutable complète incluant le script d’entraînement, consultez Affinement de LLM multinoeud avec FSDP.

Considérations relatives à la migration​

  • Remplacez module load et les étapes d'activation de l'environnement par une configuration environment. Sélectionnez une version d'environnement gérée et déclarez des packages supplémentaires dans environment.dependencies, en ligne ou via une référence -r vers un fichier requirements.txt. Pour une pile personnalisée, utilisez une image Docker personnalisée.
  • Définissez code_source.snapshot.root_path sur votre répertoire de projet local. La CLI Databricks l'upload lorsque vous soumettez la charge de travail. Référencez les fichiers upload avec $CODE_SOURCE_PATH.
  • Utilisez env_variables pour les variables d’environnement et secrets pour les références de secret Databricks.
  • Soumettez une exécution distincte pour chaque configuration, en utilisant databricks air run --override key=value pour varier les champs entre les soumissions.

Ressources supplémentaires​

Sur cette page