RFC: rollout memory and context budget for agentic RL

[Dayuxiaoshui]

Area

Rollout engine (vLLM-Omni / SGLang / train-side)

Motivation

What is blocked today

examples/deep_research/deep_research_search_judge.yaml is the only colocated recipe in the repo that
does not set enable_memory_saver. As a result the per-rollout self.rollout.sleep()
(unirl/trainer/agentic.py:175,198) returns HTTP 200 and frees zero bytes, and SGLang's own warning
about exactly this never prints — the failure is completely silent.

Because rollout memory is never returned, the KV pool has to coexist with the train state, so
mem_fraction_static is pinned at 0.3. The measured safe ceiling is 0.88, i.e. 3.64x of KV
pool is being left on the table today. With the pool that small, decode concurrency is throttled
(measured: only 6.5 of 8 requests actually running) — which is precisely the state in which one concludes
"agentic rollout needs KV offload". Measured, that conclusion holds only inside a narrow window.

A second, independent line: the agentic engine has no context budget at all. visit.py's default
max_content_chars=90000 (≈22.5k tokens) against ctx 40,960 means the second visit necessarily
overflows — measured HTTP 400 entering turn 3, silently discarding 6 of 8 turns, with the train side
seeing only a failed trajectory.

Why this is one chain, not four small things

Filed separately, each item looks too small to be worth fixing:

missing enable_memory_saver
  -> sleep() is a silent no-op (warning suppressed, HTTP 200)
    -> rollout memory must coexist with the train state
      -> mem_fraction_static pinned at 0.3 (the real ceiling is 0.88)
        -> 151k-token KV pool, decode concurrency throttled
          -> looks like "we need KV offload" (only true for R in ~1.4-2)

This round of measurement started from the question "should UniRL adopt KV offload (HiCache /
Mooncake-style tiered KV)?". The root cause turned out to be the missing flag, not a missing offload
design.

Related issues / PRs

