Distributed training in notebooks

Important

This feature is in Beta. Workspace admins can control access to this feature from the Previews page. See Manage Azure Databricks previews.

The @distributed decorator from the Serverless GPU Python API 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.

Tip

  • 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 and set gpus=8 for full multi-GPU training.

Quickstart

The serverless_gpu package is preinstalled when your notebook is connected to a serverless GPU. Decorate your training function with @distributed, then call it with .distributed():

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
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()

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 notebook 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 of the gpus you request. 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.

Coming from TorchDistributor

If you run distributed PyTorch on Spark today with TorchDistributor and your workload fits on a single node, the serverless_gpu @distributed API is the recommended replacement for new deep learning workloads. It removes the Spark cluster and gives you the same code path from single-GPU to multi-GPU.

Feature serverless_gpu @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 and set gpus=8 instead.
  • 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 AI Runtime CLI.

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 serverless_gpu library and the distributed module.

    import serverless_gpu
from serverless_gpu 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(gpus=8, gpu_type='H100')
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-parellel 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 serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
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 serverless_gpu.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 Model checkpointing.

Limitations

  • Distributed training runs across the GPUs on the single node your notebook is connected to. For full multi-GPU training, connect to an 8xH100 accelerator, which provisions one node with 8 GPUs, and set gpus=8.
  • Accelerator type must match. If you set gpu_type in @distributed, it must match the accelerator your notebook is connected to ("H100" or "A10"). A mismatch causes the workload to fail. The parameter is optional and auto-detected when omitted.
  • AI Runtime recommends 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.

Learn more

  • Last updated on