feat(nemo-gym): train context-compacted physical traces

[aroshanghias-nvd]

Dependency

This PR depends on NVIDIA-NeMo/Gym#2870 and must not merge first. After that PR merges, this branch must update the NeMo-Gym submodule pin to an upstream commit containing the unified rollout_trace_contract.

Summary

NeMo-Gym can compact the context between model calls during one logical rollout (Gym PR #2616). Training cannot flatten that rollout into one token sequence because a later model call may use a context that is not a prefix continuation of the earlier call.

This change lets NeMo-RL train those rollouts by:

• validating Gym's exact model-call token, logprob, media, and boundary evidence;
• constructing one or more physical traces for each logical rollout;
• materializing each physical trace as an independent training row;
• assigning the logical rollout's GRPO advantage to every trainable token from that rollout; and
• padding only after all physical rows have been combined, preserving data-parallel and microbatch divisibility.

Ordinary rollouts remain on the existing single-trace path. NeMo-RL enters physical-trace materialization only when Gym supplies a valid rollout-trace contract and the evidence requires it. Undeclared prefix discontinuities remain errors and cannot silently masquerade as context compaction.

Implementation

• nemo_rl/environments/nemo_gym.py validates and normalizes Gym response evidence.
• nemo_rl/environments/nemo_gym_trace.py constructs the logical-to-physical trace plan.
• nemo_rl/experience/trace_batch_materialization.py builds trainable token, mask, logprob, and multimodal rows.
• nemo_rl/algorithms/physical_trace_training.py integrates physical rows with logical advantage computation and training batch constraints.
• nemo_rl/experience/trace_replay.py keeps synchronous and asynchronous replay handling consistent.
• GRPO uses the helpers above for both sync and async training while retaining its existing path for ordinary rollouts.

Private deterministic agents and mocked model servers used for validation are intentionally excluded from this PR.

Validation

Context Compaction Validation W&B report

Focused validation completed on DFW:

• unit tests: 417 passed;
• Ruff: clean;
• Pyrefly: zero errors.

Two deterministic 20-step multimodal GRPO runs completed with W&B enabled:

• No compaction: mean token-multiplier probability error 1.000026, maximum per-step 1.000314, rewards included both 0 and 1.
• Image compaction, K=2: mean token-multiplier probability error 1.000002, maximum per-step 1.000004, rewards included both 0 and 1.

The aggregate reward remains constant across steps by design. These runs use a private manufactured deterministic environment with a fixed reward pattern so that validation isolates token, trace, and training-data correctness rather than measuring reward improvement. That environment is test scaffolding and is not included in this PR.

The artifact audit checked all 20 optimizer steps, exact trained token rows, logical/physical ownership, compaction boundaries, and transported image history with zero reported contract errors.

After rebasing onto current NeMo-RL main, syntax validation, git diff --check, and Ruff pass on the conflict-resolved files. Full CI and the exact merged-Gym integration remain required before merge.

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

[aroshanghias-nvd]

Reviewed at 7b706e9a1 by a team of specialized agents (RL/algorithms, Gym contract, bug-finding, testing, design) with an adversarial verification pass over every finding. 49 raw findings were reduced to 6 comments: 2 disputed outright, 9 downgraded on reachability, the rest deduplicated or merged.

Thanks for this — the design holds up well under scrutiny. Two things worth calling out explicitly, because they're what made this reviewable at all:

• The pure-function seam in nemo_gym_trace.py / trace_batch_materialization.py. Roughly 950 lines of the riskiest logic (boundary inference, segment identity, padding-row masking, packed-tensor normalization) take plain mappings and tensors and return plain data — no Ray actor, no policy handle, no model. Our agents could execute counterexamples against it directly on a GPU-less node, which is not true of the analogous logic in rollouts.py/grpo.py.
• parent_indices as the single source of truth for logical→physical fan-out. Expressing that mapping once and deriving every consumer from it (with -1 as a validated padding sentinel) is what prevents a family of off-by-one bugs.

Already fixed during review — verified, no reply needed:

Issue	Fixed in
Async physical-trace training raised One physical optimizer batch requires one generation policy version whenever replay mixed trajectory ages	279d65a
Advantage broadcast had no multi-rollout test — a logical_advantages[0]-everywhere regression passed the whole suite	7b706e9
7 files failed the repo's pinned ruff isort/format hooks; branch had merge conflicts	resolved by the rebase

On the first: we specifically checked that relaxing the batch-level check to per-group didn't trade a loud crash for a silent wrong gradient. It doesn't — generation_policy_versions was a function-local set never read downstream, and validate_physical_trace_training_config already hard-requires use_importance_sampling_correction=True for async before prepare_trace_batch runs, so the relaxation is gated by the exact mechanism that makes mixed-age batches sound. The added per-group invariant is stronger than what we asked for.

One item still open, and it needs a human rather than a code change: this PR has no CI:* label (labels is []), and the request job shows skipping behind the copy-pr-bot vetter gate — so no lint or test job has ever executed on this branch. That's why the ruff failures above went unnoticed. Per the cicd skill, changes touching advantage/baseline computation are CI:L2. Please attach it and get a vetter to unblock the runners.

Review caveats, stated plainly: this was run on a GPU-less node where import torch fails and pytest is unavailable, so no test in this PR was executed. Only the pure-stdlib nemo_gym_trace.py was run directly (via importlib), which is how the anti-spoofing findings in the coverage comment were confirmed. pyrefly was also skipped — your "Pyrefly: zero errors" claim is unverified here, though all four new modules are correctly registered in pyrefly.toml project-includes. The ruff results above used the repo's pinned v0.9.9 hooks with a baseline control at the merge base.

Nothing below is a blocker on the feature's design.

Generated by Claude Code

[aroshanghias-nvd]

nemo_gym.py:1364

1 action item. PR-introduced.

TL;DR — with policy.router_replay.enabled, the last token of every non-final physical trace replays a fabricated expert route instead of falling back to the sentinel.

Background: the decode route for a turn's final token is a placeholder, normally repaired from the next turn's prefill. This seen_token_ids = [] reset happens before that repair block at nemo_gym.py:1405, whose guard is if routed_experts is not None and seen_token_ids: — so at every compaction boundary the repair is skipped, for both the logical message log and previous_physical_assistant.

Skipping is correct: after compaction the new prefill contains no route for the previous trace's final token, and routed_experts[len(seen_token_ids) - 1] would index [-1]. The gap is that nothing marks the token as unknown.

The placeholder is not the sentinel. In vLLM's routed-experts packing, full is initialized to default_route = torch.arange(topk) and expected_routes = min(max(valid_length - 1, 0), padded_length), so index valid_length - 1 — the final token — is written by neither the copy (routes_to_copy = min(expected_routes, ...)) nor the sentinel fill (full[routes_to_copy:expected_routes]). It keeps arange(topk). Megatron's router-replay fallback keys on -1, so it replays that fabricated route as a genuine routing decision. backfill_missing_routed_experts fills missing fields and won't detect a present-but-stale one.

Action: in _postprocess_nemo_gym_rollout, when trace_call["starts_physical_trace"] and turn_idx > 0, stamp the previous physical trace's final assistant token with ROUTED_EXPERTS_MISSING_ROUTE_SENTINEL, and set previous_physical_assistant = None at a trace start so it can't be cross-patched later. One token per boundary; only matters when policy.router_replay.enabled.

Confidence the code misbehaves: 85. We could not execute the Megatron replay path (no GPU), so the claim that the placeholder differs from the sentinel is from reading vLLM's packing utility, not from a run.

[aroshanghias-nvd]

Good catch—agreed. I updated the physical-trace boundary handling so the previous logical and physical assistant final-token routes are replaced with ROUTED_EXPERTS_MISSING_ROUTE_SENTINEL, and previous_physical_assistant is cleared before resetting seen_token_ids. I also added regression coverage asserting that both copies use the sentinel while the new trace retains its own prefill route. The fix is prepared in local commit 8e4f87d5 and has not been pushed yet.

[aroshanghias-nvd]

nemo_gym.py:1327

2 action items. PR-introduced.

TL;DR — the highest-traffic error in the Gym path lost its diagnostic and now points users at a contract that does not exist at the pinned Gym SHA.

AI-1 — restore the tokenization diagnostic

The assert this replaces named the three real causes — token-merging when messages are concatenated, truncated chat history, and stripped reasoning — and dumped seen_token_ids, the offending prompt_token_ids, and the diverging prefix. The new message says the rollout "contains N undeclared prompt/media discontinuity boundary or boundaries. Only an exact rollout trace contract may authorize a physical trace split."

For the overwhelmingly common trigger — a Qwen3/DeepSeek-R1-style chat template stripping prior <think> blocks on re-render, or an agent truncating long tool output — that names a mechanism the user isn't using and cannot obtain: rollout_trace_contract has zero hits in NeMo-Gym at the currently pinned c3bac96314a59f28b896f597eb9845d175bb0252.

To be clear, detection is not weakened — non-strict build_rollout_trace_plan computes the same previous_context == prompt_token_ids[:len(previous_context)] predicate and converts a violation into inferred_boundaries, which this raise surfaces. Strictness is arguably up, since the old assert was strippable under python -O. Only the diagnostic regressed.

Action: when not has_rollout_trace_contract, raise with the original cause list plus the first diverging turn's turn_id, len(previous_context), and the diverging prefix slice; keep the contract-flavored wording for the has_rollout_trace_contract case. Cheapest route is to record the first inferred boundary's turn and divergence offset into trace_plan["checks"] and interpolate.

AI-2 — drop two dead projection keys

chunk_records (nemo_gym.py:84) and guard_records (nemo_gym.py:86) appear nowhere else in this repo and nowhere in Gym at the pin. Drop them, or add a comment naming the Gym PR that will introduce them so a reader can distinguish "not yet upstream" from "typo".

[aroshanghias-nvd]

Agreed with both points. The legacy no-contract path now reports the first divergent turn, token/media continuity, previous-context length, mismatch offset, and a bounded token window, while restoring the likely causes: token merging/retokenization, truncation, and stripped reasoning. The strict exact-contract path already raises its own boundary-specific error before this check. I also removed chunk_records and guard_records from the bounded full_result projection; they are real fields in Gym PR #2870, but NeMo-RL does not consume them after ingestion. Focused trace-planner coverage passes 4/4, with adapter-level assertions added as well. Prepared in local commit 237886d2; it has not been pushed yet.

[aroshanghias-nvd]

nemo_gym_trace.py:172

4 action items. PR-introduced.

TL;DR — the test suite is green regardless of what Gym actually emits, and the fail-loud ladder that is this feature's entire safety story is largely unexercised.

AI-1 — tests validate a hand-written mock, not Gym's real contract

tests/unit/trace_test_utils.py builds the whole contract from literal dicts; across all of tests/unit/ nothing imports a Gym type beyond nemo_gym.config_types. Combined with all 8 contract identifiers being absent at pin c3bac963, every field name in the contract could be misspelled and CI would not notice — so the suite currently carries no integration signal.

Action: add one contract test that constructs the trainable output item via Gym's real model (e.g. nemo_gym.openai_utils.TokenIDLogProbMixin.model_validate(...)) and asserts rollout_trace_contract is a declared field on Gym's response model. It fails loudly at the current pin and turns green the moment the pin is bumped — which also supplies the merge-order guard that today exists only as prose in the PR description. (We grepped .github/workflows/*.yml: the submodule jobs only enforce fast-forward-ness, nothing asserts the pinned Gym exposes a given symbol.)

AI-2 — the anti-spoofing check on this line has zero coverage

Gym declaring expected_append_compatible=True while the token evidence disagrees is the most security-relevant branch in this module, and nothing tests it. Three sibling branches are also uncovered. We executed all four against the real module (it is pure-stdlib) and confirmed both the raise and the message:

• turn 2 prompt_token_ids=[9] after turn 1 [1]/[2] with append_compatible=True → "append-compatibility declaration disagrees"
• append-compatible turn 2 plus boundary_events=[{"event_id":"b2","applies_to_step":2}] → "does not correspond to a rewrite"
• media_ids=["ghost"] with media_assets={} → "references an unknown media asset"
• sampled_token_ids=[2,3] with sampled_logprobs=[-0.1] → "token/logprob lengths disagree"

AI-3 — ~8 fail-loud paths in trace_batch_materialization.py are uncovered

Grepped tests/ for each message; all returned nothing: "has no messages", "one generation policy version", "Duplicate physical trace ID", "Duplicate logical rollout ID", "no eligible action tokens", "comparison groups are incomplete", "physical message logs are incomplete", "routed_experts tensors have inconsistent worker shapes". Note test_one_group_cannot_own_different_prompts passes expected_rollouts_per_group=2, which deliberately steps around the incomplete-group check. Worth covering at least the incomplete-group and multi-version cases — both are one-line mutations of the existing _prepare helper.

AI-4 — media re-send across a compaction boundary is untested end-to-end

At a boundary the new physical trace must re-own all media, not just the delta. test_nemo_gym_postprocess_builds_exact_compacted_physical_logs sets up exactly the right data (["screen-a"] → ["screen-a","screen-b"]) but its _MockSelf has no processor, so the if processor is not None: branch that calls _resolve_images_by_media_id never executes. Separately, test_declared_media_rewrite_starts_a_physical_trace uses a non-overlapping pair (["image-a"] → ["image-b"]) where a correct implementation and a buggy delta-only one differ merely as ["image-b"] vs []; an overlapping pair discriminates properly. We ran the plan layer on the realistic pattern and got new_media_ids == [['screen-a'], ['screen-a','screen-b']] — correct, but nothing pins it.

[aroshanghias-nvd]

Agreed on the missing safety coverage, with one correction: mixed policy-version rejection was already covered for both synchronous cross-group batches and asynchronous within-group mixing, so I did not duplicate it. I added: (1) a pinned-Gym compatibility guard using Gym’s real TokenIDLogProbMixin, ContextCompactedTransportResponse, RolloutTraceContract, and ModelCallMetadata; it intentionally fails at the current c3bac963 pin until a compatible Gym revision is pinned, (2) all four strict trace-planner fail-loud cases, (3) incomplete groups, empty traces, duplicate logical/physical IDs, incomplete physical logs, zero eligible tokens, and inconsistent routed-expert shapes, and (4) overlapping-media ownership at both the planner and adapter-with-processor layers. The pure trace suite passes 8/8 and Ruff/import/format checks pass. The torch/ray-dependent subsets are not runnable in this Mac environment. Prepared in local commit e25a5f4f; it has not been pushed yet.

[aroshanghias-nvd]

grpo.py:3405

4 action items — all metric-only. No gradient impact; none of these is a blocker. PR-introduced.

TL;DR — four W&B metrics are wrong, missing, or structurally incapable of being non-zero on the physical-trace path.

AI-1 — num_mask_sample_filtered is hardcoded to 0

Here and at grpo.py:5317, while the masking itself is applied — _logical_sample_masks genuinely zeroes the sample mask on mask_sample, loss_multiplier, and truncated. So the metric reports zero env-flagged drops on exactly the path where they matter. Expose the count from PreparedTraceBatch and report it on both paths. (Relatedly, "total_num_tokens": input_lengths.numpy() now includes synthetic padding rows.)

AI-2 — advantages/sum and advantages/mean change meaning between paths

grpo.py:3595 passes physical_trace_batch.logical_advantages (shape [N,1]) where the ordinary path passes the token-broadcast [B,S] tensor. _log_mixed_rewards_and_advantages_information calls .sum()/.mean() on whatever it receives, so the same W&B keys silently switch from "sum over the padded token grid" to "sum over logical rollouts" — differing by roughly a factor of sequence length — the first time a step contains a split rollout. Any dashboard comparing a compacted run against a non-compacted one is then comparing different quantities.

Action: emit the logical stats under distinct keys (e.g. physical_trace_training/logical_advantages_{mean,sum}) and keep advantages/sum/advantages/mean computed from train_data["advantages"] on both paths.

AI-3 — the sync path drops two metrics the async path emits

Async adds physical_trace_training/scheduler_step_increment and /optimizer_steps at grpo.py:5324; sync calls only metrics.update(physical_trace_batch.metrics()). scheduler_step_increment is the single number proving the row multiplication did not accelerate the LR schedule — the metric a reviewer of a sync compaction run most wants.

Action: move both keys into PreparedTraceBatch.metrics() so both call sites get them, and delete the async-only block.

AI-4 — inferred_boundary_count is dead, and its metric is a constant zero

trace_batch_materialization.py:464 initializes it to 0 and nothing ever increments it, so the guard at trace_batch_materialization.py:569 (raise ValueError("Physical training cannot consume inferred discontinuities")) is unreachable and physical_trace_training/inferred_discontinuities reads 0 forever. The real enforcement lives upstream at nemo_gym.py:1327, which reads trace_plan["checks"]["inferred_boundary_count"] — that value is never propagated into rollout_batch.

Action: either carry the per-rollout count through and populate the metric, or delete the variable, the guard, and the metric. A metric named inferred_discontinuities that structurally cannot be non-zero is worse than not publishing one.

[aroshanghias-nvd]

Agreed. I addressed all four metric issues in the current local patch:

• num_mask_sample_filtered now reports explicit logical mask_sample flags on the physical path, matching the legacy metric semantics; truncation and zero loss multipliers do not inflate it.
• total_num_tokens excludes synthetic physical padding rows.
• Standard advantage histogram/sum logging uses the materialized physical train_data["advantages"]; logical-rollout sum/mean are reported separately under physical_trace_training/logical_advantages/*.
• Shared prepared-batch metrics now provide scheduler_step_increment and optimizer_steps to both sync and async paths.
• The unreachable materialization-level inferred-boundary field/guard/constant-zero metric was removed; the upstream NeMo-Gym ingestion rejection remains intact.

Added focused materialization and sync-path metric regressions. Ruff import/check/format, Python compilation, and git diff --check pass. The focused pytest suite cannot collect in this Mac environment because torch is not installed. Local commit: 25660ca3 (not pushed yet).

[aroshanghias-nvd]

grpo.py:3402

2 action items + 1 follow-up. PR-introduced. All low severity — none of this is a bug, and we are not asking you to restructure the feature.

TL;DR — the new path is cleanly isolated in its modules but not at the call site, and the sync/async copies have already drifted three ways within a single commit.

AI-1 — drop the redundant bool and the 10 asserts it forces

materialize_physical_traces = physical_trace_batch is not None creates a bool the type checker cannot use to narrow the Optional, so 10 downstream sites re-assert assert physical_trace_batch is not None (sync: 3404, 3506, 3594, 3665, 3983; async: 5250, 5354, 5421, 5506, 5930). The file already contradicts itself on the idiom — :3569, :3993, :5390 and :5940 branch on physical_trace_batch is not None directly.

Action: delete materialize_physical_traces and branch on the Optional everywhere. All 10 asserts disappear, pyrefly narrows correctly, ~25 lines go, zero behaviour change. Worth doing independently of the follow-up below.

AI-2 — hoist the deferred import

trace_batch_materialization.py:590 defers from nemo_rl.experience.rollouts import backfill_missing_routed_experts into the function body, with a comment that explains why physical logs are normalized separately — not why the import is deferred. We verified there is no cycle: rollouts.py has zero references to trace_batch_materialization. Hoist it to module top, or keep it deferred and change the comment to name the cycle, per the linting skill's "say which".

Follow-up — tracking issue, not this PR

The integration is duplicated across grpo_train and async_grpo_train, and three drifts already exist one commit in: sync keeps metrics_logging_data["content"] while async keeps flat_messages_content; sync uses metrics.update(...) while async builds a separate physical_trace_metrics dict; and sync omits the two metrics noted in AI-3 of the metrics comment. Two of those three were independently filed as user-visible metric bugs in this same review, so this is a demonstrated hazard rather than a predicted one.

Worth a tracking issue to extract the shared prepare/override seam so the two loops cannot drift further. We are explicitly not proposing a general protocol for future algorithms — validate_physical_trace_training_config rejects everything but GRPO, so that extension is hypothetical.

[aroshanghias-nvd]

Agreed on both requested cleanups. The current local patch now branches directly on physical_trace_batch is not None throughout both sync and async loops, removing the derived boolean and all nine redundant narrowing assertions present at the current head. I also hoisted backfill_missing_routed_experts to the module imports after confirming there is no import cycle.

Ruff import/check/format, Python compilation, and git diff --check pass. This is local commit 4524e2c5 and has not been pushed yet. I have not created the suggested follow-up tracking issue; the preceding metric patch removed the demonstrated metric drift, and any broader extraction can remain separate from this PR.

[aroshanghias-nvd]

physical_trace_training.py:110

3 action items. PR-introduced.

AI-1 — the feature ships with no documentation

git diff --stat <base> HEAD -- docs/ is empty. This is a ~1500-line user-facing feature with a new Gym contract surface (rollout_trace_contract, model_call_metadata, boundary_events, media_assets), a new physical_trace_training/* metric family, and the 20-item unsupported-config matrix in this function — none of it discoverable.

Action: add docs/guides/context-compacted-trace-training.md covering the Gym-side contract fields, the exact unsupported-config list from this function, what each physical_trace_training/* metric means, and the policy.train_global_batch_size == num_prompts_per_step * num_generations_per_prompt constraint; register it in docs/index.md.

AI-2 — the deny-list is an allow-list by omission

This function names 20 rejected options, which is the right behaviour — failing at once with a bulleted list of every violated precondition is exactly what the error-handling skill asks for. The problem is the direction: when someone adds a new grpo or loss_fn field next quarter, nothing tells them to triage it against physical-trace semantics, and the run silently applies it on an unvalidated path. That reason-to-change ("GRPO gained a knob") is far more frequent than anything else in this file.

Action: add a unit test asserting set(GRPOConfig.model_fields) | set(ClippedPGLossConfig.model_fields) equals a frozenset checked into the test. A new config field then turns the test red until someone either adds it to the deny-list or records it as reviewed-and-safe. ~15 lines converts a silent-wrong into an explicit decision.

AI-3 — two specific gaps in the validation evidence

The W&B report and the token-multiplier probability error of ~1.000002 are good evidence for numerical correctness of the compacted path, and the deterministic-environment framing is sound. Two gaps remain:

1. No throughput or memory number, even though this PR adds the instrumentation to produce one — _actor_peak_rss_gib and context_compaction_transport_reduction_ratio / context_compaction_ray_env_extras_bytes. Those values are already being logged by the K=2 run, so reporting them costs nothing, and they are what a user needs to decide whether to enable compaction.
2. No evidence the ordinary path is unregressed. Since the contract is absent from Gym at the pinned SHA, the "no compaction" run is K=1 inside the new code path, not the pre-PR code path. Meanwhile this PR touches nemo_gym.py's postprocess loop for every rollout (a build_rollout_trace_plan call per rollout, plus per-call token/logprob list validation over full arrays). A short non-contract NeMo-Gym GRPO run at main vs this branch showing identical step-1 loss/advantage metrics would close it.

Action: post (1) from the existing run, and (2) as a short A/B.