SS-HSDP把运行时布局与优化单元拆开:128块GPU上MFU最高提升21.7%,峰值显存不增
核心概要
分片数据并行训练里,计算内核要的打包低精度布局和优化器要的更新单元长期被绑在一起,SS-HSDP把运行时参数与更新参数拆成两套表示,用共享数据视图和分层分片避免多余搬运。在视频与语言模型的MXFP8训练中,128块GPU上相对最快基线MFU最高提升21.7%,峰值显存与最低显存基线相当或更低。
深度剖析
训练框架SS-HSDP把用于计算的运行时参数和用于优化的更新参数分成两套表示,使计算布局与优化单元可以各自组织。 此前的分片数据并行框架把这两项选择耦合在一起,支持不同更新单元往往要重排数据或定制优化器。 在视频与语言模型的MXFP8负载上,SS-HSDP的MFU高于FSDP2、veScale-FSDP和Megatron-FSDP;128块GPU上相对最快基线最高提升21.7%,峰值显存与最低显存基线相当或更低。
两套表示靠共享存储和分层分片连起来,打包与拆包不产生额外拷贝。 数据树让形状和边界不同的张量共享同一块存储,分片链记录各层分片在原张量中的位置。 以逐注意力头的Muon更新加MXFP8的Q投影为例,各头权重是同一缓冲区的连续视图,打包后的运行时权重只是根分配的视图,无需拷贝或额外分配。
更新阶段只在需要完整矩阵时做点对点交换,主权重和动量始终保持分片。 每个更新操作声明所需的最小单元,从逐元素的一个元素到矩阵运算的完整矩阵。 示例中四个rank、ZeRO-1度为四、ZeRO-2与ZeRO-3度为二,Head0的更新跨rank 2和rank 1,rank 2取回缺失片段执行Newton–Schulz后把结果送回。
启示与展望
演示范围是视频与语言模型的MXFP8训练,GPU数从32到128,每GPU负载固定,权重与梯度分片度固定为32;ZeroSM通信的负载扫描在64块GPU、分片度64、BF16计算下进行。
ZeroSM在最小负载上吞吐比基于SM的SS-HSDP低33.5%,在最大负载上高8.8%,优势依赖负载规模;BF16SR与FP32的最终一万步平均损失差小于0.05%,但排序并非完全一致;代码计划后续发布。
