作者：胡凯文，李鹏，黄俊
一、引言
多模态大模型是近期业界关注的热点，OpenAI 的 GPT4O 以及谷歌 Gemini 等多模态大模型的出现让人机交互变得更加简单和自然，在图文检索、视觉问答等多种下游任务上表现优异。通过结合语言理解和视觉感知能力，能为用户提供更加丰富和自然的人机交互体验。
Pai-Megatron-Patch 是一款由阿里云人工智能平台PAI 研发的围绕英伟达 Megatron 的大模型训练配套工具，旨在帮助开发者快速上手大模型，打通大模型相关的高效分布式训练、有监督指令微调、下游任务评估等大模型开发链路。本文以 Qwen2-VL 为例，从易用性和训练性能优化两个方面介绍基于 Megatron 构建的 Pai-Megatron-Patch 多模态大模型训练的关键技术。
在 Megatron 的基础之上，Pai-Megatron-Patch 及时追踪社区最新需求，搭建了包括 LLaMA3、Qwen2 等在内的多种热门大语言模型，并在此基础上持续拓展新特性，如支持 Optimizer Offloading 功能的 Distributed Optimizer， 可与 Transformer Engine、MoE 以及流水并行等模块协同使用，以适应更多场景需求。Pai-Megatron-Patch 的整体技术栈如下图所示：
在易用性方面，最新的 Pai-Megatron-Patch 实现了基于 Mcore 的多模态编码器和 LLM 解码器，同时实现了支持高精度低损耗的 Huggingface 和 MCore 多模态模型权重互转转换以及并行加载，极大简化了不同框架间迁移的学习成本和技术障碍。在性能方面，实现了支持高性能的文本/图像/视频数据统一加载，实现了高性能的流水并行切分，另外通过引入 CPU 卸载技术以及 Sequence Packing 技术，进一步支持多模态长序列训练的显存优化，有效缓解了 GPU 的显存压力，提升了大模型训练的效率与稳定性。这些改进共同作用，为用户提供更加流畅且高效的多模态大模型开发体验。
二、支持模型并行的权重转换
作为开源主流高效大模型训练工具，用户在尝试使用 Megatron 训练开源模型时首先遇到的障碍是 Huggingface 与 Megatron 模型权重格式的差异。下面将以 Qwen2-VL 为例，详解从 Huggingface 到 Megatron 的模型权重转换技术。如下图所示，Qwen2-VL[2] 主要由多模态编码器和 LLM 解码器两个模块组成。Qwen2-VL 采用了专门设计的多模态编码器架构，可以同时接受文本/图像/视频作为输入，并将它们转换成统一形式的表示向量后送入 LLM 解码器。这种设计使得模型能够在同一个空间内有效地捕捉到跨模态的信息关联。与 LLaVA 等多模态模型不同，Qwen2-VL 的视觉输入不需要缩放到固定分辨率大小，而是根据图像/视频数据的原始分辨率，动态生成不同长度的视觉表示向量，这一特性被称为动态分辨率。
通过将 Huggingface 模型转换到 Megatron 格式，用户可在已经充分训练好的 Qwen2-VL 基础大模型或者微调大模型上进行二次训练。所谓支持模型并行的权重转换，其核心原理是在转换时将大型模型权重分割成多个部分，并分配给不同的 GPU 进行并行加载。通过这种方式可以显著减少单个设备上的显存占用量，加快训练速度。然而，在实际应用中这一技术面临着若干挑战。在开发过程中，我们遇到的最常见问题之一是在完成权重转换后重新加载进行训练时，step 0 loss 比预期高。造成这种情况的原因主要有两个方面：一方面是模型实现风格不一致，结构映射关系十分复杂；另一方面是模型权重并行切分算法实现容易出错。因此，针对上述两种原因，Pai-Megatron-Patch 分别进行了优化。在 Pai-Megatron-Patch 的模型转换中，针对模型结构映射关系十分复杂的问题，建模分析并构建映射表能有效对比 Huggingface 模型与 Megatron 实现的联系。下表总结了两者分别在各个算子类型上的命名对应关系。在支持模型并行的权重转换实现过程中，我们首先将 Huggingface 模型的 ckpt 赋值给 Megatron 模型，然后再在 Megatron 模型上进行 TP/PP 切分。这种先转换再切分的好处是可以提升代码的可维护性，减少转换出错的概率。
除了模块层面的实现差异，在算子层，由于3D并行的存在，Megatron 与 Huggingface 的权重及计算方式也存在很大区别，需要根据实际情况进行转换。针对模型权重并行切分算法实现容易出错的问题，逐算子精细化切分保证了自底向上的准确转换。
上图表示了一个 Huggingface 实现的 GQA 权重转换到 Megatron 的正确/错误流程，其中每个块表示一个 head 的权重，边框和内部的颜色分别代表这个 head 对应的类型及被理解成的类型。在 transformers 库中，Huggingface 采用 q_proj, k_proj, v_proj 这三个算子来计算 QKV status, 随后再调用repeat_kv()函数将同一个 Query Group 内的 KV status 数量与 Query 对齐，进而调用 attention 函数。在 Megatron 中，由于张量并行的存在，Attention Module 采用了 ColumnLinear(linear_qkv) 对 QKV 统一计算，为了降低节点间通信量，这个并行 Linear 将 Query Group 均匀切分到各个 Node 上，使每个节点内就能计算一部分 attention 结果，不需要对 QKV status 进行同步。在这个 case 下，如果直接拼接 QKV 的权重张量（如图上半部分）再进行 TP 切分，可以明显看到大部分 head 位置出现问题，因此，我们需要将 QKV 权重进行一定转换，按 Query Group 的顺序进行拼接（如图下半部分）后再切分才能保证 Attention 的正确性。具体地说，在进行 TP=2 切分时，Pai-Megatron-Patch 先将 query，key 和 value 这三个 tensor 分别 reshape 成一个 4D 的 tensor，reshape 后的 shape 是<num_query_groups, -1, head_dim, hidden_dim>。接着沿着第二个维度拼接在一起，然后沿着 num_query_groups 维度进行算子拆分。才能确保并行加载 ckpt 后的起步 loss 值误差偏移在合理范围内。
三、用户友好的多模态数据加载
相对于 LLM 预训练，在数据加载方面，由于引入了多个类型的数据，多模态模型的 DataLoader 尤其复杂。为了支持 LLaVA 等模型，英伟达的 Megatron 进一步推出了 Energon 来实现多模态数据的加载逻辑。尽管如此，由于目前多模态输入并无统一的格式，在目前的Megatron中，对于多模态数据的支持仍比较有限，具体表现在(1)仅支持部分格式简单的数据加载，例如每个样本至多包含一张图像或一个视频，应用场景有限;(2)现有 TaskEncoder 输出格式固定，灵活性不大，难以高效支持模型训练。在实际开发中，Qwen2-VL 同时遇到了上述两类障碍:与已有的 LLaVA 相比，Qwen2-VL 基于 ChatML 格式设计输入，应当支持单个样本中多张图片视频、多轮对话等复杂功能，但这类数据无法使用当前 energon 的代码进行读取;此外，现有 TaskEncoder 仅支持将图片缩放到固定大小，再传入模型训练，因此难以支持 Qwen2-VL 独有的动态分辨率特性。
为了解决上述功能性问题，我们首先对内置的 DataLoader 代码做了大量扩展，使其能基于输入的原始多模态数据构造用于动态分辨率训练的图像切片序列以及位置信息，同时支持可自定义 prompt、多轮、包含任意数量图像或视频的对话样本的数据处理。在此基础上，我们设计了一套自动化脚本，能基于用户给定的 sharegpt 格式数据集路径自动转换成 energon 可读取的 webdataset 数据集，绕过当前 energon 数据集准备的交互流程，加快数据转换效率，同时降低用户学习成本。此外，对于数量不定的图像，为了能在 energon 侧自动解码，webdataset 支持将其以 numpy array 的形式保存，然而，由于 jpg 到 ndarray 间的数据格式差异，在实际测试中，我们发现这将造成最终数据集文件体积出现几倍到几十倍的增加。为了解决这个问题，我们在运行时向 webdataset 的编解码模块增加钩子函数，使其支持图像文件列表的自动编解码，实现在按原始二进制数据保存的同时能被 energon 自动转换为图像张量的需求。通过上述优化，在 Pai-Megatron-Patch 中，经过转换的 sharegpt 格式数据相对于其原始总大小几乎没有变化。
四、视觉特征处理流程优化
与 LLM 相比，多模态模型先对视觉数据进行编码，再与文本特征拼接后送入 LLM 进行推理。与基于静态分辨率的多模态大模型不同，复杂的多模态数据格式以及动态分辨率共同使得 Qwen2-VL 的这一过程更加复杂，不仅显著影响训练效率，也与能否应用 TP Comm Overlap 等技术相关。
与静态分辨率使用固定数量视觉 token 不同，动态分辨率技术允许模型依据输入图像大小，将其编码成不定数量的视觉特征。这一改进使模型对于高分辨率图像的细节捕捉能力获得显著提升，同时也优化了低分辨率图像的推理性能。下图是 Qwen2-VL 的训练数据的简单示例，针对不同长度的视觉 token，最终它们会被嵌入到文本序列中，与同一个 batch 的数据拼合后送入到 LLM 解码器。在这个过程中主要存在两个难点:(1)如何对同一批次内的多个视觉输入高效获得特征表示;(2)如何拼合数据以支持原生性能优化开关。
由于视觉输入的分辨率具有很大不确定性，为了将同一 batch 内的不同长度视觉数据转换成推理所需的特征表示，常见的做法是将每个视觉数据填充到相同长度后统一送入视觉模型。尽管这一方式实现简洁，但填充的 token 不仅造成了显存的浪费，同时也影响了视觉编码器的吞吐，当同一批次数据中同时包含高分辨率视频及低分辨率图像，由于视觉 token 数量的显著差异，将造成极为显著的性能浪费。对此，Pai-Megatron-Patch 借鉴了 Sequence Packing 的做法，将同一批次内的所有视觉输入打包后调用 varlen attention，来避免填充操作带来的显存/性能损耗。
在生成视觉特征表示后，基于预先构造好的掩膜张量，视觉编码器的输出会被依次填入文本表示的对应位置，替换掉默认的占位符。在不使用任何性能优化技术的情况下，这一实现并不会带来什么问题。然而，当开启序列并行后，LanguageEmbedding 模块的输出会自动 reduce scatter 到各个 TP rank，导致基于原始输入构造的完整掩膜张量无法使用。为了解决这一问题，同时避免冗余通信，Pai-Megatron-Patch 针对性修改了 MCore 中实现的 LanguageEmbedding 模块，在输入序列拆分前增加利用掩膜张量替换文本表示的操作 ，不干扰序列并行特性的正常运行。在实现序列并行机制的基础上，通过进一步应用 TP Communication Overlap特性，对于Qwen2-VL-70B 我们在4机32卡a100上观察到了6%性能提升。（详见实验部分）
五、基于优化器卸载的多模态长序列显存优化
相对于大语言模型训练微调，由于视频、音频等模态存在，多模态大模型的输入序列长度将远远大于纯的语言模型。因此，在通常情况下，为了完成这类模型的训练微调，往往需要更多的计算资源。而当获取更多资源这条路径不可行时，在资源受限的情况下拉起任务成为这一场景下的首要挑战。在之前 Pai-Megatron-Patch 实现的优化器卸载特性基础上，最近，PAI 团队与英伟达 Megatron 团队深度合作，共建了一套原生基于 Megatron 的权重卸载优化器。在此基础上，通过与多种正交的卸载技术结合，Pai-Megatron-Patch 将四机32卡A100上的Qwen2-VL 的最长上下文从4K提升至32K。在本章节中，我们将以BF16训练为例，详细描述最新版权重卸载优化器的原理、实现及我们引入的最新性能优化特性--Optimizer Overlapping。如下图，与激活重计算或激活卸载等显存优化技术不同，优化器卸载的目的在于将参数更新这一计算负载较低但显存占用相对高的步骤全部或部分放置到 CPU 上进行，以达到降低显存占用的目的。对于一个参数量为M的模型，Adam 会默认分配大小为12M字节的显存以保存全精度的参数(4M)及优化器状态(8M)，当使用优化器卸载时，优化器则会将这12M大小的张量全部分配到内存上，同时额外分配4M的内存用于拷贝GPU上的全精度梯度数据。因此，当打开 Megatron 的分布式优化器时，以4机32卡训练70B为例，每张卡的显存能进一步减少约26GB。
如下图所示，在最新的原生优化器卸载特性实现中，我们设计了一个混合设备优化器类(Hybrid Device Optimizer, HDO)来支持静态优化器卸载功能。与先前的深度修改 DistributedOptimizer，难以快速迁移到其他训练代码不同，通过在顶层沿用 PyTorch 优化器的 API，HDO 能无缝替换 TorchAdam，同时仍保留一定的优化器卸载能力。我们希望这一设计能使用户在更多场景下体验到优化器卸载带来的资源需求下降*注。整个 HDO 大致分为两部分，包括针对单个参数张量更新的底层优化器集合，以及用于管理优化器参数映射关系、优化器保存/读取、以及更新过程中子优化器间同步的 HDO 类。
由于在先前版本的 Megatron-Core 中，优化器内全精度参数创建部分由 DistributedOptimizer 默认分配到 GPU 上，并不能进行卸载，导致最初 HDO 的实际显存优化仅有8M而非旧版深度定制 DO 的12M。为了解决这一问题，我们修改了 DO 的参数创建逻辑，当识别到用户启用优化器卸载特性时，DO 自动将全精度参数创建这一步骤交由 HDO 进行，避免了冗余显存分配。
注: 由于参数已经被外部创建，为了不对这些参数进行修改，这一场景下仅能将8M的 Adam/AdamW 优化器状态分配及参数更新移动到 CPU 上进行，对于 SGD 优化器，由于无额外的固定显存分配，HDO 没有卸载效果。
- 优化器初始化
与支持单一设备张量的优化器不同，为了实现位于不同设备上的多个张量的参数更新及性能优化，HDO 被实现为多个子优化器的集合，每个子优化器负责单一设备上的数个张量的参数更新，而 HDO 用于实现主优化器与子优化器间的优化状态同步、Device/Host 间的参数/梯度传递以及由此引发的 CUDA 同步问题。在优化器初始化阶段，相对于普通的 Adam，HDO 额外要求用户指定 CPU/CUDA 上的子优化器类型以及相应的卸载比例，随后基于这一比例对 param_groups 进行拆分，再重新构造对应设备上的子优化器。在训练阶段，由于 lr_scheduler 等外部因素可能会对优化器参数产生影响，在子优化器参数更新前，HDO 会将所有状态同步到子优化器，如果子优化器在 CPU 上更新，也会将 GPU 梯度复制到 CPU 上。由于子优化器也可能修改参数组的状态，在更新结束后，每个子优化器也会将这些状态同步到主优化器，从而保证与非优化器卸载训练的收敛一致性。
- 优化器保存/读取及卸载比例变换
相对于 Pai-Megatron-Patch 内的优化器卸载模块，HDO 支持完整的优化器保存/读取特性，同时也支持加载优化器时切换优化器卸载比例，具有更大的灵活性。如前文所述，HDO 内部由多个子优化器负责实际参数更新，在每个参数更新阶段开始前/结束后分别进行 HDO 及子优化器间参数的同步。在保存优化器状态时，由于 HDO 内数据与子优化器一致，HDO 的 state_dict 即包含恢复优化器所需的全部信息。在加载权重时，由于子优化器需要基于 fp32 参数进行更新，我们引入了 pre_load_state_dict_hook 以及 post_load_state_dict_hook，将 HDO 的 state 中的半精度参数数据临时用保存的 fp32 替换后同步到子优化器，再替换回来，随后调用_move_new_state_to_right_device，将每个子优化器的状态移动到其对应的设备。通过在运行时更新子优化器的对应设备，我们能将在一个卸载比例下保存的模型在另一个优化器卸载配置上拉起训练。
- Naive Optimizer Overlapping 技术
与 CUDA 上的参数更新相比，基于 CPU 的参数更新性能成为优化器卸载技术是否可用的关键因素。通常，优化器卸载的时间包含三个部分：将 CUDA 梯度同步到CPU(D2H)、在 CPU 上进行参数更新(C)、将更新后的CPU参数同步到 CUDA(H2D)。为了尽可能提升性能，我们提出了一项 Naive optimizer Overlapping 技术，通过 PyTorch 的 CUDA Stream 机制，实现通信与计算重叠，将 CPU-CUDA 间数据拷贝的时间尽可能掩盖在 CPU 计算下。在 Naive optimizer Overlapping 中，我们的一个主要改进是将 CPU 子优化器进一步打散，使一个 CPU 参数对应一个CPU优化器，同时令所有拷贝非阻塞化。这一改进允许 HDO 在单个梯度张量同步完成后立即调用 CPU 参数更新，无需等待所有梯度完成拷贝，也允许 CPU 参数更新完成后立刻非阻塞拷贝到 GPU，从而实现重叠。为了保证多个 CUDA Stream 以及 CPU 之间正确的数据同步关系，同时尽可能避免影响性能，我们仅在关键位置引入 CUDAEvent，来避免额外流同步带来的开销。下图表示了一个理想的优化器重叠场景，其中每个参数张量的更新时间与两次数据传输接近。当流水线热身结束后（time step 2），可以看到 HDO 在两个 CUDA Stream 上进行双向的数据拷贝同时，CPU 也在同时进行参数更新。这一方式最大化掩盖了冗长的数据传输时间，在我们测试的 LLaMA2-70B 训练中，我们观察到该优化能减少约1s的端到端时间。
六、多模态流水并行性能优化
多模态大模型训练微调作为英伟达 Megatron 的新特性，目前仍在持续开发中。在3D并行方面，目前以 LLaVA 为代表的多模态模型的类型均为 ModelType.encoder_and_decoder，它对 encoder（视觉编码器）与 decoder（LLM解码器）分别引入了一套 TP/PP 参数，同时支持特殊的 Encoder PP 配置：当 Encoder PP>0 时，Megatron 会将 Encoder 放在独立的 GPU 上，允许两部分模型使用不同的 TP/PP 参数初始化；当 Encoder PP=0 时，Megatron 将 Encoder 部分按照 Decoder 的 TP 切分后，与 Decoder 的第一个 PP Stage 合并，放置在同一组 GPU 上。其中，前种切分方式一般适用于编码器/解码器结构类似、参数量接近的情况；后者则适用于视觉编码器参数量较少，计算量不大的场景。考虑到 Qwen2-VL 的视觉编码器大小以及用户体验，我们最初采用了编码器与解码器部分合并的切分方式，来保持和其他 LLM 相似的切分参数配置，减少学习成本。但在实际测试中，通过与基模型 Qwen2 对比，我们观察到这一合并显著降低了 Qwen2-VL 的训练吞吐。这主要由两方面因素导致：(1)由于模型结构不一致，视觉编码器与 LLM 解码器的前向速度有差异，导致单个 micro batch 内每个 GPU 间存在负载不均衡;(2)由于视觉编码器支持动态分辨率特性，对于每个 micro batch，实际前向的token数并不一致，使得每个 mciro batch 计算量略有差异，出现各 GPU 间存在负载不均衡的情况。
为了解决上述负载不均衡的问题，进一步提升训练吞吐，在原生 Megatron 的基础上，我们从两个方向对现有框架进行改进尝试：(1)基于非均匀切分策略间的负载均衡;(2)拓展模型实现以支持虚拟流水线并行特性。
- 基于非均匀切分策略间的负载均衡：在模型总参数量较小或 PP 较大的并行配置下，pp rank 0 包含的视觉编码器带来的额外计算量相对于原有的 LLM 解码器的计算量比例有显著提升。由于在流水线中，最慢的步骤决定了整体的性能，应用 megatron 内的非均匀切分特性，将部分计算量转移到其他 pp rank 上，能有效提升整体的训练性能。为了支持基于非均匀切分的继续预训练或微调功能，我们同步更新了 Qwen2-VL 的转换模型，用户可以通过控制 MP_PP0_LAYERS 环境变量，获得非均匀切分的模型权重文件。
- 虚拟流水线并行特性：虚拟流水线并行，或 interleaved 1F1B PP 是 Megatron-LM 在1F1B的流水线并行实现上的重要改进。它通过将同一个 pp stage 内的 transformer 层分配到多个 GPU，进一步打碎计算，通过增大通信量的方式使负载更加均衡，降低 Global Batch 内的 bubble 比例。然而， 对于 encoder_and_decoder 类型的模型（即英伟达内置的所有多模态模型），MCore 不支持启用 VPP。考虑到视觉编码器的整体计算量不需要独立GPU，我们进一步改进了 Qwen2-VL 实现，移除了 Encoder PP 的支持，并将 Qwen2-VL 的模型类型设置为 LLM 的 encoder_or_decoder 来启用 VPP 选项。与官方论文的结果不同，我们观察到虚拟流水线并行也能起到训练加速效果，甚至在H20上，我们观察到相对于非均匀切分，虚拟流水线并行的加速效果更加明显。
七、实验分析
在本节中，我们对上述提到的技术对 Qwen2-VL 模型的性能及收敛情况展开了全面的分析，主要包括三方面：1. 权重转换精度分析：通过对比转换前后的模型测试集表现，我们验证了权重转换模块的正确性；2.多模态长序列显存优化分析：通过引入 HDO 对优化器进行卸载，在相同 GPU 数量的情形下，Pai-Megatron-Patch 大大提升了长序列训练的可用性；3.多模态模型训练性能优化：全面测试了各训练加速技术对于不同设备/模型大小的影响，论证了技术的有效性，同时为用户提供开关设置的指引。
权重转换精度分析
为了验证 PAI-Megatron-Patch 内多模态模型转换的正确性，我们采用 VLMEvalKit 对转换前后的 Qwen2-VL-7B 模型在多个数据集上的表现进行评估。结果如下图，其中 Ref 是 VLMEvalKit 提供的在 SEEDBench 及 MMBench 上的评估分数，Official 是官方权重在测试环境内的实测分数，Convert 为将官方原始权重进行两次转换得到的新 Huggingface 权重的分数。可以看到，转换前后的模型评估分数完全一致，且均与开源 leaderboard 指标接近，有效说明了转换代码的正确性。
基于优化器卸载的多模态长序列显存优化
在本节中，我们探索了通过 HDO 与其他显存优化技术复合，进一步提升有限机器场景下可训练上下文长度上限的可能性。我们分别测试了 Qwen2-VL 7B/72B 在单机8卡、四机32卡的无 offload 情况下的最长上下文，以及打开所有优化后的上下文长度，结果如下表。
可以看到，对于 Qwen2-VL 7B 模型，通过结合激活重计算技术及优化器卸载技术，我们将原本 64K上下文长度的训练上限进一步拓展到128K，且没有 OOM 风险。*对于72B模型，结论类似，通过结合多种正交的显存优化技术，我们也将四机32卡A100的可训练上下文长度从16K提升到128K。*注：在单独开启重计算的情况下我们也成功在 demo 数据集上实现了128K上下文长度的训练，此时显存空闲不到1G，考虑到实际场景下视觉模型占用的显存会随数据变化等原因，不推荐使用这一配置训练。
多模态模型训练性能优化
在本节中，我们针对两种大小的 Qwen2-VL 模型的训练场景，测试了 PAI-Megatron-Patch 当前支持的训练优化技术对吞吐的影响，来说明其在不同场景下的优势。具体地，我们采用预处理的 LLaVA-Pretrain 数据集(参见 Qwen2-VL 的最佳实践文档)在4K上下文长度下对 Qwen2-VL-7B/70B 进行训练，在最佳并行配置的基础上，进一步打开各类优化技术以比较模型性能的变化。其中，PP0_layers表示开启非均匀切分时，PP Rank0 上的 LLM Decode r层数，TGP内的三个字母依次表示tp-comm-overlap、overlap-grad-reduce、overlap-param-gather三个优化开关的启用状态。对于每个实验，我们记录了其训练时的每秒 token 数(TGS)及 MFU。需要注意的是，由于 megatron-core 暂不支持估计视觉编码器的 TFLOPs，记录的 MFU 仅基于 LLM 的运算量计算得到，因此略低于实际值。
overlap 开关对训练吞吐的影响
相对于 LLM，由于存在多个不同结构的 transformer，多模态模型并不原生支持 megatron-core 的某些 overlap 技术，同时差异化的结构也可能导致 overlap 的收益与 LLM 不同。在本节中，我们在A100上测试了tp-comm-overlap、overlap-grad-reduce、overlap-param-gather三个优化开关对 Qwen2-VL 7B/72B 训练吞吐的影响，结果如上表所示。主要结论如下：
- 对于 Qwen2-VL，可以观察到overlap-grad-reduce及overlap-param-gather无论在7B/72B模型规模上均导致了性能出现略微下降。
- 相对于72B模型(~8%),7B模型开启tp-comm-overlap的直接收益可以忽略不计。这主要是因为7B模型实际运行时的低TP数以及未开启 overlap 的 vision model 计算量占比相对较高。
虚拟流水并行及非均匀切分对训练吞吐的影响
基于 megatron-core 的核心能力，Pai-Megatron-Patch 为 Qwen2-VL 提供了两种性能优化方式：多模态大模型的虚拟流水并行及解码器非均匀切分技术。为了比较这两者对训练吞吐的影响差异，我们在最佳并行配置的基础上，比较了两种技术的性能上限，结果如下表。其中，VP 表示每个虚拟并行块中的解码器 Transformer 层数（区别于通常意义的 VPP Size）。
从表中我们能得出的主要结论如下：
- 尽管tp-comm-overlap仅能为 Qwen2-VL 72B 模型直接带来约8%的性能提升，然而，通过序列并行带来的显存优化，模型有进一步降低并行数来提升吞吐的潜力。例如在H20上，我们观察到，通过打开 SP 及 TP Comm overlap 开关，模型可以采用 TP4PP4 的配置进行训练，通过调整虚拟流水并行，相对于 TP8PP4 的最佳性能有额外15%的性能提升。
- 由于各 GPU 的机内、机间带宽差异，在不同的 GPU 上，最佳优化技术有所不同。
- 与 Qwen2-VL 论文的结果不同，我们的 VPP（或1F1B interleaved）实现在 A100/H20 机器上均有一定的性能提升。
八、总结
在基于 Megatron 的 Qwen2-VL 多模态大模型最佳实践开发过程中，我们围绕大模型训练测试了以下核心技术的性能：
- 支持模型并行的权重转换。
- 基于 CPU 卸载的多模态长序列显存优化的鲁棒性。
- 基于 Megatron 的多重训练加速技术的稳定性。
- 运用综合加速技术来训练 Qwen2-VL 过程中的易用性和稳定性。
后续在 Pai-Megatron-Patch 中还会陆续放出更多高质量的大模型最佳实践以及最新的训练加速技术应用示例，敬请期待。
参考文献
- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
- Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution
- MMBench: Is Your Multi-modal Model an All-around Player?
- SEED-Bench: Benchmarking Multimodal LLMs with Generative Comprehension
- https://github.com/open-compass/VLMEvalKit
