Parameter sharding: materialize only what the next operation needs

In the scoped FULL_SHARD lifecycle, parameter storage is sharded between uses. An all-gather unshards the relevant managed unit before its forward computation, and resharding follows forward. The unit is unsharded again before its backward computation and reshards afterward; gradients are synchronized and sharded through reduce-scatter, and owners update their optimizer state locally. These operations have stream and completion dependencies. Calling a submodule directly while bypassing its managing wrapper can violate the assumption that its needed parameters have first been materialized.

The managed unit matters. Wrapping the whole model as one unit can require a much larger full-parameter materialization than wrapping suitable submodules. Conversely, excessively small units can create many communication launches and harm overlap. DeepSpeed’s documented stage-three assumptions require the relevant submodule forward and backward to fit, with supported techniques needed when they do not. Parameter sharding does not divide activations automatically, and a large intermediate or one indivisible operator can still exceed device memory even when the persistent parameter shard is small.

With equal shards of a unit, each rank already holds one fraction and receives the remaining fractions to reconstruct that unit’s logical parameter data. This is a data-volume model, not an algorithm-independent wall-clock formula. Topology, protocol, dtype, padding, overlapping collectives, and reuse policy change time and actual traffic. Distinguish bytes received by one rank from bytes sent, aggregate fabric traffic, and the size of the final full buffer. Confusing those counters can produce an apparently consistent but incorrect communication budget.

Do not retain arbitrary references to temporary full-parameter views and assume they remain valid after resharding. The wrapper can replace views during forward and backward. Shared or externally accessed parameters require the framework’s supported coordination so that the right data is gathered and ownership is preserved. Build optimizers against the supported wrapped representation and preserve canonical parameter identity in checkpoints. A logical parameter existing in the model definition does not mean every rank has its complete current data resident at every instant.

Trace one managed unit through forward, backward, reduction, update, and the next iteration. Record which representation is full, which is sharded, and which stream still uses each allocation. A strategy that keeps full parameters after forward can reduce the need for another gather but has a different memory schedule from FULL_SHARD. Reentrant computation or activation rematerialization can also change how many materializations occur. Count the actual events rather than transfer a two-gather illustration to every model and configuration.

Use the cache lab only as a reuse analogy. Reset, explicitly set capacity two and stride four, predict the outcome after one Step, then Step once. A separate Reset with capacity four and stride four gives a different reuse outcome after exactly one Step: 17 misses versus four misses and 13 hits. The toy starts empty and traces a four-address cycle over 17 requests. It does not model all-gather, managed parameter units, GPU allocations, or a legal FSDP retention policy; real schedule evidence determines whether reuse is possible and affordable.