Why one dropped prompt group stops a whole context-compaction training run

Explainer for a review finding on NVIDIA-NeMo/RL #4124 (paired with NVIDIA-NeMo/Gym #3368). RL links pin c617a2c; Gym links pin b27e0db.

Before
NeMo-RL main has no context compaction (CC). On the ordinary path, a dropped prompt group can be retried or skipped within a budget.
This PR
Under CC, a failed model call now masks only its own rollout. But a whole dropped group still ends the run, while the new guide says it stays within the drop budget.
Status
Open — not fixed in this PR. Raised in the review (comment below). The Gym half belongs in Gym #3368, and retry is tracked in #4557.

Nine pages on this PR, one question each. This one: what happens when one prompt group fails.

Background. A GRPO training step needs a fixed number of prompt groups; each group is one prompt run several times (each run is a rollout). NeMo-RL normally tolerates a few lost groups: a drop budget and a minimum step size decide how many.
Why context compaction (CC) cannot absorb a drop. CC turns off both ways to absorb one. It forces the minimum to 100% (config.py L1437-1438), because each CC optimizer step must hold exactly train_global_batch_size rollouts, or it raises "CC optimizer step requires the full logical sample count" (single_controller.py L3206-3213). It also never re-sends a dropped group, because RL keeps the record of which responses the agent accepted only inside the request that ran it (guide L47). Neither the code nor the guide says why a short step would be wrong. The per-rollout advantage and the token-level loss work for any number of rollouts, so this looks like a limit of the first version, not a math requirement. Dotted words have a short explanation: hover or tap them.

The figure follows one step of 4 prompt groups with 2 rollouts each, and one bad rollout in group g1. Where the failure happens decides everything. If a model call fails, the agent still finishes and returns a result, and RL masks just that rollout (nemo_gym.py L1114-1120). If the agent itself raises, Gym answers with an HTTP 500, RL files it as an infrastructure failure, and CC drops the whole group without a retry (rollout_manager.py L2210-2214). Both ways of handling a dropped group then end the run.

How compaction is triggered, and where the length is checked

In this PR's built-in client (simple_agent_with_compaction), compaction is not set off by vLLM hitting its limit. It runs on a schedule. With a context guard set (guards.max_total_tokens, off by default), Gym also counts the prompt before every call and compacts early when it would be too long. The guard's refusal (a "guard rejection") happens only when compacting early still is not enough. The guard checks the prompt, not the reply: a reply can still run into the limit, and is then cut off (bottom box).

Before every model call, both lanes: scheduled compaction drops old images and old reasoning blocks, every call (rolling_recency) or every K actions (turn_chunked_recency). It is a plain rewrite: it never calls the model. context_management/config.py L52-94 Guard on: Gym checks the prompt length before sending Guard off (default): nobody checks first Count the prompt tokens. Is prompt + reserved output (default 0) under guards.max_total_tokens? Send the call to vLLM as is. fits too long Send the call. ✅ The normal case. Compact early: close the chunk now, then count again. client.py L167-177 fits too long (400) vLLM runs it. ✅ The normal case. No compaction happens. Gym returns an empty reply with finish_reason "length"; the episode ends. vllm_model/app.py L1113-1124 fits still too long Send the call. ✅ Compaction was triggered by the length. Refuse the call. The agent crashes: HTTP 500, group dropped. Lane B. That one rollout is masked; the rest of its group trains. Lane A. Once a call is sent, both lanes: RL caps the reply to the room left (max_model_len minus the prompt). The guard does not check the reply unless reserved_generation_tokens is set. A reply that reaches the cap is cut off with finish_reason "length": a normal truncated reply, not an error. RL then marks the rollout truncated; with overlong_filtering the whole rollout is masked. vllm_worker_async.py L1018-1037 "Still too long" happens when what is left is text, tool output or user turns: this client never drops those. Early compaction exists only for turn_chunked_recency; with rolling_recency the history is already compacted, so a guard hit refuses at once. Harnesses that summarize with the model (Claude Code, OpenCode) work differently, but this PR does not support them (page 1).

Same bad rollout, two places it can fail

one training step · 4 prompt groups (g0–g3) · 2 rollouts each · rollout g1_0's history grows too long for the model A · no guard: vLLM notices, after the call is sent The call is sent; vLLM refuses it (400, prompt too long). Gym's model server catches the 400 and returns an empty reply with finish_reason="length", so the agent ends normally. ledger: c3 failed, no tokens reward: from verifier Gym vllm_model/app.py L1113-1124 RL sees a capture that never finished and poisons only this rollout. capture_failed=True → receipt.capture_poisoned=True nemo_gym.py L1114-1143 Group g1 keeps its slot: g1_0 is masked (no gradient), g1_1 trains as usual. g1_0: sample_mask=0 g1_1: sample_mask=1 Step trains with 4 of 4 groups. The run continues. B · guard on: Gym notices, before the call is sent The context guard (optional, off by default) counts the prompt before each call and refuses to send it if it is too long. That refusal is an error the agent never catches → HTTP 500. 500 "ContextGuardRejected('Context guard …')" Gym client.py L176server_utils.py L1297-1303 RL reads a 500 as an infrastructure failure. Under CC it does not retry; it drops all of group g1 (g1_1 too). GymTransportError("NeMo-Gym /run failed with HTTP 500: …") nemo_gym.py L287-289rollout_manager.py L2210-2214 budget 0 (default) budget = 5 Drop 1 is more than the budget of 0, so it raises: RolloutRedispatchExhausted "… exceeding max_consecutive_dropped_ prompts=0 …" rollout_manager.py L2240-2248 The step keeps 3 of 4 groups; the floor is 4, so it raises: RuntimeError "training step … lost 1 of 4 prompt group(s), leaving 3, below … 4" single_controller.py L2638-2647 The whole training run stops. Also dropped the same way, so also run-ending: · a timeout or transport error before the agent returns · any other agent error (resample budget, call limit) · a group where every rollout failed: single_controller.py L2062-2071 only one rollout masked; training continues the run stops a step on the way Numbers in the error text are this example's: 4 groups, 1 lost, 3 left, floor 4 (= every group, because CC requires 100%).

