Megatron-LM 后端 =================== 最后更新：2025 年 06 月 24 日。 我们通过实现 actor、critic、reference、rollout 和 reward 模型的各种 worker 来支持 Megatron 后端。我们还在 `megatron_vllm.py `_ 和 `megatron_sglang.py `_ 中使用 Megatron-LM 和 vLLM/SGLang 实现 ``3DHybridEngine``。 **优点** - 支持 5D 并行（TP、EP、CP、DP、PP）和序列并行，以获得最佳的可扩展性和吞吐量。 - 3D HybridEngine 可以显著降低峰值内存使用量，并减少 actor 和 rollout 之间的权重同步开销。 **缺点** - Huggingface 模型和 Megatron 检查点需要转换工具。 开发进度 ---------- 请注意，[Deprecated] 表示该功能在 verl 的最新版本中不受支持。 [To-Optimize] 表示该功能已实现但尚未优化。 [WIP] 表示该功能正在开发中。 [In-Release] 表示该功能已准备就绪并处于审查过程中，随时可能发布。 +---------------+-----------------------------------------------------------+ | [Deprecated] | 使用自定义模型的 Megatron 3D 并行 | +---------------+-----------------------------------------------------------+ | [Done] | Megatron 0.11.0 ``GPTModel`` 支持 | +---------------+-----------------------------------------------------------+ | [Done] | Megatron GRPO 支持 | +---------------+-----------------------------------------------------------+ | [Done] |"),(Megatron + vLLM 0.8.2，支持每个张量的权重加载 | +---------------+-----------------------------------------------------------+ | [Done] | Megatron 支持上下文并行 (Context Parallel) | +---------------+-----------------------------------------------------------+ | [Done] | Qwen2MoE 模型支持 | +---------------+-----------------------------------------------------------+ | [To-Optimize] | Megatron 分布式检查点 (Dist Checkpoint) | +---------------+-----------------------------------------------------------+ | [To-Optimize] | Huggingface 和 Megatron 检查点转换器 | +---------------+-----------------------------------------------------------+ | [To-Optimize] | 高效的融合线性层、熵和交叉熵计算 | +---------------+-----------------------------------------------------------+ | [Done] | Megatron 卸载 (参数、梯度、优化器) | +---------------+-----------------------------------------------------------+ | [Done] | Megatron Profiler | +---------------+-----------------------------------------------------------+ | [In-Release] | Megatron 0.12.0, TE 2.2 w/ vLLM 0.8.3 and Fused Attn | +---------------+-----------------------------------------------------------+ | [WIP] | Moonlight/DeepSeek-V3 模型支持 | +---------------+-----------------------------------------------------------+ | [WIP] | Expert Parallel 支持 | +---------------+-----------------------------------------------------------+ | [WIP] | Megatron 支持动态批次大小 | +---------------+-----------------------------------------------------------+ | [To-Do] | 性能调优 | +---------------+-----------------------------------------------------------+ | [MileStone] | 可运行 DeepSeek-V3 671B 训练后模型 | +---------------+-----------------------------------------------------------+ Megatron Worker 工具 ------------------------- MegatronWorker ^^^^^^^^^^^^^^ ``MegatronWorker`` 是不同 megatron worker 类的基类。在此类中，``get_megatron_global_info`` 和 ``get_megatron_rank_info`` 函数用于检索运行在特定 GPU 上的每个 ``Worker`` 的 3D 并行世界大小和 rank。这些信息将用于 Megatron 后端的传输协议。 以下用于不同模型的 ``Worker`` 类将用于构建 ``WorkerGroup``。 我们为每个 ``Worker`` 类实现了由 ``@register(dispatch_mode=)`` 装饰的各种 API。这些 API 可以被 ray driver 进程调用。数据可以根据每个函数上的 ``dispatch_mode`` 被正确地收集和分发。支持的 ``dispatch_model``（即传输协议）可以在 `decorator.py `_ 中找到。 ActorRolloutRefWorker ^^^^^^^^^^^^^^^^^^^^^ 此类是为 Actor/Rollout 混合引擎或引用模型实现其模型初始化和执行计算而实现的。 Actor/Rollout 混合引擎 '''''''''''''''''''''''''' 1. 混合引擎、Actor 和 Rollout 初始化 API。 .. code:: python @register(dispatch_mode=Dispatch.ONE_TO_ALL) def init_model(self): ``ONE_TO_ALL``：当从 driver 进程调用 ``init_model`` 函数时，每个 worker（在 GPU 上）将执行以下模型初始化过程。 混合引擎、Actor 和 Rollout 的初始化细节如下： 1. ``MegatronPPOActor`` 在使用 Megatron 构建模型时，实现了简单的 PPO 计算逻辑，包括计算 log prob 和模型更新。 2. ``vLLMRollout`` 支持使用 vLLM 进行生成。我们修改了 vLLM Engine，使其在 SPMD 下运行，以适应我们的 ``WorkerGroup`` 设计。 3. ``MegatronVLLMShardingManager`` 是一个上下文管理器，用于在 actor 和 rollout 之间执行实际的权重重分片。 有关更多信息，请参阅 `源代码 `_。 .. code:: python # 构建 actor 模型 self.actor = MegatronPPOActor(config=self.config.actor, model_config=self.actor_model_config, megatron_config=megatron_config, actor_module=self.actor_module, actor_optimizer=self.actor_optimizer, actor_optimizer_config=self.actor_optim_config) # 构建 rollout # rollout 初始化 rollout = vLLMRollout(actor_module=params, config=self.config.rollout, tokenizer=self.tokenizer, model_hf_config=self.actor_model_config, train_tp=mpu.get_tensor_model_parallel_world_size()) # 在 actor 和 rollout 之间执行权重重分片 sharding_manager = MegatronVLLMShardingManager(module=self.hybrid_engine, inference_engine=rollout.inference_engine, model_config=self.actor_model_config, layer_name_mapping=layer_name_mapping) ... 1. 生成序列并重新计算 log prob .. code:: python @register(dispatch_mode=Dispatch.MEGATRON_PP_AS_DP_PROTO) def generate_sequences(self, prompts: DataProto): - ``Dispatch.MEGATRON_PP_AS_DP_PROTO``：actor 模型的 PP 维度将被视为 DP 维度。然后，driver 进程将根据此重组来分发和收集数据。这是因为，在混合引擎中，通常应用较大 3D 并行尺寸的 actor 权重将沿着 PP 维度和 TP 维度进行聚合。因此，相应的数据应该通过 rollout 模型的 3D 并行组而不是 actor 模型进行分发和收集。但是，world_size 和 rank 信息只能从 ``get_megatron_global_info`` 和 ``get_megatron_rank_info`` 中检索，它们记录了 actor 模型的 3D 信息。此外，TP 维度内的重分片将在混合引擎中处理。 - 在此函数中，rollout 模型将执行自回归生成，actor 模型将重新计算旧 log prob 以用于生成的响应。 3. 更新 actor 模型 .. code:: python @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO) def update_actor(self, data: DataProto): - ``Dispatch.MEGATRON_COMPUTE_PROTO``：用户传入按 DP 维度分区的数据。数据将被分发到同一 dp 组内的所有 tp/pp rank，最终只收集 tp=0 和最后一个 pp 的输出数据。 - 使用 PPO 和熵损失更新 actor 模型权重。 ..note:: 目前，训练 Tensor Parallel Size 可以不同于推理 Tensor Parallel Size。 ReferenceModel '''''''''''''' 1. 引用模型初始化 引用模型使用与 actor 模型相同的函数进行初始化，但不初始化混合引擎和优化器。然后 actor 模型也由 ``MegatronPPOActor`` 包装。 2. 计算引用 log prob .. code:: python @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO) def compute_ref_log_prob(self, data: DataProto): - 在此函数中，引用模型将调用 ``MegatronPPOActor`` 中的计算 log prob 函数来计算引用 log prob。 CriticWorker 和 RewardWorker ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 1. 模型初始化 与引用模型非常相似。CriticWorker 将为优化器执行额外的初始化。 2. 计算 CriticWorker 的值 .. code:: python @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO) def compute_values(self, data: DataProto): 3. 更新 Critic .. code:: python @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO) def update_critic(self, data: DataProto): 4. 计算 Reward .. code:: python @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO) def compute_rm_score(self, data: DataProto): 训练优化工具 --------------------------- Offload (卸载) ^^^^^^^^^^^ 当资源紧张时，卸载方法可以降低 GPU 内存使用量，帮助训练和推理框架在 verl 下良好运行。它将参数、梯度和优化器移至 CPU 内存，仅在需要时将其加载回 GPU。 如果要使用卸载，可以分别为 actor 和 ref 添加以下参数。 .. code:: python # 对于 actor actor_rollout_ref.actor.megatron.param_offload=True \ actor_rollout_ref.actor.megatron.grad_offload=True \ actor_rollout_ref.actor.megatron.optimizer_offload=True \ # 对于 ref（不带梯度和优化器） actor_rollout_ref.ref.megatron.param_offload=True \ 对于 critic，可以包含这些参数。 .. code:: python # 对于 critic critic.megatron.param_offload=True \ critic.megatron.grad_offload=True \ critic.megatron.optimizer_offload=True \ 相关的 MCore 文档 ---------------------- 此外，还有一份关于使用 MCore 训练不同类型模型的详细文档，请参考 `MCore 文档 `_。