SS-HSDP splits runtime layout from optimizer units: up to 21.7% higher MFU at 128 GPUs with no memory penalty
Synopsis
In sharded data-parallel training, the packed low-precision layouts that compute kernels want and the update units that optimizers need have been tied together; SS-HSDP separates runtime parameters from update parameters and uses shared data views and hierarchical sharding to avoid redundant movement. On MXFP8 video and language-model workloads at 128 GPUs, MFU rises by up to 21.7% over the fastest baseline, with peak memory comparable to or lower than the lowest-memory baseline.
Interpretation
The training framework SS-HSDP splits parameters into runtime parameters for computation and update parameters for optimization, so computation layouts and optimization units can be organized independently. Earlier sharded data-parallel frameworks coupled these two choices, so supporting different update units often required reorganizing data or customizing the optimizer. On MXFP8 video and language-model workloads, SS-HSDP reaches higher MFU than FSDP2, veScale-FSDP and Megatron-FSDP; at 128 GPUs the largest gain over the fastest baseline is 21.7%, with peak memory comparable to or lower than the lowest-memory baseline.
Shared storage and hierarchical sharding connect the two representations, so packing and unpacking add no extra copies. A data tree lets tensors with different shapes and boundaries share one allocation, and a shard chain records where each level's shard sits in the original tensor. In the example of per-head Muon updates with an MXFP8 Q projection, each head's weights are consecutive views of one buffer, and the packed runtime weight is only a view of the root allocation, needing no copy or extra allocation.
The update phase exchanges data point-to-point only when a complete matrix is required, while main weights and momentum stay sharded throughout. Each update operation declares the smallest unit it needs, from a single element for elementwise work to a full matrix for matrix operations. In the four-rank example with ZeRO-1 degree four and ZeRO-2 and ZeRO-3 degree two, Head0's update spans ranks 2 and 1; rank 2 fetches the missing fragments, runs Newton–Schulz and returns the result.
Perspective
The demonstrated scope is MXFP8 training of video and language models from 32 to 128 GPUs with fixed per-GPU workloads and weight and gradient sharding degree fixed at 32; the ZeroSM communication sweep runs at 64 GPUs with sharding degree 64 under BF16 computation.
ZeroSM delivers 33.5% lower throughput than SM-based SS-HSDP at the smallest workload and gains 8.8% at the largest, so its advantage depends on workload size; BF16SR and FP32 differ by less than 0.05% in final-10k-step mean loss, but their ranking is not identical; code release is planned for the future.