feat: Support for dtensor ppo by fujial-code · Pull Request #2837 · NVIDIA-NeMo/RL

fujial-code · 2026-06-16T08:23:23Z

What does this PR do ?

Adds DTensor/FSDP2 support for the PPO value model path and aligns the DTensor PPO 1.5B GSM8K recipe with the existing Megatron-Core PPO baseline.
The main correctness fix is a value temporal alignment issue: DTensor value inference/training was not using the same V(s_t) convention as the Megatron PPO path. After right-shifting token values to match Megatron's temporal semantics, the DTensor reward curve aligns with the Megatron baseline on the Qwen2.5-1.5B GSM8K 1-node recipe.

Issues

close [#2047]

Summary of Changes

DTensor PPO Value Model

Adds DTensorValueWorkerV2 for PPO value model training and inference on the DTensor/FSDP2 backend.
Wires DTensor value worker selection through LMValue.
Supports PPO value inference, value training, and checkpoint save/load paths for the DTensor backend.
Adds lifecycle compatibility with the PPO loop, including inference/training finish hooks where needed.

Value Temporal Alignment

Fixes the DTensor value convention to match the Megatron PPO path.
Applies right-shifted values so token-level value predictions correspond to V(s_t).
This was the key fix for the DTensor vs Megatron reward mismatch observed during GSM8K validation.

Recipe and Config Alignment

Adds the DTensor 1-node 8-GPU Qwen2.5-1.5B GSM8K PPO recipe:
- examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-dtensor.yaml
Aligns the DTensor recipe with the Megatron recipe for A/B comparison:
- ppo_epochs: 1
- math_verify_impl: hf_math_verify
- matching value loss config structure
Removes the unused GSM8K verifier-specific path and uses the existing HF math verifier for the GSM8K recipe.

Runtime Compatibility Fixes

Updates DTensor worker call sites to match current Automodel API signatures.
Removes stale kwargs that no longer exist in the current forward/checkpoint APIs.
Passes checkpointing config when saving the PPO value model checkpoint.
Adds missing imports and no-op lifecycle methods needed by the PPO loop.

Tests

Unit tests: Added DTensor PPO value temporal alignment tests for _right_shift_values and _RightShiftLossWrapper.
Functional test: Added a 2-GPU DTensor PPO smoke test using Qwen2.5-0.5B on GSM8K for 2 steps with metric sanity checks.
Nightly test: Added a 1-node 8-GPU Qwen2.5-1.5B GSM8K DTensor PPO test, checking train/reward > 0.75 and validation/accuracy > 0.65 at step 40.
Local checks: Ran Python compile, shell syntax checks, nightly dry-run, and git diff --check.

Validation

Validated the DTensor PPO path against the Megatron-Core PPO baseline on the Qwen2.5-1.5B GSM8K 1-node 8-GPU setup.

Ran 100-step PPO training for both backends with matched recipe settings:

Model: Qwen2.5-1.5B-Instruct
Dataset: GSM8K
Cluster: 1 node, 8 GPUs
Backends:
- DTensor/FSDP2 Automodel backend
- Megatron-Core backend
Matched settings:
- ppo_epochs: 1
- math_verify_impl: hf_math_verify
- matched value loss config
- DTensor train/logprob micro batch size aligned to 4 for the final run

Results:

train/reward: DTensor and Megatron follow the same learning curve and both converge to around 0.88-0.90 by step 100.
validation/accuracy: DTensor and Megatron track closely across training and reach around 0.78-0.80 by step 100.
The reward and validation accuracy curves are aligned after fixing the value temporal shift and aligning the recipe settings.
These results indicate that the DTensor PPO value path now matches the Megatron-Core PPO baseline behavior on the 1.5B GSM8K recipe.

copy-pr-bot · 2026-06-16T08:23:27Z

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.

fujial-code · 2026-06-16T08:26:11Z

/ok to test 980ef04

fujial-code · 2026-06-16T18:34:06Z

/ok to test 980ef04

copy-pr-bot · 2026-06-16T18:34:10Z

/ok to test 980ef04

@fujial-code, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

fujial-code · 2026-06-16T18:34:46Z

/ok to test 980ef04

copy-pr-bot · 2026-06-16T18:34:49Z

/ok to test 980ef04

@fujial-code, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

fujial-code · 2026-06-16T18:39:44Z

/ok to test 8681202

fujial-code · 2026-06-16T18:43:12Z

/ok to test 138ae84

fujial-code · 2026-06-17T03:33:28Z

/ok to test a9d9ea8

fujial-code · 2026-06-17T07:36:30Z

/ok to test aa7a690

Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 · 2026-06-26T10:24:44Z

/ok to test 51cad68

Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 · 2026-06-26T11:32:32Z

/ok to test c569d70

yuki-97

Two cleanup nits: dead _v2=False branch in Value.save_checkpoint and an unused-arg wrapper in the new DTensor value worker. Neither is blocking.

… actor, tighten value-backend assert Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 · 2026-06-26T12:22:59Z

/ok to test 8403d74

The PPO value worker (MegatronValueWorkerImpl) landed in main after this branch was first written and runs a Megatron model on GPU like the policy worker, but did not call bind_to_gpu_numa(). Add the bind early in __init__ for parity with the Megatron/DTensor policy workers and vLLM workers. Note: the DTensor PPO value worker (PR #2837, feat: Support for dtensor ppo) will need the same call; whichever of the two PRs merges second must add it. Signed-off-by: Terry Kong <terryk@nvidia.com>

qiaochuz-nv · 2026-06-29T17:24:45Z

POR Test Summary — NVIDIA-NeMo/RL PR #2837 (DTensor PPO support)

Result: PASS ✅

What PR #2837 adds

DTensor (FSDP2) backend support for the value model in PPO. Prior to this,
the value/critic path was Megatron-only (cf. PR #2530). The test drives the real
shipped recipe on the DTensor path with megatron_cfg.enabled: false.

It now states the test covers:
- PPO algorithm value/critic path: nemo_rl/algorithms/ppo.py
- DTensor value worker: nemo_rl/models/value/workers/dtensor_value_worker_v2.py
- DTensor policy worker: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
- Automodel DTensor setup/checkpoint path: nemo_rl/models/automodel/{setup.py,checkpoint.py}
- Ray worker registry path: nemo_rl/distributed/ray_actor_environment_registry.py
- Shipped DTensor PPO recipe: ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-valuetp2sp.yaml

Test

ID: test_ppo_dtensor_qwen1.5b_1n8g
Tier: function (1 node × 8 GPU, submit_mode: srun)
Wrapper: nemo_llm/test_suite/rl/testcases/algorithms/test_ppo_dtensor_qwen1.5b_1n8g.sh
Entrypoint: examples/run_ppo.py (real product recipe — NOT a re-wrap of PR unit tests)
Config: examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-automodel-valuetp2sp.yaml
Backend: policy.dtensor_cfg.enabled: true, value.dtensor_cfg.enabled: true, megatron_cfg.enabled: false
Model: Qwen/Qwen2.5-1.5B-Instruct; value TP2/SP
Shrink: step counts only (ppo.max_num_steps=2, max_num_epochs=1, val_period=999999); model/parallelism/batch untouched
Tags: [codecoverage, rc_p0, por]

Run

Cluster: EOS (login-eos02), account coreai_dlalgo_ci, partition batch
JOBID: 5528566 (scoring srun); driver run-id 20260629T062002Z-2cbd7c20
State: COMPLETED, duration 268s
Container: nemo-rl-nightly-20260629.sqsh
Workdir (isolated): /lustre/fsw/coreai_dlalgo_ci/qiaochuz/nemo_llm_por_2837

Metric gate (all PASS — from real stdout.log)

Assertion	Value
max `train/critic/loss` < 6.0	3.090
min `train/critic/loss` >= 0	1.982
max `train/critic/explained_var` <= 1.0001	-13.370
max `train/critic/grad_norm` < 350	105.634
max `train/token_mult_prob_error` < 1.05	1.0158
`train/probs_ratio_clamped_min` in (0.79, 1.21)	1.0
`train/probs_ratio_clamped_max` in (0.79, 1.29)	1.0

DTensor backend confirmed exercised in logs: DTensorValueWorkerV2 +
DTensorPolicyWorkerV2 workers; value model weights loaded (338 shards) and
critic loop emitted train/critic/{loss,grad_norm,explained_var}.

Status flip

testcase_list.json → function.test_ppo_dtensor_qwen1.5b_1n8g.status: untested → pass

Artifacts

~/nemo/nmfw_tests_por_2837/nemo_llm/test_suite/rl/logs/20260629T062002Z-2cbd7c20/

tests/test_ppo_dtensor_qwen1.5b_1n8g/summary.json (status: passed, job_id 5528566)
tests/test_ppo_dtensor_qwen1.5b_1n8g/stdout.log

Branch / worktree

Worktree: ~/nemo/nmfw_tests_por_2837
Branch: qiaochuz/por_2837_ppo_dtensor (off main)
slurm_config.py host/identity/workdir edits are local launch-env overrides — keep out of the committed test contract.

fujial-code added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jun 16, 2026

copy-pr-bot Bot temporarily deployed to public June 16, 2026 08:26 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 08:27 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 08:30 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 18:40 Inactive

fujial-code force-pushed the ppo-dtensor-path-a branch from 8681202 to 138ae84 Compare June 16, 2026 18:41

copy-pr-bot Bot temporarily deployed to public June 16, 2026 18:43 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 18:44 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 18:49 Inactive

fujial-code force-pushed the ppo-dtensor-path-a branch 2 times, most recently from c5dc3ae to a9d9ea8 Compare June 17, 2026 03:25

copy-pr-bot Bot temporarily deployed to public June 17, 2026 03:33 Inactive

copy-pr-bot Bot temporarily deployed to public June 17, 2026 03:34 Inactive

copy-pr-bot Bot temporarily deployed to public June 17, 2026 03:37 Inactive

copy-pr-bot Bot temporarily deployed to public June 17, 2026 07:36 Inactive

yuki-97 force-pushed the ppo-dtensor-path-a branch 5 times, most recently from 473ddaf to 6f4ddc6 Compare June 26, 2026 07:02

yuki-97 added 2 commits June 26, 2026 00:21

refactor: tighten worker finish_* lifecycle signatures

2dc053b

Signed-off-by: Yuki Huang <yukih@nvidia.com>

mirror mcore unit test

a20f96f

Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 mentioned this pull request Jun 26, 2026

PPO Dtensor value model support sequence packing + CP #2951

Open

2 tasks

yuki-97 added 5 commits June 26, 2026 00:59

assert sequence packing and CP

942ab8d

Signed-off-by: Yuki Huang <yukih@nvidia.com>

nit changes

71e5b9f

Signed-off-by: Yuki Huang <yukih@nvidia.com>

refactor: inline temporal-shift helpers into DTensor value worker

16d0c1c

Signed-off-by: Yuki Huang <yukih@nvidia.com>

update functional test

4746783

Signed-off-by: Yuki Huang <yukih@nvidia.com>

update nightly test to tpsp

51cad68

Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 added CI:L1 Run doctests, unit tests, and functional tests and removed CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) labels Jun 26, 2026

yuki-97 mentioned this pull request Jun 26, 2026

PPO DTensor value-head initialization is nondeterministic across Ray worker launches #2953

Open

fix unit test

c569d70

Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 reviewed Jun 26, 2026

View reviewed changes

Comment thread nemo_rl/models/value/lm_value.py Outdated

Comment thread nemo_rl/models/value/workers/dtensor_value_worker_v2.py Outdated

refactor(value): drop V1 DTensor branch, split worker into Impl + Ray…

8403d74

… actor, tighten value-backend assert Signed-off-by: Yuki Huang <yukih@nvidia.com>

terrykong approved these changes Jun 26, 2026

View reviewed changes

terrykong merged commit 4c0c268 into main Jun 26, 2026
164 of 167 checks passed

terrykong deleted the ppo-dtensor-path-a branch June 26, 2026 18:08

terrykong added the Ready for QA label Jun 26, 2026

qiaochuz-nv added the QA:Verified label Jun 29, 2026

fujial-code mentioned this pull request Jul 23, 2026

feat: support sequence packing and context parallelism for DTensor value models #3320

Open

Uh oh!

Conversation

fujial-code commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Summary of Changes

DTensor PPO Value Model

Value Temporal Alignment

Recipe and Config Alignment

Runtime Compatibility Fixes

Tests

Validation

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

fujial-code commented Jun 16, 2026

Uh oh!

fujial-code commented Jun 16, 2026

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

fujial-code commented Jun 16, 2026

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

fujial-code commented Jun 16, 2026

Uh oh!

fujial-code commented Jun 16, 2026

Uh oh!

fujial-code commented Jun 17, 2026

Uh oh!

fujial-code commented Jun 17, 2026

Uh oh!

yuki-97 commented Jun 26, 2026

Uh oh!

yuki-97 commented Jun 26, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yuki-97 commented Jun 26, 2026

Uh oh!

Uh oh!

qiaochuz-nv commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

POR Test Summary — NVIDIA-NeMo/RL PR #2837 (DTensor PPO support)

What PR #2837 adds

Test

Run

Metric gate (all PASS — from real stdout.log)

Status flip

Artifacts

Branch / worktree

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

fujial-code commented Jun 16, 2026 •

edited

Loading

qiaochuz-nv commented Jun 29, 2026 •

edited

Loading