Megatron-LM 原生强化学习后训练框架 Megatron-RL架构设计与 Off-Policy Generation 实战指南【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMegatron-RL 是 Megatron-LM 仓库中新增的原生强化学习RL后训练模块它把环境Environment、智能体Agent与底层 RL 算法实现解耦并在训练循环中引入 Off-Policy Generation 机制来重叠推理与训练、最大化硬件利用率。本文以仓库内 megatron/rl/README.md 为骨架结合 megatron/rl 模块源码、参数定义 与 examples/rl 实例完整讲解 Megatron-RL 的组件划分、collection lag 概念、生成粒度模式以及一套可运行的 GRPO 训练配置帮助读者理解并上手这套面向研究团队的 LLM RL 后训练方案。Megatron-RL 概览与项目状态Megatron-RL 的目标是在 Megatron-LM 内原生支持基于强化学习的后训练。官方 README 明确指出megatron/rl/README.md新库的主体代码位于megatron/rl目录对 Megatron Core 的推理inference代码做了显著修改对 Megatron-LM 训练循环做了少量改动用于支持 RL 步骤示例环境位于examples/rl目录。它面向的是“在最新 NVIDIA 硬件上、对前沿基础模型进行大规模 RL 后训练”的研究团队其定位是研究框架而非企业级框架——README 明确说明它不保证对任意开源模型的即开即用支持此类能力由 NVIDIA 的 NeMo RL 提供。需要特别提醒的是项目状态截至 2025-08-27Megatron-RL 仍处于活跃开发阶段在 NVIDIA 内部可用但并非所有代码都已发布到本仓库现有代码与示例会随开发进度变化。因此使用时应以当前仓库实际内容为准并关注其 roadmap 的更新。设计哲学Agent/Environment 与 RL 实现的解耦Megatron-RL 的核心设计理念是尽可能让智能体/环境Agent/Environment设计与底层 RL 实现解耦。整体职责被划分为三个概念组件Agent and Environment合称 Agent持有InferenceInterface的句柄返回Rollout或EvaluationResponse对象负责采样参数、自定义生成参数如停止条件 stop conditions、内联评估 inline evaluation等。从源码看Agent 的抽象定义在 megatron/rl/agent/api.py 中Agent是抽象基类通过get_rollout_response获取单个 rollouts 的模型响应其下派生多种生成器接口RolloutGenerator产生包含 rollout 字符串与对应奖励的RolloutContrastiveRolloutGenerator产生包含 chosen/rejected 两条轨迹的对比数据用于偏好类训练TokenizedRolloutGenerator产生包含 token id 与奖励的TokenRollout可选携带生成掩码GroupedRolloutGenerator面向 GRPO 等需要按组训练的算法通过prepare_group_rollout生成一组 rollouts是RolloutPipeline消费的核心契约EvaluationAgent接收推理接口并返回基准分数。GroupedRolloutRequest同一文件中携带num_groups、rollouts_per_group、submission_granularity、consumption_granularity等字段是粒度控制传入 Agent 的载体。Trainer/Evaluator管理 rollout 生成与评估的控制流协调或创建InferenceInterface和 Agent。在 megatron/rl/agent/rollout_pipeline.py 中RolloutPipeline实现了“调用方持有的分组 rollout 生成编排”它把提交/消费粒度策略封装为_GranularityConfig从GroupedRolloutRequest中解析并校验粒度组合。Inference Interface为 Agent 环境提供.generate(prompt, **generation_args)端点可以有多种形态如 Megatron、OpenAI、HuggingFace并支持多种配置选项。源码定义位于 megatron/rl/inference/inference_interface.pyInferenceInterface提供同步generate与异步agenerate两个入口prepare_request负责把字符串 prompt 包装成LLMChatMessage。此外还有三个能力标记 Mix-InReturnsRaw返回完整字符串、ReturnsTokens返回完整 token 列表、ReturnsLogProbs返回 logprobs——后两者是实现重要性采样校正importance sampling correction等高级特性的接口基础。Off-Policy Generation让推理与训练重叠为什么需要 Off-Policy纯 on-policy RL 的 rollout 生成与训练是串行的每一步训练都必须等待当前策略产生的最新 rollout。这会导致硬件利用率显著不足——训练时推理引擎空转推理时训练卡空闲。Off-policy generation 通过两种维度把二者重叠起来时间维度类似 Python asyncio在训练与推理之间切换资源空间维度类似 Python threading把资源划分为训练与推理两部分也可以两者兼用。无论采用哪种重叠方式核心行为几乎相同当训练忙于第N步时推理引擎持续生成将被第N1、N2… 步消费的 rollout。由于这些 rollout 来自略旧的策略权重它们存在一定程度的“陈旧性staleness”README 称之为collection lag滞后量L表示当前训练步消费的 rollout 是在L个训练步之前生成的即由一个落后当前训练策略L次更新的策略产生。更高的 lag 意味着更多的重叠更高吞吐代价是更高的陈旧度样本质量/稳定性下降。找到这个权衡点上合适的位置正是这个旋钮存在的意义。开启方式与 lag 取值Off-policy generation 需要显式开启--rl-partial-rollouts滞后量用--rl-generation-lag设置--rl-partial-rollouts --rl-generation-lag 2在 参数定义 中可以看到--rl-generation-lag是 float 类型允许小数或负数--rl-partial-rollouts默认关闭开启后“允许推理在训练更新策略权重的同时继续生成 rollout”。参数校验逻辑同文件 L548-L576 附近还强制了若干约束--rl-generation-lag与--rl-max-inflight-requests互斥当 lag 大于 0 时必须开启--rl-partial-rollouts--rl-generation-lag取值必须 ≥ -1--rl-submission-granularity同样要求--rl-partial-rollouts--rl-consumption-granularity目前不支持R提交粒度B与消费粒度G的组合不被允许。有用取值README 原文L -1最小值——同时只有一个单位unit的生成工作在途实际上是串行的主要用于测试只有在提交粒度G或R下才与L 0可区分。L 0恰好填满一个训练步所需的生成容量推理引擎永远不会领先于训练器。L 1, 2, …引擎提前运行这么多步。L 越大吞吐越高、陈旧度越大。允许小数。自动调优Autotuning与资源公式如果不提供--rl-generation-lag系统会自动设置 lag 并在启动时打印所选择的值若自行提供则启动警告会告知你的--rl-generation-lag是否**过度订阅oversubscribe**了推理引擎——即你要求在途的工作量超过了引擎实际可服务的能力。这正是允许 L 取负值的原因当推理资源过少、即使是完全同步的 RL 也会过度订阅引擎时L 会变为负数。基于此用户可用如下公式合理缩放推理资源max_effective_lag DP * engine.max_requests / (G * P) - 1其中DP推理数据并行大小G--grpo-group-size每个 GRPO 组的采样数P--grpo-prompts-per-step每步的提示组数即论文中的组数 Gengine.max_requests推理引擎每 rank 的请求容量由 KV-cache 内存推导也可通过--inference-dynamic-batching-max-requests覆盖。提交粒度与消费粒度Submission and Consumption Granularity两个进一步的控制开关通过选择流水线的工作单位来塑造重叠形态。工作单位有三档按粗细排列为单个 rolloutR→ 提示组G一个 prompt 的--grpo-group-size个样本→ 整个训练批次B一步中全部P个组。该粗细阶梯在 megatron/rl/rollout_granularity.py 中以GRANULARITY_RANK {R: 0, G: 1, B: 2}的形式显式编码。--rl-submission-granularity {R,G,B}设置引擎容量按什么单位切分。一个单位从被受理到完成期间占据其份额的引擎容量完成后该容量才释放给新任务。--rl-consumption-granularity {G,B}设置训练器等待的单位。G按完成顺序消费提示组B按提交顺序消费整个批次。消费粒度不能细于提交粒度B提交强制B消费且在 GRPO 下不可能按 rollout 消费——因为 advantage 是相对组内奖励计算的组是训练的最小原子单位。以“提交/消费”模式书写支持的模式为B/B默认、G/G、G/B、R/G、R/B。它们的行为由一条事实支配一个单位只有在其最慢的 rollout 完成时才算完成。在长尾生成长度下单位预期完成时间会超过平均 rollout 时间超额部分用尾部因子τ E[单位中最慢 rollout] / E[rollout]刻画且随单位增大而增大τ_rollout 1 ≤ τ_group ≤ τ_batch。对一台恰好按 lag 门控配置autotuning 所做之事大小的引擎由此导出两条定律利用率只由提交粒度决定U 1 / τ_submit。粗单位会让 straggler 占用引擎槽位使引擎空闲1 − 1/τ_submit的时间比例更细的提交能立即回填释放的槽位R提交能让引擎完全忙起来。陈旧度由比值决定平均首 token 陈旧度≈ (τ_consume / τ_submit) · (1 L)。消费粒度粗于提交粒度会把中间跨越的每个尾部因子都乘进陈旧度却买不到任何利用率提升——因此在固定提交粒度下匹配的模式G/G优于G/BR/G优于R/B占优。消费粒度还决定陈旧度的分布在G消费下一个组的陈旧度与其自身完成时间成比例——长组总是按比例训练得更“旧”。B/B是一个特殊角落只在消费时释放容量使生成与训练时钟相位锁定因此每一步中每个组都以恰好L版策略的陈旧度训练——确定性最强且比“完成即释放”模式支付的(1 L)少一个版本。在L 0时B/B就是完全同步的 on-policy 训练代价是引擎空闲1 − 1/τ_batch——这是唯一一种不把空闲转化为超前陈旧度的模式。模式速查规则总结模式特点与适用场景B/B默认on-policy 角落陈旧度最小且确定吞吐最低最适合调试、可复现实验与消融lag 方差会混淆结论时G/G相比B/B获得τ_batch/τ_group倍的利用率代价是额外一版本平均陈旧度呈长度相关的散布R/G唯一能完全利用引擎的路径陈旧度为G/G的τ_group倍单独看是等价交换配合重要性采样校正--rl-inference-logprobs-is-correction则成为净收益——校正后的陈旧度是统计成本而空闲引擎是不可挽回的墙钟时间G/B、R/B被支配利用率与对应匹配模式相同陈旧度严格更高仅当确实需要“流式提交 批次有序消费”时才使用内部旋钮映射lag 如何转化为在途任务底层实现中lag 控制同时在途的提交单位tasks数量tasks max(1, round((L 1) * P * G / unit)) # 在途的生成工作单位数 rollouts_in_flight tasks * unit # 并发推理请求总数其中unit是每个提交单位包含的 rollout 数B为P * GG为GR为1。通常无需手动推算——选定或让 autotune 选定一个L或直接用--rl-max-inflight-requests直接限定在途 rollout 请求总数等价于(L 1) * P * G与--rl-generation-lag互斥——然后在启动日志中核对门控容量确认任务数符合预期即可。相关 RL 参数速查除粒度与 lag 外megatron/training/arguments.py 的_add_rl_argsL2826 起还定义了以下常用参数含默认值参数默认值说明--perform-rl-stepFalse使用 RL 训练步--rl-prompts-per-eval32每个 RL 任务的评估 prompt 数passk 类评估代价高默认取较小值--grpo-prompts-per-step32每步 GRPO 组数论文中的 G--grpo-group-size2每个 GRPO 组的采样数--grpo-iterations2GRPO 实现的迭代次数--grpo-clamp-eps-lower0.01GRPO 下界裁剪--grpo-clamp-eps-upper0.01GRPO 上界裁剪如 DAPO 般上下界可不同vanilla GRPO 令二者相等--grpo-kl-beta0.001GRPO loss 中 KL 项权重--grpo-entropy-term-weight0.0熵项权重--grpo-filter-groups-with-same-rewardFalse过滤奖励相同的组--langrl-env-configNoneRL 环境配置 YAML 路径--rl-default-temperature/--rl-default-top-p/--rl-default-top-k1.0 / 0 / -1推理默认采样参数--rl-offload-optimizer-during-inferenceFalse推理/rollout 阶段将优化器状态卸载到 CPU 以省显存--rl-kv-cache-management-modepersistKV-cache 管理persist驻留显存/offload卸载到 CPU/recompute释放并逐周期重算--rl-persist-cuda-graphsFalse推理引擎挂起时是否持久化 CUDA graphs--rl-inference-logprobs-is-correctionFalse用推理 logprobs 做 loss 的重要性采样校正--rl-importance-sampling-truncation-coefNone配合上式对 IS 校正做截断--rl-use-sequence-packingFalse启用序列打包--rl-durable-rollout-bankFalse把完成的 rollout 组持久化到可靠账本SIGKILL如 SLURM 时限后重启可恢复--rl-rollout-bank-dirNone持久化 rollout bank 目录默认save/rollout_bank--rl-rollout-bank-max-bytes0bank 软上限字节0 表示不限制--rl-training-cuda-graphsFalse训练阶段不反复切换 CUDA graphs 开关实战示例GRPO 训练 Qwen2.5-32B仓库在 examples/rl/README.md 中提供了一个完整可参考的 GRPO 实验在 DAPO-17k如qwen_2p5_32b.sh、common.sh环境配置位于 examples/rl/environment_configs。环境准备使用nvcr.io/nvidia/pytorch:25.06-py3容器额外安装依赖pip install flask-restful uvloop datasets evaluate设置环境变量并创建目录export CUDA_DEVICE_MAX_CONNECTIONS1 CHECKPOINT # 指定基础模型 checkpoint 路径 RUN_DIR # 指定记账/工作目录 WANDB_PROJECT # 指定 WANDB_EXP_NAME # 指定 LOG_DIR$RUN_DIR/logs DATA_CACHE_DIR$RUN_DIR/data_cache CHECKPOINT_DIR$RUN_DIR/checkpoints TB_DIR$RUN_DIR/tensorboard将 HuggingFace checkpoint 转换为 Megatron 格式使用 tools/checkpoint/convert.py 将 HF Qwen checkpoint 转为 Megatron-LM 格式TP8 HF_FORMAT_DIRPATH_TO_HF_FORMAT_DIR MEGATRON_FORMAT_DIRPATH_TO_MEGATRON_FORMAT_DIR TOKENIZER_MODELHF_FORMAT_DIR python ./tools/checkpoint/convert.py \ --bf16 \ --model-type GPT \ --loader llama_mistral \ --saver core \ --target-tensor-parallel-size ${TP} \ --checkpoint-type hf \ --load-dir ${HF_FORMAT_DIR} \ --save-dir ${MEGATRON_FORMAT_DIR} \ --tokenizer-model ${TOKENIZER_MODEL} \ --model-size qwen2.5 \ --loader-transformer-impl transformer_engine \ --make-vocab-size-divisible-by 128 \训练命令与关键参数以下命令片段来自 examples/rl/README.md实际运行时需按环境微调COMMON_OPTIONS\ --tensor-model-parallel-size $TP \ --pipeline-model-parallel-size $PP \ --use-mcore-models \ --transformer-impl transformer_engine \ --bf16 \ --te-rng-tracker \ --cuda-graph-impl local \ --inference-dynamic-batching-num-cuda-graphs 1 \ --inference-dynamic-batching-buffer-size-gb 20 \ --data-parallel-random-init \ --attention-backend flash \ --timing-log-level 1 \ --log-timers-to-tensorboard \ --initialize-socket-comms \ GRPO_CLAMP_EPS_LOWER0.2 GRPO_CLAMP_EPS_UPPER0.28 MAX_INFERENCE_BS32 GRPO_GROUP_SIZE16 GRPO_PROMPTS_PER_STEP64 GRPO_ITERATIONS1 GRPO_KL_BETA0.0 TRAINING_BATCH_SIZE1024 MICRO_BATCH_SIZE1 MAX_SEQ_LENGTH11999 MODEL_OPTIONS\ --ckpt-format torch \ --seq-length $MAX_SEQ_LENGTH \ --inference-max-seq-length $MAX_SEQ_LENGTH \ --inference-max-requests $MAX_INFERENCE_BS \ --pretrained-checkpoint $CHECKPOINT \ --untie-embeddings-and-output-weights \ --disable-bias-linear \ --add-qkv-bias \ --normalization RMSNorm \ --norm-epsilon 1e-5 \ --group-query-attention \ --num-query-groups 8 \ --no-masked-softmax-fusion \ --attention-softmax-in-fp32 \ --attention-dropout 0.0 \ --hidden-dropout 0.0 \ --weight-decay 0.0 \ --position-embedding-type rope \ --rotary-percent 1.0 \ --rotary-base 1000000 \ --use-rotary-position-embeddings \ --swiglu \ --num-layers 64 \ --hidden-size 5120 \ --ffn-hidden-size 27648 \ --num-attention-heads 40 \ --max-position-embeddings 131072 \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model unsloth/Qwen2.5-32B \ --lr 1e-6 \ --lr-warmup-samples 0 \ --make-vocab-size-divisible-by 128 \ --clip-grad 1.0 \ --recompute-granularity selective \ --recompute-activations ENV_DEPENDENT\ --langrl-env-config examples/rl/environment_configs/dapo.yaml \ --micro-batch-size $MICRO_BATCH_SIZE \ --global-batch-size $TRAINING_BATCH_SIZE \ --grpo-group-size $GRPO_GROUP_SIZE \ --grpo-prompts-per-step $GRPO_PROMPTS_PER_STEP \ --grpo-iterations $GRPO_ITERATIONS \ --grpo-clamp-eps-lower $GRPO_CLAMP_EPS_LOWER \ --grpo-clamp-eps-upper $GRPO_CLAMP_EPS_UPPER \ --grpo-kl-beta $GRPO_KL_BETA \ --env-config $ENV_CONFIG torchrun \ --nproc-per-node8 \ --nnodes8 \ train_rl.py \ --mock-data \ --distributed-timeout-minutes 60 \ --train-samples 48828125 \ --log-interval 10 \ --log-progress \ --timing-log-option minmax \ --log-params-norm \ --log-num-zeros-in-grad \ --log-throughput \ --adam-beta1 0.9 \ --adam-beta2 0.95 \ --adam-eps 1e-8 \ --no-create-attention-mask-in-dataloader \ --accumulate-allreduce-grads-in-fp32 \ --calculate-per-token-loss \ --log-straggler \ --disable-straggler-on-startup \ --perform-rl-step \ --use-distributed-optimizer \ --straggler-minmax-count 16 \ --eval-interval 20 \ --rl-prompts-per-eval 32 \ --tensorboard-log-interval 1 \ --empty-unused-memory-level 2 \ --data-cache-path ${DATA_CACHE_DIR} \ --save $CHECKPOINT_DIR \ --load $CHECKPOINT_DIR \ --tensorboard-dir $TB_DIR \ --seed $SEED \ --sequence-parallel \ --finetune \ --save-interval 20 \ --wandb-project $WANDB_PROJECT \ --wandb-exp-name $WANDB_EXP_NAME \ ${MODEL_OPTIONS} \ ${COMMON_OPTIONS} \ ${ENV_DEPENDENT} $要点解读--perform-rl-step是启用 RL 训练步的总开关对应_add_rl_args中--perform-rl-step参数--inference-max-requests示例中 32对应上文的engine.max_requests直接决定推理引擎的并发请求容量是评估max_effective_lag公式的关键输入--grpo-clamp-eps-lower/upper采用 DAPO 式非对称裁剪0.2/0.28如要做 vanilla GRPO 令二者相等即可--grpo-kl-beta 0.0表示关闭 KL 惩罚项评估间隔--eval-interval 20、每次评估 prompt 数--rl-prompts-per-eval 32与 README 中“评估可能非常昂贵”的提醒一致。环境配置示例--langrl-env-config指向的环境 YAML 声明了参与训练/评估的 Agent 及其权重。以 examples/rl/environment_configs/gsm8k.yaml 为例- agent_type: GSM8KAgent agent_args: answer_format: boxed format_reward: 0.5 weight: 1.0 evaluation_only: falseDAPO 配置 则展示了多 Agent 组合DAPOAgent参与训练weight 1.0AIMEAgent仅用于评估evaluation_only: trueweight 0.0- agent_type: DAPOAgent agent_args: format_reward: 0.0 weight: 1.0 - agent_type: AIMEAgent agent_args: format_reward: 0.0 weight: 0.0 evaluation_only: true这些配置中的agent_type对应 examples/rl/environments 目录下的具体实现如math/dapo_agent.py、math/gsm8k_agent.py、countdown/countdown_agent.py。Agent 的实际分配逻辑由GroupedRolloutGenerator.rollout_allocationsmegatron/rl/agent/api.py定义返回每个环境在每训练批中的固定组数分配实现多任务混合训练weight即各任务占比。值得关注的工程化特性围绕 RL 训练循环仓库还提供了若干值得了解的工程化模块Durable Rollout Bankmegatron/rl/rollout_bank.py把已完成的RolloutGroup在组装完成瞬间持久化到磁盘默认save/rollout_bank使 SLURM 4 小时时限导致的 SIGKILL 不再摧毁已完成的生成工作。其布局为单写者仅 rank 0面向 LustreMANIFEST.json原子地选择一个完整 generation生成数据以 JSONL 索引 二进制 sidecartokens/logprobs/masks形式存储——由于每个生成 token 同时携带 token id 与 logprobJSON 文本形式约 28 B/token而 sidecar 二进制存储可大幅压缩。重启时通过restore重放活动 generation按 checkpoint 步号T决定每个组的去留marker ≤ T 丢弃、marker T 恢复、无 marker 恢复。Rollout Bank 恢复语义同上文件注释追加写入按组fsync直写kill 最多损坏最后一条记录且校验和会在读取时丢弃checkpoint 会先写完整的新 generation含压缩后的 marker 日志再原子翻转 manifest。这两个特性与 megatron/rl/rollout_granularity.py、megatron/rl/agent/rollout_pipeline.py 共同构成了 Megatron-RL 面向大规模、长时间 RL 训练的基础设施。小结Megatron-RL 通过三组件解耦Agent/Environment、Trainer/Evaluator、Inference Interface提供了一套灵活、可扩展的 RL 后训练框架而 Off-Policy Generation 机制--rl-partial-rollouts--rl-generation-lag是其吞吐优化的核心理解 collection lag 的权衡、掌握max_effective_lag资源公式、选择合适的提交/消费粒度B/B、G/G、R/G…就能在陈旧度与硬件利用率之间找到正确平衡点。配合 examples/rl 中完整的 GRPO 实验配置与 megatron/rl 的模块化实现研究团队可以在最新的 NVIDIA 硬件上、以 Megatron-LM 的规模和效率对前沿模型进行 RL 后训练。需要再次强调该模块仍处于活跃开发中代码与接口可能持续变化使用时请以当前仓库内容为准。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考