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

nemo_rl/algorithms/distillation.py

@@ -503,6 +503,9 @@ def init_nemo_gym():
                     "invalid_tool_call_patterns", None
                 )
                 thinking_tags = nemo_gym_dict.pop("thinking_tags", None)
+                truncate_noncontiguous_episodes = nemo_gym_dict.pop(
+                    "truncate_noncontiguous_episodes", False
+                )
                 # Pass prebuilt cache + venv dirs through the global config so the
                 # gym reuses image-baked venvs instead of rebuilding them.
                 uv_cache_dir = get_nemo_gym_uv_cache_dir()
@@ -516,6 +519,7 @@ def init_nemo_gym():
                     base_urls=deferred_vllm.dp_openai_server_base_urls,
                     invalid_tool_call_patterns=invalid_tool_call_patterns,
                     thinking_tags=thinking_tags,
+                    truncate_noncontiguous_episodes=truncate_noncontiguous_episodes,
                     initial_global_config_dict=nemo_gym_dict,
                 )
                 nemo_gym_opts = {

nemo_rl/algorithms/grpo.py

@@ -483,6 +483,9 @@ def _spinup_nemo_gym(base_urls, model_name):
             "invalid_tool_call_patterns", None
         )
         thinking_tags = nemo_gym_dict.pop("thinking_tags", None)
+        truncate_noncontiguous_episodes = nemo_gym_dict.pop(
+            "truncate_noncontiguous_episodes", False
+        )
         # Pass prebuilt cache + venv dirs through the global config so the gym reuses
         # image-baked venvs instead of rebuilding them.
         uv_cache_dir = get_nemo_gym_uv_cache_dir()
@@ -497,6 +500,7 @@ def _spinup_nemo_gym(base_urls, model_name):
             invalid_tool_call_patterns=invalid_tool_call_patterns,
             thinking_tags=thinking_tags,
             require_routed_experts=router_replay_enabled(policy_config),
+            truncate_noncontiguous_episodes=truncate_noncontiguous_episodes,
             initial_global_config_dict=nemo_gym_dict,
         )
         nemo_gym_opts = {}

nemo_rl/environments/nemo_gym.py

@@ -80,6 +80,14 @@ class NemoGymConfig(TypedDict):
     require_routed_experts: NotRequired[
         bool
     ]  # Require Gym output items to carry R3 routed_experts
+    # When true, an episode whose turn breaks token-prefix contiguity (e.g. a
+    # rare tokenization/re-render edge case in a long multi-turn rollout) is
+    # truncated at the last contiguous turn: the corrupted tail is dropped and
+    # the valid prefix stays trainable (same philosophy as overlong filtering).
+    # When absent/false (default), such an episode raises the contiguity
+    # assertion below, which kills the rollout task; under async GRPO a single
+    # such episode can then stall the training step indefinitely.
+    truncate_noncontiguous_episodes: NotRequired[bool]
 
 
 def _detect_invalid_tool_call_and_malformed_thinking(
@@ -335,10 +343,23 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
             if "generation_token_ids" not in output_item_dict:
                 continue
 
-            assert (
+            is_contiguous = (
                 seen_token_ids
                 == output_item_dict["prompt_token_ids"][: len(seen_token_ids)]
-            ), f"""Non-contiguous messages found! This may be a tokenization issue where certain tokens are combined when messages are concatenated, or it may be due to part of the chat history being truncated (like if super long history is truncated or if reasoning is stripped out).
+            )
+            if not is_contiguous and self.cfg.get("truncate_noncontiguous_episodes"):
+                # Opt-in resilience: keep the contiguous prefix trainable and
+                # drop the corrupted tail (same philosophy as overlong
+                # filtering) instead of killing the rollout task — under async
+                # GRPO a single asserting episode can stall the step
+                # indefinitely.
+                print(
+                    "[nemo_gym] WARNING: non-contiguous turn; truncating episode "
+                    f"at {len(nemo_rl_message_log)} messages "
+                    f"(seen={len(seen_token_ids)} tokens); dropping corrupted tail."
+                )
+                break
+            assert is_contiguous, f"""Non-contiguous messages found! This may be a tokenization issue where certain tokens are combined when messages are concatenated, or it may be due to part of the chat history being truncated (like if super long history is truncated or if reasoning is stripped out).
 Seen token IDs: {seen_token_ids}
 Output prompt token IDs: {output_item_dict["prompt_token_ids"]}
 """

tests/unit/environments/test_nemo_gym.py

@@ -223,6 +223,64 @@ class _MockSelf:
     assert nemo_gym_result["response"]["output"][1]["generation_str"] == "6 7"
 
 
+def _make_noncontiguous_nemo_gym_result():
+    """Two assistant turns where turn 2's prompt breaks token-prefix contiguity.
+
+    After turn 1, seen_token_ids == [1, 2, 3]; turn 2's prompt starts with
+    [1, 99, ...] (99 != 2), e.g. a tokenization/re-render edge case in a long
+    multi-turn rollout.
+    """
+    return {
+        "response": {
+            "output": [
+                {
+                    "prompt_token_ids": [1, 2],
+                    "generation_token_ids": [3],
+                    "generation_log_probs": [-0.1],
+                },
+                {
+                    "prompt_token_ids": [1, 99, 3, 4],
+                    "generation_token_ids": [6],
+                    "generation_log_probs": [-0.2],
+                },
+            ]
+        },
+        "responses_create_params": {"input": []},
+    }
+
+
+class _JoinTokenizer:
+    def batch_decode(self, batch):
+        return [" ".join(map(str, token_ids)) for token_ids in batch]
+
+
+def test_nemo_gym_postprocess_noncontiguous_asserts_by_default():
+    class _MockSelf:
+        cfg = {}
+
+    with pytest.raises(AssertionError, match="Non-contiguous messages found"):
+        NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result(
+            _MockSelf(), _make_noncontiguous_nemo_gym_result(), _JoinTokenizer()
+        )
+
+
+def test_nemo_gym_postprocess_noncontiguous_truncates_when_enabled():
+    class _MockSelf:
+        cfg = {"truncate_noncontiguous_episodes": True}
+
+    result = (
+        NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result(
+            _MockSelf(), _make_noncontiguous_nemo_gym_result(), _JoinTokenizer()
+        )
+    )
+
+    # The corrupted second turn is dropped; the contiguous first turn survives
+    # as a trainable (user prompt, assistant generation) pair.
+    assert len(result["message_log"]) == 2
+    assert result["message_log"][0]["token_ids"].tolist() == [1, 2]
+    assert result["message_log"][1]["token_ids"].tolist() == [3]
+
+
 @pytest.mark.nemo_gym
 def test_nemo_gym_sanity(
     nemo_gym,