Multi-GPU distributed training

Preview

The @distributed decorator used in the notebook examples is in Public Preview.

These examples scale model training across multiple GPUs and nodes on AI Runtime. The notebook examples use the serverless_gpu Python API on attached GPU compute. The Databricks CLI examples provide workload configurations and source code for single-node and multi-node training.

note

Notebook-based multi-GPU distributed training is supported on H100 GPUs.

Choose your parallelism technique

Choose a parallelism technique based on your model size, available GPU memory, and performance requirements. The following table compares the options.

Technique	When to use
DDP (Distributed Data Parallel)	Full model fits in single GPU memory; need to scale data throughput
FSDP (Fully Sharded Data Parallel)	Very large models that don't fit in single GPU memory
DeepSpeed ZeRO	Large models with advanced memory optimization needs

Technique	When to use
DDP (Distributed Data Parallel)	Full model fits in single GPU memory; need to scale data throughput
FSDP (Fully Sharded Data Parallel)	Very large models that don't fit in single GPU memory
DeepSpeed ZeRO	Large models with advanced memory optimization needs

For detailed information about each technique, see DDP, FSDP, and DeepSpeed.

For advice and patterns to make your training pipeline more efficient and resilient, see the performance and resiliency guide.

Notebook examples by technique and framework

The following table organizes notebook examples by the framework or library and the parallelism technique applied. Multiple notebooks might appear in a single cell.

Framework/Library	DDP examples	FSDP examples	DeepSpeed examples
PyTorch (native)	Simple MLP neural network RetinaNet image detection	10M parameter transformer	—
Huggingface TRL	Fine-tune Gpt OSS 20B	Fine-tune Gpt OSS 120B	Fine-tune Llama 3.2 1B
Unsloth	Fine-tune Llama 3.2 3B	—	—
Axolotl	Fine-tune Olmo3 7B	—	—
Lightning	Two-tower recommender system	—	—

Framework/Library	DDP examples	FSDP examples	DeepSpeed examples
PyTorch (native)	Simple MLP neural network RetinaNet image detection	10M parameter transformer	—
Huggingface TRL	Fine-tune Gpt OSS 20B	Fine-tune Gpt OSS 120B	Fine-tune Llama 3.2 1B
Unsloth	Fine-tune Llama 3.2 3B	—	—
Axolotl	Fine-tune Olmo3 7B	—	—
Lightning	Two-tower recommender system	—	—

CLI examples

The following examples link to runnable recipes in the AI Runtime Cookbook.

Example	Technique	Description
ResNet-50 fine-tuning with Ray Train	DDP	Fine-tune ResNet-50 with one Ray Train worker per A10 GPU across two nodes.
Qwen2.5-3B fine-tuning with Ray Train	DDP	Fine-tune Qwen2.5-3B with one Ray Train worker per GPU on an 8x H100 node.
Multi-node fine-tuning of Llama 3.1 8B	FSDP	Fine-tune Llama 3.1 8B with PyTorch FSDP across two 8x H100 nodes.
Multi-node fine-tuning of Qwen3.8-27B	FSDP2	Fine-tune Qwen3.8-27B with PyTorch FSDP2 across four 8x H100 nodes.
Qwen3-8B fine-tuning for invoice extraction	DeepSpeed ZeRO-3	Fine-tune Qwen3-8B for invoice and receipt entity extraction with TRL on one 8x H100 node.
ModernBERT fine-tuning with DDP	DDP	Fine-tune ModernBERT for topic classification with PyTorch DDP on an 8x H100 node.
LoRA fine-tuning of Nemotron 3.5 Lightning	FSDP2 and expert parallelism	Fine-tune Nemotron 3.5 Lightning with NeMo AutoModel across eight H100 GPUs.
LoRA fine-tuning of Nemotron 3.5 Super VL	FSDP2 and expert parallelism	Fine-tune the Nemotron 3.5 Super VL vision-language model with NeMo AutoModel across eight H100 GPUs.
Deep learning recommendation model with TorchRec	Embedding sharding	Train a click-through-rate model while TorchRec shards its embedding tables across four A10 nodes.

Example	Technique	Description
ResNet-50 fine-tuning with Ray Train	DDP	Fine-tune ResNet-50 with one Ray Train worker per A10 GPU across two nodes.
Qwen2.5-3B fine-tuning with Ray Train	DDP	Fine-tune Qwen2.5-3B with one Ray Train worker per GPU on an 8x H100 node.
Multi-node fine-tuning of Llama 3.1 8B	FSDP	Fine-tune Llama 3.1 8B with PyTorch FSDP across two 8x H100 nodes.
Multi-node fine-tuning of Qwen3.8-27B	FSDP2	Fine-tune Qwen3.8-27B with PyTorch FSDP2 across four 8x H100 nodes.
Qwen3-8B fine-tuning for invoice extraction	DeepSpeed ZeRO-3	Fine-tune Qwen3-8B for invoice and receipt entity extraction with TRL on one 8x H100 node.
ModernBERT fine-tuning with DDP	DDP	Fine-tune ModernBERT for topic classification with PyTorch DDP on an 8x H100 node.
LoRA fine-tuning of Nemotron 3.5 Lightning	FSDP2 and expert parallelism	Fine-tune Nemotron 3.5 Lightning with NeMo AutoModel across eight H100 GPUs.
LoRA fine-tuning of Nemotron 3.5 Super VL	FSDP2 and expert parallelism	Fine-tune the Nemotron 3.5 Super VL vision-language model with NeMo AutoModel across eight H100 GPUs.
Deep learning recommendation model with TorchRec	Embedding sharding	Train a click-through-rate model while TorchRec shards its embedding tables across four A10 nodes.

Get started in a notebook

Use the following tutorials to get started with the serverless GPU Python library for distributed training:

Tutorial	Description
AI Runtime with H100 GPUs	Learn how to use Databricks AI Runtime with H100 accelerators to run distributed GPU workloads using the serverless_gpu Python library.

Tutorial	Description
AI Runtime with H100 GPUs	Learn how to use Databricks AI Runtime with H100 accelerators to run distributed GPU workloads using the serverless_gpu Python library.