megatron_critic.py — Megatron 并行 Critic

文件概述

基于 Megatron-LM 的 MegatronPPOCritic 实现。与 megatron_actor.py 类似，支持 TP + PP 并行策略。

核心特点

与 FSDP 版本的主要区别： - 使用 Megatron 的流水线调度器执行前向/反向传播 - 模型输出在最后一个 PP stage 上，需要跨 stage 广播结果 - 支持张量并行下的梯度聚合

与其他模块的关系

• 继承自 base.py 的 BasePPOCritic
• 被 megatron_workers.py 的 CriticWorker 使用

小结

Megatron Critic 为大规模模型提供了流水线并行支持的价值网络训练。