Migrate from Slurm

Preview

This guide shows how to migrate distributed training workloads from Slurm to AI Runtime. It explains the core concepts, shows how to monitor and control runs, and walks through translating a Slurm batch script.

In AI Runtime, you define the GPU type, total GPU count, environment, and command to run on each node in a YAML workload configuration. When you submit that configuration with databricks air run --file train.yaml, AI Runtime handles scheduling, provisioning, and resource cleanup for the run.

Before you start, install the Databricks CLI, configure authentication, and follow the quickstart.

Slurm and AI Runtime concepts

GPU resources

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 run. By default, the service provisions that compute on demand.

Set compute.num_accelerators to the total GPU count. The compute.accelerator_type determines how many GPUs each node has. For example, GPU_8xH100 provides 8 H100 GPUs per node, so requesting 16 GPUs gives you 2 nodes. The total must be a multiple of the GPUs per node.

Training processes

In Slurm, sbatch submits a job that requests nodes and GPUs, and srun launches tasks on the allocated nodes. For distributed PyTorch training, a common pattern is to use srun to start one torchrun launcher per node. Each launcher starts one training process per GPU.

AI Runtime runs the same command one time on each node. Use torchrun in that command to launch the training processes. AI Runtime supplies each node's rank and the connection information the nodes need to coordinate. torchrun assigns each training process its RANK, LOCAL_RANK, and WORLD_SIZE. With 2 nodes and 8 GPUs per node, this gives you 2 launchers and 16 training processes.

Storage and retries

Use Unity Catalog volumes (/Volumes/<catalog>/<schema>/<volume>/...) for shared datasets and checkpoints that must survive a run. Treat each node's local disk as temporary working space.

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

Monitor and control runs

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 run with logs and automatically collected system metrics. Use MLflow in your training code to log parameters, training metrics, and artifacts. See Track runs with MLflow and the Jobs run page.

Example: translate an `sbatch` script

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
'

The equivalent 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

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.

Migration considerations

Replace module load and environment activation steps with an environment configuration. Select a managed environment version and declare additional packages in environment.dependencies, inline or through a -r reference to a requirements.txt file. For a custom stack, use a custom Docker image.
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 environment variables and secrets for Databricks secret references.
Submit a separate run for each configuration, using databricks air run --override key=value to vary fields between submissions.

Slurm and AI Runtime concepts​

GPU resources​

Training processes​

Storage and retries​

Monitor and control runs​

Example: translate an sbatch script​

Migration considerations​

Additional resources​