fix(sc): preserve DatumSpec loss_multiplier through rollouts

[tianyi-zhang-02]

What does this PR do?

Preserves DatumSpec.loss_multiplier across the SingleController rollout boundary.

The legacy GRPO path turns that field into sample_mask, including fractional and zero weights. SingleController previously dropped it while building PromptGroupRecord, so record_to_train_batch always started from ones.

This PR only adds the missing producer-to-consumer plumbing:

• PromptGroupRecord stores the prompt-level loss multiplier;
• both rollout implementations copy it from the input DatumSpec;
• record_to_train_batch fills sample_mask from that value.

Environment masking, completion truncation, filtering semantics, and related metrics are deliberately left to the earlier #3766. This branch no longer overlaps that implementation.

Validation

Final SHA: d85b2696593ccb3ba6fe69e5f02baf20e623cb27, based on upstream main at ccbcd4cc5.

Runpod Secure Cloud, nvcr.io/nvidia/nemo-rl:v0.7.0, Python 3.13.14, PyTorch 2.11.0+cu130, vLLM 0.25.1:

Check	Result
payload/record unit tests	6 passed
real vLLM caller-level tests with Qwen2.5-1.5B	2 passed
real NeMo Gym caller and legacy-parity tests with Qwen3-0.6B	2 passed

The vLLM caller tests used GPUs 0/1 and verified that both rollout callers preserve loss_multiplier=0.25 in their produced records.

The NeMo Gym supplement used 2×H100 80 GB. Each node ran in a fresh pytest/Ray process with separate locked vLLM and NeMo Gym actor environments. The direct manager test passed in 101.92 s and preserved loss_multiplier=0.25 through two real Gym rollouts; the legacy-vs-manager parity test passed in 75.21 s after both paths completed two rollouts against the same resource, agent, and model servers.

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

[tianyi-zhang-02]

Closed a gap in my own description: grpo.py has three rules here, not two.

Besides the env flag and truncated completions, processors.py sets DatumSpec.loss_multiplier to 0.0 when the prompt itself exceeded max_seq_length. That one is unconditional — nothing to do with overlong_filtering, which is about completions — and on the batched path it reaches the loss as sample_mask = repeated_batch["loss_multiplier"].

build_sample_mask started from ones, so an over-long prompt still trained at full weight on SC. And it was not recoverable from the completions: Completion carries only truncated, env_extras, reward and message_log, so nothing about the rollout records that decision.

It now rides on PromptGroupRecord. Both construction sites in RolloutManager already have the DatumSpec in scope, so it is two lines there, and build_sample_mask uses it as the starting weight exactly as grpo.py does. Defaults to 1.0, so a record built without it behaves as before.

The updated table:

	grpo.py	SC before	SC after
starting value	DatumSpec.loss_multiplier	1.0	DatumSpec.loss_multiplier
env-flagged (mask_sample)	0	1.0	0
truncated, overlong_filtering: true	0	1.0	0

Tests go red under either mutation — starting from ones again, or not threading it from the record. tests/unit/experience/, tests/unit/single_controller/: 616 passed. Also merged current main in :)

[tianyi-zhang-02]

Heads-up on merge order, with the resolution written out.

Four of my open PRs edit the same twelve-line block — the unsupported knob list in _validate_algo_settings:

• fix(sc): preserve DatumSpec loss_multiplier through rollouts #3786 removes overlong_filtering from it (build_sample_mask now applies it)
• fix(sc): apply advantage clipping and reward scaling in the advantage stage #3787 removes reward_scaling from it (_advantage_stage now applies it)
• feat(sc): run the distillation teacher in the SingleController train pump #3846 / test(sc): add the distillation recipe and functional test for SingleController #3849 wrap the list in an is_distillation_run guard, because DistillationConfig defines none of those knobs and the comprehension would AttributeError

I checked every pair by actually merging them. #3786 × #3787 is clean now — they only ever conflicted on the comment wording and on both adding a test before the same anchor, and both are fixed. The remaining four pairs are a genuine textual conflict on that block, and no restructuring avoids it: extracting or reordering the list conflicts just as hard.

