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
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
| Stage | What it shards | When to use |
|---|---|---|
| ZeRO-2 | Optimizer states + gradients | 2–4 GPUs, 7B–13B models |
| ZeRO-3 | Everything incl. parameters | 4+ GPUs, 30B+ models |
| FSDP2 | Fully sharded (PyTorch native) | Alternative to ZeRO-3 |
Config for Llama 3.1 70B on 8× A100
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/llama70bSharding is not a config key. There is no
training.distributedblock. DeepSpeed and FSDP are selected on the command line, so the same YAML runs on one GPU or eight.
Launch on 8 GPUs
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
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: trueis validated whensoup trainstarts, 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:
pip install "soup-cli[ring-attn]"data:
max_length: 131072
training:
use_ring_attention: true # validated at launch; nothing applies it at v0.75.2For 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: truefor 70B+ models - CPU offload trades speed for VRAM — use only if OOM
- Profile first:
soup profile --config llama70b.yamlestimates memory + throughput before you spend GPU hours
Related
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.