feat: PPO with MCore

[bg51717]

What does this PR do ?

Adds full Proximal Policy Optimization (PPO) support to NeMo-RL with an actor-critic architecture, using the Megatron-Core (mcore) backend for both the policy and value models. The policy (actor) and value function (critic) are jointly trained using Generalized Advantage Estimation (GAE). Both models run on Megatron-Core with GPU/CPU offloading for colocated execution on the same set of GPUs as vLLM generation.

Backend support: This PR implements PPO on the Megatron-Core backend only. The DTensor/FSDP2 backend is not yet supported for the value model. All recipes and tests use megatron_cfg.enabled: true.

Issues

close #2048

Summary of Changes

PPO Training Algorithm

• Complete PPO training loop with critic-before-actor update order
• Multiple training steps per rollout (steps_per_epoch)
• Configurable critic warmup (policy_training_start_step) — trains value model alone before policy updates begin
• Dynamic sampling support
• Colocated architecture with GPU memory management via model offloading

Generalized Advantage Estimation (GAE)

• Token-level GAE with carry-forward masking for correct multi-turn/padding handling
• Token-level KL penalty in rewards (configurable coefficient and KL type: k1/k3)
• VAPO decoupled GAE: separate lambda for value returns vs. policy advantages
• Length-adaptive lambda: lambda_policy = 1 - 1/(alpha * response_length)
• Reward whitening

Value Model (Megatron-Core backend)

• LM backbone + scalar value head, reusing Megatron-Core policy infrastructure
• Supports TP/PP/DP parallelism, distributed optimizer, sequence packing
• GPU/CPU offloading for colocated execution with policy and vLLM
• Full checkpoint save/load including value head weights
• Clipped MSE value loss with configurable loss scale and clip range
• VAPO NLL auxiliary loss on correct samples

Shared Algorithm Improvements

• Refactored clipped PG loss to support both GRPO and PPO
• Added Reinforce++ and raw-reward advantage estimators
• GSM8K answer extraction and verification environment

Configuration and Recipes

• Base config: examples/configs/ppo_math_1B_megatron.yaml (DAPO-style PPO: no KL penalty, asymmetric clipping, dual-clip, reward scaling)
• ppo-dsr1-7b-math-8n8g-megatron — DeepSeek-R1-7B on DAPOMath-17K, 8 nodes, KL penalty + importance sampling
• ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron — Qwen2.5-1.5B-Instruct on GSM8K, 1 node, VAPO decoupled GAE

Tests

• Unit tests: 17 tests for GAE computation, value loss, advantage estimator factory; 78 tests for Megatron model setup
• Functional test: End-to-end PPO training on 2 GPUs with metric assertions on ratio clipping
• Nightly tests (2 recipes):
  • 8-node DeepSeek-R1-7B on DAPOMath, 40 steps, checks reward > 0.3 and accuracy > 0.42 at step 40
  • 1-node Qwen2.5-1.5B on GSM8K, 100 steps, checks reward > 0.85 and accuracy > 0.7 at step 100
• Reference config snapshot tests for all algorithms

Documentation

• Algorithm overview: key differences from GRPO
• In-depth guide: value model, GAE, VAPO decoupled GAE, training loop, loss functions, configuration, and metrics

Architecture

PPO Training Loop (mcore)

1. Generate responses (vLLM, colocated)
2. Score with environment (math verification)
3. Value inference → per-token V(s_t)
4. Policy logprobs → π_θ(a|s)
5. GAE → advantages A_t, returns R_t
6. Train critic (MSE on returns)
7. Train actor (clipped surrogate objective)
   8.Steps 6-7 repeat steps_per_epoch times

Experimental Results

GSM8K: Qwen2.5-1.5B-Instruct, 1 node x 8 GPUs

1. val:accuracy over steps — shows convergence on GSM8K test set
2. train/reward over steps — shows reward progression

DAPOMath-17K: DeepSeek-R1-7B, 8 nodes x 8 GPUs

Metrics to screenshot from wandb (project: nemo-rl, run: ppo-dsr1-7b-math-8n8g-megatron):

1. val:accuracy (AIME 2024) over steps — shows convergence on competition math
2. train/reward over steps — shows reward progression

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[bg51717]

/ok to test 50e878e

[bg51717]

/ok to test 47d84e6

[bg51717]

/ok to test a53ce1a

[bg51717]

/ok to test 47d5e37

[bg51717]

/ok to test 4d9ae3f

[bg51717]

/ok to test cc0b381

[bg51717]

/ok to test 40a67cc

[bg51717]

/ok to test a567bfb

[yuki-97]

@bg51717 LGTM, thanks so much for the efforts!

[qiaochuz-nv]

Test Summary — NVIDIA-NeMo/RL PR #2530 ("feat: PPO with MCore")

POR RL coverage for the new PPO + Megatron-Core actor-critic path. MR:
gitlab-master.nvidia.com/dl/DLQASH/nmfw_tests/-/merge_requests/265
Branch: qiaochuz/por_2530_ppo_mcore · Suite: nemo_llm/test_suite/rl

All cases drive the real product entrypoint examples/run_ppo.py (no mock /
no 2-GPU "shrink" for the recipe cases). Per NeMo recipe-test convention, only
ppo.max_num_steps is reduced for CI; model / parallelism / sequence / batch
dimensions stay at recipe scale so the real code path is exercised.

