Gradient sharding: reduce globally and retain the owner slice

Start with one logical gradient vector per worker evaluated at compatible parameters. Split corresponding vectors into compatible slices and assign each output slice to an owner. A SUM reduce-scatter adds corresponding input coordinates and returns only the assigned reduced slice at each rank. If the intended aggregate is a worker average, a separate division by the participating group size is needed somewhere in the controlled protocol. A framework may implement that division differently for numerical or scheduling reasons; never add a manual division without checking what the wrapper already does.

ZeRO stage two partitions reduced gradient storage along with optimizer-owned state while the simple stage-two compute-parameter view remains replicated. A full local backward gradient or communication bucket can still exist transiently before reduction and release. Therefore the persistent gradient-shard size does not prove that every instant uses only that much memory. Track when a slice becomes ready, when communication completes, when its owner updates state, and when replicated parameter views become consistent. Ownership and buffer lifetime are correctness properties as well as memory-budget inputs.

Clipping each shard independently generally changes the update because a shard norm is not the whole-model norm. For disjoint owner slices of one already reduced gradient, sum their squared norms across owners and take one square root. Compute a common clipping coefficient and apply it consistently. Do not first average those squared norms merely because the operation uses an all-reduce: the mathematical norm requires a sum. Replicated or tied coordinates also need careful counting so that the logical vector is not included twice.

The same distinction appears in logging and optimizer diagnostics. A small local norm can coexist with a large global norm, and averaging local norm values is not equivalent to taking the square root of summed squares. Mixed precision adds another ordering constraint: evaluate the intended unscaled gradient, account for overflow, and avoid applying an owner update before the shared clipping decision is available. State whether the comparison uses exact arithmetic or a tolerance for the reduction order and precision actually chosen.

Build a small fixture with distinct coordinate values and uneven parameter shapes. Reconstruct the logical reduced gradient from owner slices and compare it with an unsharded reference before testing the optimizer. Include clipping, zero gradients, tied parameters if supported, and a partially filled final bucket. A successful reduce-scatter on the wrong slice layout can produce plausible numbers in the wrong parameter locations. Preserve a canonical coordinate map and check both values and identities rather than compare only one aggregate norm.

The function lab isolates the squared-contribution arithmetic. Reset, explicitly set coefficient one and position two, predict y, then Play or Step. For an independent run Reset and explicitly set coefficient one and position one before predicting again. The results four and one are the squared contributions of a two-coordinate example; their sum is five. The instrument computes one parabola at a time. It does not contain shards, reduce-scatter, clipping, distributed norm reduction, or optimizer ownership, so those protocols must be checked with the actual fixture.