Skip to content

[Fix] retract-mode flush_cache no-op crash - #1750

Merged
Zhichenzzz merged 2 commits into
mainfrom
zhichen/retract-flush-cache-fix-megatron
Jul 23, 2026
Merged

[Fix] retract-mode flush_cache no-op crash#1750
Zhichenzzz merged 2 commits into
mainfrom
zhichen/retract-flush-cache-fix-megatron

Conversation

@Zhichenzzz

@Zhichenzzz Zhichenzzz commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Summary

Companion to sgl-project/sglang#31962. sglang's flush_cache() refuses to run while
waiting_queue is non-empty — but pause_generation(mode="retract") re-queues every retracted
request into exactly that queue, and under high concurrency the queue also always holds
never-scheduled requests. So the flush that weight updates depend on could never succeed, and
retract mode was effectively unusable for RL training. The real fix is sglang-side (exempt the
paused-engine waiting_queue from the idle check); this PR carries the small miles-side
companions:

  • sglang_engine.py flush_cache(): back off on 400 responses too — "Cache not flushed" is
    a normal 400 response, not an exception, so without the sleep the 60 retries burned through in
    under a second instead of giving in-flight generation ~60s to drain. Also surface the last
    response body in the eventual TimeoutError instead of discarding it (Timeout while flushing cache: Flush cache failed. instead of a bare timeout).
  • update_weight_from_distributed/{delta.py,mixin.py}: if mode not in ("in_place"): is a
    parenthesized string, not a tuple — a substring check, not membership. Harmless today (no valid
    mode is a proper substring of "in_place") but a latent footgun; fixed to mode != "in_place".

For the megatron backend the sglang bug was a crash, not a silent no-op: the bucket-transfer
weight update leaves flush_cache=True at its default, hits sglang's internal
assert flush_cache_success, and the whole scheduler process dies with SIGQUIT.

A third equivalent fix (torchtitan_utils hardcodes mode="retract" instead of reading
args.pause_generation_mode) is tracked separately on feat/torchtitan_exp since that file
doesn't exist on main yet.

Test plan

  • Baseline reproduction (megatron backend, fully-async, Qwen3.5-4B on 8xH200): without the
    sglang fix, the first deep-queue update_weights under --pause-generation-mode retract
    dies exactly as described (Cache not flushed because there are pending requests retried
    until TimeoutError); with [sglang-miles] Fix flush_cache() no-op after pause_generation in retract sgl-project/sglang#31962 applied, the same run completes.
  • Production-scale stress (megatron, fully-async, 512 in-flight requests, global-batch 128
    → 4 weight broadcasts per rollout, 20 rollouts ≈ 80 retract-mode broadcasts, checkpoint
    persistence + pause-the-world eval): zero flush warnings, eval version pinning intact,
    training curve statistically identical to an in_place control run.
  • MoE + TP4/EP4 (Qwen3-30B-A3B, megatron actor TP4/EP4 + sglang --tp 4 --ep-size 4):
    retract-mode weight updates and pause-the-world eval green end to end.
  • sglang-side regression tests live in the paired PR (5 scheduler cases incl. the
    never-scheduled-backlog scenario).

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@Zhichenzzz Zhichenzzz changed the title Fix retract-mode flush_cache no-op crash for megatron backend [Fix] retract-mode flush_cache no-op crash Jul 22, 2026
…ubstring-check footgun

Companion to sgl-project/sglang#31962 (exempt the paused-engine
waiting_queue from the flush idle-check). miles side: back off on 400
flush responses so the 60 retries give real drain time, surface the
last response body in the flush TimeoutError, and fix the
'mode not in ("in_place")' substring check in delta.py/mixin.py.
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/retract-flush-cache-fix-megatron branch from 6acad47 to 09aebbf Compare July 23, 2026 03:25
@Zhichenzzz
Zhichenzzz merged commit fe281ad into main Jul 23, 2026
37 checks passed
@Zhichenzzz
Zhichenzzz deleted the zhichen/retract-flush-cache-fix-megatron branch July 23, 2026 05:11
fzyzcjy added a commit that referenced this pull request Aug 27, 2026
pause-generation-mode=in_place freezes in-flight requests across a weight
update so their KV survives, which also means the engine-side prefix-cache
flush cannot run (flush_cache refuses while requests hold KV, and the
engine asserts on it). The update path therefore skips the flush in this
mode -- but nothing ever invalidates the radix tree, so prefixes cached
under pre-update weights keep serving every later request that shares
them. Rollouts silently mix old-policy KV into new-policy generations,
the recorded rollout log probs describe that hybrid, and the
train-vs-rollout log-prob gap grows with every update. The per-request
weight_version bookkeeping (mixed_version_ratio) only tracks requests
alive across an update, so it stays near zero while the prefix reuse
poisons all subsequent traffic.

retract does not have this problem: every retracted request recomputes
its KV under the new weights and the update path flushes the prefix
cache, which actually succeeds under load since #1750 and
sgl-project/sglang#31962. Flip every recipe that defaulted or hardcoded
in_place over to retract, and document why in the fully-async guide.
examples/multi_lora keeps in_place deliberately: adapter upsert relies on
it (an unload would deadlock behind paused requests) and its --use-tis
corrects the training-side bias.
fzyzcjy added a commit that referenced this pull request Aug 27, 2026
in_place was chosen to dodge the fully-async flush_cache deadlock, but that
deadlock is fixed (miles #1750, sgl-project/sglang#31962) and in_place has
a worse cost: the update path cannot flush the prefix cache in that mode,
so radix-cached prefixes computed under pre-update weights keep serving
later requests and the rollout behavior policy silently lags the trainer.
fzyzcjy added a commit that referenced this pull request Aug 27, 2026
in_place was chosen to dodge the fully-async flush_cache deadlock, but that
deadlock is fixed (miles #1750, sgl-project/sglang#31962) and in_place has
a worse cost: the update path cannot flush the prefix cache in that mode,
so radix-cached prefixes computed under pre-update weights keep serving
later requests and the rollout behavior policy silently lags the trainer.
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.

2 participants