NVIDIA-NeMo / NVIDIA-NeMo/RL

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

Open
#2,953 0 comments 0 reactions 1 assignee Claimed by @fujial-code View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Description

The original report observed a much higher initial critic loss for DTensor PPO with value dynamic batching, while reward remained aligned: [original PR discussion](https://github.com/NVIDIA-NeMo/RL/pull/2837#discussion_r3480775913).

Original DB-OFF/DB-ON comparison

Follow-up experiments show that the mismatch was not caused by dynamic batching. It came from nondeterministic initialization of the DTensor value head across Ray worker launches.

## Root cause

The PPO controller seeded its own process, but each DTensor value worker runs in a separate Ray process. The base causal-LM checkpoint does not contain a regression `score` head, so the head is initialized when the worker constructs the value model. Before the fix, the controller seed did not reach this initialization, and separate DB-OFF/DB-ON launches received different value heads.

MCore did not show the same cross-launch symptom because its worker setup already seeded model initialization. A shared numeric seed does not make fresh MCore and Automodel heads identical, because the two backends use different head implementations and initialization schemes.

## Resolution

PR #3321 makes fresh DTensor value-model initialization reproducible and enables value dynamic batching:

- add `seed` to the value configuration and default it to `${ppo.seed}`;
- seed each DTensor value worker after distributed setup and before checkpoint-manager/model construction;
- add fresh-initialization and checkpoint-pinned TP2+SP equivalence coverage for DB-OFF versus DB-ON.

## Validation

- Two independent Automodel launches with `seed=42` saved the value model before the first optimizer step. All 584 checkpoint tensors matched exactly, including `score.weight` and `score.bias`:
- [repeat 1](https://wandb.ai/nvidia/nemo-rl-ppo-port-v1/runs/pr3321-automodel-seed42-r1)
- [repeat 2](https://wandb.ai/nvidia/nemo-rl-ppo-port-v1/runs/pr3321-automodel-seed42-r2)
- Fresh TP2+SP 100-step DB-OFF/DB-ON runs with `seed=42` closely align without sharing a checkpoint:
- [DB OFF](https://wandb.ai/nvidia/nemo-rl-ppo-port-v1/runs/seed42-db-off-100)
- [DB ON](https://wandb.ai/nvidia/nemo-rl-ppo-port-v1/runs/seed42-db-on-100)
- critic-loss correlation: `0.9945`
- explained-variance correlation: `0.9956`
- reward correlation: `0.9880`

These results isolate the original mismatch to value-head initialization rather than dynamic batching.

## Related follow-up

Issue #3412 remains a separate follow-up. Fixing the value seed removes the fresh-head randomness confound, but does not eliminate the remaining critic-metric regression associated with the vLLM rollout-distribution change.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.