Course 4b: Training Framework¶
Paper coverage: Section 3.4 Four places where the training machine had to be reshaped because the architecture broke a standard assumption.
Each subsection follows the same shape: what the standard tool assumes → what V4's design violates → the fix.
Evidence tags: PAPER stated by DeepSeek · DERIVED mathematically implied · IMPLEMENTED reproduced in this project · MEASURED observed in our experiments · INFERRED our interpretation · UNKNOWN not established. Full definitions: evidence ledger
1 · Muon under ZeRO: when updates stop being element-wise¶
ZeRO sharding works because AdamW's update for parameter i depends only on gradient statistics for element i: split tensors anywhere, update independently, results identical.
Muon breaks that contract. Its update is a matrix-wide object, the Newton-Schulz orthogonalization mixes every row and column. Slice a matrix across ranks and there is no longer "an update" for any piece of it.
V4's answers, per parameter class:
| class | strategy |
|---|---|
| dense matrices | cap ZeRO width; knapsack-pack whole matrices into ≤5 per rank (<10% padding); beyond the data-parallel limit, recompute redundantly rather than shard |
| MoE expert matrices | flatten all down-projections across experts/layers (then up, then gate) so vectors split evenly without cutting any single matrix |
| same-shape neighbors | merge consecutive identical shapes → one batched Newton-Schulz call |
Two numerics choices complete the picture: NS iterations run in BF16 (stable, verified), and cross-rank gradient sync uses BF16 stochastic rounding through all-to-all followed by FP32 local sums, avoiding the accumulation error a low-precision ring reduce would add [PAPER Section 3.4.1].
2 · mHC made affordable: the widened stream has a storage bill¶
Hyper-connections multiply residual-lane activations by n_hc=4 and add pipeline communication. Left alone, that is a memory tax on every layer. Three reductions bring measured wall-time overhead to 6.7%:
- fused kernels for the lane mixing;
- selective recomputation: recompute cheap things (hidden states, normalized inputs), never recompute expensive things (the F-layer itself); checkpointing here is a budgeting decision;
- adjusted DualPipe scheduling so mHC's extra pipeline traffic overlaps.
The transferable idea: activation memory is not uniform: annotate which tensors are worth keeping, not just which modules to checkpoint (see 4 below).
3 · Contextual parallelism: compression vs sequence partitioning¶
Context parallelism splits one long sequence across ranks. Compression fights it twice: an m-token pooling window can straddle a rank boundary, and pooled lengths vary with packing. V4's two-stage protocol makes both problems disappear:
- rank i ships its last m raw entries to rank i+1; each rank then compresses locally to a fixed s/m+1 length (padding absorbed locally);
- all-gather + fused select-and-pad reassemble exactly cp_size·s/m entries globally.
Causality survives because visibility rules are precomputed per query (HCA/indexer) or carried explicitly by top-k indices (CSA sparse).
4 · Tensor-level activation checkpointing¶
Standard checkpointing chooses whole modules: keep everything or recompute everything. V4 annotates individual tensors; TorchFX then traces backward from each annotated tensor to find its minimal recomputation subgraph, inserts it before the needing gradient, reuses storage pointers (no copies), and deduplicates shared-storage tensors automatically.
The programming-model win matters as much as the memory win: developers declare what to forget, not how to recompute it.