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. Since v0.74.0 your own file is resolved the way a preset is, and only for keys that are provably invalid for the run: zero_hpz_partition_size when the world size is not divisible by it, and the fp16 quantiser keys against a bf16 run. That matters because the documented way to customise ZeRO++ is to copy the preset JSON, which copies both defects, so an unresolved user file inherited a crash the presets are already protected from. A repair is printed and written to a temp copy: your file on disk is never modified, a config using none of those keys is passed through byte-identical, and malformed JSON is passed through untouched because DeepSpeed reports a bad config better than a traceback would. 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 MiB 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, which repaired the supervised trainer; v0.74.0 then attached the same repair to every other DeepSpeed-capable trainer. 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

Not applied at v0.75.2. training.use_ring_attention: true is validated when soup train starts, and a failed check exits 1, but no trainer applies ring attention afterwards: the run proceeds with ordinary attention, so the flag does not split a long sequence across GPUs and saves no memory.

The checks are CUDA, two or more GPUs, ring-flash-attn importable (that is all the [ring-attn] extra is needed for) and data.max_length of at least 8192:

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

training:
  use_ring_attention: true   # validated at launch; nothing applies it at v0.75.2

For long context today, gradient_checkpointing, --deepspeed zero3 and --fsdp full_shard reduce the memory each GPU holds (recomputed activations, sharded parameters and optimizer state); none of them splits one sequence across GPUs.

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.