Prefetching: hide communication only inside a live-memory budget

Draw a timeline for persistent shards, current full parameters, next full parameters, unreleased gradients, activations, workspaces, and communication buffers. Count shared storage once: two views into one allocation do not consume two complete allocations, while two independent buffers do. Use measured allocation lifetimes to decide whether the persistent shard remains alongside a full buffer or whether storage is reused. A simple sum is meaningful only when the listed allocations are disjoint and simultaneously live. Otherwise it can either double-count aliases or miss a short but decisive overlap.

Reserved allocator memory and live tensor memory are also different observations. Fragmentation, caches, asynchronous release, and library workspaces can make a process fail below a naive sum of named tensors. Leave a measured margin and record the exact memory statistic used. Activations change with sequence length, microbatch size, and rematerialization policy; a parameter-only budget does not capture that variation. Test a representative worst-case input and the complete update lifecycle, including first-iteration initialization and checkpoint operations if they are part of the claimed capacity envelope.

The pinned FSDP wrapper’s BACKWARD_PRE fetches the next parameter set before the current gradient computation, increasing overlap opportunity while potentially holding current parameters, next parameters, and current gradients together. BACKWARD_POST changes that issue and release order to reduce the simultaneous parameter footprint with less overlap. Its forward-prefetch option follows the learned execution order and is documented for static-graph cases. These are schedule contracts, not interchangeable “faster” switches; dynamic execution and other wrappers require their own analysis.

The limit_all_gathers option throttles CPU issue so that parameter all-gather allocation stays bounded to the relevant consecutive instances in its documented schedule. A deliberate CPU gap can therefore be a memory protection mechanism rather than an idle-GPU bug. It does not bound activations, optimizer state, workspaces, or every allocation in the process. Removing it can issue work earlier while causing an out-of-memory failure. Measure device execution and the actual critical path instead of treating the absence of a CPU enqueue gap as proof of improvement.

Change one prefetch or wrapping decision while holding the objective and input workload fixed. Check parameter availability and update equivalence, then compare step time, communication completion, peak live memory, and reserved memory. An apparent speedup from a short run can disappear under long sequences or an accumulation window that retains full gradients. Include recovery after a failed allocation in the operational test plan, but do not continue a partially failed distributed step as though the optimizer state were known consistent. Memory safety and collective consistency remain constraints on any throughput experiment.

The queue lab offers only an average outstanding-work analogy. Reset, explicitly set arrival four and service three, predict growth from four items, then Play or Step. Compare a separate Reset with arrival two and service three explicitly restored before predicting drainage. Edits retain elapsed time and recompute the initial-value formula. Items have no byte sizes or lifetimes, and the model contains no FSDP rate limiter, all-gather, stream, or allocator. A live-allocation timeline, not this analogy, establishes the actual prefetch memory peak.