Migrate from Slurm

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

Cada carga de trabalho tem uma execução do MLflow com logs e métricas do sistema coletadas automaticamente. Use o MLflow no seu código de treinamento para registrar parâmetros, métricas de treinamento e artefatos. Consulte Acompanhamento de experimentos e observabilidade.

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

Para obter uma versão completa executável, incluindo o script de treinamento, consulte Ajuste fino FSDP multinó de Llama 3.1 8B no Cookbook do AI Runtime.

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.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​