🔴 review comment — context-compaction.md:47

What would fix it — proposed, not in the PR

Nothing below is implemented or has been run.

Where the fix lands: RL code and docs (this PR), the Gym agent (Gym #3368), and a follow-up issue for retry. The ordinary path has two safety nets for a failed group: retry it, then train without it. CC turned both off. The fix puts them back, one at a time.

1 · RL, in this PR: train without the dropped group

Let a CC step train on the groups it has, the way the ordinary path does. The loss is divided by the token count and the advantage is per rollout, so nothing in the math needs the full count (page 7).

@@ config.py L1437-L1438 · CC config check
             raise ValueError(
                 "context compaction requires direct routed-expert assembly"
             )
-        if async_config.rollout_failure.min_step_batch_fraction != 1:
-            raise ValueError("context compaction requires min_step_batch_fraction=1")
         validate_cc_objective(master_config.grpo, master_config.loss_fn)
@@ single_controller.py L3206-L3213 · _train_pump, before finish_train_step
                 if not self._is_ppo:
                     if context_compaction and (
-                        groups_dispatched != self._algo_cfg.num_prompts_per_step
-                        or len(logical_owners)
-                        != self._master_config.policy["train_global_batch_size"]
+                        # A dropped group shrinks the step, as on the ordinary path.
+                        groups_dispatched != self._target_groups_for_step(version_during_step)
+                        or len(logical_owners)
+                        != groups_dispatched * self._algo_cfg.num_generations_per_prompt
                     ):
                         raise ValueError(
-                            "CC optimizer step requires the full logical sample count"
+                            "CC optimizer step requires every rollout of every kept group"
                         )

_target_groups_for_step already returns the number of groups minus the dropped ones, and still enforces min_step_batch_fraction. The "gbs" case in test_cc_optimizer_batch.py needs its expected message updated. Also add one test where a group drops and the step trains with the rest. With this change, a dropped group in lane B costs one group of data instead of the whole run. That holds once the user sets a drop budget above 0, exactly as on the ordinary path.

2 · Gym, in #3368: a guard rejection should not crash the agent

simple_agent_with_compaction should catch the guard's refusal, end the episode there, and return what it has, the way lane A already does. RL then masks that one rollout instead of dropping the whole group.

3 · Follow-up: retry the failed group (#4557)

The ordinary path re-runs a failed group with backoff (rollout_manager.py L2215-2219). CC skips that (L2210-2214), because RL keeps the record of which responses the agent accepted only for that one request. This can land after this PR, but it matters: at moderate scale, timeouts and transport errors are routine. Without retry, each one throws away a group the ordinary path would have recovered.

4 · Docs: say what the code does

@@ docs/guides/context-compaction.md L47 · Measurement and lifecycle
- Set `async_rl.rollout_failure.nemo_gym.rollout_timeout_s` as a backstop for a hung prompt-group stream; it is not a per-model-call deadline or an acknowledgement that the remote engine stopped. A timeout or transport failure that prevents an agent result drops the group under the existing consecutive-drop budget. This baseline does not redispatch CC groups because selected-response metadata is request-local; checkpoint/retry support is a separate change. Data-integrity errors still raise, and operator cancellation still propagates.
+ Set `async_rl.rollout_failure.nemo_gym.rollout_timeout_s` as a backstop for a hung prompt-group stream; it is not a per-model-call deadline or an acknowledgement that the remote engine stopped. A timeout, transport failure or agent exception that prevents an agent result drops the group without a retry, and so does a group with no valid owner. A dropped group shrinks that training step, within `max_consecutive_dropped_prompts` and `min_step_batch_fraction`, the same as on the ordinary path; with the default budget of 0, the first drop still stops the run. CC groups are not redispatched because selected-response metadata is request-local; retry is tracked in NVIDIA-NeMo/RL#4557. Data-integrity errors still raise, and operator cancellation still propagates.
So what. Today a context guard makes a CC run less robust, not more: without a guard, an oversized request masks one rollout (lane A); with a guard, the same situation stops the run (lane B). Until fix 1 lands, any rollout error that keeps the agent from returning a result ends a context-compaction training run. With fix 1, it costs one group. With retry (#4557), most such errors would cost nothing.