Multi-GPU training with DeepSpeed ZeRO-3

Use Soup CLI with DeepSpeed ZeRO-3 to train large models across multiple GPUs. This guide shows how to fine-tune a 70B model across 4–8 GPUs.

Install

bash
pip install "soup-cli[train,deepspeed]"

[deepspeed] adds ZeRO on top of the training stack; it does not pull it in on its own, because the base package is PyTorch-free.

When to use which ZeRO stage

StageWhat it shardsWhen to use
ZeRO-2Optimizer states + gradients2–4 GPUs, 7B–13B models
ZeRO-3Everything incl. parameters4+ GPUs, 30B+ models
FSDP2Fully sharded (PyTorch native)Alternative to ZeRO-3

Config for Llama 3.1 70B on 8× A100

yaml
base: meta-llama/Llama-3.1-70B-Instruct
task: sft
backend: transformers       # root-level, NOT under training

data:
  train: train.json
  format: alpaca
  max_length: 4096          # sequence length lives under data

training:
  epochs: 2
  lr: 1.0e-4                # the key is lr, not learning_rate
  batch_size: 1
  gradient_accumulation_steps: 16
  gradient_checkpointing: true
  auto_mixed_precision: true
  lora:
    r: 32                   # LoRA turns on when r > 0; there is no enabled flag
    alpha: 64

output: ./runs/llama70b

Sharding is not a config key. There is no training.distributed block. DeepSpeed and FSDP are selected on the command line, so the same YAML runs on one GPU or eight.

Launch on 8 GPUs

bash
soup train --config llama70b.yaml --gpus 8 --deepspeed zero3

--deepspeed accepts zero2, zero3, zero2_offload, zero3_offload, zero++, or a path to a DeepSpeed config JSON. zero3_offload (v0.73.0) is stage 3 with CPU parameter offload, which is what a run short of VRAM actually wants and which could not be named on the command line before: zero3 set no offload at all and the only offload preset was stage 2, optimizer-only. Measured on one H100 with Llama-3.1-8B in bf16: 21.65 tok/s at a 38,135 MB peak. Its offload_optimizer deliberately stays off, because turning it on makes DeepSpeed build a CPU Adam kernel against a matching CUDA toolkit and fail without one; copy the emitted JSON and flip it if you have nvcc. Note that DeepSpeed with LoRA was broken on every stage until v0.73.0, and the repair covers the supervised trainer only. Use zero2_offload when the optimizer state is what does not fit. Soup handles the torchrun / deepspeed launcher config automatically.

Alternative: FSDP2

bash
soup train --config llama70b.yaml --gpus 8 --fsdp full_shard

--fsdp accepts full_shard, shard_grad or full_offload. FSDP2 is PyTorch-native and often simpler for LoRA workloads.

Ring FlashAttention for 128k+ context

For very long sequences:

bash
pip install "soup-cli[ring-attn]"
yaml
data:
  max_length: 131072

training:
  use_ring_attention: true

Tips

  • Always enable gradient_checkpointing: true for 70B+ models
  • CPU offload trades speed for VRAM — use only if OOM
  • Profile first: soup profile --config llama70b.yaml estimates memory + throughput before you spend GPU hours

Soup is free and Apache-2.0. If it saved you a training run, starring the repo costs nothing and helps most. You can also fund the GPU time behind the work a 4 GB laptop cannot reach.