fix(grpo): pass rewards into the GRPO train batch for VAPO's NLL term

[khazic]

Problem

ClippedPGLossFn's VAPO positive-example NLL term is gated like this:

nll_loss = torch.tensor(0.0, device=mask.device)
if self.positive_example_nll_weight > 0 and "rewards" in data:
    ...
loss = actor_loss + kl + self.positive_example_nll_weight * nll_loss

If "rewards" is absent from the data dict passed to the loss, the term is silently zero regardless of positive_example_nll_weight, with no warning.

ppo.py already carries this key: "rewards": repeated_batch["total_reward"] is part of its train_data dict. GRPO's two in-process training paths do not:

• grpo_train() builds train_data inline without a rewards key.
• the async no-TQ path's _build_async_grpo_train_data() builds the same shape of dict, also without it.

So setting loss_fn.positive_example_nll_weight on a GRPO run is a silent no-op. The config accepts the value, training runs normally, and the loss is bit-identical to leaving the weight at 0.

Fix

Add "rewards" to both dict-construction sites, mirroring PPO's existing pattern exactly:

• In _build_async_grpo_train_data(): "rewards": repeated_batch["total_reward"].
• In grpo_train()'s inline train_data dict: "rewards": rewards, where rewards is the same batch-aligned reward tensor already computed a few lines above (repeated_batch["total_reward"], or repeated_batch["filtered_reward"] under DAPO dynamic sampling) and used to build advantages. It stays row-aligned with repeated_batch through this section: dynamic sampling has already filtered repeated_batch to match by the time train_data is built, and neither overlong filtering nor the mask-sample filter that follow drop rows, they only zero loss_multiplier.

Evidence

Three tests, from the API boundary up to the actual loss activating:

1. TestBuildAsyncGrpoTrainData::test_train_data_carries_rewards calls _build_async_grpo_train_data() directly and asserts "rewards" is present and equals repeated_batch["total_reward"].
2. TestBuildAsyncGrpoTrainData::test_rewards_row_matches_reward_hungry_loss_term feeds the resulting train_data straight into ClippedPGLossFn with positive_example_nll_weight=0.5 and asserts metrics["positive_nll_loss"] != 0, i.e. the fix actually turns the feature on end-to-end.
3. test_grpo_train_sync_train_data_carries_rewards locks the synchronous path's inline dict literal via source inspection (grpo_train() builds this dict inline rather than through an extracted, directly callable helper, so invoking the real function would require standing up rollout, policy, and generation actors).

On the current main, all three fail:

FAILED tests/unit/algorithms/test_grpo.py::TestBuildAsyncGrpoTrainData::test_train_data_carries_rewards
FAILED tests/unit/algorithms/test_grpo.py::TestBuildAsyncGrpoTrainData::test_rewards_row_matches_reward_hungry_loss_term
FAILED tests/unit/algorithms/test_grpo.py::test_grpo_train_sync_train_data_carries_rewards

AssertionError: assert 'rewards' in {'input_ids': ..., 'sample_mask': tensor([1., 1., 1., 1.])}
...
AssertionError: grpo_train's train_data dict lost the 'rewards' key; loss_fn.positive_example_nll_weight would silently no-op again

3 failed

With the fix applied, all three pass:

3 passed

The full test_grpo.py suite (144 tests, covering dynamic sampling, mask/reward-penalty filtering, config validation, refit handshakes, and rollout-metric aggregation) passes on the fix branch, so the change does not regress any other GRPO path:

144 passed

Known related gap, out of scope here

The data-plane (TQ) synchronous path (grpo_sync.py, a documented sibling of grpo_train() for data_plane.enabled=true) has the identical gap, traced end to end:

• DP_TRAIN_FIELDS (nemo_rl/data_plane/schema.py), the column list workers fetch for train_from_meta, does not include a rewards-shaped field.
• The rollout actor's TQ bulk write (nemo_rl/experience/sync_rollout_actor.py, bulk_batch = BatchedDataDict[Any](...) built from DP_TRAIN_FIELDS plus multimodal extras) does not write one either, even though total_reward is already one of the PROMOTE_1D_FIELDS the schema knows how to carry through the Mooncake adapter, so the reward is available in the data plane, just not under a name or field list this fetch reads.
• The driver's own write_to_dataplane call in grpo_sync.py writes only {"advantages", "sample_mask"}.

So "rewards" in data is always False in ClippedPGLossFn on this path too. Fixing it means adding a fetched column and aliasing it to rewards for the loss call, in a materially different fetch mechanism (a distributed column store, not a literal dict) that this PR does not touch and that deserves its own focused change and testing.

Scope

This only adds the rewards key to the two in-process GRPO training paths. Everything else about how train_data is built, and the loss function itself, is unchanged.

[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.

[khazic]

Verified end to end on a GRPO run. The observable is positive_nll_loss, which ClippedPGLossFn already reports as its own metric, so the term can be watched directly instead of being inferred from the total loss. That also makes the check immune to rollout sampling noise.

The term is inert on main

GRPO, Qwen3-0.6B, gsm8k, loss_fn.positive_example_nll_weight: 0.1, 3 steps.

build	positive_nll_loss, steps 1 to 3
main	0.000000 / 0.000000 / 0.000000
with this PR	1.077010 / 2.346028 / 2.476855

On main the configured weight has no effect whatsoever. Training runs to completion, no warning is emitted, and the term contributes exactly zero at every step. Anyone who sets positive_example_nll_weight on a GRPO run today gets a run that is bit for bit a run with the weight left at 0.

As a control, two independent runs of the unpatched configuration reproduce each other exactly at step 1 (train/loss 0.022877, train/reward 0.093750), so the comparison is not being carried by run to run variation.

Please merge this together with #3791

This is the part I would ask reviewers to weigh, because merging this PR alone makes the situation worse rather than better.

The term is currently dead, so the normalization bug in #3791 is latent. Merging this PR activates the term, and it comes back carrying that bug. Measured on the same run:

build	positive_nll_loss at step 1	train/loss
main, term inert	0.000000	0.022877
this PR only	1.077010	0.130577
this PR plus #3791	0.033657	0.026242

1.077010 / 0.033657 = 32.000, exactly train_global_batch_size, which is the number of microbatches the global batch is split into. That is the signature of #3791: the term is normalized by each microbatch's own correct-token count, so every microbatch contributes a full local mean instead of its share of one global mean.

So the three possible states are:

• main today: the feature does not work, silently.
• this PR alone: the feature works, with a 32x inflated weight.
• this PR with fix(loss): normalize the VAPO positive-example NLL by the global token count #3791: the feature works correctly.

The middle state is worse than the first, because a wrong contribution to the gradient is harder to notice than no contribution at all. I would suggest landing them together, or landing #3791 first.

Note on the shape of the fix

The change mirrors what ppo.py already does ("rewards": repeated_batch["total_reward"] in its train_data), so after this PR all three paths that build a train batch carry the key the loss function gates on. That removes the silent-no-op failure mode rather than papering over it.

A separate improvement worth considering, outside this PR: ClippedPGLossFn silently skips the term when positive_example_nll_weight > 0 but "rewards" is absent. A warning there would have surfaced this immediately. Happy to open a follow-up if that seems useful.