Convert train loops and RayTrainGroup to async - #849
Merged
Conversation
The field uses mpu.get_data_parallel_src_rank(with_context_parallel=True) and is used with dp_cp_group_gloo, so it belongs in the dp_cp namespace.
…rlier - Add set_parallel_state()/get_parallel_state() global accessors - create_megatron_parallel_state no longer requires model; VPP fields derived from args (defaults to pipeline_model_parallel_size) - verify_megatron_parallel_state checks against model config post-creation - set_parallel_state called in initialize.init() right after _initialize_distributed, so ParallelState is available everywhere - Both Megatron and FSDP actors updated
Address PR review comment.
fzyzcjy
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the training pipeline to utilize Python's asyncio framework, replacing synchronous ray.get calls with async/await semantics. Key changes include converting RayTrainGroup methods to asynchronous functions, implementing a _broadcast helper for concurrent actor communication via asyncio.gather, and updating the training loops in train.py and train_async.py to support asynchronous execution and task-based parallelism. I have no feedback to provide as the implementation correctly transitions the logic to an asynchronous model.
fzyzcjy
changed the title
…an hook Migrate all scattered mpu.get_data_parallel_* calls to use get_parallel_state() accessors. Add pygrep pre-commit hook to prevent future direct mpu.get_data_parallel_* usage. Allowed exceptions: parallel.py (constructs ParallelState), tools/convert_to_hf.py (standalone script without actor init).
…an hook Migrate all scattered mpu.get_data_parallel_* calls to use get_parallel_state() accessors. Add pygrep pre-commit hook to prevent future direct mpu.get_data_parallel_* usage. Allowed exceptions: parallel.py (constructs ParallelState), tools/convert_to_hf.py (standalone script without actor init).
Prefix all data-parallel fields with intra_ to distinguish from independent DP (cross-cell) fields that will be added later.
Prefix all data-parallel fields with intra_ to distinguish from independent DP (cross-cell) fields that will be added later.
Extract repeating (rank, size, group) pattern into GroupInfo dataclass. ParallelState fields become: - intra_dp: GroupInfo - intra_dp_cp: GroupInfo (with gloo_group and src_rank) - cp: GroupInfo - tp: GroupInfo Access pattern changes: e.g. parallel_state.intra_dp_rank -> parallel_state.intra_dp.rank
…ParallelState fields)
…imize diff with /6
… behavioral difference
…ix lint warnings - Use create_mode string param with inline if/else instead of async wrapper functions that could spuriously yield - Add stacklevel=2 to warnings.warn calls (B028)
…f relying on GC timing" This reverts commit 244b8e0.
…h on unawaited coroutine
warnings.filterwarnings("error") alone doesn't crash: the warning fires
inside __del__, so the exception is swallowed by sys.unraisablehook.
Also: remove destroyed-pending-task logic (asyncio.run cancels pending
tasks on shutdown, so the warning is never triggered in practice).
Use shared _FATAL_ASYNC_PATTERN string to avoid regex duplication.
- Add pytest-asyncio install to CI workflow (both j2 template and generated yml) - Configure asyncio_mode = "auto" in pyproject.toml - Register asyncio marker to avoid strict-markers warning - Replace deprecated asyncio.get_event_loop() with asyncio.run() - Restore sys.unraisablehook in test fixture for proper test isolation
Revert the test extras approach — add pytest-asyncio directly to requirements.txt so it is installed by default with pip install -e .
fzyzcjy
force-pushed
the
feat/refactor_dp/7
branch
from
April 1, 2026 10:28
68799b4 to
35eb213
Compare
maocheng23
approved these changes
Apr 3, 2026
Collaborator
Author
|
ci status 
|
fzyzcjy
changed the title
Resolve conflicts: - actor_group.py: keep async _broadcast version (HEAD) - placement_group.py: keep both imports (eager_create_task + compute_ray_pin_head_options) - async_utils.py: keep HEAD + add main's logger and AsyncioGatherUtils - pyproject.toml: keep asyncio_mode=auto, drop removed markers per main - update_weight_from_distributed.py: accept deletion
Resolve conflicts combining async conversion (HEAD) with #839's role/with_ref-at-init pattern (main): - actor_group.py: init() now async + uses self.args/self.role/self.with_ref - placement_group.py: call sites use parameterless init()
Shi-Dong
pushed a commit
that referenced
this pull request
Apr 5, 2026
Closed
GuanxingLu
pushed a commit
to GuanxingLu/miles
that referenced
this pull request
Apr 21, 2026
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
2 participants
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.
No description provided.