megatron_workers.py — Megatron 后端 PPO Workers

文件概述

基于 Megatron-LM 后端的 PPO 训练 Worker 实现。与 fsdp_workers.py 功能对等，但使用 Megatron 的张量并行（TP）、流水线并行（PP）、专家并行（EP）等高级并行策略。

核心类

ActorRolloutRefWorker

与 FSDP 版本类似，但内部使用 MegatronPPOActor 替代 DataParallelPPOActor。

class ActorRolloutRefWorker:
    def init_model(self):
        # 初始化 Megatron 模型
        # 支持张量并行、流水线并行
        self.actor = MegatronPPOActor(...)

CriticWorker

使用 MegatronPPOCritic 实现价值网络。

与 FSDP 版本的主要区别

特性	FSDP Workers	Megatron Workers
并行策略	FSDP 分片	TP + PP + EP
模型加载	HuggingFace	Megatron-Bridge
权重同步	state_dict	export_weights
流水线并行	不支持	支持
适用场景	中小规模	大规模训练

与其他模块的关系

• 依赖 actor/megatron_actor.py 中的 MegatronPPOActor
• 依赖 critic/megatron_critic.py 中的 MegatronPPOCritic
• 被 verl/trainer/ppo/ray_trainer.py 调度

小结

Megatron Workers 提供了更强大的并行能力，适合超大规模模型训练，但配置和使用复杂度更高。