• [Tracking] SGLang AR rollout arg cleanup (#40 item 4) #261 [Tracking] SGLang AR rollout arg cleanup (open) — the nearest neighbour. It covers warning on
  silently-dropped engine_kwargs keys and aligning AR recipe presets. Two differences: (i) [Tracking] SGLang AR rollout arg cleanup (#40 item 4) #261's scope
  is AR recipes and misspelled/dropped keys, whereas this is the agentic/deep_research recipe
  and a missing key; (ii) [Tracking] SGLang AR rollout arg cleanup (#40 item 4) #261 plans a doc blurb on "which knobs matter for colocate memory" — this
  issue supplies exactly the measured numbers that blurb needs. If maintainers prefer, I can submit the
  first two items as sub-items of [Tracking] SGLang AR rollout arg cleanup (#40 item 4) #261 rather than opening a duplicate PR.
• feat(weight-sync): support SGLang TP/EP full-weight sync #226 feat(weight-sync): support SGLang TP/EP full-weight sync (merged) — does touch
  enable_memory_saver, but in the opposite direction: it fails closed when such fields are explicitly
  requested and the installed SGLang build does not support them. This issue is about the field being
  absent, which degrades silently; feat(weight-sync): support SGLang TP/EP full-weight sync #226's check never fires on that path.
• [rollout, ckpt] RFC: Multi-GPU SGLang Rollouts with TP, EP, and Checkpoint-Engine Weight Sync #217 RFC: Multi-GPU SGLang Rollouts with TP, EP, and Checkpoint-Engine Weight Sync — its
  motivation also notes that training and rollout memory peaks compete on one device, but the subject
  differs.
• Agentic RL Roadmap #93 Agentic RL Roadmap — the context-budget item belongs to this roadmap's scope.
• fix(rollout): make sglang_diffusion engine.sleep() idempotent #67 (closed) — made sglang_diffusion's sleep() idempotent; same theme, different defect.

Duplicate check: searched memory_saver (0), hicache (0), hierarchical cache (0),
deep_research memory (0), context overflow (0), mem_fraction_static (8),
release_memory_occupation (1, #67 closed), context length agentic (1, #94 closed), and read the
bodies of #261 / #226 / #217 / #67. No duplicate.

Feature request

Four items, in causal order, each with measured numbers.

1. Missing enable_memory_saver makes sleep() free zero bytes, silently

unirl/rollout/engine/sglang/config.py:114-115 declares both flags as Optional[bool] = None, and
:221-224 forwards them only when explicitly set. On the SGLang side enable_memory_saver defaults
to False, and TorchMemorySaverAdapter.create(enable=False) returns a Noop adapter — so
release_memory_occupation succeeds, returns 200, and frees nothing, and the adapter's own
will not save memory because torch_memory_saver is not enabled warning never fires (grep count: 0).

All seven other colocated recipes set both flags:

recipe	lines
examples/ar/qwen3_cppo_30b_a3b_base_dapo_sglang_lora.yaml	118-119
examples/ar/qwen3_cppo_30b_a3b_base_dapo_sglang_full_weight_sync.yaml	107-108
examples/ar/qwen3_5_moe_grpo_35b_a3b_base_dapo_sglang.yaml	100-101
examples/ar/qwen3_5_moe_grpo_35b_a3b_geo3k_mc_sglang.yaml	98-99
examples/pe/pe_sglang_full_wise.yaml	239-240
examples/pe/pe_sglang_full_pickscore.yaml	266-267
examples/pe/pe_sglang_lora_pickscore.yaml	249-250
examples/deep_research/deep_research_search_judge.yaml:111-117	neither

One H20 (97,871 MiB), Qwen3-4B, mem_fraction_static: 0.6, POSTing /release_memory_occupation and then
/resume_memory_occupation:

	baseline	after release	after resume
saver OFF (the recipe today)	60,016 MiB	60,056 MiB — frees nothing	60,056 MiB
saver ON	60,068 MiB	1,910 MiB — frees 58.2 GiB	60,108 MiB

Both return HTTP 200. Turning the flag on costs 0.1% of pool size.

Ask: add enable_memory_saver: true + enable_weights_cpu_backup: true to that recipe, and warn or
fail closed when a colocated layout does not receive the flag.

2. The real mem_fraction_static ceiling is 0.88, not 0.3

BaseBackend.offload() (unirl/train/backend/base_backend.py:562) moves params + grads + optimizer to
CPU and calls empty_cache(). Measured, Qwen3-4B fp32 + Adam, world_size 1:

	GPU
after optimizer.step() (params + grads + Adam m,v)	65,163 MiB
after offload()	383 MiB — the GPU is genuinely empty
after onload()	62,195 MiB

So during rollout only the server should be resident, and the pool ceiling should be a purely
rollout-phase question — but only if sleep() really releases before onload(). With the saver off,
the two must coexist: mf 0.88 (81 GiB) + ~32 GiB of sharded train state at num_devices: 2 = 113 GiB

96 GiB → OOM, while mf 0.3 fits. That is the entire reason for the 0.3.

Pool vs mf (Qwen3-4B tp=1, cell_size 147,456 B/token):

mf	KV pool	free GPU after boot
0.30 (recipe today)	151,276 tok	—
0.60	358,146 tok	—
0.80	495,691 tok	17.45 GB
0.88 (ceiling)	550,856 tok (3.64x)	9.99 GB
0.93	585,335 tok	5.21 GB (too tight)

Ask: once item 1 lands, raise mem_fraction_static to the measured safe value.

3. The KV-offload benefit is unimodal in R — do not recommend it unconditionally

Define R = working set / KV pool, where working set = concurrency x per-trajectory context. Seven
HiCache on/off pairs across N=8/16/32/64 and three independent pool sizes (151k / 358k / 551k)
collapse onto a single unimodal curve when sorted by R:

R	N	pool	hc off (s)	hc on (s)	HiCache marginal
0.72	8	358,146	296.2	295.2	−0.3%
0.94	16	550,856	394.2	393.8	−0.1%
1.45	16	358,146	497.3	462.0	−7.1%
1.71	8	151,276	344.6	304.3	−11.7% (peak)
1.88	32	550,856	1083.8	989.4	−8.7%
2.89	32	358,146	1233.2	1210.3	−1.9%
3.76	64	550,856	2394.2	2390.2	−0.2%

Below R=1 there is nothing to spill. Far above ~2.5 the host tier thrashes too (restored blocks are
evicted again before reuse, and the bottleneck has moved from prefill to decode/retraction), so it buys
nothing again. Raising the pool wins at every R (N=8: −14.0%, N=16: −20.7%, N=32: −12.1%) and lowers
R, which is also the way out of the window. Stacking both at N=32 is −19.8%.

Mechanically, a small pool throttles decode concurrency; it is not wasting prefill:

pool	peak token usage	mean #running-req	decode iters
151k, N=8	1.00 (full)	6.73 / 6.52 of 8	972 / 1004
358k, N=8	0.68	8.00	818

HiCache restores prefix reuse but the device pool stays full, so the throttle remains — that gap is
exactly 304.3 vs 296.2. This is why prefix-reuse percentage is a bad proxy for wall clock: at R=3.76
HiCache lifted reuse from 9.3% to 13.2% and bought 0.2%.

Two hard constraints any host-tier design will hit:

1. HiCache's host tier must strictly exceed the device pool (AssertionError: The host memory should be larger than the device memory with the current protocol). At mf 0.88 the device KV is 75.7 GiB, so
   --hicache-size 64 will not boot and 96 is required. On an 8-GPU colocated node that is ~600 GB of
   host RAM — raising the pool and adding a host tier are in direct tension.
2. See item 4: the long-observation regime is not an offload opportunity, it is a crash.

Ask: document the R criterion in the rollout-side README.md under ## Gotchas, replacing the
unconditional "agentic rollout needs KV offload". Compute R = (concurrency x per-trajectory context) /
(pool at the max safe mf). R<1 → raise the pool only, a host tier buys ~0. R in [1.4, 2] → raise the pool
and a host tier is worth 9–12%. R>2.5 → neither helps much; cut concurrency or shard differently.
Always check token usage and mean #running-req before proposing offload.

4. No context budget in the agentic engine — the default config dies on turn 3

unirl/rollout/env/tools/visit.py:55 has max_content_chars: int = 90000 (≈22.5k tokens), and :125
skips summarization when SUMMARY_URL is unset, giving ctx/obs = 40,960 / 22,530 = 1.82 — so
the second visit necessarily overflows, regardless of generation length (verified at 300 / 800 /
2048 gen-tok per turn).

Measured against a live server with an 8-turn budget (max_turns: 8):

turn 1: wire=14     status=200
turn 2: wire=22528  status=200
turn 3: wire=45042  status=400
        {'error': {'message': "The input (45042 tokens) is longer than the model's context length (40960 tokens)."}}

Nothing in unirl/rollout/engine/agentic/engine.py budgets context; the per-trajectory net at :85-92
marks the trajectory failed, so 6 of 8 turns are silently lost and the train side has no indication
that the cause was context overflow. At 2048 gen-tok/turn the largest observation that survives 8 turns
is about 3,510 tokens (≈14k chars) — 6.4x smaller than the default.

Two related problems:

• HTTPBackend._post_generate (unirl/rollout/engine/sglang/backends/http.py:343,
  max_retries: int = 60) retries this deterministic 400 sixty times at 1 s intervals — one wasted
  minute per overflowing trajectory.
• The recipe sets no context length, and nothing anywhere validates
  max_turns x (max_new_tokens + observation budget) <= ctx.

Ask: clip observations against the remaining context; report overflow as an explicit outcome rather
than a generic failed; stop retrying deterministic 4xx sixty times.

Method and honest limits

• Hardware: one H20 (97,871 MiB), Qwen3-4B (tp=1, cell_size 147,456 B/token).
• Engine args copied from the recipe's engine_kwargs (attention_backend: triton, cuda_graph_max_bs,
  skip_server_warmup); only max_running_requests was raised so that concurrency, not the scheduler
  cap, is the independent variable.
• This is a single-GPU rollout-side proxy, not an end-to-end run of the num_devices: 2 recipe. The
  113 GiB > 96 GiB OOM in item 2 is computed from the sharded train state, not observed as a crash.
• Trajectories run in ignore_eos lockstep (a reproducible worst case; real runs stagger), and
  observations are synthetic and fixed-length. The absolute seconds in item 3 are therefore not an
  end-to-end prediction; the relative differences between paired arms are trustworthy.
• I did not run real deep_research end-to-end (missing SERPER_KEY_ID / JINA_API_KEYS, judge server
  down, and the token budget needs authorization). The item-4 overflow was reproduced with real tokenizer
  lengths against a real server.
• Harness: serve.sh / drive.py / overflow.py / offload_mem.py, archived together with 12
  server/probe logs. Happy to attach them or post a gist.

Your contribution

Yes, I can send PRs. Split into four independently reviewable pieces:

1. Recipe fix (minimal, mergeable immediately): add enable_memory_saver: true +
   enable_weights_cpu_backup: true to examples/deep_research/deep_research_search_judge.yaml and
   raise mem_fraction_static from 0.3 to the measured safe value, with the before/after memory numbers
   from items 1 and 2 in the Test Plan.
2. Warn / fail closed: in a colocated layout, warn or fail closed when the rollout engine receives no
   enable_memory_saver (unirl/rollout/engine/sglang/config.py). This overlaps [Tracking] SGLang AR rollout arg cleanup (#40 item 4) #261 the most, so
   maintainers' call — I can retarget it as a sub-PR to [Tracking] SGLang AR rollout arg cleanup (#40 item 4) #261, or let @yhl48's planned PR absorb it.
3. Context budget: clip observations against the remaining context in the agentic engine /
   ToolEnvironment, report overflow as an explicit outcome, and stop retrying deterministic 4xx sixty
   times.
4. Docs: put the R criterion and the two hard constraints from item 3 into the rollout-side
   README.md under ## Gotchas.

Per CLAUDE.md §5, the one-off verification harnesses are not committed; their commands and results are
quoted in the PR Test Plan. I am happy to open each item as a separate PR, and I will not open a
duplicate PR if maintainers decide items 1 and 2 belong in #261.

[CjhHa1]

Thanks for the measurements — this is convincing. Please go ahead and send the PRs.

The “agentic rollout needs KV offload” conclusion looks like an artifact of the missing enable_memory_saver, not a missing HiCache / Mooncake-style host-tier. Do not start a Mooncake / HiCache integration on this path until someone re-measures R after the recipe fix. If R stays below 1, raising the pool is the whole fix.

Split is good; don’t wait on #261:

1. Recipe fix first, independently mergeable. Add enable_memory_saver: true + enable_weights_cpu_backup: true to deep_research_search_judge.yaml and raise mem_fraction_static to the measured safe value. Put the before/after GPU numbers from items 1–2 in the Test Plan. Fine that the 113 GiB colocated OOM is computed rather than observed — just say so there.
2. Warn / fail closed when a colocated layout does not receive enable_memory_saver. Own PR is fine. [Tracking] SGLang AR rollout arg cleanup (#40 item 4) #261 can absorb AR-side docs later; that should not block this.
3. Context budget next. Clip observations against remaining context, report overflow as an explicit outcome (not a generic failed), and stop retrying deterministic 4xx sixty times. This is a real training-side hole, independent of the memory chain.
4. Docs with the R criterion after (1) lands, so the numbers match the recipe. Replace any unconditional “agentic rollout needs KV offload” with the unimodal window.

Per item, one PR, measured numbers in the Test Plan, no one-off harnesses in tree. Looking forward to (1).