Optimizer-state sharding: assign update ownership explicitly

Separate the low-precision model parameters used for compute, the stored gradients, any full-precision master parameter copy, and optimizer moments. Whether these allocations exist, share storage, or change precision depends on the implementation. An illustrative mixed-precision Adam-style layout uses two bytes for the compute parameter, two for its gradient, and twelve for a four-byte master parameter plus two four-byte moments. Activations, communication buffers, allocator reservations, and temporary kernels are not included. A different layout needs a different inventory rather than the same memorized multiplier.

Stage-one reasoning partitions optimizer-owned state while leaving the compute parameters and, in the basic model here, gradients replicated. Each owner receives the intended synchronized gradient for its assigned parameters and advances the associated optimizer state. Updated parameters then become available consistently to all replicas before their next use. PyTorch’s documented ZeroRedundancyOptimizer broadcasts owner-updated parameters; that specific mechanism should not be confused with a claim that every sharded optimizer uses the same communication schedule or splits every tensor into equal pieces.

A partition map associates each parameter or parameter slice with the right state, hyperparameter group, and checkpoint identity. The map cannot be reconstructed safely by assuming that rank order equals module traversal order. PyTorch’s pinned ZeroRedundancyOptimizer description uses a sorted-greedy packing of whole parameters, so one large tensor can limit balance; the round-robin scalar map used below is only a teaching model. Tied parameters and parameter groups need deliberate handling so that one logical state is neither duplicated inconsistently nor updated twice.

State sharding is mathematically straightforward for coordinatewise optimizer state only after all global operations are specified. Global gradient clipping requires the norm of the intended full gradient, and any optimizer rule coupling coordinates needs its corresponding distributed operation. Learning-rate schedules, step counters, overflow decisions, and decoupled weight decay must also match the reference update. A state-memory reduction does not authorize different update frequency. Inspect multiple steps, since freshly initialized moments can hide a missing-state restoration bug on the first step.

After local owners update their state, future computation needs the correct updated parameter values. Communication and any required synchronization therefore move rather than disappear. Checkpoint operations may consolidate or repartition state and can have very different peak memory from steady training; all required ranks must participate according to the API. Construct the optimizer against the wrapper’s supported parameter representation and retain a canonical logical mapping for tests and restoration. A checkpoint containing only one owner’s local state is not automatically a complete training checkpoint.

The modular lab makes one ownership-enumeration pitfall visible. Reset, explicitly set modulus four and additive step one, predict the residues over four Steps, then Step four times to visit 1,2,3,0 from the initial zero. For a fresh comparison Reset and explicitly set modulus four and step two before predicting and stepping: only zero and two are visited. This deterministic orbit is not a sharded optimizer; it illustrates how a stride that is not coprime can fail to enumerate every owner in a proposed toy traversal.