针对大语言模型与多模态模型训练中数据并行(Data Parallel,DP)负载不均的行业难题,百度百舸近日正式发布 LoongForge DP 负载均衡优化方案。该方案通过精准计算开销建模与自适应数据重分配技术,在超大规模分布式训练场景下实现近 10% 的性能提升,有效破解制约模型训练效率的核心瓶颈。

行业痛点:DP负载不均成为训练效率"拦路虎"

当前,大语言模型与多模态大模型训练普遍采用数据并行机制,将训练数据分发至各计算节点并行计算后,通过 AllReduce 完成梯度同步。然而,分布式同步机制的特性决定了任意单个节点的计算延迟都会被全局放大,形成全员等待的冗余开销。

传统固定长度 packing 策略仅从 Token 数量维度实现基础负载均衡,却忽视了 Transformer 注意力机制二次方计算复杂度的特性。百度百舸技术团队发现,即便两个节点处理的 Token 总数相同,只要样本长度分布存在差异,实际计算开销就会出现巨大差距。”

这一问题在多模态模型训练中更为严峻。图像分辨率、图片张数、视频帧数等差异,导致视觉编码器与文本解码器双重模块均存在负载不均,进一步加剧分布式训练瓶颈。

核心创新:两大阶段破解负载均衡难题

LoongForge DP 负载均衡方案采用两阶段优化策略,全程嵌入原生训练流程,无需离线预处理:

热身建模阶段:通过在线实时性能探测机制,动态采集各 DP 节点的真实计算执行耗时与样本特征数据,自适应构建与模型计算特性相匹配的开销估计模型。该模型能够同时刻画 Attention 的二次复杂度开销、线性层开销及固定开销,精准拟合不同样本组合的真实计算成本。

在线自适应重分配阶段:基于构建的计算开销模型,实时评估各 DP 节点待训练样本的计算压力,动态完成跨节点样本重分配调度,最小化所有节点单迭代最大总计算开销,抹平节点间计算耗时差异。

四大特性:易用性与性能兼顾

该方案具备四大核心特性:

多模态双重负载均衡:同时适配 LLM 文本解码器与 ViT 视觉编码器,实现全维度负载均衡优化

• 迭代级全局均衡:支持跨微批次负载持续追踪,达成完整迭代周期内的全局负载均衡

• 智能触发机制:自动跳过无效重排操作,避免无效通信与资源浪费

• 零额外时延:采用异步流水线设计,数据重排开销完全隐藏,无额外训练时延增量

方案支持开箱即用,全面兼容 InternVL、Qwen2-VL/2.5-VL/3-VL 等主流多模态模型,覆盖图像、视频全场景训练,无需修改模型训练代码,仅通过简单命令行参数即可一键启用。

实验验证:超大规模场景收益显著

在固定未开启 All-Reduce 通信重叠优化的实验条件下,针对不同 DP 并行规模开展对照实验:

未开启 DP 负载均衡机制前:随着并行规模从 DP32 扩展至 DP512,模型训练整体吞吐性能(TGS)持续下降,在 DP256 向 DP512 扩展阶段性能退化尤为显著。

开启 LoongForge DP 负载均衡机制后:在所有 DP 并行规模下训练吞吐水平均明显提升,且并行规模越大优化收益越显著。在 DP256 规模下性能提升约 3.3%,在 DP512 超大规模场景下性能提升接近 10%。

应用价值:推动大规模模型训练降本增效

实验结果表明,LoongForge 的 DP 负载均衡方案通过对计算负载进行精细化建模并实施自适应动态数据重分配,从根本上缓解了负载不均问题,显著减少梯度同步阶段的无效等待时间,整体提升分布式训练吞吐率及 GPU 资源利用效率,尤其适用于超大规模集群训练场景。

该方案的发布,为业界提供了一套行之有效的大规模模型训练优化工具,有望显著降低训练成本、提升算力利用效率,推动大模型技术发展与产业化应用。

关于 LoongForge

LoongForge 是百度百舸开源的全模态训练框架,为原生多模态时代提供一套统一、高效、易用的训练加速解决方案。已在 GPU 与昆仑芯 XPU 两大平台、数千卡规模集群上完成长期生产验证,覆盖 LLM 到 VLM、VLA 等多种业务场景。更多信息请访问:

• 官方网页:https://baidu-baige.github.io/LoongForge/

• GitHub:https://github.com/baidu-baige/LoongForge