Distributed training in notebooks

Important

This feature is in Beta.

The @distributed decorator from the Databricks AI Runtime SDK (databricks-sdk-air) is the most convenient way to run distributed training from a Databricks notebook. Decorate your training function, call it, and AI Runtime runs it across all GPUs on the node your notebook is connected to. The same code scales from single-GPU to multi-GPU with no cluster to provision and no distributed launcher to configure.

Key points on this page:

• The @distributed decorator runs a training function across every GPU on your node from inside a notebook.
• It supports PyTorch DDP, FSDP, and DeepSpeed, and moves single-GPU code to multi-GPU with minimal changes.
• Connect your notebook to an 8xH100 accelerator for full multi-GPU training. By default, @distributed uses every GPU on the node.

Note

This page covers distributed training from Databricks notebooks with the Databricks AI Runtime SDK. To submit distributed training workloads from your local machine, use the Databricks CLI commands for AI Runtime, which are in Public Preview. See Use the Databricks CLI with AI Runtime.

Quickstart

The examples on this page use the Databricks AI Runtime SDK, which is preinstalled in GPU environment version 6 and above. On version 5 and below, use the serverless_gpu package instead. See Migrate from serverless_gpu. To connect a notebook to an accelerator, see Connect to a notebook.

Decorate your training function with @distributed, then call it with .distributed():

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

Distributed training requires an 8xH100 accelerator, which provisions a single node with 8 GPUs. By default, @distributed uses all eight GPUs. Set num_accelerators to use fewer. The accelerator_type parameter is also optional and automatically detected from the accelerator your notebook is connected to.

Each .distributed() call creates an MLflow run (or a nested child run if one is already active) and prints a run link in the cell output. For a complete, runnable walkthrough, see Full example.

Supported frameworks

The @distributed API integrates with major distributed training libraries:

• PyTorch Distributed Data Parallel (DDP): Standard multi-GPU data parallelism.
• Fully Sharded Data Parallel (FSDP): Memory-efficient training for large models.
• DeepSpeed: Microsoft's optimization library for large model training.

For real training scenarios that use each library, see distributed training examples.

How the @distributed decorator works

When you call a decorated function with .distributed(), AI Runtime handles the mechanics that you would otherwise configure by hand with a distributed launcher:

• Serialization and fan-out: The function is serialized and launched on each selected GPU. Every GPU runs a copy of the function with the same arguments.
• Environment synchronization: The Python environment and dependencies are replicated across all ranks, so every process runs the same code.
• Rank environment variables: Standard variables such as LOCAL_RANK are populated for each process. Read them in your function to place the model and data on the correct device.
• Result collection: Return values are collected from all ranks and returned to the caller.
• MLflow tracking: Each .distributed() call creates an MLflow run, or a nested child run if one is already active, so metrics logged from your function land on the same run.
• Lifecycle and timeout: Distributed execution runs within the lifecycle of the notebook. Terminating the notebook terminates the run. The decorator has a default timeout of 3 hours. Pass timeout in seconds to change it, or timeout=None to disable it. Custom timeouts require GPU environment v5 and above.

The API builds on the standard PyTorch libraries: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), and DeepSpeed.

Migrate from serverless_gpu

The serverless_gpu package is deprecated and no longer receives new features. Starting in GPU environment version 6, the Databricks AI Runtime SDK (databricks-sdk-air) replaces it and provides equivalent APIs, with renamed arguments, under the databricks.air namespace. The package you use depends on the GPU environment version:

• Version 5 and below: Include only serverless_gpu. Continue to use serverless_gpu on these versions.
• Version 6: Includes both packages, but only the AI Runtime SDK receives updates. Version 6 is the last GPU environment version that includes serverless_gpu.

The AI Runtime SDK is a separate package from the Databricks SDK for Python (databricks-sdk).

	serverless_gpu (deprecated)	AI Runtime SDK (databricks.air)
GPU environment versions	6 and below	6 and above
Import	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
GPU count argument	gpus	num_accelerators
GPU type argument	gpu_type (GPUType)	accelerator_type (AcceleratorType)
GPU type values	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Data loading and checkpointing module	serverless_gpu.data	databricks.air.data. See Load data on AI Runtime.

To migrate a notebook:

1. Connect the notebook to GPU environment version 6 or above, which includes the AI Runtime SDK.
2. Change imports from serverless_gpu to databricks.air.
3. Rename gpus to num_accelerators and gpu_type to accelerator_type, and use the new accelerator type values. The old argument names and GPU type values still work in databricks.air but emit a FutureWarning. Support will end in GPU environment version 7.

Tip

Tell Genie Code to do this for you:

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

For Databricks CLI workloads that use databricks.air.data, add databricks-sdk-air[data] to the workload's environment.dependencies.

Coming from TorchDistributor

If you run distributed training using TorchDistributor on Spark clusters and your workload fits on a single node, the @distributed API is the recommended replacement on AI Runtime. Your DDP, FSDP, or DeepSpeed model code doesn't change.

Feature	@distributed API	TorchDistributor
Infrastructure	Fully serverless, no cluster management	Requires a Spark cluster with GPU workers
Setup	Single decorator, minimal configuration	Requires Spark cluster and TorchDistributor setup
Framework support	PyTorch DDP, FSDP, DeepSpeed	Primarily PyTorch DDP
Data loading	Inside the decorator, uses Unity Catalog volumes (UCVolumeDataset for streaming file data)	Via Spark or filesystem

