多模态 大语言模型 训练的异构并行性

基金会 模型训练 正在变得多模式，从 后训练 管道到大规模预训练。随着模态覆盖范围的扩大、上下文窗口的增长以及编码器 LLM 规模的发散，以 LLM 为中心的单个 TP/CP/PP/DP/EP 布局越来越限制吞吐量。这种耦合迫使编码器继承 LLM 驱动的分片和放置选择，这些选择可以添加通信、限制编码器并行性或约束 LLM 调度；这种不匹配在长上下文中最为明显，其中融合多模态序列需要 LLM 上下文并行性，但编码器输入仍然有限。我们提出了多模式 大语言模型 训练的异构并行性，这是一种抽象，允许一个端到端图中的模块使用独立的布局和等级放置，支持共享 GPU 上的共置执行和不相交的等级集上的非共置执行。关键的挑战是在独立布局中保留边界张量语义：必须为目标布局具体化前向激活，而后向梯度必须路由回源布局。我们通过边界通信器来解决这个问题，该边界通信器实现了前向和后向布局转换，以及两种布局模式的调度扩展。我们跨多模式工作负载和 GPU 规模评估优化的同构、共置异构和非共置异构配置，以表征何时增加的布局和放置自由度会带来更好的操作点。在此过程中，共置异构性将 TFLOPS/GPU 提高了高达 49.3%，而非共置异构性将总词元吞吐量提高了高达 13.0%，将 TFLOPS/GPU 提高了高达 9.6%。我们根据同质基线验证损耗收敛奇偶性，并将该系统作为开源 Megatron-LM 扩展发布。