Repository navigation
[RL] Keep pause_generation and weight updates from deadlocking each other - #40779
Conversation
…ther Three fixes for the pause -> update_weights -> continue cycle RL drivers run: - The idle sleeper's periodic empty_cache() ran while the engine was paused, racing the in-progress weight update on the allocator. IdleSleeper and RustServerIdleSleeper take a can_empty_cache gate that the scheduler binds to `not self._engine_paused`. - flush_cache() and release_memory_occupation() gate on is_fully_idle(), which counts waiting_queue. After pause_generation(mode="retract") the retracted requests (and never-scheduled ones under high concurrency) sit in waiting_queue holding no KV state, so the flush no-ops and the weight update's cache flush retries until the caller times out. is_fully_idle() gains ignore_waiting, set while the engine is paused. Squash of sglang-miles #22754, #22623, #31962. Co-authored-by: maocheng23 <maocheng@berkeley.edu> Co-authored-by: Zhichenzzz <zczeng@uw.edu>
|
/tag-run-ci-label |
|
/rerun-test test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py test/registered/unit/managers/test_scheduler_pause_generation.py test/registered/unit/managers/test_scheduler_flush_cache.py test/registered/rl/test_pause_generation_tensor_consistency.py test/registered/unit/managers/test_scheduler_weight_version_tracking.py test/registered/rl/test_update_weights_from_distributed.py test/registered/rl/test_update_weights_from_tensor.py test/registered/rl/test_multi_instance_release_memory_occupation.py test/registered/scheduler/test_scheduler_control.py |
|
Results for 🚀 🚀 🚀 🚀 |
|
/rerun-failed-ci |
…ther (sgl-project#40779) Co-authored-by: maocheng23 <maocheng@berkeley.edu> Co-authored-by: Zhichenzzz <zczeng@uw.edu>
…ract pause_generation(mode="retract") parks in-flight requests in waiting_queue, so flush_cache()'s is_fully_idle() guard refused every flush during a weight sync. slime therefore skipped the flush in retract mode, and the radix cache kept serving old-weight prefix KV (MambaRadixCache.cache_unfinished_req re-points each request at cached prefix pages after prefill, regardless of logprob_start_len): the shared system-prompt/tool-def KV stayed at the first rollout's weights for a whole run. Port of upstream sgl-project#40779 (squash of sglang-miles sgl-project#22754/sgl-project#22623/sgl-project#31962): - is_fully_idle(ignore_waiting=False); flush_cache() and release_memory_occupation() pass ignore_waiting=self._engine_paused. Queued requests hold no KV/tree/mamba state while paused (retract releases with is_insert=False; get_new_batch_prefill frees a rejected request's COW mamba slot), so ignoring them is safe. in_place pauses keep requests in running_batch, so their flush is still refused. - IdleSleeper(can_empty_cache=...) skips its periodic empty_cache() while paused, so it can't race an in-progress weight update. - test_scheduler_flush_cache_after_retract.py: upstream test adapted to this fork's scheduler attributes, plus an in_place-refusal case. Verified on GPU (Qwen3.6-35B-A3B, EP4/DeepEP): with 24 requests retracted, /flush_cache returns 200 (400 before), and post-sync decodes no longer read stale prefix KV (tool-name log p -7.81 vs -0.71 unpatched).
Motivation
RL drivers run
pause_generation->update_weights_*->continue_generationevery step. Two things in that cycle break on main:empty_cache()keeps firing while the engine is paused, racing the in-progress weight update on the allocator.flush_cache()andrelease_memory_occupation()gate onis_fully_idle(), which countswaiting_queue. Afterpause_generation(mode="retract")the retracted requests, plus never-scheduled ones under high concurrency, sit inwaiting_queueholding no KV state, so the flush no-ops and the weight update's cache flush retries until the caller times out (reproduced in RL training at 512 in-flight requests).Modifications
IdleSleeper/RustServerIdleSleepertake acan_empty_cachegate; the scheduler binds it tonot self._engine_paused.Scheduler.is_fully_idle()gainsignore_waiting.flush_cache()andrelease_memory_occupation()passignore_waiting=self._engine_paused.test_scheduler_flush_cache_after_retract.pycovers the retract case.Squash of sglang-miles #22754, #22623, #31962 (authors credited as co-authors).
Testing
pre-commit clean. CPU unit:
test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py,test_scheduler_pause_generation.py. E2E on GPU via CI:test/registered/rl/test_update_weights_from_distributed.py(pause modesin_place/retract).CI States
Latest PR Test (Base): ❌ Run #35787077698
Latest PR Test (Extra): ❌ Run #35787077454
Latest PR Test (AMD ROCm 10): ❌ Run #35787077689