Skip to content

feat(nemo-gym): opt-in truncation of non-contiguous rollout episodes - #3063

Draft
Kh4L wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
Kh4L:nemo-gym-noncontig-truncate
Draft

feat(nemo-gym): opt-in truncation of non-contiguous rollout episodes#3063
Kh4L wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
Kh4L:nemo-gym-noncontig-truncate

Conversation

@Kh4L

@Kh4L Kh4L commented Jul 4, 2026

Copy link
Copy Markdown

Problem

_postprocess_nemo_gym_to_nemo_rl_result hard-asserts that each assistant
turn's accumulated token IDs form a prefix of the next turn's prompt. In long
multi-turn agentic rollouts (100+ turns), rare tokenization/re-render edge
cases can break this contiguity. The assert then kills the rollout task — and
under async GRPO, the step stalls indefinitely waiting for a trajectory that
will never arrive ("Need N trajectories, only N-1 ready"), until walltime
kills the job.

Change

New opt-in NemoGymConfig knob, truncate_noncontiguous_episodes
(default false = today's assert, unchanged). When enabled, a
non-contiguous turn truncates the episode at the last contiguous turn with a
warning: the corrupted tail is dropped, the valid prefix stays trainable —
same philosophy as overlong filtering. Wired through GRPO + distillation
setup identically to invalid_tool_call_patterns.

Evidence

Validated at scale on a 16-node Qwen3-30B (MoE) SWE-bench async-GRPO run with
an SGLang generation backend: 63 training steps across four clean
checkpoint+resume segments, ~200 graceful truncations, zero stalls, reward
converged (0.10 -> 0.22 plateau). Without the flag, single corrupted episodes
reproducibly stalled the step until walltime death.

Tests

Two unit tests in tests/unit/environments/test_nemo_gym.py:

  • default path still raises the contiguity assertion;
  • enabled path returns the truncated, valid message log.

Run against clean main: 2 passed.

_postprocess_nemo_gym_to_nemo_rl_result hard-asserts that each
assistant turn's accumulated token IDs form a prefix of the next
turn's prompt. In long multi-turn agentic rollouts, rare
tokenization/re-render edge cases can break this contiguity; the
assert then kills the rollout task, and under async GRPO the step
stalls indefinitely waiting for a trajectory that never arrives.

Add an opt-in NemoGymConfig knob, truncate_noncontiguous_episodes
(default false = current assert, unchanged). When enabled, a
non-contiguous turn truncates the episode at the last contiguous
turn with a warning: the corrupted tail is dropped and the valid
prefix stays trainable, analogous to overlong filtering. Wired
through GRPO and distillation setup identically to
invalid_tool_call_patterns.

Signed-off-by: Serge Panev <spanev@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 4, 2026

Copy link
Copy Markdown

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.

@Kh4L

Kh4L commented Jul 4, 2026

Copy link
Copy Markdown
Author

/ok to test c0075c4

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant