Skip to content

Distributed Training (PyTorch DDP)

Run large-scale distributed training with PyTorch DDP or DeepSpeed on a multi-GPU bare-metal instance.

Recommended hardware

Instance type overview

Spheron GPU offerings are classified by two criteria: interruptibility and hardware isolation.

All Spot instances are VM-based and can be reclaimed by the provider at any time. Use Spot only for fault-tolerant jobs with checkpointing. Dedicated instances carry a 99.95% SLA and are not reclaimed after deployment.

Within Dedicated, two hardware isolation options are available:

  • VM: Runs in an isolated virtual machine on shared physical hardware. The default across most providers and GPU offers.
  • Bare Metal: Full physical server with no hypervisor, no shared tenants. GPU count varies by offer and provider, from single-GPU up to multi-GPU servers. On the dashboard, identified by the BAREMETAL suffix in the GPU type name.

For multi-GPU distributed training, use a Dedicated Bare Metal offer with multiple GPUs on a single host. Bare metal removes the hypervisor layer and gives training processes direct access to every GPU and to the interconnect between them, which is what makes gradient synchronization efficient.

Offer: Look for the BAREMETAL suffix in the GPU type name and a GPU count of 4 or 8

GPU form factor and interconnect

Gradient synchronization speed depends on how the GPUs on the host talk to each other, which is determined by the GPU form factor in the offer you select:

Form factorIntra-node interconnectNotes
SXM (B200 SXM6, H200 SXM5, H100 SXM5)NVLink / NVSwitchHighest GPU-to-GPU bandwidth; optimal for all-reduce-heavy DDP and ZeRO-3 runs
PCIe (H100 PCIE, A100 PCIE)PCIe lanesLower cost; sufficient when gradient synchronization is not the bottleneck

Choose an SXM offer for large model training where gradient synchronization dominates step time. PCIe offers are adequate for smaller models or when cost is the priority.

SXM multi-GPU offers are available on Spheron ES (B200 SXM6, H200 SXM5, H100 SXM5), Spheron AI (H200 SXM, H100 SXM5), and Verda (H100). See Regions and Providers for current inventory per provider.

Deploy the instance

Deploy a multi-GPU bare-metal instance from the dashboard. On the Deploy GPUs page, select a Dedicated offer with the BAREMETAL suffix and the GPU count your run needs. Choose Ubuntu 22.04 as the operating system and attach your SSH key.

Running distributed training with torchrun

Once SSH'd into the instance, launch your training script with torchrun:

torchrun \
  --nproc_per_node=8 \
  --nnodes=1 \
  train.py \
  --batch_size 32 \
  --gradient_checkpointing

--nproc_per_node=8 uses all 8 H100 GPUs. For a 4-GPU offer, use --nproc_per_node=4.

PyTorch DDP training script

Minimal example of a DDP-compatible training loop:

import argparse
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
 
def setup():
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
def cleanup():
    dist.destroy_process_group()
 
def train():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--gradient_checkpointing", action="store_true")
    args = parser.parse_args()
 
    setup()
    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
 
    model = YourModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])
 
    # Enable gradient checkpointing to reduce VRAM usage
    if args.gradient_checkpointing:
        model.module.gradient_checkpointing_enable()
 
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
 
    # DistributedSampler ensures each worker sees a disjoint shard of the data
    dataset = YourDataset()  # replace with your dataset
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler)
 
    num_epochs = 3
    for epoch in range(num_epochs):
        # Reshuffle the dataset differently for each epoch across all workers
        sampler.set_epoch(epoch)
 
        for step, batch in enumerate(dataloader):
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                loss = model(**batch).loss
 
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
 
            # Save checkpoint every 100 steps
            if step % 100 == 0 and rank == 0:
                torch.save({
                    'step': step,
                    'epoch': epoch,
                    'model_state_dict': model.module.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                }, f'/checkpoints/checkpoint_epoch{epoch}_step{step}.pt')
 
    cleanup()
 
if __name__ == '__main__':
    train()

DeepSpeed ZeRO-3 for models >30B

For models too large to fit in a single GPU's memory, use DeepSpeed ZeRO-3 to shard parameters, gradients, and optimizer states across all GPUs.

ds_config.json:

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": { "device": "cpu", "pin_memory": true },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8,
    "stage3_prefetch_bucket_size": 5e7,
    "stage3_param_persistence_threshold": 1e6
  },
  "bf16": { "enabled": true },
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": true
  },
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 8
}

Launch with DeepSpeed:

deepspeed --num_gpus=8 train.py \
  --deepspeed ds_config.json \
  --model_name_or_path meta-llama/Meta-Llama-3.1-70B

Mixed precision (BF16)

H100, H200, and B200 GPUs have native BF16 support. Always use BF16 for training on these GPUs; it is faster and more numerically stable than FP16:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    outputs = model(**batch)

Checkpoint persistence

Mount a persistent volume at /checkpoints before your training run to protect checkpoints across deployments:

  1. Create a volume: see Volume Mounting
  2. Mount it at /checkpoints in your cloud-init script
  3. Save checkpoints to /checkpoints/ in your training loop (example above)

Dataset storage

For large datasets, stage the data on the instance's local NVMe disk rather than reading it from a network volume during training. Local disks deliver much higher I/O bandwidth than NFS or virtiofs volumes. Keep the authoritative copy on a persistent volume and copy the working set to local disk at the start of the run.

GPU monitoring

Watch per-GPU utilization during training:

nvidia-smi dmon -s u

Check NVLink health and bandwidth (SXM offers only):

nvidia-smi nvlink --status
nvidia-smi nvlink --capabilities

Monitor GPU memory:

nvidia-smi --query-gpu=memory.used,memory.free --format=csv -l 1

What's next