What the PR adds and these tests gate: the value model (critic) + GAE
machinery. GRPO has no value head, so every critic/GAE/warmup/decoupled-lambda
metric below only exists because the PR's PPO path is live.

---

1. Function / code-path coverage matrix

#	Testcase	Tier	GPUs	Recipe / config	PPO-specific code path exercised
1	test_ppo_megatron_basic	function	2	ppo_math_1B_megatron.yaml (Qwen2.5-0.5B ×2)	Core actor-critic PPO loop: clipped surrogate (ratio clamp), value head MSE loss, explained_var, critic grad_norm. Standard coupled GAE.
2	test_ppo_megatron_critic_warmup	function	2	same + policy_training_start_step=1, 3 steps	Critic warmup branch (PPO-unique): critic trains with policy frozen on step 0, policy starts at step 1. Asserts ≥3 critic-loss points + clamping active post-warmup.
3	test_ppo_megatron_kl_penalty	function	2	same + reference_policy_kl_penalty=0.01, skip_reference_policy_logprobs_calculation=false, use_kl_in_reward=false	Reference-policy KL penalty branch in PPOLossFn (loss_functions.py:325,403). Default DAPO recipe short-circuits this; test flips it ON, initializes ref model, asserts train/kl_penalty populated, finite, sane.
4	test_ppo_megatron_vapo_gae	function	2	same + adv_estimator.gae_lambda_value=1.0, length_adaptive_alpha=0.05	VAPO decoupled GAE (arXiv:2504.05118): separate lambda for policy advantages vs value returns. MC-like value targets; asserts value-loss & grad_norm DECREASE step1→step2 (real value-learning signal).
5	test_ppo_megatron_qwen1.5b_1n8g	function	8 (1n)	ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-dynbatch.yaml	Shipped 1-node recipe at native parallelism: policy TP=2/CP=2/SP, value TP=2, colocated vLLM, dynamic batching, gsm8k. Single-node → srun.
6	test_ppo_megatron_dsr1_7b_8n8g	e2e	64 (8n)	ppo-dsr1-7b-math-8n8g-megatron.yaml	Shipped multi-node recipe at full scale: DeepSeek-R1-Distill-Qwen-7B, policy TP=4/CP=2/SP, value TP=2, GAE + KL(0.001) + truncated importance sampling (tis ratio 5), seqlen 10240 / max_new 8192. sbatch via ray.sub batched-job.

Coverage spread: cases 1–4 are fast 2-GPU function smokes isolating one PPO
sub-feature each (core loop, warmup, KL, decoupled-GAE); cases 5–6 are the real
shipped recipes at native single-node and multi-node scale.

2. Metric assertions (what "pass" actually verifies)

Common to all (clipped-surrogate / actor sanity):

• train/token_mult_prob_error < 1.05 — logprob recompute matches generation
• train/probs_ratio_clamped_min > 0.79, ..._max < 1.29 — PPO ratio clamping live

Critic / value-head (PPO-unique — proves the PR's value model ran):

• train/critic/loss >= 0, magnitude-bounded (<6 basic/warmup/kl, <12 vapo)
• train/critic/explained_var <= 1.0001
• train/critic/grad_norm bounded (<150, <500 vapo)

Branch-specific teeth:

• warmup: len(train/critic/loss) >= 3
• kl_penalty: len(train/kl_penalty) >= 1, >= 0, < 1000
• vapo: mean(loss,last) < mean(loss,first) AND same for grad_norm (value learning progresses)

3. Run results

Testcase	Result	Evidence	Job
test_ppo_megatron_dsr1_7b_8n8g	PASS	summary.json on disk: status=passed, exit_code=0, duration 1044s, run_id 20260628T034758Z-0b1f468b	5524727
test_ppo_megatron_qwen1.5b_1n8g	PASS	END marker exit 0; all 5 metric assertions PASS in stdout	5524807
test_ppo_megatron_basic	PASS (reported)	per author; assertion set verified in script	—
test_ppo_megatron_critic_warmup	PASS (reported)	per author; assertion set verified in script	—
test_ppo_megatron_kl_penalty	PASS (reported)	per author; assertion set verified in script	—
test_ppo_megatron_vapo_gae	PASS (reported)	per author; assertion set verified in script	—

Evidence note: dsr1 has an on-disk summary.json (status=passed). qwen is cited
via stdout END-marker exit 0 + 5/5 assertions (job 5524807) — no summary.json was
keyed at that lustre run_id. Cases 1–4 are reported passing by the author; their
metric-assertion contents are verified from the committed scripts, but separate
run artifacts were not re-read for this summary.

group_raw_exit_code=1 in dsr1's notes is benign Ray-teardown noise, not a test
failure — the test exit_code is 0 and status=passed.

4. Coverage gaps / notes

• Only the Megatron-Core backend is exercised — that's the only backend PR
  feat: PPO with MCore #2530 supports, so this is complete for the PR.
• All cases use max_num_steps=2 (3 for warmup): they gate that the path RUNS and
  produces sane metrics, not convergence quality. Convergence is out of scope for
  POR functional/e2e tiers.
• File hashes verified byte-identical to the launched/passing scripts before commit
  (qwen 86cb2a05…, dsr1 c26e6b17…, + 4 others matched staging).