Skip to content

fix(fsdp): apply pause mode during non-colocated weight sync - #2494

Open
shangrilar wants to merge 1 commit into
radixark:mainfrom
shangrilar:fix/fsdp-pause-generation-mode
Open

fix(fsdp): apply pause mode during non-colocated weight sync#2494
shangrilar wants to merge 1 commit into
radixark:mainfrom
shangrilar:fix/fsdp-pause-generation-mode

Conversation

@shangrilar

Copy link
Copy Markdown

Summary

  • forward pause_generation_mode to every rollout engine in the FSDP weight-update lifecycle
  • wait for all engines to pause before conditionally flushing their caches
  • skip the cache flush for in_place updates
  • cover abort, retract, and in_place with a parameterized lifecycle regression test

Problem

The FSDP updater called pause_generation() without forwarding the configured mode and always called
flush_cache(). Because the SGLang wrapper defaults to retract, a fully async FSDP run configured with
--pause-generation-mode in_place behaved like retract + flush instead of preserving its in-flight
requests and KV cache.

The updated lifecycle is:

abort/retract: pause(mode) -> flush -> begin weight update
in_place:      pause(mode) -> begin weight update

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_place support.

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

  • focused tests/fast/backends/test_fsdp_update_weight.py: 5 passed in a local Python 3.12 harness that
    loads 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.py
  • git diff --check

The full Miles dependency suite is left to upstream CI.

@shangrilar
shangrilar marked this pull request as ready for review August 13, 2026 00:36
@shangrilar

Copy link
Copy Markdown
Author

Ready for review. The fork workflows are currently awaiting approval. Could a maintainer apply a run-ci label when convenient? Thank you!

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.

1 participant