Fill in the DeepSpeed config keys that select a ZeRO stage

**DeepSpeed** is a library that adds memory-saving training tricks on top of PyTorch. The most-used trick is called **ZeRO**, which shards the training state across multiple GPUs so each GPU holds only a slice. The knob that controls how aggressively ZeRO shards lives in the DeepSpeed JSON config under the key `zero_optimization.stage`. You set its value to 1, 2, or 3 depending on how much memory you need to save. Stage 1 shards the optimizer state across the GPUs in your data-parallel group. This is the biggest single saving for Adam, which keeps two fp32 moments per parameter. Stage 2 also shards the gradients. Each GPU holds only the gradient slices that match its optimizer state slice, then all-gathers what it needs for the backward pass. Stage 3 also shards the parameters themselves. The model weights are split across GPUs, and the missing slices are pulled in on demand during the forward pass. Each step up saves more memory per GPU. The cost is more communication between GPUs, because the sharded pieces have to be gathered together when the math needs them. On fast NVLink interconnects this is cheap; on slower PCIe or network links it can bottleneck training.

**The selector key is `zero_optimization.stage`** with values 1, 2, or 3. The DeepSpeed JSON config nests it as `{"zero_optimization": {"stage": 2, ...}}`. Each stage shards more of the training state across the data-parallel ranks at the cost of more cross-GPU communication. Stage 1 shards optimizer state only. For Adam this is the dominant memory term, 8 bytes per parameter for the two fp32 moments. Each rank owns its slice of the moments and applies its slice of the optimizer step. Stage 2 adds gradient sharding. Each rank holds only the gradients that match its optimizer slice, reducing peak training memory further. The reduce-scatter operation that produces these sharded gradients overlaps with the backward pass. Stage 3 adds parameter sharding. The model weights themselves are split, and a parameter is gathered just before any layer that needs it, then dropped after. This is what makes 70B-class training viable on commodity 80 GB GPUs. The practical pick depends on memory pressure and interconnect speed. Stage 2 is the common default for mid-size fine-tuning where it fits. Stage 3 unlocks the largest models but pays the most communication tax, so it benefits from NVLink or NVSwitch fabrics. On slower interconnects, stage 3 can be slower than stage 2 even when memory allows it.