From 303126a5972f9acb604231bcdf69d961269e4e58 Mon Sep 17 00:00:00 2001 From: Elad Sarafian Date: Thu, 10 Sep 2026 03:55:06 -0700 Subject: [PATCH 1/6] feat(grpo): train all Gym traces per logical rollout Signed-off-by: Elad Sarafian --- docs/guides/gym-multi-trace-grpo.md | 95 +++++ examples/configs/grpo_math_1B.yaml | 2 + nemo_rl/algorithms/grpo.py | 203 +++++++++- nemo_rl/environments/gym_traces.py | 267 ++++++++++++++ nemo_rl/environments/nemo_gym.py | 67 ++++ nemo_rl/experience/rollouts.py | 50 ++- nemo_rl/models/policy/interfaces.py | 3 + nemo_rl/models/policy/lm_policy.py | 14 + .../policy/workers/megatron_policy_worker.py | 12 +- pyrefly.toml | 1 + tests/unit/environments/test_gym_traces.py | 346 ++++++++++++++++++ .../test_nemo_gym_token_capture.py | 82 +++++ .../unit/reference_configs/grpo_math_1B.yaml | 2 + 13 files changed, 1122 insertions(+), 22 deletions(-) create mode 100644 docs/guides/gym-multi-trace-grpo.md create mode 100644 nemo_rl/environments/gym_traces.py create mode 100644 tests/unit/environments/test_gym_traces.py diff --git a/docs/guides/gym-multi-trace-grpo.md b/docs/guides/gym-multi-trace-grpo.md new file mode 100644 index 00000000000..f5b4dcb5a72 --- /dev/null +++ b/docs/guides/gym-multi-trace-grpo.md @@ -0,0 +1,95 @@ +# Train multiple Gym traces per rollout + +`grpo.gym_multi_trace: true` trains all independently conditioned sequences +returned by Gym's `training_traces` v1 envelope. This supports agent forks, +interleaved conversations, and context compaction without flattening unrelated +chats into one causal sequence. The default remains the existing single-response +Gym path. + +## Configuration + +Apply these overrides to a working synchronous NeMo Gym Megatron recipe: + +```yaml +grpo: + gym_multi_trace: true + num_prompts_per_step: 2 + num_generations_per_prompt: 4 + async_grpo: + enabled: false +data_plane: null +policy: + train_global_batch_size: 8 # logical rollouts, before trace expansion + train_micro_batch_size: 1 + megatron_cfg: + enabled: true +loss_fn: + token_level_loss: true + sequence_level_importance_ratios: false +env: + should_use_nemo_gym: true + should_mask_flagged_samples: true + nemo_gym: + token_id_capture: + enabled: true + all_agents: true + delivery: all_traces + builder: per_request # or prefix_merging + dir: /shared/gym-capture +``` + +The Gym source must support `delivery: all_traces` and the v1 trace envelope. +Capture writers and the Gym actor must see the same capture directory. +RL calls Gym's finalizer after each completed rollout and retains the capture +files for inspection; it does not retire them before durable training handoff. + +## Identity, reward, and loss + +A task group is the original input example together with its independent +rollouts. Group membership is assigned before invoking the agent and survives +prompt rewriting. Gym's unique `rollout_id` identifies one such rollout; +`trace_id` identifies a physical training row within it. + +GRPO computes a scalar advantage once per logical rollout using its task group +and outcome reward, then broadcasts that advantage to its eligible sampled +tokens. Trace count does not create extra reward observations. A masked rollout +or a rollout with no in-limit eligible trace is excluded from its group's +baseline and standard deviation. An overlong trace is individually replaced by +an inert row; valid sibling traces keep training. An entirely ineligible batch +fails before an optimizer update. + +The objective is the mean over unique eligible sampled tokens in the complete +logical batch. It does not apply inverse-trace-count or equal-session weights. +Consequently longer sampled responses carry more token weight. Outcome reward +broadcast is not a process reward or a claim of causal credit assignment. + +Each trace carries full `token_ids`, aligned `generation_logprobs` and a binary +`loss_mask`. `sampled_spans` identifies the model call owning each trainable +span. RL checks vector lengths, finiteness, masked initial conditioning tokens, +span coverage, and unique ownership. It never retokenizes training IDs. Shared +ancestor copies and intervening context remain masked. Prefixes and masks are +not inferred from rendered text or from a terminal `response.output`. + +All physical rows form one optimizer update. Zero-loss rows pad the batch to +the data-parallel size times microbatch size; no trailing traces are dropped. +The Megatron learning-rate scheduler advances by the logical rollout count, +while checkpoint progress retains its existing count of original input prompts. +Metrics record logical rollouts, physical rows, padding, overlong traces, and +invalid rollouts. Training JSONL includes task-group, rollout, trace, and logical +row identifiers when training-data logging is enabled. + +## Supported scope + +This initial path requires synchronous GRPO, a text-only Megatron policy, and +token-level loss and importance ratios. Setup rejects asynchronous/TQ/staged +capture, custom capture sinks, router replay, multimodal models, alternative +advantage estimators, dynamic sampling, legacy episode-length filtering/shaping, +message-level penalties, and post-advantage sequence-logprob-error masking. +`train_global_batch_size` must equal the logical rollout count per step. +Use `data_plane: null` or a complete disabled data-plane configuration. + +The regression suite tests exact token custody, logical grouping and baseline +masking, unequal trace lengths/counts, individual overlength masking, physical +padding, configuration guards, and equivalent split/merged loss and gradients. +The Gym-marked tests also exercise real capture-store records through Gym's +finalizer and the RL actor for both builders. diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index a9364d5ccee..f952b8901ab 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -1,5 +1,7 @@ # GRPO Algorithm Configuration grpo: + # Opt in to Gym all_traces delivery; synchronous Megatron only. + gym_multi_trace: false num_prompts_per_step: 32 num_generations_per_prompt: 16 max_rollout_turns: 1 # for multi-turn rollouts. Math Environments just have 1 turn (answering the question) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index f31b01007f4..c9cdeb17c77 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -85,6 +85,7 @@ prepare_segment_topology, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.gym_traces import prepare_gym_trace_batch from nemo_rl.environments.nemo_gym import should_use_nemo_gym, spinup_nemo_gym_actor from nemo_rl.experience.interfaces import ( FRONTIER_ORDINAL_KEY, @@ -306,6 +307,8 @@ def _require_unwanted_token_ids_when_penalized(self) -> "RewardPenaltyConfig": class GRPOConfig(BaseModel, extra="allow"): + # Train all Gym-owned traces per logical rollout (synchronous Megatron only). + gym_multi_trace: bool = False num_prompts_per_step: int = 32 num_generations_per_prompt: int = 16 max_num_epochs: int = 1 @@ -449,6 +452,100 @@ class MasterConfig(BaseModel, extra="allow"): # =============================================================================== +def _validate_gym_multi_trace_capability(master_config: MasterConfig) -> None: + """Fail before worker allocation for unsupported trace-training paths.""" + grpo = master_config.grpo + gym = master_config.env.get("nemo_gym") or {} + capture = gym.get("token_id_capture") or {} + if not grpo.gym_multi_trace: + if capture.get("delivery") == "all_traces": + raise ValueError( + "Gym all_traces delivery requires grpo.gym_multi_trace=true" + ) + return + if not master_config.env.get("should_use_nemo_gym"): + raise ValueError("gym_multi_trace requires env.should_use_nemo_gym=true") + if not capture.get("enabled") or capture.get("delivery") != "all_traces": + raise ValueError( + "gym_multi_trace requires enabled Gym token_id_capture with delivery=all_traces" + ) + if not capture.get("all_agents"): + raise ValueError("gym_multi_trace requires token_id_capture.all_agents=true") + if ( + capture.get("sink") + or capture.get("external_staging") + or (master_config.data_plane or {}).get("enabled") + ): + raise NotImplementedError( + "gym_multi_trace requires local capture and data_plane.enabled=false; staged/TQ capture is unsupported" + ) + if grpo.async_grpo is not None and grpo.async_grpo.enabled: + raise NotImplementedError( + "gym_multi_trace currently supports synchronous GRPO only" + ) + policy = master_config.policy + if not (policy.get("megatron_cfg") or {}).get("enabled"): + raise NotImplementedError( + "gym_multi_trace currently supports the Megatron policy only" + ) + if policy.get("is_vlm") or router_replay_enabled(policy): + raise NotImplementedError( + "gym_multi_trace does not support multimodal or router-replay payloads" + ) + if ( + grpo.adv_estimator.name != "grpo" + or master_config.on_policy_distillation is not None + ): + raise NotImplementedError( + "gym_multi_trace requires the standard GRPO advantage estimator" + ) + if ( + not master_config.loss_fn.token_level_loss + or master_config.loss_fn.sequence_level_importance_ratios + ): + raise NotImplementedError( + "gym_multi_trace requires token-level loss and token-level importance ratios" + ) + if ( + grpo.use_dynamic_sampling + or grpo.overlong_filtering + or grpo.reward_shaping.enabled + ): + raise NotImplementedError( + "gym_multi_trace does not support dynamic sampling or episode-length filtering/shaping; trace lengths are filtered independently" + ) + if grpo.seq_logprob_error_threshold is not None: + raise NotImplementedError( + "gym_multi_trace does not yet support post-advantage sequence error masking" + ) + if ( + grpo.invalid_tool_call_advantage is not None + or grpo.malformed_thinking_advantage is not None + ): + raise NotImplementedError( + "gym_multi_trace does not support message-level advantage penalties" + ) + if any( + master_config.reward_penalties.model_dump()[flag] + for flag in _REWARD_PENALTY_FLAGS + ): + raise NotImplementedError( + "gym_multi_trace does not support legacy response-based reward penalties" + ) + effort = _get_effort_config(master_config) + if effort is not None and effort.low_weight and effort.low_string: + raise NotImplementedError( + "gym_multi_trace does not support legacy response-length effort shaping" + ) + if master_config.env.get("should_mask_flagged_samples") is False: + raise ValueError("gym_multi_trace requires masking invalid Gym rollouts") + logical_gbs = grpo.num_prompts_per_step * grpo.num_generations_per_prompt + if policy["train_global_batch_size"] != logical_gbs: + raise ValueError( + "gym_multi_trace requires train_global_batch_size to equal the logical rollout count per step" + ) + + def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None: """Reject configurations whose media transfer path is not qualified.""" if not master_config.grpo.deduplicate_multimodal_data: @@ -554,6 +651,7 @@ def setup( """ # Start timing the entire setup process setup_start_time = time.perf_counter() + _validate_gym_multi_trace_capability(master_config) # Extract individual configs for easier access policy_config = master_config.policy @@ -3086,6 +3184,13 @@ def grpo_train( enabled=master_config.grpo.debug_payload_metrics, ) ) + if master_config.grpo.gym_multi_trace: + # Stable input-example groups survive harness prompt rewrites. + repeated_batch["gym_task_group_id"] = torch.arange( + batch.size + ).repeat_interleave( + master_config.grpo.num_generations_per_prompt + ) # Convert LLMMessageLogType to FlatMessagesType for generation batched_flat, input_lengths = batched_message_log_to_flat_message( repeated_batch["message_log"], @@ -3207,7 +3312,13 @@ def grpo_train( master_config.grpo.debug_payload_metrics ), ) - input_ids = nemo_gym_rollout_result.input_ids + input_ids = ( + nemo_gym_rollout_result.final_batch[ + "gym_task_group_id" + ].unsqueeze(-1) + if master_config.grpo.gym_multi_trace + else nemo_gym_rollout_result.input_ids + ) repeated_batch = nemo_gym_rollout_result.final_batch rollout_metrics = nemo_gym_rollout_result.rollout_metrics del nemo_gym_rollout_result @@ -3397,6 +3508,29 @@ def grpo_train( num_mask_sample_filtered = _apply_mask_sample_filter(repeated_batch) metrics["num_mask_sample_filtered"] = num_mask_sample_filtered + trace_batch = None + if master_config.grpo.gym_multi_trace: + trace_batch = prepare_gym_trace_batch( + repeated_batch, + estimator=adv_estimator, + pad_token_id=tokenizer.pad_token_id, + max_sequence_length=master_config.policy[ + "max_total_sequence_length" + ], + row_multiple=policy.data_parallel_size + * master_config.policy["train_micro_batch_size"], + ) + repeated_batch = trace_batch.batch + metrics.update( + { + "gym/logical_rollouts": trace_batch.logical_count, + "gym/physical_rows": trace_batch.physical_count, + "gym/padding_rows": trace_batch.padding_count, + "gym/overlong_traces": trace_batch.overlong_trace_count, + "gym/invalid_rollouts": trace_batch.invalid_rollout_count, + } + ) + add_grpo_token_loss_masks_and_generation_logprobs( repeated_batch["message_log"] ) @@ -3531,7 +3665,11 @@ def grpo_train( else: seq_error_result = compute_and_apply_seq_logprob_error_masking( train_data=train_data, - rewards=rewards, + rewards=( + rewards[trace_batch.logical_indices] + if trace_batch is not None + else rewards + ), seq_logprob_error_threshold=seq_logprob_error_threshold, ) seq_logprob_error_metrics = seq_error_result @@ -3555,14 +3693,21 @@ def grpo_train( sample_mask = train_data["sample_mask"] mask = token_mask * sample_mask.unsqueeze(-1) - train_data["advantages"] = adv_estimator.compute_advantage( - prompt_ids=prompt_ids_for_adv, - rewards=rewards, - mask=mask, - repeated_batch=repeated_batch, - logprobs_policy=train_data["prev_logprobs"], - logprobs_reference=train_data.get("reference_policy_logprobs"), - ) + if trace_batch is not None: + train_data["advantages"] = trace_batch.advantages.unsqueeze( + -1 + ).expand_as(mask) + else: + train_data["advantages"] = adv_estimator.compute_advantage( + prompt_ids=prompt_ids_for_adv, + rewards=rewards, + mask=mask, + repeated_batch=repeated_batch, + logprobs_policy=train_data["prev_logprobs"], + logprobs_reference=train_data.get( + "reference_policy_logprobs" + ), + ) del prompt_ids_for_adv # Log rewards and advantages information @@ -3602,11 +3747,22 @@ def grpo_train( **{"rl.iteration": total_steps + 1}, ), ): - train_results = policy.train( - train_data, - loss_fn, - timer=timer, - ) + if trace_batch is not None: + # All physical rows form one optimizer update. The scheduler + # advances by logical rollouts, independent of trace counts. + train_results = policy.train( + train_data, + loss_fn, + timer=timer, + gbs=train_data.size, + scheduler_step_samples=trace_batch.logical_count, + ) + else: + train_results = policy.train( + train_data, + loss_fn, + timer=timer, + ) # Recompute KV scales after policy training if needed if sync_kv_scales: @@ -3913,7 +4069,22 @@ def grpo_train( if "agent_ref" in repeated_batch: log_data["agent_ref"] = repeated_batch["agent_ref"] log_data["content"] = flat_messages["content"] - log_data["rewards"] = rewards.tolist() + log_data["rewards"] = ( + repeated_batch["total_reward"].tolist() + if trace_batch is not None + else rewards.tolist() + ) + if trace_batch is not None: + for key in ( + "gym_task_group_id", + "gym_logical_index", + "gym_rollout_id", + "gym_trace_id", + ): + value = repeated_batch[key] + log_data[key] = ( + value.tolist() if isinstance(value, torch.Tensor) else value + ) if master_config.grpo.use_dynamic_sampling: log_data["filtered_rewards"] = rewards.tolist() log_data["rewards"] = repeated_batch["total_reward"].tolist() diff --git a/nemo_rl/environments/gym_traces.py b/nemo_rl/environments/gym_traces.py new file mode 100644 index 00000000000..e265dbc9f6f --- /dev/null +++ b/nemo_rl/environments/gym_traces.py @@ -0,0 +1,267 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Gym trace custody and logical-rollout GRPO batch expansion. + +No tokenizer is used to construct training IDs. Each physical row has its own +causal context; only sampled spans owned by that row receive loss. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch + +from nemo_rl.data.interfaces import LLMMessageLogType +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + +if TYPE_CHECKING: + from nemo_rl.algorithms.advantage_estimator import GRPOAdvantageEstimator + + +@dataclass(frozen=True) +class GymTrainingTrace: + rollout_id: str + trace_id: str + token_ids: tuple[int, ...] + generation_logprobs: tuple[float, ...] + loss_mask: tuple[int, ...] + + +def parse_gym_training_traces(envelope: dict[str, Any]) -> list[GymTrainingTrace]: + """Validate the public Gym v1 wire contract, including sampled-span ownership. + + Invalid custody is an error, never a reason to retokenize or silently fall + back to a legacy terminal response. + """ + if ( + type(envelope.get("schema_version")) is not int + or envelope["schema_version"] != 1 + ): + raise ValueError("Gym training_traces requires schema_version=1") + rollout_id = envelope.get("rollout_id") + if not isinstance(rollout_id, str) or not rollout_id: + raise ValueError("Gym training_traces requires a nonempty rollout_id") + if envelope.get("builder") not in {"per_request", "prefix_merging"}: + raise ValueError("Unsupported Gym training_traces builder") + rows = envelope.get("traces") + if not isinstance(rows, list) or not rows: + raise ValueError("Gym training_traces.traces must be a list") + trace_ids: set[str] = set() + owned_call_ids: set[str] = set() + traces = [] + for row in rows: + trace_id = row.get("trace_id") + if not isinstance(trace_id, str) or not trace_id or trace_id in trace_ids: + raise ValueError("Gym trace_id must be nonempty and unique per rollout") + trace_ids.add(trace_id) + tokens = row.get("token_ids") + logprobs = row.get("generation_logprobs") + mask = row.get("loss_mask") + if not all(isinstance(value, list) for value in (tokens, logprobs, mask)): + raise ValueError("Gym trace vectors must be lists") + if not tokens or len(tokens) != len(logprobs) or len(tokens) != len(mask): + raise ValueError("Gym trace token/logprob/mask vectors must align") + if any(type(token) is not int or token < 0 for token in tokens): + raise ValueError("Gym trace token IDs must be nonnegative integers") + if any(type(bit) is not int or bit not in (0, 1) for bit in mask) or mask[0]: + raise ValueError("Gym loss_mask must be binary with first token masked") + if any( + type(value) not in (int, float) or not math.isfinite(value) + for value in logprobs + ): + raise ValueError("Gym generation_logprobs must be finite") + call_ids = row.get("model_call_ids") + spans = row.get("sampled_spans") + if not isinstance(call_ids, list) or not isinstance(spans, list): + raise ValueError("Gym trace requires model_call_ids and sampled_spans") + if any(not isinstance(call_id, str) or not call_id for call_id in call_ids): + raise ValueError("Gym model_call_ids must be nonempty strings") + covered = [0] * len(tokens) + for span in spans: + call_id, start, end = ( + span.get("model_call_id"), + span.get("start"), + span.get("end"), + ) + if call_id not in call_ids or call_id in owned_call_ids: + raise ValueError( + "Gym sampled model call must have exactly one trace owner" + ) + if ( + type(start) is not int + or type(end) is not int + or not 0 < start < end <= len(tokens) + ): + raise ValueError("Gym sampled span is outside the trace") + owned_call_ids.add(call_id) + for position in range(start, end): + if covered[position]: + raise ValueError("Gym sampled spans overlap") + covered[position] = 1 + if len(set(call_ids)) != len(call_ids): + raise ValueError("Gym model_call_ids must be unique within a trace") + if covered != mask or not any(mask): + raise ValueError("Gym sampled_spans must exactly cover loss_mask") + traces.append( + GymTrainingTrace( + rollout_id, trace_id, tuple(tokens), tuple(logprobs), tuple(mask) + ) + ) + return traces + + +def gym_trace_message_log(trace: GymTrainingTrace) -> LLMMessageLogType: + """Represent one independently conditioned trace with exact sampled masks.""" + messages = [] + start = 0 + while start < len(trace.token_ids): + trainable = trace.loss_mask[start] + end = start + 1 + while end < len(trace.token_ids) and trace.loss_mask[end] == trainable: + end += 1 + message = { + "role": "assistant" if trainable else "user", + "content": "", + "token_ids": torch.tensor(trace.token_ids[start:end], dtype=torch.long), + } + if trainable: + message["generation_logprobs"] = torch.tensor( + trace.generation_logprobs[start:end], dtype=torch.float32 + ) + messages.append(message) + start = end + return messages + + +def gym_masked_message_log(pad_token_id: int) -> LLMMessageLogType: + """Return an explicit inert row for invalid rollouts and batch padding.""" + return [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([pad_token_id, pad_token_id], dtype=torch.long), + } + ] + + +@dataclass +class GymTraceBatch: + batch: BatchedDataDict + advantages: torch.Tensor + logical_indices: torch.Tensor + logical_count: int + physical_count: int + padding_count: int + overlong_trace_count: int + invalid_rollout_count: int + + +def prepare_gym_trace_batch( + logical_batch: BatchedDataDict, + *, + estimator: GRPOAdvantageEstimator, + pad_token_id: int, + max_sequence_length: int, + row_multiple: int, +) -> GymTraceBatch: + """Compute logical GRPO advantages, then expand and zero-pad physical rows. + + A rollout votes in its prompt group's baseline only if its episode mask is + positive and it has a trainable trace within the policy context limit. + An overlong trace is replaced by an inert row, preserving valid siblings. + Padding and invalid rows never contribute reward votes, loss, or tokens. + """ + if row_multiple < 1 or max_sequence_length < 2: + raise ValueError("Gym trace row multiple and context limit must be positive") + logical_count = logical_batch.size + traces_by_rollout = logical_batch["gym_training_traces"] + groups = logical_batch["gym_task_group_id"] + rewards = logical_batch["total_reward"] + if not torch.isfinite(rewards).all(): + raise ValueError("Gym logical rewards must be finite") + valid_rollouts = torch.zeros(logical_count, dtype=torch.float32) + indices, message_logs, trace_ids, rollout_ids, row_validity = [], [], [], [], [] + overlong_count = 0 + known_rollout_ids: set[str] = set() + for logical_index, traces in enumerate(traces_by_rollout): + rollout_id = logical_batch["gym_rollout_id"][logical_index] + if rollout_id in known_rollout_ids: + raise ValueError("Duplicate Gym rollout_id in a logical batch") + known_rollout_ids.add(rollout_id) + if any(trace.rollout_id != rollout_id for trace in traces): + raise ValueError("Gym trace does not belong to its logical rollout") + for trace in traces: + overlong = len(trace.token_ids) > max_sequence_length + overlong_count += int(overlong) + valid = ( + bool(sum(trace.loss_mask)) + and not overlong + and bool(logical_batch["loss_multiplier"][logical_index] > 0) + ) + if valid: + valid_rollouts[logical_index] = 1.0 + indices.append(logical_index) + message_logs.append( + gym_trace_message_log(trace) + if valid + else gym_masked_message_log(pad_token_id) + ) + row_validity.append(float(valid)) + trace_ids.append(trace.trace_id) + rollout_ids.append(trace.rollout_id) + if not traces: + indices.append(logical_index) + message_logs.append(gym_masked_message_log(pad_token_id)) + row_validity.append(0.0) + trace_ids.append("") + rollout_ids.append(rollout_id) + if not valid_rollouts.any(): + raise ValueError("Gym logical batch has no eligible training tokens") + logical_advantages = estimator.compute_advantage( + prompt_ids=groups.unsqueeze(-1), + rewards=rewards, + mask=valid_rollouts.unsqueeze(-1), + valid_mask=valid_rollouts, + ).squeeze(-1) + physical_count = len(indices) + padding_count = (-physical_count) % row_multiple + for _ in range(padding_count): + indices.append(0) + message_logs.append(gym_masked_message_log(pad_token_id)) + row_validity.append(0.0) + trace_ids.append("") + rollout_ids.append("") + index_tensor = torch.tensor(indices, dtype=torch.long) + batch = logical_batch.select_indices(index_tensor) + batch["message_log"] = message_logs + batch["length"] = torch.tensor( + [len(messages[0]["token_ids"]) for messages in message_logs] + ) + batch["loss_multiplier"] = batch["loss_multiplier"] * torch.tensor(row_validity) + batch["gym_trace_id"] = trace_ids + batch["gym_rollout_id"] = rollout_ids + batch["gym_logical_index"] = index_tensor + return GymTraceBatch( + batch, + logical_advantages[index_tensor], + index_tensor, + logical_count, + physical_count, + padding_count, + overlong_count, + int((valid_rollouts == 0).sum()), + ) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 2c7c2235853..f31b3340eea 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -39,6 +39,11 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.gym_traces import ( + gym_masked_message_log, + gym_trace_message_log, + parse_gym_training_traces, +) from nemo_rl.environments.nemo_gym_multimodal import ( _index_per_turn_images, _is_trainable_output_item, @@ -649,6 +654,15 @@ async def run_rollouts( nemo_rl_result = await self._postprocess_receipt_mode( nemo_gym_row, nemo_gym_result ) + elif ( + (self.cfg.get("initial_global_config_dict") or {}).get( + "token_id_capture" + ) + or {} + ).get("delivery") == "all_traces": + nemo_rl_result = await self._postprocess_all_gym_traces( + nemo_gym_row, nemo_gym_result, tokenizer + ) else: nemo_rl_result = self._postprocess_nemo_gym_to_nemo_rl_result( nemo_gym_row, @@ -696,6 +710,59 @@ async def run_rollouts( timing_metrics, ) + async def _postprocess_all_gym_traces( + self, + nemo_gym_row: dict, + nemo_gym_result: dict, + tokenizer: PreTrainedTokenizerBase, + ) -> dict: + """Finalize the local Gym capture and preserve every owned training trace.""" + # Gym is installed in the actor environment, not the trainer environment. + from nemo_gym.global_config import ROLLOUT_ID_KEY_NAME + from nemo_gym.token_id_capture.config import TokenIdCaptureConfig + from nemo_gym.token_id_capture.delivery import finalize_rollout_token_capture + from nemo_gym.token_id_capture.store import TokenCaptureStore + + capture_config = TokenIdCaptureConfig.model_validate( + self.cfg["initial_global_config_dict"] + ) + rollout_id = nemo_gym_row.get(ROLLOUT_ID_KEY_NAME) + if not rollout_id: + raise ValueError("Gym all_traces result is missing its capture rollout ID") + nemo_gym_result[ROLLOUT_ID_KEY_NAME] = rollout_id + store = TokenCaptureStore(capture_config.resolved_dir()) + await finalize_rollout_token_capture( + nemo_gym_result, + store, + builder=capture_config.token_id_capture.builder, + delivery=capture_config.token_id_capture.delivery, + ) + envelope = nemo_gym_result.get("training_traces") + if envelope is None: + if not nemo_gym_result.get("mask_sample"): + raise ValueError( + "Gym all_traces delivery returned no training envelope" + ) + traces = [] + else: + traces = parse_gym_training_traces(envelope) + if envelope["rollout_id"] != rollout_id: + raise ValueError("Gym trace envelope belongs to a different rollout") + logs = [gym_trace_message_log(trace) for trace in traces] + # The existing logical rollout protocol needs a representative prompt. + # Training expands gym_training_traces before flattening any sequence. + message_log = ( + logs[0] if logs else gym_masked_message_log(tokenizer.pad_token_id) + ) + return { + "message_log": message_log, + "input_message_log": message_log[:1], + "gym_training_traces": traces, + "gym_rollout_id": rollout_id, + "gym_metrics_message_log": [message for log in logs for message in log], + "full_result": nemo_gym_result, + } + async def _postprocess_receipt_mode( self, nemo_gym_row: dict, nemo_gym_result: dict ) -> dict: diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index a0688be3620..371c0094588 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -2827,18 +2827,48 @@ def _postprocess_single_nemo_gym_group( "total_reward": r["full_result"]["reward"], "assistant_tokens": sum( len(m["token_ids"]) - for m in r["message_log"] + for m in ( + r["gym_metrics_message_log"] + if "gym_metrics_message_log" in r + else r["message_log"] + ) if m["role"] == "assistant" ), - "total_tokens": sum(len(m["token_ids"]) for m in r["message_log"]), - "turn_count": sum(1 for m in r["message_log"] if m["role"] == "user"), - "hit_max_tokens": sum(len(m["token_ids"]) for m in r["message_log"]) + "total_tokens": sum( + len(m["token_ids"]) + for m in ( + r["gym_metrics_message_log"] + if "gym_metrics_message_log" in r + else r["message_log"] + ) + ), + "turn_count": sum( + 1 + for m in ( + r["gym_metrics_message_log"] + if "gym_metrics_message_log" in r + else r["message_log"] + ) + if m["role"] == "user" + ), + "hit_max_tokens": sum( + len(m["token_ids"]) + for m in ( + r["gym_metrics_message_log"] + if "gym_metrics_message_log" in r + else r["message_log"] + ) + ) == max_total_tokens_per_sample, # max_gen_tokens_per_turn: Diagnostic for long single generations "max_gen_tokens_per_turn": max( ( len(m["token_ids"]) - for m in r["message_log"] + for m in ( + r["gym_metrics_message_log"] + if "gym_metrics_message_log" in r + else r["message_log"] + ) if m["role"] == "assistant" ), default=0, @@ -2971,6 +3001,16 @@ def _postprocess_single_nemo_gym_group( ), } ) + if any("gym_training_traces" in result for result in results): + if not all("gym_training_traces" in result for result in results): + raise ValueError("Gym batch mixes legacy and multi-trace rollout results") + final_batch["gym_training_traces"] = [ + result["gym_training_traces"] for result in results + ] + final_batch["gym_rollout_id"] = [result["gym_rollout_id"] for result in results] + if "gym_task_group_id" in input_batch: + final_batch["gym_task_group_id"] = input_batch["gym_task_group_id"] + # Carry the raw env/agent flag downstream; the advantage stage composes it # into sample_mask. env.should_mask_flagged_samples=false skips this. if mask_env_flagged_samples: diff --git a/nemo_rl/models/policy/interfaces.py b/nemo_rl/models/policy/interfaces.py index 9fc9aa0b5ce..8f2f3fd9590 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -126,6 +126,7 @@ def train( gbs: Optional[int] = None, mbs: Optional[int] = None, timer: Optional[Timer] = None, + scheduler_step_samples: Optional[int] = None, ) -> dict[str, Any]: """Train the policy on a global batch of data. @@ -135,6 +136,8 @@ def train( eval_mode: Whether to run in evaluation mode (no gradient updates) gbs: Global batch size override (if None, uses config default) mbs: Micro batch size override (if None, uses config default) + scheduler_step_samples: Logical sample increment for Megatron when + physical training rows outnumber logical rollouts. None uses gbs. """ pass diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 8ee689a5783..ded8ea36f4d 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -856,16 +856,25 @@ def train( mbs: Optional[int] = None, timer: Optional[Timer] = None, check_dim_skip_keys: Optional[Iterable[str]] = None, + scheduler_step_samples: Optional[int] = None, ) -> dict[str, Any]: """Train the policy on a batch of data with a given loss function. Args: + scheduler_step_samples: Megatron-only logical sample increment when + physical rows differ from the configured logical batch size. + None preserves the existing global-batch scheduler increment. check_dim_skip_keys: Keys whose tensors are not student-sequence-aligned at dim 1 and must be excluded from the worker's sequence-dim pre-flight check. Used by cross-tokenizer distillation to pass through teacher / alignment auxiliaries that ride on the same data dict. """ + if scheduler_step_samples is not None: + if scheduler_step_samples <= 0 or not self.cfg["megatron_cfg"]["enabled"]: + raise ValueError( + "scheduler_step_samples requires a positive logical batch size and Megatron" + ) batch_size = gbs or self.cfg["train_global_batch_size"] micro_batch_size = mbs or self.cfg["train_micro_batch_size"] # Shard and replicate the batch @@ -905,6 +914,11 @@ def train( "gbs": batch_size, "mbs": micro_batch_size, "check_dim_skip_keys": check_dim_skip_keys, + **( + {"scheduler_step_samples": scheduler_step_samples} + if scheduler_step_samples is not None + else {} + ), }, ) results = self.worker_group.get_all_worker_results(futures) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 651cd9e5a35..48f5f970ad6 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -816,9 +816,13 @@ def train( gbs: Optional[int] = None, mbs: Optional[int] = None, check_dim_skip_keys: Optional[Iterable[str]] = None, + scheduler_step_samples: Optional[int] = None, ) -> dict[str, Any]: """Train the policy on a batch of data with a given loss function. + ``scheduler_step_samples`` optionally preserves the logical rollout + clock when a batch expands into several independently conditioned rows. + ``check_dim_skip_keys`` is accepted for parity with the v1/v2 DTensor workers (cross-tokenizer ride-along tensors whose dim 1 is not the student sequence axis). Megatron doesn't run cross-tokenizer, so it @@ -828,6 +832,8 @@ def train( "check_dim_skip_keys is only supported by the v2 DTensor worker; " "Megatron does not run cross-tokenizer distillation." ) + if scheduler_step_samples is not None and scheduler_step_samples <= 0: + raise ValueError("scheduler_step_samples must be positive") self.timer.start("train") # Note: zero_grad_buffer is called at the start of each global batch iteration # in the loop below, so we don't need to call it here. @@ -1106,7 +1112,11 @@ def train( # samples: NeMo init scales lr_warmup_steps by gbs internally, so # passing increment=gbs cancels that scaling and one tick == one # train() call regardless of batch size. - self.scheduler.step(increment=gbs) + self.scheduler.step( + increment=gbs + if scheduler_step_samples is None + else scheduler_step_samples + ) # Aggregate metrics across all microbatches mb_metrics, global_loss = aggregate_training_statistics( diff --git a/pyrefly.toml b/pyrefly.toml index e4ffa4a3673..650184dd76f 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -167,6 +167,7 @@ project-includes = [ "nemo_rl/environments/code_jaccard_environment.py", "nemo_rl/environments/games/sliding_puzzle.py", "nemo_rl/environments/interfaces.py", + "nemo_rl/environments/gym_traces.py", "nemo_rl/environments/math_environment.py", "nemo_rl/environments/metrics.py", "nemo_rl/environments/nemo_gym_multimodal.py", diff --git a/tests/unit/environments/test_gym_traces.py b/tests/unit/environments/test_gym_traces.py new file mode 100644 index 00000000000..2297a79c2cc --- /dev/null +++ b/tests/unit/environments/test_gym_traces.py @@ -0,0 +1,346 @@ +"""CPU tests of the Gym wire contract and the actual GRPO loss boundary.""" + +from copy import deepcopy + +import pytest +import torch + +from nemo_rl.algorithms.advantage_estimator import ( + AdvEstimatorConfig, + GRPOAdvantageEstimator, +) +from nemo_rl.algorithms.grpo import ( + GRPOConfig, + MasterConfig, + RewardPenaltyConfig, + _validate_gym_multi_trace_capability, + add_grpo_token_loss_masks_and_generation_logprobs, +) +from nemo_rl.algorithms.loss import ClippedPGLossConfig, ClippedPGLossFn +from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.environments.gym_traces import ( + gym_trace_message_log, + parse_gym_training_traces, + prepare_gym_trace_batch, +) + + +def _row(trace_id, tokens, spans): + mask = [0] * len(tokens) + for _, start, end in spans: + mask[start:end] = [1] * (end - start) + return { + "trace_id": trace_id, + "model_call_ids": [call for call, _, _ in spans], + "token_ids": tokens, + "generation_logprobs": [-0.125 if bit else 0.0 for bit in mask], + "loss_mask": mask, + "sampled_spans": [ + {"model_call_id": call, "start": start, "end": end} + for call, start, end in spans + ], + } + + +def _envelope(rollout, rows, builder="prefix_merging"): + return { + "schema_version": 1, + "rollout_id": rollout, + "builder": builder, + "traces": rows, + } + + +def _batch(envelopes, *, rewards, groups, masks=None): + traces = [ + parse_gym_training_traces(envelope) if envelope else [] + for envelope in envelopes + ] + return BatchedDataDict( + { + "gym_training_traces": traces, + "gym_rollout_id": [ + envelope["rollout_id"] if envelope else f"masked-{index}" + for index, envelope in enumerate(envelopes) + ], + "gym_task_group_id": torch.tensor(groups), + "total_reward": torch.tensor(rewards, dtype=torch.float32), + "loss_multiplier": torch.tensor( + masks if masks is not None else [1.0] * len(envelopes) + ), + "message_log": [ + gym_trace_message_log(rows[0]) if rows else [] for rows in traces + ], + "length": torch.ones(len(envelopes), dtype=torch.long), + } + ) + + +def _prepare(batch, *, max_length=32, multiple=4): + return prepare_gym_trace_batch( + batch, + estimator=GRPOAdvantageEstimator( + AdvEstimatorConfig(normalize_rewards=False), ClippedPGLossConfig() + ), + pad_token_id=0, + max_sequence_length=max_length, + row_multiple=multiple, + ) + + +def _flatten(prepared): + add_grpo_token_loss_masks_and_generation_logprobs(prepared.batch["message_log"]) + flat, lengths = batched_message_log_to_flat_message( + prepared.batch["message_log"], pad_value_dict={"token_ids": 0} + ) + return BatchedDataDict( + { + "input_ids": flat["token_ids"], + "input_lengths": lengths, + "token_mask": flat["token_loss_mask"], + "generation_logprobs": flat["generation_logprobs"], + "sample_mask": prepared.batch["loss_multiplier"], + "advantages": prepared.advantages.unsqueeze(-1).expand_as( + flat["token_ids"] + ), + } + ) + + +def test_exact_tokens_logprobs_masks_and_context_edits(): + envelope = _envelope( + "rollout", + [ + _row("a", [1, 2, 3, 9, 4, 5], [("a1", 2, 3), ("a2", 4, 6)]), + _row("b", [7, 8], [("b1", 1, 2)]), + ], + ) + prepared = _prepare(_batch([envelope], rewards=[1], groups=[17])) + flat = _flatten(prepared) + assert flat["input_ids"][0].tolist() == [1, 2, 3, 9, 4, 5] + assert flat["input_lengths"][0].item() == 6 + assert prepared.batch["length"][0].item() == 2 + assert flat["token_mask"][0].tolist() == [0, 0, 1, 0, 1, 1] + assert flat["generation_logprobs"][0].tolist() == [0, 0, -0.125, 0, -0.125, -0.125] + assert flat["sample_mask"].tolist() == [1, 1, 0, 0] + assert ( + prepared.logical_count == 1 + and prepared.physical_count == 2 + and prepared.padding_count == 2 + ) + + +def test_task_groups_and_masked_rollouts_vote_before_expansion(): + envelopes = [ + _envelope( + "a", + [_row("a1", [1, 2], [("a1", 1, 2)]), _row("a2", [3, 4, 5], [("a2", 1, 3)])], + ), + _envelope("b", [_row("b1", [1, 3], [("b1", 1, 2)])]), + _envelope("c", [_row("c1", [1, 4], [("c1", 1, 2)])]), + _envelope("d", [_row("d1", [1, 5], [("d1", 1, 2)])]), + ] + prepared = _prepare( + _batch( + envelopes, rewards=[0, 100, 1, 0], groups=[7, 7, 8, 8], masks=[1, 0, 1, 1] + ) + ) + assert prepared.logical_indices.tolist() == [0, 0, 1, 2, 3, 0, 0, 0] + assert prepared.advantages[:5].tolist() == [0, 0, 0, 1, -1] + assert prepared.batch["loss_multiplier"].tolist() == [1, 1, 0, 1, 1, 0, 0, 0] + assert prepared.invalid_rollout_count == 1 + + +def test_overlong_trace_masks_only_its_row_and_empty_rollout_is_inert(): + envelope = _envelope( + "a", + [ + _row("long", [1, 2, 3, 4, 5], [("long", 1, 5)]), + _row("short", [1, 2], [("short", 1, 2)]), + ], + ) + prepared = _prepare( + _batch([envelope, None], rewards=[1, 99], groups=[5, 5]), max_length=3 + ) + assert prepared.batch["loss_multiplier"].tolist() == [0, 1, 0, 0] + assert prepared.overlong_trace_count == 1 and prepared.invalid_rollout_count == 1 + assert prepared.batch["gym_rollout_id"] == ["a", "a", "masked-1", ""] + assert prepared.advantages[1].item() == 0 + assert _flatten(prepared)["input_lengths"].max().item() <= 3 + + +@pytest.mark.parametrize( + "mutation", + [ + lambda env: env.update(schema_version=2), + lambda env: env["traces"][0]["token_ids"].append(4), + lambda env: env["traces"][0]["generation_logprobs"].__setitem__( + 1, float("nan") + ), + lambda env: env["traces"][0]["loss_mask"].__setitem__(0, 1), + lambda env: env["traces"][0]["sampled_spans"][0].update(end=99), + lambda env: env["traces"].append(dict(env["traces"][0], trace_id="other")), + ], +) +def test_bad_custody_is_rejected(mutation): + envelope = _envelope("a", [_row("trace", [1, 2], [("call", 1, 2)])]) + mutation(envelope) + with pytest.raises(ValueError): + parse_gym_training_traces(envelope) + + +def test_duplicate_rollout_and_all_invalid_batch_are_rejected(): + envelope = _envelope("a", [_row("trace", [1, 2], [("call", 1, 2)])]) + with pytest.raises(ValueError, match="Duplicate Gym rollout_id"): + _prepare(_batch([envelope, envelope], rewards=[1, 0], groups=[0, 0])) + with pytest.raises(ValueError, match="no eligible training tokens"): + _prepare(_batch([None], rewards=[1], groups=[0])) + + +def test_equivalent_split_and_merged_conditioning_has_identical_loss_and_gradient(): + merged = _envelope( + "a", [_row("merged", [1, 2, 3, 4, 5, 6], [("a1", 2, 3), ("a2", 4, 6)])] + ) + split = _envelope( + "a", + [ + _row("first", [1, 2, 3], [("a1", 2, 3)]), + _row("second", [1, 2, 3, 4, 5, 6], [("a2", 4, 6)]), + ], + builder="per_request", + ) + sibling = _envelope("b", [_row("other", [7, 8], [("b1", 1, 2)])]) + torch.manual_seed(11) + base_model = torch.nn.Sequential(torch.nn.Embedding(16, 5), torch.nn.Linear(5, 16)) + loss_fn = ClippedPGLossFn(ClippedPGLossConfig(reference_policy_kl_penalty=0.0)) + outcomes = [] + for envelope in (merged, split): + model = deepcopy(base_model) + prepared = _prepare( + _batch([envelope, sibling], rewards=[1, 0], groups=[42, 42]) + ) + data = _flatten(prepared) + # A deterministic causal model: the same complete prefix gives the same logits. + hidden = model[0](data["input_ids"][:, :-1]).cumsum(dim=1) + logits = model[1](hidden) + logprobs = ( + logits.log_softmax(dim=-1) + .gather(-1, data["input_ids"][:, 1:].unsqueeze(-1)) + .squeeze(-1) + ) + data["prev_logprobs"] = torch.nn.functional.pad(logprobs.detach(), (1, 0)) + valid_tokens = ( + data["token_mask"][:, 1:] * data["sample_mask"].unsqueeze(-1) + ).sum() + loss, _ = loss_fn( + logprobs, + data, + global_valid_seqs=data["sample_mask"].sum(), + global_valid_toks=valid_tokens, + ) + loss.backward() + outcomes.append( + ( + loss.detach(), + torch.cat( + [parameter.grad.flatten() for parameter in model.parameters()] + ), + valid_tokens, + ) + ) + torch.testing.assert_close(outcomes[0][0], outcomes[1][0]) + torch.testing.assert_close(outcomes[0][1], outcomes[1][1]) + assert outcomes[0][2] == outcomes[1][2] == 4 + + +def _config(): + return MasterConfig.model_construct( + grpo=GRPOConfig( + gym_multi_trace=True, num_prompts_per_step=2, num_generations_per_prompt=2 + ), + policy={"megatron_cfg": {"enabled": True}, "train_global_batch_size": 4}, + env={ + "should_use_nemo_gym": True, + "nemo_gym": { + "token_id_capture": { + "enabled": True, + "all_agents": True, + "delivery": "all_traces", + } + }, + }, + loss_fn=ClippedPGLossConfig(), + reward_penalties=RewardPenaltyConfig(), + ) + + +@pytest.mark.parametrize( + "configure", + [ + lambda cfg: setattr(cfg.grpo.async_grpo, "enabled", True), + lambda cfg: setattr(cfg, "data_plane", {"enabled": True}), + lambda cfg: cfg.policy.update(is_vlm=True), + lambda cfg: setattr(cfg.loss_fn, "token_level_loss", False), + lambda cfg: setattr(cfg.grpo, "use_dynamic_sampling", True), + lambda cfg: setattr(cfg.grpo, "seq_logprob_error_threshold", 2.0), + lambda cfg: setattr(cfg.grpo.adv_estimator, "name", "gdpo"), + lambda cfg: cfg.env["nemo_gym"]["token_id_capture"].update(sink="custom:Sink"), + lambda cfg: cfg.policy.update(train_global_batch_size=8), + ], +) +def test_unsupported_configs_fail_before_worker_setup(configure): + config = _config() + _validate_gym_multi_trace_capability(config) + configure(config) + with pytest.raises((NotImplementedError, ValueError)): + _validate_gym_multi_trace_capability(config) + + +def test_both_sides_must_opt_in(): + config = _config() + config.grpo.gym_multi_trace = False + with pytest.raises(ValueError, match="requires grpo.gym_multi_trace"): + _validate_gym_multi_trace_capability(config) + config.env["nemo_gym"]["token_id_capture"]["delivery"] = "main_chain" + _validate_gym_multi_trace_capability(config) + + +@pytest.mark.parametrize("logical_count", [None, 8]) +def test_policy_forwards_logical_scheduler_count_only_when_requested(logical_count): + from types import SimpleNamespace + from unittest.mock import Mock + + from nemo_rl.models.policy.lm_policy import Policy + + worker_group = SimpleNamespace( + run_all_workers_sharded_data=Mock(return_value=[]), + get_all_worker_results=Mock( + return_value=[ + { + "global_loss": torch.tensor(0.0), + "grad_norm": torch.tensor(1.0), + "all_mb_metrics": {}, + } + ] + ), + ) + policy = SimpleNamespace( + cfg={ + "megatron_cfg": {"enabled": True}, + "train_global_batch_size": 8, + "train_micro_batch_size": 1, + }, + _shard_for_train=Mock(return_value=[]), + _report_sharded_payload=Mock(), + flops_tracker=None, + worker_group=worker_group, + ) + data = BatchedDataDict({"input_ids": torch.zeros(12, 2, dtype=torch.long)}) + Policy.train(policy, data, Mock(), gbs=12, scheduler_step_samples=logical_count) + kwargs = worker_group.run_all_workers_sharded_data.call_args.kwargs["common_kwargs"] + assert kwargs["gbs"] == 12 + if logical_count is None: + assert "scheduler_step_samples" not in kwargs + else: + assert kwargs["scheduler_step_samples"] == 8 diff --git a/tests/unit/environments/test_nemo_gym_token_capture.py b/tests/unit/environments/test_nemo_gym_token_capture.py index 0925fdcbe28..f7bc80f4d52 100644 --- a/tests/unit/environments/test_nemo_gym_token_capture.py +++ b/tests/unit/environments/test_nemo_gym_token_capture.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio +from copy import deepcopy import hashlib from unittest.mock import AsyncMock @@ -398,3 +399,84 @@ def test_postprocess_passes_the_scored_response_to_attribution() -> None: receipt = result["receipt"] assert receipt["terminal_model_call_id"] == "c2" assert receipt["terminal_selection"] == "response_id" + + +@pytest.mark.parametrize("builder", ["per_request", "prefix_merging"]) +def test_local_all_traces_finalizer_preserves_branches_and_compaction( + tmp_path, builder +): + """Exercise the actual Gym store, public schema and RL actor handoff.""" + from types import SimpleNamespace + + from nemo_gym.global_config import ROLLOUT_ID_KEY_NAME + from nemo_gym.token_id_capture.records import ( + ParentResolutionStatus, + TokenEntry, + stamp_lineage, + ) + from nemo_gym.token_id_capture.store import TokenCaptureStore + from nemo_gym.token_id_capture.training_traces import TrainingTraceBatch + + store = TokenCaptureStore(tmp_path) + for call, prompt, generated, parent in [ + ("a", [1, 2], [10, 11], None), + ("b", [1, 2, 10, 11, 3], [12], "a"), + ("c", [1, 2, 10, 11, 4], [13, 14], "a"), + ("compact", [5, 6], [15], None), + ]: + entry = TokenEntry( + rollout_id="r1", + model_call_id=call, + model="policy", + prompt_token_ids=prompt, + generation_token_ids=generated, + generation_log_probs=[-token / 100 for token in generated], + response_id=f"response-{call}", + ) + store.append( + stamp_lineage( + entry, + parent, + parent_resolution=ParentResolutionStatus.RESOLVED + if parent + else ParentResolutionStatus.ROOT, + ) + ) + env = _capture_env() + env.cfg = { + "initial_global_config_dict": { + "token_id_capture": { + "enabled": True, + "all_agents": True, + "dir": str(tmp_path), + "delivery": "all_traces", + "builder": builder, + } + } + } + original_response = { + "id": "response-c", + "output": [ + { + "type": "message", + "content": [{"type": "output_text", "text": "scored answer"}], + } + ], + } + original_result = {"reward": 0.75, "response": deepcopy(original_response)} + result = asyncio.run( + env._postprocess_all_gym_traces( + {ROLLOUT_ID_KEY_NAME: "r1"}, + original_result, + SimpleNamespace(pad_token_id=0), + ) + ) + public = TrainingTraceBatch.model_validate(result["full_result"]["training_traces"]) + assert len(public.traces) == (4 if builder == "per_request" else 3) + assert len(result["gym_training_traces"]) == len(public.traces) + assert sum(sum(trace.loss_mask) for trace in result["gym_training_traces"]) == 6 + assert result["full_result"]["reward"] == 0.75 + assert result["full_result"]["response"] == original_response + assert { + span.model_call_id for trace in public.traces for span in trace.sampled_spans + } == {"a", "b", "c", "compact"} diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 6501ad190e6..e3a48b9e575 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -1,5 +1,7 @@ # GRPO Algorithm Configuration grpo: + # Opt in to Gym all_traces delivery; synchronous Megatron only. + gym_multi_trace: false num_prompts_per_step: 32 num_generations_per_prompt: 16 max_rollout_turns: 1 # for multi-turn rollouts. Math Environments just have 1 turn (answering the question) From ca263193c7baca9a94c5ca72631491507f044d25 Mon Sep 17 00:00:00 2001 From: Elad Sarafian Date: Thu, 10 Sep 2026 04:13:33 -0700 Subject: [PATCH 2/6] fix: guard trace-dependent correction and align baseline diagnostics Signed-off-by: Elad Sarafian --- docs/guides/gym-multi-trace-grpo.md | 4 +- nemo_rl/algorithms/grpo.py | 13 ++++++ nemo_rl/environments/gym_traces.py | 2 + tests/unit/environments/test_gym_traces.py | 51 ++++++++++++++++++++++ 4 files changed, 69 insertions(+), 1 deletion(-) diff --git a/docs/guides/gym-multi-trace-grpo.md b/docs/guides/gym-multi-trace-grpo.md index f5b4dcb5a72..4c9eb005b44 100644 --- a/docs/guides/gym-multi-trace-grpo.md +++ b/docs/guides/gym-multi-trace-grpo.md @@ -84,7 +84,9 @@ This initial path requires synchronous GRPO, a text-only Megatron policy, and token-level loss and importance ratios. Setup rejects asynchronous/TQ/staged capture, custom capture sinks, router replay, multimodal models, alternative advantage estimators, dynamic sampling, legacy episode-length filtering/shaping, -message-level penalties, and post-advantage sequence-logprob-error masking. +message-level penalties, active `seq-mask-tis` importance correction, and +post-advantage sequence-logprob-error masking. Sequence-level correction masks +depend on physical trace boundaries and cannot preserve split/merged equivalence. `train_global_batch_size` must equal the logical rollout count per step. Use `data_plane: null` or a complete disabled data-plane configuration. diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index c9cdeb17c77..4ce16968dac 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -506,6 +506,13 @@ def _validate_gym_multi_trace_capability(master_config: MasterConfig) -> None: raise NotImplementedError( "gym_multi_trace requires token-level loss and token-level importance ratios" ) + if ( + master_config.loss_fn.use_importance_sampling_correction + and master_config.loss_fn.truncated_importance_sampling_type == "seq-mask-tis" + ): + raise NotImplementedError( + "gym_multi_trace does not support seq-mask-tis: sequence-level masks depend on physical trace boundaries" + ) if ( grpo.use_dynamic_sampling or grpo.overlong_filtering @@ -3520,6 +3527,12 @@ def grpo_train( row_multiple=policy.data_parallel_size * master_config.policy["train_micro_batch_size"], ) + baseline_for_log, _ = calculate_baseline_and_std_per_prompt( + repeated_batch["gym_task_group_id"].unsqueeze(-1), + rewards, + trace_batch.logical_valid_mask, + leave_one_out_baseline=adv_estimator.use_leave_one_out_baseline, + ) repeated_batch = trace_batch.batch metrics.update( { diff --git a/nemo_rl/environments/gym_traces.py b/nemo_rl/environments/gym_traces.py index e265dbc9f6f..0c737fc7db1 100644 --- a/nemo_rl/environments/gym_traces.py +++ b/nemo_rl/environments/gym_traces.py @@ -168,6 +168,7 @@ class GymTraceBatch: padding_count: int overlong_trace_count: int invalid_rollout_count: int + logical_valid_mask: torch.Tensor def prepare_gym_trace_batch( @@ -264,4 +265,5 @@ def prepare_gym_trace_batch( padding_count, overlong_count, int((valid_rollouts == 0).sum()), + valid_rollouts, ) diff --git a/tests/unit/environments/test_gym_traces.py b/tests/unit/environments/test_gym_traces.py index 2297a79c2cc..e88d6caf4ea 100644 --- a/tests/unit/environments/test_gym_traces.py +++ b/tests/unit/environments/test_gym_traces.py @@ -17,6 +17,7 @@ add_grpo_token_loss_masks_and_generation_logprobs, ) from nemo_rl.algorithms.loss import ClippedPGLossConfig, ClippedPGLossFn +from nemo_rl.algorithms.utils import calculate_baseline_and_std_per_prompt from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.gym_traces import ( @@ -150,6 +151,7 @@ def test_task_groups_and_masked_rollouts_vote_before_expansion(): assert prepared.advantages[:5].tolist() == [0, 0, 0, 1, -1] assert prepared.batch["loss_multiplier"].tolist() == [1, 1, 0, 1, 1, 0, 0, 0] assert prepared.invalid_rollout_count == 1 + assert prepared.logical_valid_mask.tolist() == [1, 0, 1, 1] def test_overlong_trace_masks_only_its_row_and_empty_rollout_is_inert(): @@ -170,6 +172,43 @@ def test_overlong_trace_masks_only_its_row_and_empty_rollout_is_inert(): assert _flatten(prepared)["input_lengths"].max().item() <= 3 +@pytest.mark.parametrize("leave_one_out", [False, True]) +def test_baseline_diagnostics_use_the_same_logical_validity_as_advantages( + leave_one_out, +): + envelopes = [ + _envelope("valid-a", [_row("short", [1, 2], [("call", 1, 2)])]), + _envelope("valid-b", [_row("short", [1, 3], [("call", 1, 2)])]), + _envelope("overlong", [_row("long", [1, 2, 3, 4], [("call", 1, 4)])]), + None, + ] + logical_batch = _batch(envelopes, rewards=[1, 0, 100, 99], groups=[7, 7, 7, 7]) + prepared = prepare_gym_trace_batch( + logical_batch, + estimator=GRPOAdvantageEstimator( + AdvEstimatorConfig( + normalize_rewards=False, use_leave_one_out_baseline=leave_one_out + ), + ClippedPGLossConfig(), + ), + pad_token_id=0, + max_sequence_length=3, + row_multiple=4, + ) + assert prepared.logical_valid_mask.tolist() == [1, 1, 0, 0] + baseline, _ = calculate_baseline_and_std_per_prompt( + logical_batch["gym_task_group_id"].unsqueeze(-1), + logical_batch["total_reward"], + prepared.logical_valid_mask, + leave_one_out_baseline=leave_one_out, + ) + # Invalid rewards cannot influence the surviving rollouts or their diagnostics. + assert baseline[:2].tolist() == ([0.0, 1.0] if leave_one_out else [0.5, 0.5]) + torch.testing.assert_close( + prepared.advantages[:2], logical_batch["total_reward"][:2] - baseline[:2] + ) + + @pytest.mark.parametrize( "mutation", [ @@ -297,6 +336,18 @@ def test_unsupported_configs_fail_before_worker_setup(configure): _validate_gym_multi_trace_capability(config) +def test_sequence_mask_tis_is_rejected_only_when_correction_is_enabled(): + config = _config() + config.loss_fn.truncated_importance_sampling_type = "seq-mask-tis" + config.loss_fn.truncated_importance_sampling_ratio = 1.1 + config.loss_fn.truncated_importance_sampling_ratio_min = 0.9 + config.loss_fn.use_importance_sampling_correction = True + with pytest.raises(NotImplementedError, match="seq-mask-tis"): + _validate_gym_multi_trace_capability(config) + config.loss_fn.use_importance_sampling_correction = False + _validate_gym_multi_trace_capability(config) + + def test_both_sides_must_opt_in(): config = _config() config.grpo.gym_multi_trace = False From 98d5449eacdccfe6ed91297e322881c2d1994fc9 Mon Sep 17 00:00:00 2001 From: Elad Sarafian Date: Thu, 10 Sep 2026 04:19:57 -0700 Subject: [PATCH 3/6] build: pin Gym trace delivery and refresh its dependency lock Signed-off-by: Elad Sarafian --- 3rdparty/Gym-workspace/Gym | 2 +- uv.lock | 36 +++++++++++++++++++++++++++++------- 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index fd5e84d6b1c..2ce5e44ef8d 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit fd5e84d6b1c485c80e7ae61553bbd485611c03b4 +Subproject commit 2ce5e44ef8dbf44410c8ca1914134a0a9ec2e9f8 diff --git a/uv.lock b/uv.lock index c5235dd97a9..efb3a199b40 100644 --- a/uv.lock +++ b/uv.lock @@ -2096,14 +2096,34 @@ wheels = [ name = "gitpython" version = "3.1.59" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] dependencies = [ - { name = "gitdb" }, + { name = "gitdb", marker = "extra == 'extra-7-nemo-rl-sglang' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ca/dc/126b28e76b24a9268ba931ad3e012f71ebdadf62fd9f17758f7074bb0b20/gitpython-3.1.59.tar.gz", hash = "sha256:0a1475cfdc38a5bfba1a3e9a4a9da52a39749ecec322b772915c019f94e5b7e4", size = 230445, upload-time = "2026-08-10T12:03:20.271Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ef/ed/ae57eb7d344f43f87b74b3a281ead6ec7d6394eef72a7b1dcb28dd089550/gitpython-3.1.59-py3-none-any.whl", hash = "sha256:67a82f537384578643624c8b2c531938a9b82be431663e575dcf638526631d4c", size = 220996, upload-time = "2026-08-10T12:03:18.804Z" }, ] +[[package]] +name = "gitpython" +version = "3.1.62" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] +dependencies = [ + { name = "gitdb", marker = "extra == 'extra-7-nemo-rl-automodel' or extra == 'extra-7-nemo-rl-fsdp' or extra == 'extra-7-nemo-rl-mcore' or extra == 'extra-7-nemo-rl-nemo-gym' or extra != 'extra-7-nemo-rl-sglang' or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e0/db/3ca813cbacb23ab6fe46ff38a9b5ef8e73e970c8051f2ce903aacafe0446/gitpython-3.1.62.tar.gz", hash = "sha256:1791de66309bc0c7cfca40bf8d2e3de7ca091cbf94e6051be1ad0722c61062af", size = 231728, upload-time = "2026-09-07T02:57:21.155Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d6/0b/29d7965215f8ef830a7ca1f42997fe13e5693d85e9edb18f938d063ef5f2/gitpython-3.1.62-py3-none-any.whl", hash = "sha256:7002251225e10e29d2e1f49e6532613fe5d5d9f0b6f1f02997a52b38fe56899e", size = 222753, upload-time = "2026-09-07T02:57:19.762Z" }, +] + [[package]] name = "google-api-core" version = "2.32.0" @@ -3514,7 +3534,8 @@ dependencies = [ { name = "cloudpickle" }, { name = "databricks-sdk" }, { name = "fastapi", extra = ["standard"] }, - { name = "gitpython" }, + { name = "gitpython", version = "3.1.59", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-sglang' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "gitpython", version = "3.1.62", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or extra == 'extra-7-nemo-rl-fsdp' or extra == 'extra-7-nemo-rl-mcore' or extra == 'extra-7-nemo-rl-nemo-gym' or extra != 'extra-7-nemo-rl-sglang' or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "importlib-metadata" }, { name = "opentelemetry-api" }, { name = "opentelemetry-proto" }, @@ -3929,7 +3950,7 @@ dependencies = [ { name = "devtools" }, { name = "fastapi", extra = ["standard"], marker = "extra == 'extra-7-nemo-rl-automodel' or extra == 'extra-7-nemo-rl-fsdp' or extra == 'extra-7-nemo-rl-mcore' or extra == 'extra-7-nemo-rl-nemo-gym' or (extra == 'extra-7-nemo-rl-modelopt' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nvrx' and extra != 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-trtllm')" }, { name = "fonttools" }, - { name = "gitpython" }, + { name = "gitpython", version = "3.1.62", source = { registry = "https://pypi.org/simple" } }, { name = "httptools" }, { name = "hydra-core", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "hydra-core", version = "1.3.4", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or (extra == 'extra-7-nemo-rl-fsdp' and extra != 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-nemo-gym') or (extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-modelopt' and extra != 'extra-7-nemo-rl-sglang') or (extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-nvrx' and extra != 'extra-7-nemo-rl-sglang') or (extra != 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra != 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra != 'extra-7-nemo-rl-mcore' and extra != 'extra-7-nemo-rl-sglang' and extra != 'extra-7-nemo-rl-trtllm') or (extra != 'extra-7-nemo-rl-modelopt' and extra != 'extra-7-nemo-rl-nemo-gym' and extra != 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, @@ -3974,15 +3995,15 @@ requires-dist = [ { name = "fastapi" }, { name = "flashinfer-python", marker = "extra == 'vllm'", specifier = "==0.6.13" }, { name = "fonttools", specifier = ">=4.60.2" }, - { name = "gitpython", specifier = ">=3.1.57" }, + { name = "gitpython", specifier = ">=3.1.61" }, { name = "gprof2dot", marker = "extra == 'dev'" }, { name = "httptools" }, { name = "httpx-aiohttp", marker = "extra == 'sandbox'", specifier = ">=0.2.0" }, { name = "hydra-core" }, { name = "itsdangerous" }, { name = "mcp", specifier = ">=1.28.1,<2" }, - { name = "mlflow", specifier = ">=3.15.2" }, - { name = "mlflow-skinny", specifier = ">=3.15.2" }, + { name = "mlflow", specifier = ">=3.16.0" }, + { name = "mlflow-skinny", specifier = ">=3.16.0" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" }, { name = "nemo-gym", extras = ["dev", "sandbox"], marker = "extra == 'all'" }, { name = "nemo-lens", extras = ["sdk"], marker = "extra == 'telemetry'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, @@ -4205,7 +4226,8 @@ dev = [ { name = "types-requests" }, ] docs = [ - { name = "gitpython" }, + { name = "gitpython", version = "3.1.59", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-sglang' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "gitpython", version = "3.1.62", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or extra == 'extra-7-nemo-rl-fsdp' or extra == 'extra-7-nemo-rl-mcore' or extra == 'extra-7-nemo-rl-nemo-gym' or extra != 'extra-7-nemo-rl-sglang' or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "myst-parser" }, { name = "nvidia-sphinx-theme" }, { name = "python-dotenv" }, From accc9cdc8a4f4c0134881457515b708137168518 Mon Sep 17 00:00:00 2001 From: Elad Sarafian Date: Thu, 10 Sep 2026 04:30:01 -0700 Subject: [PATCH 4/6] fix: add copyright header to Gym trace tests Signed-off-by: Elad Sarafian --- tests/unit/environments/test_gym_traces.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/unit/environments/test_gym_traces.py b/tests/unit/environments/test_gym_traces.py index e88d6caf4ea..a4ea0a93b44 100644 --- a/tests/unit/environments/test_gym_traces.py +++ b/tests/unit/environments/test_gym_traces.py @@ -1,3 +1,16 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """CPU tests of the Gym wire contract and the actual GRPO loss boundary.""" from copy import deepcopy From bff7e4ab70bd7f3c3cc30ec1e68364310879ad6c Mon Sep 17 00:00:00 2001 From: Elad Sarafian Date: Thu, 10 Sep 2026 04:42:53 -0700 Subject: [PATCH 5/6] fix: correlate synchronous Gym capture before dispatch Signed-off-by: Elad Sarafian --- nemo_rl/environments/nemo_gym.py | 57 ++++++-- .../test_nemo_gym_token_capture.py | 138 +++++++++++++++++- 2 files changed, 181 insertions(+), 14 deletions(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index f31b3340eea..717e540834d 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -616,6 +616,30 @@ async def run_rollouts( # Megatron's HTTP backend consumes the same normalized Responses payload. normalize_media_in_examples(nemo_gym_examples) + all_traces_delivery = ( + (self.cfg.get("initial_global_config_dict") or {}).get("token_id_capture") + or {} + ).get("delivery") == "all_traces" + if all_traces_delivery: + from uuid import uuid4 + + from nemo_gym.global_config import ROLLOUT_ID_KEY_NAME + from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body + + # Synchronous rows need not carry Gym task indices. Assign capture + # identity before /run so writers and finalization use the same key. + # Fresh dispatches must not collide when task numbering restarts. + capture_ids = set() + for row in nemo_gym_examples: + if ROLLOUT_ID_KEY_NAME not in row: + row[ROLLOUT_ID_KEY_NAME] = uuid4().hex + capture_id = maybe_rollout_id_from_run_body(row) + if not capture_id or capture_id in capture_ids: + raise ValueError( + "Gym all_traces dispatch requires unique valid capture rollout IDs" + ) + capture_ids.add(capture_id) + timer = Timer() timer.start("_run_rollouts_total") nemo_gym_result_iterator = self.rch.run_examples( @@ -654,12 +678,7 @@ async def run_rollouts( nemo_rl_result = await self._postprocess_receipt_mode( nemo_gym_row, nemo_gym_result ) - elif ( - (self.cfg.get("initial_global_config_dict") or {}).get( - "token_id_capture" - ) - or {} - ).get("delivery") == "all_traces": + elif all_traces_delivery: nemo_rl_result = await self._postprocess_all_gym_traces( nemo_gym_row, nemo_gym_result, tokenizer ) @@ -718,7 +737,13 @@ async def _postprocess_all_gym_traces( ) -> dict: """Finalize the local Gym capture and preserve every owned training trace.""" # Gym is installed in the actor environment, not the trainer environment. - from nemo_gym.global_config import ROLLOUT_ID_KEY_NAME + from nemo_gym.global_config import ( + ATTEMPT_INDEX_KEY_NAME, + ROLLOUT_ID_KEY_NAME, + ROLLOUT_INDEX_KEY_NAME, + TASK_INDEX_KEY_NAME, + ) + from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body from nemo_gym.token_id_capture.config import TokenIdCaptureConfig from nemo_gym.token_id_capture.delivery import finalize_rollout_token_capture from nemo_gym.token_id_capture.store import TokenCaptureStore @@ -726,10 +751,24 @@ async def _postprocess_all_gym_traces( capture_config = TokenIdCaptureConfig.model_validate( self.cfg["initial_global_config_dict"] ) - rollout_id = nemo_gym_row.get(ROLLOUT_ID_KEY_NAME) + # Ordinary RL dispatch identifies requests with task/rollout indices; + # explicit IDs are optional. Use the same resolver as Gym's capture writer. + rollout_id = maybe_rollout_id_from_run_body(nemo_gym_row) if not rollout_id: raise ValueError("Gym all_traces result is missing its capture rollout ID") - nemo_gym_result[ROLLOUT_ID_KEY_NAME] = rollout_id + # Preserve the original fields, not an already suffixed canonical ID: + # the finalizer applies the attempt suffix exactly once. The dispatched + # row is authoritative even if an agent echoes stale correlation fields. + for key in ( + ROLLOUT_ID_KEY_NAME, + TASK_INDEX_KEY_NAME, + ROLLOUT_INDEX_KEY_NAME, + ATTEMPT_INDEX_KEY_NAME, + ): + if key in nemo_gym_row: + nemo_gym_result[key] = nemo_gym_row[key] + else: + nemo_gym_result.pop(key, None) store = TokenCaptureStore(capture_config.resolved_dir()) await finalize_rollout_token_capture( nemo_gym_result, diff --git a/tests/unit/environments/test_nemo_gym_token_capture.py b/tests/unit/environments/test_nemo_gym_token_capture.py index f7bc80f4d52..a57e1e25f29 100644 --- a/tests/unit/environments/test_nemo_gym_token_capture.py +++ b/tests/unit/environments/test_nemo_gym_token_capture.py @@ -402,13 +402,24 @@ def test_postprocess_passes_the_scored_response_to_attribution() -> None: @pytest.mark.parametrize("builder", ["per_request", "prefix_merging"]) +@pytest.mark.parametrize( + "row,rollout_id", + [ + ({"_ng_rollout_id": "r1"}, "r1"), + ({"_ng_task_index": "group", "_ng_rollout_index": 0}, "group-0"), + ({"_ng_rollout_id": "r1", "_ng_attempt_index": 2}, "r1-a2"), + ( + {"_ng_task_index": "group", "_ng_rollout_index": 1, "_ng_attempt_index": 2}, + "group-1-a2", + ), + ], +) def test_local_all_traces_finalizer_preserves_branches_and_compaction( - tmp_path, builder + tmp_path, builder, row, rollout_id ): """Exercise the actual Gym store, public schema and RL actor handoff.""" from types import SimpleNamespace - from nemo_gym.global_config import ROLLOUT_ID_KEY_NAME from nemo_gym.token_id_capture.records import ( ParentResolutionStatus, TokenEntry, @@ -425,7 +436,7 @@ def test_local_all_traces_finalizer_preserves_branches_and_compaction( ("compact", [5, 6], [15], None), ]: entry = TokenEntry( - rollout_id="r1", + rollout_id=rollout_id, model_call_id=call, model="policy", prompt_token_ids=prompt, @@ -463,15 +474,21 @@ def test_local_all_traces_finalizer_preserves_branches_and_compaction( } ], } - original_result = {"reward": 0.75, "response": deepcopy(original_response)} + original_result = { + "reward": 0.75, + "response": deepcopy(original_response), + "_ng_rollout_id": "stale-agent-id", + "_ng_attempt_index": 99, + } result = asyncio.run( env._postprocess_all_gym_traces( - {ROLLOUT_ID_KEY_NAME: "r1"}, + row, original_result, SimpleNamespace(pad_token_id=0), ) ) public = TrainingTraceBatch.model_validate(result["full_result"]["training_traces"]) + assert public.rollout_id == result["gym_rollout_id"] == rollout_id assert len(public.traces) == (4 if builder == "per_request" else 3) assert len(result["gym_training_traces"]) == len(public.traces) assert sum(sum(trace.loss_mask) for trace in result["gym_training_traces"]) == 6 @@ -480,3 +497,114 @@ def test_local_all_traces_finalizer_preserves_branches_and_compaction( assert { span.model_call_id for trace in public.traces for span in trace.sampled_spans } == {"a", "b", "c", "compact"} + + +@pytest.mark.parametrize("already_flagged", [False, True]) +def test_no_model_calls_with_normal_dispatch_identity_becomes_inert( + tmp_path, already_flagged +): + from types import SimpleNamespace + from unittest.mock import Mock + + from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body + + from nemo_rl.experience.rollouts import _prepare_nemo_gym_rows + + env = _capture_env() + env._require_spinup = Mock() + env._tokenizer = SimpleNamespace(pad_token_id=0) + env._token_capture_enabled = False + env.head_server_config = None + env.cfg = { + "initial_global_config_dict": { + "token_id_capture": { + "enabled": True, + "all_agents": True, + "dir": str(tmp_path), + "delivery": "all_traces", + "builder": "per_request", + } + } + } + dispatched_ids = [] + + async def complete(row): + return row, { + "reward": 0.0, + "response": {"output": []}, + "mask_sample": already_flagged, + } + + def run_examples(*, examples, head_server_config): + # Observe the exact row the real actor sends to Gym, before any model call. + for row in examples: + assert row.get("_ng_rollout_id") + dispatched_ids.append(maybe_rollout_id_from_run_body(row)) + return [complete(row) for row in examples] + + env.rch = SimpleNamespace(run_examples=run_examples) + + async def collect(rows): + return [item async for item in env.run_rollouts(rows, "test")] + + results = [] + for batch_index in range(3): + rows = [ + {"responses_create_params": {}, "agent_ref": {"name": "test"}} + for _ in range(2) + ] + # Task numbering can restart; synchronous rows may also omit it entirely. + if batch_index == 0: + rows[0]["_ng_rollout_id"] = "supplied" + rows[0]["_ng_attempt_index"] = 2 + else: + for row in rows: + row["_ng_task_index"] = 0 + _prepare_nemo_gym_rows( + rows, + {"max_new_tokens": 16}, + SimpleNamespace(temperature=0.7, top_p=0.95), + num_generations=2, + ) + with pytest.warns(UserWarning, match="capture contains no token records"): + collected = asyncio.run(collect(rows)) + if batch_index == 0: + assert rows[0]["_ng_rollout_id"] == "supplied" + for _, _, result, _ in collected: + assert result["gym_training_traces"] == [] + assert result["gym_metrics_message_log"] == [] + assert result["message_log"][0]["role"] == "user" + assert result["message_log"][0]["token_ids"].tolist() == [0, 0] + assert result["full_result"]["mask_sample"] is True + assert ( + result["full_result"]["_ng_token_capture"]["error"] + == "capture contains no token records" + ) + results.append(result) + assert dispatched_ids[0] == "supplied-a2" + assert len(set(dispatched_ids)) == 6 + assert [result["gym_rollout_id"] for result in results] == dispatched_ids + + +@pytest.mark.parametrize("row", [{}, {"_ng_rollout_id": "invalid/path"}]) +def test_flagged_rollout_without_valid_correlation_still_fails(tmp_path, row): + from types import SimpleNamespace + + env = _capture_env() + env.cfg = { + "initial_global_config_dict": { + "token_id_capture": { + "enabled": True, + "all_agents": True, + "dir": str(tmp_path), + "delivery": "all_traces", + "builder": "per_request", + } + } + } + with pytest.raises(ValueError): + asyncio.run( + env._postprocess_all_gym_traces( + row, {"mask_sample": True}, SimpleNamespace(pad_token_id=0) + ) + ) From a425d8b3c0bb991caeb40cf31ad6dd995d373df6 Mon Sep 17 00:00:00 2001 From: Elad Sarafian Date: Thu, 10 Sep 2026 05:39:01 -0700 Subject: [PATCH 6/6] fix(gym): pin complete-prompt trace normalization Signed-off-by: Elad Sarafian --- 3rdparty/Gym-workspace/Gym | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 2ce5e44ef8d..407fcf9f55e 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 2ce5e44ef8dbf44410c8ca1914134a0a9ec2e9f8 +Subproject commit 407fcf9f55e31bcd62290e6a3050f01ef8dbba3f