Repository navigation
swift_megatron: full-parameter DiLoCo sync core and ms-swift Megatron hook (1/3) - #340
Merged
Merged
Conversation
…quorum timeout Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…al to serial Independent reductions (tensors, learners, HeLoCo's three norms) run concurrently; element-wise passes (accumulation, rescale, Nesterov and Delayed Nesterov steps, finiteness checks) are chunked. No reduction is split, so every thread count gives the exact bits of the serial path. --merge-threads / YETO_SYNCER_MERGE_THREADS selects the count (default: cores capped at 32); 1 runs the historical code unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…e full roster Releasing the slot of a refused stale-base push let an exact or streaming fold close one push short and stop the syncer with 'fold incomplete'. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…und) With quorum 1 a round whose only delta from the fast island arrived as a late carry (its anchor spent) waited for the stalled island, and with --pipeline 2 so did the next round. The new flag counts expected carried members as responders for quorum and the all-answered check. Off keeps the historical RL async rule; the flag joins the semantic profile hash only when on. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
build_layout recomputed every bin's numel for each placed tensor, which is quadratic in the tensor count: the unfused 120B MoE shape list (41,643 tensors, 32 fragments) took 107 s per learner process. Track the totals instead; the placement (first smallest bin wins) is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… in parallel Every INIT_PARAMS payload is hashed and decoded on the connection's read loop. On a Kunpeng 920 the vendored sha2 runs its portable code (its aarch64 backend needs the asm feature and another crate), about 0.25 GB/s, and the f32 decode plus the finiteness check are one more single-threaded pass, so a 120B model's 483 GB of f32 INIT took about 40 minutes to admit. - sha256.rs: SHA-256 with the ARMv8 SHA-2 instructions when the CPU has them (runtime detected), else sha2; byte-identical digests (padding boundaries, split updates and known vectors tested against sha2). Used for INIT digests and the downlink digest. - decode_tensor splits large f32 / bf16 payloads over the merge threads (disjoint chunks, same per-value conversion, same error for a non-finite value), and INIT computes its digest while it decodes. Measured on the production syncer host with the same 120B INIT stream: 0.21 -> 0.5 GB/s; the digest (now ~1.6 GB/s, one sequential pass) is the remaining bound. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
append_tensor_bytes sliced every value by the chunk's byte range and appended it on its own, about 0.4 GB/s on a Kunpeng 920. Every f32 broadcast of a 120B fragment (15 GB) goes through it, one 4 MiB chunk at a time on the forwarder, so the syncer's downlink to one island ran at ~0.42 GB/s on a 50 Gb/s link: the initial cut alone took ~20 minutes. Only the values cut by the range ends now go byte by byte; the bytes are unchanged (tested against the value-by-value encoding, NaN, -0.0, subnormal and every cut position). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…cheduler A --checkpoint-every checkpoint (other than per-round ones) was written synchronously in publish_round: the scheduler merged and broadcast nothing until the whole state was on disk. For the 120B state (966 GB of f32 params and momentum) that is a 13.7 minute stall per checkpoint on an NVMe that sustains ~1.2 GB/s of buffered writes (measured on the production syncer host). With --background-checkpoint the cut is taken as the existing copy-free CheckpointSnapshot (later commits copy only a fragment the writer has not reached) and saved on its own thread while rounds go on. At most one such write runs: the next periodic checkpoint, and every synchronous save (terminal, budget cutoff, shutdown), waits for it, so two writers never share <path>.tmp and the newest cut lands last. A failed write surfaces from that wait; the previous file stays intact. The flag needs a checkpoint path and --checkpoint-every >= 2 outside policy sweeps (per-round checkpoints still persist before broadcast). Off by default; the files are byte-identical to the synchronous path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…oad copy A reassembled INIT_PARAMS was copied out of its reassembly buffer, then hashed and decoded on the connection's read loop, so that stream stopped reading for each fragment's ~10 s digest (one sequential pass) and the learner's round-robin chunks backed up behind it. A 120B f32 INIT (483 GB) arrived at ~0.5 GB/s on a 50 Gb/s link. The reassembled frame now keeps its buffer (no copy), and an INIT_PARAMS is validated, hashed and decoded on the blocking pool while the read loop goes on; up to four per connection are in flight (each holds one fragment's payload and values) before the loop waits for the oldest. A failed decode is reported to the learner and ends the connection as before; the coordinator applies INITs in any order, from the current generation only. Direct (unchunked) INIT frames keep the inline path. Measured on the production syncer host with a synthetic 120B learner: the same INIT stream now takes 322 s (1.5 GB/s, bound by the Python client's send loop) instead of 16.9 min. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_send_large_parts filled each 4 MiB CHUNK envelope by allocating a zeroed bytearray and assigning slices into it. On a Kunpeng 920 that loop sends at most ~2 GB/s and measured 1.5 GB/s for a 120B INIT (483 GB, 322 s), the bound once the syncer keeps up. The envelope is now b"".join of its header and the part slices: one copy, ~5.9 GB/s in the same micro-benchmark. Bytes and validation are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
When --total-steps is a multiple of --checkpoint-every, the periodic checkpoint of the last step and the final checkpoint wrote the same cut back to back. For the 120B state that is two ~14 minute writes of 966 GB at the end of a run (measured on the production syncer host: 10:44 -> 10:58 -> 11:12). The periodic write now skips the last step; the final checkpoint, which is always written when a checkpoint path is set, persists it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A late push waiting for its fragment's next round held a decoded f32 outer gradient (about 4P bytes per slow island at 120B). Queue the original compressed payload instead and decode it when the push is attached to the next round. The decode is deterministic, so merged results are bit-identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Drop the wire copy of a push once it joins a round. The compressed bytes exist so a late push can wait in the late-carry queue without its decoded gradient; a push accepted by its own round, or attached to a later one, kept that copy until the round was committed (up to 2P more bytes per bf16 island, on top of the gradient the streaming fold had just released). - Release the slot of a stale-base push under a streaming fold as well. Only the exact fold fixes its order and weights for the full roster; a streaming fold normalises over what it folded, and async carry (where stale bases happen) needs one, so the earlier guard left those rounds waiting for the quorum timeout again. - Decode PUSH payloads on the calling thread, as before the parallel INIT decode: a large push no longer spawns a thread pool per connection that competes with the merge threads. INIT_PARAMS asks for threads explicitly. - Assert the SHA-2 feature on the hardware hasher path in debug builds. - Document the worst-case memory of --background-checkpoint. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Drop a leftover `+ 0` and give the speed check headroom on a slow CI runner; the rescanning greedy it guards against took minutes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MegatronDiLoCoSync speaks the torch learner's DiLoCo protocol (layout, INIT, raw-anchor PUSH deltas, BCAST apply before pulls, finalization, HMAC HELLO) to the unchanged syncer from a TP/EP/PP Megatron island: - global trainable set named by global Megatron names (global layer and expert indices), frozen parameters and buffers excluded; - chunked reads of FP32 optimizer masters (DistributedOptimizer shards, HybridDeviceOptimizer host masters) gathered over DP and TP and sent point-to-point to rank 0; per-rank anchors keep deltas local; - writes install the merged base into the FP32 master and the bf16 parameter, keeping the Adam moments; - collectives on host buffers over gloo mirrors by default, so the sync allocates nothing on the accelerator; - rank 0's per-boundary plan is broadcast so every rank runs the same collectives and rank-0 errors are raised everywhere. write_masters also updates a distinct HybridDeviceOptimizer host copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ms-swift's `megatron sft` has no DiLoCo boundary of its own. The hook loads through its official `--external_plugins` / `--callbacks` mechanism and registers `yeto_diloco`, which ms-swift instantiates after the model, optimizer and resume checkpoint are loaded. Every rank then calls the sync core once per optimizer step (skipped steps included, with zero target tokens) from `on_step_end`, before ms-swift's own save decision, so a checkpoint written at that step already holds the merged base. Configuration comes from YETO_DILOCO_* environment variables because the ms-swift argument parser refuses unknown flags; YETO_DILOCO_AUTO_ATTACH=1 appends the callback without touching the command line. The plugin adds its own repository root to sys.path, since ms-swift imports it by file path on every node, and composes with other plugins that patch `train_step` at class level. The hook passes update_successful, the final-cut payload mode and the collective transport through to the sync core. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…tron sync Move client connection, wire dtype names, pull readiness, push delta encoding and boundary metrics from learner.py / sync_diloco_boundary into public yeto.diloco_sync helpers (torch learner behavior unchanged) and drive MegatronDiLoCoSync's counters, pulls and pushes through them and DiLoCoSyncState. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…or other islands A synchronous boundary stalls every island on the slowest one and on the syncer, which defeats DiLoCo across heterogeneous or WAN-linked islands. - PUSH bytes go to a background sender thread on rank 0; while a push is still streaming, pulls stay pending instead of blocking the step. - One fragment is gathered per boundary; a pull is not answered from an anchor already spent on another round (late-carry anchor_reused), and a dropped push releases its anchor. - wire_dtype accepts the RL full-parameter push names (int4, int8, int4-log, push-bf16) with tensor-block quantization. - Test: quorum 1 / max_base_lag 1 / int4 with one island stalled, against the real syncer with --carry-answers-round; the other island keeps stepping and merging without blocking boundaries, the stalled island rejoins, and both install the syncer's final state. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…re 0.18 Megatron-Core 0.18 sets local_expert_indices on MoELayer only, not on its experts container, so every expert-parallel rank named its experts with the local indices and the island synced only EP rank 0's experts (the others aliased them). Fall back to the parent MoE layer and refuse an expert parameter that still has no global index. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ncers, docs - Check broadcast payload sizes while rank 0 builds the boundary plan, and the final cut's versions and sizes before the finalization outcome is broadcast. Raised later, on rank 0 alone, they left the other ranks waiting in the chunk broadcast for good. - End training at the step that installed the final cut (ms-swift saves at train_iters), instead of training past the global result and saving local steps over it. - Refuse a comma-separated (sharded) syncer address on every rank before connecting: chunked INIT and push streaming talk to one syncer. - Report push_bytes per boundary as the torch learner does, and leave out push_delta_norm, which this sync never computes. - Docs: list FINAL_CUT_PAYLOAD and TRANSPORT and every wire dtype, drop a reference to a file outside the repository, and correct where the callback is added and where the initial base comes from. Drop a debug print from a test and a comment that tripped the secret scanner. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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 #335 (stacked on
pr/syncer-capacity; the decoupled test runs the real syncer with--carry-answers-round).First of three PRs adding a full-parameter DiLoCo learner for ms-swift Megatron (
megatron sft) islands: this one is the sync core and the ms-swift hook; #341 adds what a multi-node 120B island needs, #342 restart / rejoin correctness.Summary
yeto/swift_megatron/lets an ms-swift Megatron SFT island (TP/EP/PP, distributed or host-offloaded optimizer) take part in yeto's DiLoCo rounds against the unchanged syncer, through ms-swift's own plugin mechanism.Motivation
The torch learner covers single-process / DDP / FSDP models. A 120B MoE needs Megatron parallelism, and ms-swift is how that recipe is trained; patching ms-swift is not an option for production islands.
Changes
sync.py(MegatronDiLoCoSync): speaks the torch learner's DiLoCo protocol (layout, INIT, raw-anchor PUSH deltas, BCAST before pulls, finalization, HMAC HELLO). The trainable set is named by global Megatron names (global layer and expert indices); FP32 masters are read in bounded chunks (DistributedOptimizer shards, HybridDeviceOptimizer host masters), gathered over DP/TP to rank 0; writes install the merged base into the FP32 master and the bf16 parameter, keeping Adam moments. Collectives run on host buffers over gloo mirrors by default. Rank 0's per-boundary plan is broadcast so every rank runs the same collectives and rank-0 errors are raised everywhere.swift_plugin.py:--external_plugins .../swift_plugin.py --callbacks yeto_diloco(orYETO_DILOCO_AUTO_ATTACH=1); configured byYETO_DILOCO_*environment variables; one sync call per optimizer step fromon_step_end, before ms-swift's save decision.yeto/diloco_sync.py: session helpers moved out oflearner.pyand shared by both learners (torch learner behavior unchanged).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_learner_units.py tests/test_hybrid_offload_masters.py tests/test_integration.py: 99 passed. The sync tests run TP2xDP2 and two TP2xPP2 islands as gloo process groups on CPU, against a scripted client and against the real syncer (including a connection drop and a stalled island).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
push_bytesreported per boundary; env-var docs completed.False(link down) is treated as a drop only from swift_megatron, syncer: resume / rejoin correctness and graceful syncer stop (3/3) #342 on.🤖 Generated with Claude Code