1M上下文训练狂飙7.59倍！块并行让扩散语言模型长文本训练不再卡通信
Block Parallelism For Efficient Distributed Long-Context Diffusion Language Model Training
📖 AI 总结
本文针对块扩散语言模型（BDLM）长上下文训练中分布式注意力通信开销大、激活内存受限的问题，提出了一种新的分布式并行维度——块并行（BP）。作者观察到BDLM的训练目标在各目标块之间是可分离的，因此BP将每个受损块的计算分配到单个rank上，使受损K/V及其梯度保持本地化。为扩展到更长上下文，进一步提出上下文分片块并行（CSBP），在BP基础上将共享的干净序列也分片到各rank，避免复制干净前缀，同时保持BDLM训练语义不变。实验表明，在16块H200 GPU、256K上下文下，CSBP相比最优基线在监督微调上吞吐提升1.18至1.45倍，在自回归模型转BDLM任务上提升1.27至1.33倍，峰值显存持平或更低；512K上下文下全模型加速达1.61倍。在8块H100上，CSBP使DFlash2投机解码器训练在512K和1M上下文分别加速2.48倍和7.59倍。在12小时DiffusionGemma 26B-A4B微调对比中，CSBP在各检查点的SWE-bench Verified和Terminal-Bench Lite通过率均更高。该工作为长上下文扩散语言模型的高效分布式训练提供了新的并行范式。
🔑 关键词速览
| Block Diffusion Language Models (BDLMs) | 块扩散语言模型，一种结合自回归跨块依赖与块内并行去噪的生成模型。 |
| Context Parallelism (CP) | 上下文并行，一种按位置将序列分片到不同设备以处理长上下文的分布式训练技术。 |
| Block Parallelism (BP) | 块并行，本文提出的新并行维度，将每个损坏块的计算分配给一个独立 rank。 |
| Context-Sharded Block Parallelism (CSBP) | 上下文分片块并行，在块并行基础上进一步分片共享干净序列，以扩展至长上下文。 |
| Speculative Decoder Training | 推测解码器训练，训练用于加速自回归模型推理的草稿模型。 |
📱 每天一份 AI 前沿日报
关注公众号，每天 09:00 推送 · 不错过任何重磅