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] 表示该功能已准备就绪并处于审查过程中，随时可能发布。

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 混合引擎

混合引擎、Actor 和 Rollout 初始化 API。

@register(dispatch_mode=Dispatch.ONE_TO_ALL)
def init_model(self):

ONE_TO_ALL：当从 driver 进程调用 init_model 函数时，每个 worker（在 GPU 上）将执行以下模型初始化过程。

混合引擎、Actor 和 Rollout 的初始化细节如下：

MegatronPPOActor 在使用 Megatron 构建模型时，实现了简单的 PPO 计算逻辑，包括计算 log prob 和模型更新。
vLLMRollout 支持使用 vLLM 进行生成。我们修改了 vLLM Engine，使其在 SPMD 下运行，以适应我们的 WorkerGroup 设计。
MegatronVLLMShardingManager 是一个上下文管理器，用于在 actor 和 rollout 之间执行实际的权重重分片。

有关更多信息，请参阅源代码。

# 构建 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)
...

生成序列并重新计算 log prob

@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 以用于生成的响应。

更新 actor 模型

@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

引用模型初始化

引用模型使用与 actor 模型相同的函数进行初始化，但不初始化混合引擎和优化器。然后 actor 模型也由 MegatronPPOActor 包装。

计算引用 log prob

@register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)
def compute_ref_log_prob(self, data: DataProto):

在此函数中，引用模型将调用 MegatronPPOActor 中的计算 log prob 函数来计算引用 log prob。

CriticWorker 和 RewardWorker

模型初始化

与引用模型非常相似。CriticWorker 将为优化器执行额外的初始化。

计算 CriticWorker 的值

@register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)
def compute_values(self, data: DataProto):

更新 Critic

@register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)
def update_critic(self, data: DataProto):

计算 Reward

@register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)
def compute_rm_score(self, data: DataProto):

训练优化工具

Offload (卸载)

当资源紧张时，卸载方法可以降低 GPU 内存使用量，帮助训练和推理框架在 verl 下良好运行。它将参数、梯度和优化器移至 CPU 内存，仅在需要时将其加载回 GPU。

如果要使用卸载，可以分别为 actor 和 ref 添加以下参数。

# 对于 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，可以包含这些参数。

# 对于 critic
critic.megatron.param_offload=True \
critic.megatron.grad_offload=True \
critic.megatron.optimizer_offload=True \