fix(fsdp): apply pause mode during non-colocated weight sync - #2494
Open
shangrilar wants to merge 1 commit into
Open
fix(fsdp): apply pause mode during non-colocated weight sync#2494shangrilar wants to merge 1 commit into
shangrilar wants to merge 1 commit into
Conversation
shangrilar
marked this pull request as ready for review
August 13, 2026 00:36
shangrilar
requested review from
Rockdu,
Zhichenzzz,
fzyzcjy,
guapisolo and
yueming-yuan
as code owners
August 13, 2026 00:36
Author
|
Ready for review. The fork workflows are currently awaiting approval. Could a maintainer apply a run-ci label when convenient? Thank you! |
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
pause_generation_modeto every rollout engine in the FSDP weight-update lifecyclein_placeupdatesabort,retract, andin_placewith a parameterized lifecycle regression testProblem
The FSDP updater called
pause_generation()without forwarding the configured mode and always calledflush_cache(). Because the SGLang wrapper defaults toretract, a fully async FSDP run configured with--pause-generation-mode in_placebehaved likeretract + flushinstead of preserving its in-flightrequests and KV cache.
The updated lifecycle is:
The separate waits also ensure every rollout engine has completed its pause before a cache flush or
weight-update session begins.
Scope
Fully async training rejects colocation, so FSDP uses the non-colocated distributed weight-sync path fixed
here. This change does not remove the colocated tensor updater's separate per-bucket cache flush and does
not claim universal colocated
in_placesupport.PR #1861 rewrites the same FSDP lifecycle as part of a larger SGLang client refactor, but currently retains
the no-argument pause and unconditional flush behavior.
Validation
tests/fast/backends/test_fsdp_update_weight.py: 5 passed in a local Python 3.12 harness thatloads the production updater and test file while stubbing unavailable Ray/SGLang imports
pre-commit run --files miles/backends/fsdp_utils/update_weight_utils.py tests/fast/backends/test_fsdp_update_weight.pygit diff --checkThe full Miles dependency suite is left to upstream CI.