SS-HSDP: Efficient and Flexible Sharded Data Parallelism via Runtime-Update Decoupling
Chendong Xiang, Pengle Zhang, Jinliang Li, Zichen Liang, Jintao Zhang, Fan Bao, Jun Zhu, Jianfei Chen
原文摘要
Sharded data parallelism enables memory-efficient training of large models. While efficient training kernels favor packed, kernel-specific data layouts, optimizers may require variant data structures as optimization units. Existing sharded data-parallel frameworks couple these two choices, limiting runtime efficiency and optimization flexibility. To address this, we present \textbf{\our}, a training framework that separates runtime parameters from update parameters, allowing computation layouts and optimization units to be organized independently. Through shared data views and hierarchical sharding, \our supports this separation while avoiding redundant data movement. This design supports a broad range of optimizers, including matrix-based methods such as Muon, and performance optimizations such as fully zero-copy execution. Experiments on video and language model workloads with MXFP8 training show that \our achieves higher Model FLOPs Utilization (MFU) than FSDP2, veScale-FSDP, and Megatron-FSDP. At 128 GPUs, MFU improves by up to 21.7\% over the fastest evaluated baseline, with comparable or lower peak memory usage.
cs.DC