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. 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
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:
pip install "soup-cli[ring-attn]"data:
max_length: 131072
training:
use_ring_attention: trueTips
- 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.