[Fix] retract-mode flush_cache no-op crash - #1750
Merged
Merged
Conversation
Zhichenzzz
requested review from
Shi-Dong,
fzyzcjy,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
July 21, 2026 18:03
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
…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
force-pushed
the
zhichen/retract-flush-cache-fix-megatron
branch
from
July 23, 2026 03:25
6acad47 to
09aebbf
Compare
yueming-yuan
approved these changes
Jul 23, 2026
Zhichenzzz
added a commit
that referenced
this pull request
Jul 23, 2026
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Companion to sgl-project/sglang#31962. sglang's
flush_cache()refuses to run whilewaiting_queueis non-empty — butpause_generation(mode="retract")re-queues every retractedrequest 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_queuefrom the idle check); this PR carries the small miles-sidecompanions:
sglang_engine.pyflush_cache(): back off on 400 responses too — "Cache not flushed" isa 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
TimeoutErrorinstead 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 aparenthesized 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 tomode != "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=Trueat its default, hits sglang's internalassert flush_cache_success, and the whole scheduler process dies with SIGQUIT.A third equivalent fix (
torchtitan_utilshardcodesmode="retract"instead of readingargs.pause_generation_mode) is tracked separately onfeat/torchtitan_expsince that filedoesn't exist on
mainyet.Test plan
sglang fix, the first deep-queue
update_weightsunder--pause-generation-mode retractdies exactly as described (
Cache not flushed because there are pending requestsretrieduntil
TimeoutError); with [sglang-miles] Fix flush_cache() no-op after pause_generation in retract sgl-project/sglang#31962 applied, the same run completes.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_placecontrol run.--tp 4 --ep-size 4):retract-mode weight updates and pause-the-world eval green end to end.
never-scheduled-backlog scenario).