Migrate from Slurm

info

Visualização

Este guia mostra como migrar cargas de trabalho de treinamento distribuído do Slurm para o AI Runtime. It explains the core concepts, shows how to monitor and control execuções, and walks through translating a Slurm lotes script.

No AI Runtime, você define o tipo de GPU, a contagem total de GPUs, o ambiente e o comando a ser executado em cada nó em uma configuração de carga de trabalho YAML. Ao enviar essa configuração com databricks air run --file train.yaml, o AI Runtime lida com a programação, o provisionamento e a limpeza de recursos para a execução.

Antes de começar, instale a CLI do Databricks, configure a autenticação e siga o guia de início rápido.

Conceitos do Slurm e do AI Runtime

Recursos de GPU

In Slurm, a job receives an allocation of nodes and GPUs from a partition. With AI Runtime, you request a GPU type and total GPU count for each execução. By default, the service provisions that compute on demand.

Defina compute.num_accelerators como a contagem total de GPUs. O compute.accelerator_type determina quantas GPUs cada nó possui. Por exemplo, o GPU_8xH100 fornece 8 GPUs H100 por nó, portanto, solicitar 16 GPUs resulta em 2 nós. O total deve ser um múltiplo das GPUs por nó.

Processos de treinamento

No Slurm, sbatch envia um job que solicita nós e GPUs, e srun inicia tarefas nos nós alocados. Para treinamento distribuído com o PyTorch, um padrão comum é usar srun para começar um torchrun lançador por nó. Cada lançador inicia um processo de treinamento por GPU.

O AI Runtime executa o mesmo command uma vez em cada nó. Use torchrun nesse comando para iniciar os processos de treinamento. O AI Runtime fornece a classificação de cada nó e as informações de conexão de que os nós precisam para se coordenar. torchrun atribui a cada processo de treinamento seu RANK, LOCAL_RANK e WORLD_SIZE. Com 2 nós e 8 GPUs por nó, isso fornece 2 executores de inicialização e 16 processos de treinamento.

Armazenamento e novas tentativas

Use volumes do Unity Catalog (/Volumes/<catalog>/<schema>/<volume>/...) para datasets compartilhados e pontos de verificação que devem sobreviver a uma execução. Trate o disco local de cada nó como espaço de trabalho temporário.

timeout_minutes limits each attempt, and max_retries controls how many times a failed workload is retried. Each retry starts the comando again. To continue treinamento from a checkpoint, your treinamento code must load the saved state. See Improve treinamento desempenho and resiliency on AI Runtime for checkpoint and recovery patterns.

Monitorar e controlar execuções

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

Each workload has an MLflow execução with Logs and automatically collected system métricas. Use MLflow in your código de treinamento to log parameters, métricas de treinamento, and artifacts. See Track execuções with MLflow and the Jobs run page.

Exemplo: traduzir um script `sbatch`

A Slurm launch script for 2 nodes with 8 GPUs each (16 GPUs 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
'

O train.yaml de AI Runtime equivalente:

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

Submit and follow it:

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

For a complete runnable version, including the training script, see Multi-node LLM fine-tuning with FSDP.

Considerações sobre a migração

Substitua module load e os passos de ativação do ambiente por uma configuração de environment. Selecione uma versão de ambiente gerenciada e declare pacotes adicionais em environment.dependencies, de forma embutida ou por meio de uma referência de -r a um arquivo requirements.txt. Para uma pilha personalizada, use uma Docker Image personalizada.
Set code_source.snapshot.root_path to your local project directory. The Databricks CLI uploads it when you submit the workload. Reference the uploaded files with $CODE_SOURCE_PATH.
Use env_variables for variáveis de ambiente and secrets for Databricks secret references.
Envie uma execução separada para cada configuração, usando databricks air run --override key=value para variar os campos entre os envios.

Conceitos do Slurm e do AI Runtime​

Recursos de GPU​

Processos de treinamento​

Armazenamento e novas tentativas​

Monitorar e controlar execuções​

Exemplo: traduzir um script sbatch​

Considerações sobre a migração​

Recursos adicionais​