用于高效分布式长上下文扩散语言模型训练的块并行性

arXiv:2609.19242v1 公告类型：新 摘要：块扩散语言模型（BDLM）将块间的自回归依赖性与块内的并行去噪相结合，但长上下文训练受到分布式注意力通信和激活记忆的限制。传统的上下文并行性 (CP) 按位置对组合的干净加损坏序列进行分片，将共享的干净 K/V 与特定于块的损坏 K/V 及其梯度进行通信。我们观察到 BDLM 目标在目标块上进行了分离。我们引入了块并行性（BP），这是一种新的分布式并行性维度，它将每个损坏的块计算分配给一个等级。为了将 BP 扩展到长上下文，我们引入了上下文分片块并行性（CSBP），它也跨这些等级分片了共享的干净序列。 CSBP 将损坏的 K/V 和梯度保留在本地，避免复制干净的前缀，并保留 BDLM 训练语义。在 256K 环境下的 16 个 H200 GPU 上，CSBP 将监督微调的吞吐量比最佳基线提高了 1.18-1.45 倍，将自回归模型转换为 BDLM 的吞吐量提高了 1.27-1.33 倍，同时匹配或降低了峰值 HBM。 512K 时，全模型加速达到 1.61 倍。在 8 个 H100 GPU 上，CSBP 将 DFlash2 推测解码器训练速度在 512K 时加速 2.48 倍，在 1M 时加速 7.59 倍。在匹配的 12 小时 DiffusionGemma 26B-A4B SFT 运行中，CSBP 在 SWE-bench Verified 和 Terminal-Bench Lite 上的每个经过训练的检查点都实现了更高的通过率。代码：https://github.com/ScalingIntelligence/Turbo-dLLM