To migrate a single-node workload:

• Replace the TorchDistributor(...).run(train_fn, ...) call with the @distributed decorator on train_fn, then launch with train_fn.distributed(...).
• Remove the Spark cluster and GPU worker configuration. Connect your notebook to an 8xH100 accelerator instead. @distributed uses all eight GPUs by default.
• Move data loading inside the decorated function. See Data loading.
• Keep your existing DDP, FSDP, or DeepSpeed model code. The decorator supports all three.

@distributed runs on a single node (see Limitations), so it does not replace every TorchDistributor workload. Keep workloads that depend on Spark integration on TorchDistributor. To run distributed training from your local machine or across multiple nodes, use the AI Runtime CLI instead, which is in Public Preview. See Use the Databricks CLI with AI Runtime.

Full example

The following example trains a multilayer perceptron (MLP) model on 8 H100 GPUs from a notebook.

1. Set up your model and define utility functions.

   
   # Define the model
   import os
   import torch
   import torch.distributed as dist
   import torch.nn as nn
   
   def setup():
       torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
       dist.init_process_group("nccl")
   
   def cleanup():
       dist.destroy_process_group()
   
   class SimpleMLP(nn.Module):
       def __init__(self, input_dim=10, hidden_dim=64, output_dim=1):
           super().__init__()
           self.net = nn.Sequential(
               nn.Linear(input_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, output_dim)
           )
   
       def forward(self, x):
           return self.net(x)
   

2. Import the distributed decorator from the AI Runtime SDK.

   from databricks.air import distributed
   

3. Wrap the model training code in a function and decorate the function with the @distributed decorator. The decorated function is the entrypoint for distributed execution, so define all training logic, data loading, and model initialization inside it.

   @distributed()
   def run_train(num_epochs: int, batch_size: int) -> None:
       import mlflow
       import torch.optim as optim
       from torch.nn.parallel import DistributedDataParallel as DDP
       from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
   
       # 1. Set up multi-GPU environment
       setup()
       device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
   
       # 2. Apply the Torch distributed data parallel (DDP) library for data-parallel training.
       model = SimpleMLP().to(device)
       model = DDP(model, device_ids=[device])
   
       # 3. Create and load dataset.
       x = torch.randn(5000, 10)
       y = torch.randn(5000, 1)
   
       dataset = TensorDataset(x, y)
       sampler = DistributedSampler(dataset)
       dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
   
       # 4. Define the training loop.
       optimizer = optim.Adam(model.parameters(), lr=0.001)
       loss_fn = nn.MSELoss()
   
       for epoch in range(num_epochs):
           sampler.set_epoch(epoch)
           model.train()
           total_loss = 0.0
           for step, (xb, yb) in enumerate(dataloader):
               xb, yb = xb.to(device), yb.to(device)
               optimizer.zero_grad()
               loss = loss_fn(model(xb), yb)
               # Log loss to MLflow metric
               mlflow.log_metric("loss", loss.item(), step=step)
   
               loss.backward()
               optimizer.step()
               total_loss += loss.item() * xb.size(0)
   
           mlflow.log_metric("total_loss", total_loss)
           print(f"Total loss for epoch {epoch}: {total_loss}")
   
       cleanup()
   

4. Run the distributed training by calling the distributed function with user-defined arguments.

   run_train.distributed(num_epochs=3, batch_size=1)
   

5. When executed, an MLflow run link is generated in the notebook cell output. Click the MLflow run link or find it in the Experiment panel to see the run results. For details on customizing experiment names, tracking metrics, and resuming runs, see Experiment tracking and observability.

Data loading

Place data loading code inside the @distributed function. A dataset can exceed the maximum size allowed by pickle, so generating or loading it inside the decorator avoids serialization errors:

from databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

For file-based data stored in Unity Catalog volumes, use UCVolumeDataset from databricks.air.data, which streams files with local caching and partitions them across ranks and workers automatically. To checkpoint distributed training to a volume, use UCVolumeWriter and UCVolumeReader. See Load data on AI Runtime and Checkpoint with Distributed Checkpoint (DCP).

Limitations

• @distributed supports only a single node: it runs across the GPUs on the node your notebook is connected to. To train across multiple nodes, use the AI Runtime Databricks CLI. See Use the Databricks CLI with AI Runtime.
• If you set accelerator_type, it must match the accelerator your notebook is connected to (for example, "GPU_8xH100"). A mismatch causes the workload to fail. The parameter is optional and auto-detected when omitted.
• AI Runtime supports GPU environment v4 and above. Custom timeouts (the timeout parameter) require GPU environment v5 and above.
• The decorator times out after 3 hours by default. Pass timeout in seconds to change it, or timeout=None to disable it.
• Execution runs within the lifecycle of the notebook. Terminating the notebook terminates the run.

Additional resources

• For the @distributed decorator, AcceleratorType, and Ray APIs, see the Databricks AI Runtime SDK API reference.
• For patterns that make your training pipeline more efficient and resilient, see the performance and resiliency guide.
• For end-to-end training scenarios, see distributed training examples.