Optimizers and ZeRO / FSDP sharding
Optimizer state bytes per parameter and how ZeRO stages and FSDP shard training state across GPUs.
On this page · 4 sections
For a full fine-tune, the weights are the smallest part of training memory. The FP32 master copy, the gradients and above all the optimizer state take up to seven times more. Two choices move that number: which optimizer keeps how many bytes per parameter, and whether that state is replicated on every GPU or sharded across them.
Optimizer state per parameter
| Optimizer | State kept | Bytes per trainable parameter |
|---|---|---|
| AdamW | Momentum and variance, both FP32 | 8 |
| SGD-momentum | Momentum, FP32 | 4 |
| Lion | Momentum, FP32 | 4 |
| 8-bit Adam | Two moments in 8 bits (bitsandbytes) | 2 |
The state exists only for trainable parameters. In a full fine-tune that is every weight; in LoRA and QLoRA it is the adapters only, so the optimizer choice barely moves a LoRA run. On a full fine-tune of Llama 3.1 8B, AdamW holds 64 GB of state; 8-bit Adam holds 16 GB and brings the run from about 132 GB to about 84 GB per GPU.
The full-fine-tune total with 8-bit Adam keeps the separate FP32 master copy (4 bytes per parameter). Some bitsandbytes BF16 setups skip it, so for those the estimate is about 4 bytes per parameter conservative — deliberately, in the out-of-memory-safe direction.
Sharding strategies
The Sharding control has five options. They differ in what each GPU keeps a full copy of:
| Strategy | Sharded across N GPUs | Replicated on every GPU |
|---|---|---|
| Single | — | Everything; one GPU |
| DDP | — | Everything, plus an all-reduce gradient bucket |
| ZeRO-1 | Optimizer state and the FP32 master copy | Weights, gradients, activations |
| ZeRO-2 | + gradients | Weights, activations |
| ZeRO-3 / FSDP | + weights and adapters | Activations and overhead only |
The FP32 master copy belongs to the optimizer partition in the ZeRO paper (arXiv:1910.02054), so it shards from ZeRO-1, not only at ZeRO-3. Activations and the per-GPU overhead never shard: they belong to the micro-batch each GPU processes.
per GPU = weights·(1/N if ZeRO-3) + master·(1/N if ZeRO-1+) + grads·(1/N if ZeRO-2+)
+ optimizer·(1/N if ZeRO-1+) + activations + 1.5 GB
weights 2 B · master 4 B · grads 2 B · AdamW 8 B per parameter
Who picks the GPU count
- Single evaluates one GPU of each card.
- DDP evaluates the GPU count you type (2 by default). Every GPU holds a full copy, so DDP adds throughput, not capacity — the calculator says so — and the gradient bucket it keeps live during the backward pass adds about one more copy of the gradients.
- ZeRO-1, ZeRO-2 and ZeRO-3 / FSDP pick the count for you: for each card they climb 1, 2, 4, 8, 16, 32, 64 GPUs and take the first count whose per-GPU memory fits within 95 %. PCIe cards stop at one node of 8; rack-scale systems start above 8.
Example: a full fine-tune of Llama 3.3 70B with AdamW and ZeRO-3 is about 146 GB per GPU at 8 GPUs — a fit for a 180 GB B200 node — and about 76 GB per GPU at 16 GPUs, which brings 94–96 GB cards into the list.
Choosing
- The model fits on one card with room to spare: Single, or DDP to train faster on several.
- Optimizer state is the problem (full fine-tune, mid-size model): ZeRO-1 or ZeRO-2 keeps weights local and communication light.
- The weights themselves do not fit: ZeRO-3 / FSDP, at the price of gathering weights every layer — more inter-GPU traffic, which the time estimate does not penalize (see time and cost to train).
- Before sharding a LoRA run, check that the frozen base is the problem; QLoRA often removes the need for sharding altogether.
CPU and NVMe offload (ZeRO-Offload, paged optimizers) would trade GPU memory for host RAM; they are not modeled.
- Method
- Full
- Optimizer
- AdamW
- Sharding
- Single
- Seq len
- 2,048