Repository navigation
swift_megatron: multi-node island support (cross-node gloo, per-rank scatter, INIT overlap, step log) (2/3) - #341
Merged
Conversation
This was referenced Nov 7, 2026
michaellchung
changed the base branch from
pr/swift-megatron-learner-core
to
main
November 11, 2026 00:50
…t staging A pipeline-parallel island can span two Ascend nodes without RoCE: groups whose ranks cross nodes are created with gloo, and collectives / P2P on them (and on WORLD when the job spans nodes) run on host copies of the device tensors. Intra-node TP/EP/DP groups keep HCCL. Enabled by YETO_CROSS_NODE_GLOO=1 when loaded as an ms-swift external plugin. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_apply_chunk broadcast every chunk of a received fragment to the whole world group, so each rank received the full global tensors and kept only its TP slice. On the 24-rank 120B island the gloo world broadcast of one 4-fragment payload (about 120 GB) ran for more than 20 minutes and stalled the step; the cost is world copies of the fragment. Gather every rank's (name, TP partition) list once at construction and have rank 0 send each peer only its slices of the chunk, point to point. About DP copies of the chunk cross the wire instead of world copies. Values, anchors and write-back are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…al cut before training INIT gathered a whole fragment on every rank and only then streamed it, so the island's gather and the link took turns. Rank 0 now streams fragment k on a background thread while the island gathers fragment k + 1 (at most two fragments' wire parts on rank 0). The syncer's initial cut (every fragment's version 0, or the current versions on a reconnect) was drained only at the first boundary, one optimizer step after INIT: rank 0's client queued all of it meanwhile, 4 N bytes on an f32-bulk session (~480 GB for 120B; ~400 GB of rank 0 RSS was observed). With initial_cut_timeout > 0 (YETO_DILOCO_INITIAL_CUT_TIMEOUT) the island installs each fragment as it lands, collectively, before training starts, and frees it; after the timeout the rest is applied at the boundaries as before. Off by default. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The sync core now accumulates the wall time spent gathering the fp32 master for a push (export_s) and writing a received broadcast back (apply_s). With YETO_DILOCO_STEP_LOG set, sync rank 0 appends one JSON line per step with the time inside the sync boundary, those two splits, and the push/apply counters, so per-step overhead can be read against the trainer's own step timings. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
When the syncer reopens a round after a quorum timeout it sends a pull with a higher round_attempt. An island that still held the old attempt's pull answered it, and the syncer rejected the push as out of round. take_ready_pulls now keeps only the newest attempt per (fragment, step). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
After a redial the syncer may resend the pull of an attempt the learner already holds; both copies were answered, pushing the fragment twice for one round (the second is refused as a duplicate, after the bytes went out). Keep one copy per attempt. Also document the cross-node gloo module's ranks-per-node override. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
michaellchung
force-pushed
the
pr/swift-megatron-learner-scale
branch
from
November 11, 2026 00:51
6e2603d to
0a2d562
Compare
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.
Depends on #340 (stacked on
pr/swift-megatron-learner-core).Summary
What a multi-node 120B ms-swift Megatron island needed on top of the sync core: cross-node process groups without RoCE, per-rank scatter of applied fragments, INIT overlapped with the send, the initial cut installed before training, a per-step timing log, and pull de-duplication in the shared learner helpers.
Changes
cross_node_gloo.py(YETO_CROSS_NODE_GLOO=1, loaded as an extra ms-swift plugin): groups whose ranks span nodes are created with gloo and their collectives / P2P run on host copies; intra-node TP/EP/DP groups keep HCCL. For PP islands across nodes without RoCE.YETO_DILOCO_INITIAL_CUT_TIMEOUTthe syncer's initial cut is installed fragment by fragment before training instead of queueing on rank 0 until the first boundary (~400 GB of rank-0 RSS observed at 120B).YETO_DILOCO_STEP_LOG: one JSON line per step from sync rank 0 (time inside the boundary, export / apply split, counters).take_ready_pullskeeps only the newest attempt of a round, and one copy of a resent pull of the same attempt.Testing
pytest tests/test_swift_megatron_sync.py tests/test_swift_megatron_plugin.py tests/test_swift_megatron_expert_index.py tests/test_diloco_sync.py tests/test_cross_node_gloo.py tests/test_learner_units.py tests/test_protocol_finalization.py: 105 passed.pr/swift-megatron-host-memory, which contains this branch): 8997 passed, 181 skipped, 2 failed while three suites ran in parallel;test_dispense_data_mix_samplers::test_engine_samplers_claim_from_the_dispenser...passes on rerun (load), andtest_rl_full_parameter_restart_rejoin::test_heartbeats_keep_the_lease_while_the_pool_waits_for_initexposed the graceful-stop wait fixed in the last commit of the resume PR (passes after it). For reference,mainunder the same load: 8920 passed, 1 failed (test_integration::test_learner_budget_restart_freezes_exactly_and_marks_complete_checkpoint, timing).scripts/precommit_check.py: clean.Run locally on macOS arm64, Python 3.11 with the CI install (
torch==2.8.0CPU,-r pyproject.toml --extra test),PYTHONPATH=. python -m pytest ... --timeout=600. Optional stacks that are not installed (Miles, harbor, lark-oapi, matplotlib) skip as on CI.Notes
🤖 Generated with Claude Code