Your weights fit on one device but Adam's moments do not — how do you shard the optimizer state?

## Why the optimizer state is the thing that overflows A parameter in a mixed-precision training loop is not one number in memory, it is several. There is the half-precision copy the forward and backward passes actually use, its gradient, and then the optimizer's bookkeeping: Adam maintains an exponential moving average of the gradient (the first moment) and of the squared gradient (the second moment), both per parameter, and mixed-precision training also keeps a full-precision master copy of the weight so that tiny updates are not lost to rounding. Those three full-precision buffers are the bulk. Roughly speaking the optimizer-related state is several times the size of the weight copy itself — which is why a model whose weights would comfortably fit can still fail to train on the same device. The second moment is what makes Adam Adam: the update divides the (bias-corrected) first moment by the square root of the (bias-corrected) second moment plus a small constant, giving each parameter its own effective step size. You cannot drop either buffer without changing the algorithm. ## The redundancy to exploit Run the same model on `N` devices with different data on each, and the gradients are combined so that every device ends the backward pass holding the *same* averaged gradient. Every device then feeds that same gradient into the same optimizer holding the same moments, and computes bit-for-bit the same update. `N` devices are doing one device's worth of work and storing `N` copies of one device's worth of state. That duplication is pure waste, and it is the thing to remove. ## The scheme Partition the parameter vector into `N` disjoint slices and assign slice `i` to device `i`. Then: 1. Each device runs the full forward and backward pass on its own data and produces a full gradient — nothing about the model changes. 2. Instead of an all-reduce, do a **reduce-scatter** of the gradients: device `i` ends up with the averaged gradient for *only* its own slice, and no one holds the whole averaged gradient. 3. Device `i` holds the first moment, second moment and master weights for its slice alone, and applies the optimizer update to that slice. 4. **All-gather** the updated half-precision weights so every device again has the complete model for the next forward pass. Optimizer-state memory per device falls by a factor of `N`, exactly. ## Why the communication is nearly free A ring all-reduce is not a primitive — it is implemented as a reduce-scatter followed by an all-gather, each moving roughly one parameter-vector's worth of data per device. The scheme above uses precisely those two collectives, with the optimizer step inserted between them. The volume on the wire is essentially the same as the replicated version; what changes is that the all-gather now carries updated weights rather than reduced gradients. This is the reason this technique is close to a free win, and it is why it should usually be the first thing you try when memory is short. ## What it is not This is **not** model parallelism. No weight matrix is split, no layer is assigned to a particular device, and every device still executes the entire network on every step. What is partitioned is the optimizer's private bookkeeping, plus, momentarily, the gradient. A full copy of the weights is still materialised on every device before the forward pass. So the technique relieves state pressure and nothing else — if a single layer's weights, or the activations one device must retain, still do not fit, sharding the moments will not save you and you have to split the model itself. The idea does extend further along the same axis. You can also keep the gradients partitioned rather than materialising a full gradient buffer, and, more aggressively, keep the *weights* partitioned too and gather each layer's parameters only just before that layer is used, discarding them immediately afterwards. That last step removes almost all of the per-device weight copy, but it adds a collective for every layer in both the forward and the backward pass — a real memory-for-bandwidth trade rather than a free one. ## Operational sharp edges **Gradient clipping.** Clipping by global norm needs the norm over *all* parameters, but no device holds them all. Each device computes the sum of squares over its own slice and those scalars are all-reduced before the square root is taken and the scale applied. Skip that and each device clips against its own partial norm, which silently changes the effective threshold and can quietly destabilise training. **Checkpointing.** The state now lives in `N` pieces. Either gather it on save to produce a device-count-independent checkpoint, or accept a sharded format and pay a resharding step whenever the device count changes. **Partition granularity.** Slicing at arbitrary byte offsets splits individual tensors across owners, which complicates any per-tensor logic; partitioning whole tensors is simpler but can leave the load uneven when one tensor dominates the parameter count. **Determinism.** The updated weights now arrive through a collective, so reduction order affects the last bits. That is normally harmless, but it makes bitwise run-to-run reproducibility harder to guarantee.