So rather than force a stack, here is the resolution. Whichever order they land in, this is the merged form:

    # An enabled one here describes shaping this run does not do. An entry
    # leaves this list when the SC path starts implementing it -- rejecting a
    # knob is only right while nobody honours it.
    #
    # DistillationConfig defines none of them: there is no reward to shape or
    # filter on, so the list cannot even be evaluated on that path.
    unsupported = (
        []
        if is_distillation_run(master_config)
        else [
            name
            for name, enabled in (
                ("use_dynamic_sampling", algo_cfg.use_dynamic_sampling),
                ("reward_scaling", algo_cfg.reward_scaling.enabled),
                ("reward_shaping", algo_cfg.reward_shaping.enabled),
            )
            if enabled
        ]
    )

I built that merge locally and ran the SC and config suites against it — 1084 passed. Happy to rebase whichever ones are left once the first lands; just say which order you want.

[tianyi-zhang-02]

Stepping back from most of this — #3766 got here first and covers more of it.

I went looking for duplicate work across the open PRs and found @macandro96's #3766 ("fix(grpo): support environment sample masking in single controller"), opened two days before this one. It solves the same problem and its coverage is a superset of mine:

• it handles overlong_filtering as well as the env mask_sample flag;
• it reports num_mask_sample_filtered, the metric grpo.py emits and which this PR does not — a gap in mine;
• it carries mask_sample and truncated as data-plane columns rather than collapsing them, so grpo.reward_shaping (which needs truncated) can reuse them later.

The one design difference I did check rather than assume: #3766 keeps a separate advantage_mask over the full prompt group. That is not a correctness difference for GRPO — GRPOAdvantageEstimator passes torch.ones_like(rewards) as the valid-mask to calculate_baseline_and_std_per_prompt and uses the mask argument only for advantages.expand(mask.shape), so narrowing sample_mask first changes no advantage value. Both give the same numbers; theirs is the more defensive arrangement.

One thing #3766 does not cover, and it is why I am not just closing this: the prompt-level DatumSpec.loss_multiplier. That is set to 0.0 when the prompt exceeded the length budget before generation ever ran, and record_to_train_batch on main drops it entirely — git grep loss_multiplier nemo_rl/experience/payload.py returns nothing, and #3766 does not add it. So an over-long prompt still trains at full weight on SC even after #3766 lands. That is completion-independent and orthogonal to the env/truncation flags.

So: I'll cut this down to just the PromptGroupRecord.loss_multiplier → build_sample_mask plumbing and drop everything that overlaps, once #3766 lands or if @macandro96 would rather I rebase on top of it now. Happy either way — say which you prefer and I'll reshape it. :)

[tianyi-zhang-02]

Correction to what I wrote above — I said #3766's coverage is a superset of this PR's. It isn't, and the difference matters in your favour, @macandro96.

I checked #3766's config.py hunk properly afterwards. It adds two AdvantageConfig field names and nothing else:

+    mask_sample_field: str = "mask_sample"
+    truncated_field: str = "truncated"

("overlong_filtering", algo_cfg.overlong_filtering) is still in the unsupported list at single_controller_utils/config.py:739 on main, and #3766 does not remove it. So on that branch a config with grpo.overlong_filtering: true raises NotImplementedError at setup and its truncation-filter path is never reached. Its env-flagged mask_sample path works; the truncated half is unreachable until something takes that entry off the list, which is what this PR does.

So the accurate split is:

	env mask_sample	truncated / overlong_filtering	num_mask_sample_filtered	prompt-level loss_multiplier
#3766	yes	implemented but gated off	yes	no
#3786 (this)	yes	yes, and unblocks the knob	no	yes

Which doesn't change my conclusion — #3766 came first, its column-based shape is the better one, and I'd rather not have two mechanisms for one rule. But the right outcome isn't "close this"; it's that #3766 needs the unsupported-list removal to make its own truncation filter live, and the prompt-level loss_multiplier is still dropped by both.

Happy to send those two as a small PR on top of #3766 once it lands, and close this. Or, if it's easier, take the two lines from here directly — no attribution needed. :)