Migrazione da Slurm

Questa guida mostra come migrare carichi di lavoro di addestramento distribuiti da Slurm a AI Runtime. Spiega i concetti fondamentali, mostra come monitorare e controllare le run e guida la traduzione di uno script batch Slurm.

In AI Runtime, si definisce il tipo di GPU, il numero totale di GPU, l'ambiente e il comando da eseguire su ogni nodo in una configurazione di carico di lavoro YAML. Quando invii quella configurazione con databricks air run --file train.yaml, AI Runtime gestisce la pianificazione, il provisioning e la pulizia delle risorse per l'esecuzione.

Prima di iniziare, installa la CI di Databricks, configura l'autenticazione e segui il quickstart.

Concetti di Slurm e AI Runtime

Risorse GPU

In Slurm, un lavoro riceve un'allocazione di nodi e GPU da una partizione. Con AI Runtime, richiedi un tipo di GPU e il conteggio totale della GPU per ogni esecuzione. Di default, le disposizioni di servizio che si calcolano su richiesta.

Imposta compute.num_accelerators sul numero totale di GPU. Questo compute.accelerator_type determina quante GPU ha ogni nodo. Ad esempio, GPU_8xH100 fornisce 8 GPU H100 per nodo, quindi richiedendo 16 GPU hai 2 nodi. Il totale deve essere un multiplo delle GPU per nodo.

Processi di formazione

In Slurm, sbatch invia un lavoro che richiede nodi e GPU e srun avvia i compiti sui nodi assegnati. Per l'addestramento distribuito con PyTorch, uno schema comune consiste nell'usare srun per avviare un torchrun avviatore per nodo. Ogni launcher avvia un processo di addestramento per ogni GPU.

Il runtime dell'IA esegue lo stesso command una volta su ciascun nodo. Usa torchrun in quel comando per avviare i processi di addestramento. AI Runtime fornisce il rango di ogni nodo e le informazioni di connessione necessarie per coordinarsi. torchrun assegna a ogni processo di addestramento i suoi RANK, LOCAL_RANK, e WORLD_SIZE. Con 2 nodi e 8 GPU per nodo, questo ti dà 2 launcher e 16 processi di addestramento.

Archiviazione e ritentativi

Usa i volumi di Unity Catalog (/Volumes/<catalog>/<schema>/<volume>/...) per dataset e checkpoint condivisi che devono persistere oltre un'esecuzione. Considera il disco locale di ogni nodo come uno spazio di lavoro temporaneo.

timeout_minutes limita ogni tentativo e max_retries controlla quante volte viene riprovato un carico di lavoro fallito. Ogni nuovo tentativo riavvia il comando. Per continuare l'addestramento da un checkpoint, il tuo codice di addestramento deve caricare lo stato salvato. Vedi Miglioramento delle prestazioni e della resilienza dell'addestramento su AI Runtime per informazioni sui checkpoint e sui modelli di ripristino.

Monitoraggio e controllo

Ogni carico di lavoro ha un flusso ML eseguito con log e metriche di sistema raccolte automaticamente. Usa MLflow nel tuo codice di addestramento per registrare parametri, metriche di addestramento e artefatti. Vedi Tracciamento e osservabilità degli esperimenti.

Esempio: tradurre uno sbatch script

Uno script di avvio Slurm per 2 nodi con 8 GPU ciascuno (16 GPU in totale):

#!/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'equivalente runtime di train.yaml IA:

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:
  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

Invia e monitoralo:

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

Per una versione completa eseguibile, incluso lo script di addestramento, vedi Multi-node FSDP fine-tuning di Llama 3.1 8B nel AI Runtime Cookbook.

Considerazioni sulla migrazione

• Sostituisci module load e i passaggi di attivazione dell'ambiente con una configurazione environment. Seleziona una versione dell'ambiente gestito e dichiara pacchetti aggiuntivi in environment.dependencies, in linea o tramite un -r riferimento a un requirements.txt file. Per uno stack personalizzato, usa un'immagine Docker personalizzata.
• Imposta code_source.root_path sulla directory locale del progetto. La CLI Databricks lo carica quando invii il carico di lavoro. Fai riferimento ai file caricati con $CODE_SOURCE_PATH.
• Usa env_variables per variabili di ambiente e secrets per riferimenti segreti di Databricks.
• Invia una run separata per ogni configurazione, usando databricks air run --override key=value per variare i campi tra una sottomissione e l'altra.

Risorse aggiuntive

• Guida introduttiva ad AI Runtime
• Informazioni di riferimento su YAML del carico di lavoro
• Messa a punto FSDP multi-nodo del Llama 3.1 8B
• Formazione distribuita nei quaderni