Repository navigation
swift_megatron, syncer: resume / rejoin correctness and graceful syncer stop (3/3) - #342
Merged
Merged
Conversation
This was referenced Nov 7, 2026
michaellchung
force-pushed
the
pr/swift-megatron-resume-rejoin
branch
from
November 11, 2026 00:51
01e4573 to
f10af40
Compare
michaellchung
force-pushed
the
pr/swift-megatron-learner-scale
branch
from
November 11, 2026 00:51
6e2603d to
0a2d562
Compare
michaellchung
changed the base branch from
pr/swift-megatron-learner-scale
to
main
November 11, 2026 00:51
…atch-up A restarted island took the syncer's catch-up as a fresh initial cut and overwrote every fragment's FP32 master and bf16 copy with the global, so an island resumed from its own checkpoint lost all local steps since its anchors. Each checkpoint now gets the sync's record (on_save -> save_state): every fragment's anchor version, H-gate and token counters and spent-anchor guard, optionally every rank's anchor shards (YETO_DILOCO_SAVE_ANCHORS). On resume the first broadcast of a fragment is installed by resume_fragment_mode: the same version only sets the anchor and keeps the local values bit for bit; a newer one becomes global_new + (local - anchor_old) when the old anchors were saved, else the global; no record (or a foreign one) is the old cold-start overwrite. A kept fragment keeps its counters, so the H gate is not restarted by the island's restart. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Without a handler a stopped syncer lost every merge since its last --checkpoint-every write, while the islands already held those versions. The first signal stops launching rounds, gives the rounds still gathering --shutdown-timeout-s (default 600) to close and then drops them (they have not touched the state), commits every computed round, writes --checkpoint-path synchronously after any background write and exits 0. The checkpoint's global step is the last outer step on the tape, so --resume continues the tape at the next step. A signal before training starts exits at once; a second signal exits without a checkpoint. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
An island checkpoint written before the sync records existed can only be resumed without the catch-up overwrite through a record written by hand, and its session contract hash is not logged anywhere. The contract is still checked whenever the record has one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… a redial An island whose syncer went away kept refusing that syncer's (or its successor's) pulls forever. The background sender released a spent anchor only when the push raised ConnectionError / OSError, but push_fragment_parts returns False when the link is down (nothing queued) and raises PartialMessageGenerationLost (a RuntimeError, which failed the island) when the connection dies mid-message. A push that did go out on the old connection held its anchor too, though a restarted syncer no longer has the round it was gathered or carried into. Every later pull of that fragment for another round was refused by _pushable, so a syncer restart stalled every island that had pushed shortly before it. - _PushSender treats False and PartialMessageGenerationLost as a drop (anchor released) and records the link (connection generation) each push goes out on; the q4 push goes through push_fragment_parts too. - At each boundary with the sender idle, a spent anchor whose push went out on a link that has since been replaced is released (a record restored from the previous process counts as such), so the next pull is answered from that anchor. - syncer LateCarry: the one-delta-per-anchor guard (anchor_reused) is kept per island instead of per connection, and covers deltas gathered by their own round as well as carried ones. A redialed island answering from an anchor whose delta this syncer already holds is refused instead of merged twice; a restarted syncer holds nothing and merges it once; a round restart still frees the anchors of the pushes it dropped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ms-swift loads the OptimizerParamScheduler state only together with the optimizer (load_mcore_checkpoint: not finetune and not no_load_optim), and a --no_save_optim checkpoint holds no scheduler state at all. An island that saves parameters only and resumes with --no_load_optim true --finetune false gets its iteration and data position back but restarts the schedule at step 0: warmup again and the whole decay shifted by the resumed iteration (resumed lr at i = the uninterrupted lr at i - k), with no error. The callback now steps a fresh scheduler (num_steps == 0) at a resumed iteration i > 0 by i * global_batch_size samples (the scheduler counts samples, one global batch per successful update), which sets every parameter group's lr and weight decay as i uninterrupted updates would. A scheduler restored from the checkpoint and a cold start are left alone; YETO_DILOCO_LR_RESUME=off disables it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…a rewound one Validation showed the re-HELLO catch-up does not overwrite a running island: SyncerClient drops a broadcast of a version it already holds with the same bytes, so unchanged fragments stay bit for bit. Two gaps remained: - A syncer that crashed (or an old one stopped without a final checkpoint) resumes from its last periodic checkpoint, older than what the islands hold. The client dropped the older catch-up, the island kept bases the syncer no longer has, and every push was refused as a future base: the run stalled for good. - A catch-up newer than the island's anchor (a merge it missed while away) replaced the island's local progress even when the island never pushed from that anchor, i.e. the progress is in no global. SyncerClient flags the first broadcast of each fragment after a redial (BcastFragment.catch_up) and, with accept_rewound_catch_up (the swift sync turns it on), takes it even when its version went backwards or kept its number with other bytes; any other older or conflicting broadcast is handled as before. The sync installs a broadcast of an anchored fragment by live_install_mode: a replay of the anchor version is only confirmed (nothing written), a catch-up becomes global_new + (local - anchor) with the fragment's counters kept (a shard whose global did not change keeps its values bit for bit), and a newer catch-up of a fragment pushed from its current anchor is the ordinary install, since that delta may already be merged into it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The in-memory receive path reassembled a chunked BCAST_FRAGMENT and only then took one sequential SHA-256 over the whole payload, under the client lock, on the receiver thread that completed it: on the 120B island that put the catch-up / initial cut (483 GB of f32) at about 0.5 GB/s. Each 4 MiB CHUNK frame was also copied twice more on its way in (bytes() of the receive buffer, then the slice past its header). - The digest of that path is now SHA-256 over the SHA-256 of every 32 MiB block of the payload (bcast_payload_digest). A _BlockDigest hashes each block on a 4-thread pool as soon as all its bytes arrived (hashlib releases the GIL), so the digest is ready right after the last chunk. Only digests of this receive path are compared with each other: a replay of the version held still has the same digest, other bytes under it still conflict; a streamed sink's payload hash is kept apart by kind. - CHUNK frames are read into a bytearray that _reassemble slices through a memoryview (one copy less each); overlap checks use a sorted range list. Local micro-benchmark (4 MiB chunks over socketpairs into the client's own receive threads): 4 x 1.07 GB over 4 streams 2.7 -> 8-9 GB/s, 3 x 3.2 GB 2.5 -> 7.0 GB/s, one stream 2.1 -> 6.9 GB/s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…nk, visible logs A restarted learner 0 gathered and sent INIT_PARAMS of every fragment (~480 GB of f32 for 120B) before it installed anything, while the syncer, already initialized, ignored the INIT and sent its catch-up, which queued in the client meanwhile: rank 0's RSS went from 150 to 616 GB over ~20 min. - With a resume record and an initial-cut wait, learner 0 first waits for the catch-up and installs it fragment by fragment; INIT is sent only if nothing arrives within YETO_DILOCO_INIT_PROBE_S (default 1800 s: a syncer that lost its state and waits for INIT), or if the wait ends first. A cold start sends INIT first as before. - INIT is streamed by chunk: each exported chunk goes to the sender thread and is written into its fragment's INIT_PARAMS frame while the next one is gathered, so rank 0 holds a few chunks instead of up to three fragments. The bytes are unchanged. - The module's INFO lines (resume modes per fragment, installs) were dropped: ms-swift configures only its own logger. Without a configured root the module gets a stderr handler at INFO, and rank 0 logs a WARNING summary once every fragment is resumed (anchor / shift / overwrite counts) and at the end of the initial-cut wait. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…round With --checkpoint-every 16 the 120B run had no syncer checkpoint for its first 16 rounds (~1.5 h): a restart in that window lost the INIT and every merge, and could not resume at all. --checkpoint-on-start writes --checkpoint-path once before the first round (after INIT), on the background writer with --background-checkpoint; it is skipped when --resume loaded that checkpoint (the state just read) or the run is already at its last step. It needs --checkpoint-path. Per-round checkpoint modes keep their own version-zero write. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…y default - swift plugin: YETO_DILOCO_INITIAL_CUT_TIMEOUT defaults to 3600 s (was 0), the value every 120B island already ran with. At 0 the whole initial cut or catch-up (4 bytes per parameter) queued on rank 0 until the first boundary (400-600 GB of RSS), and a resumed island installed its first broadcasts only after one more optimizer step. 0 still turns it off. - yeto launcher, generic SFT syncer: --background-checkpoint (its periodic checkpoints, every 8 rounds by default, no longer stop the scheduler for the write) and --checkpoint-on-start. The syncer's own defaults are unchanged: the RL syncers write per-round checkpoints that persist before their broadcast, which a background write cannot provide. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The swift learner sent nothing between ~10 min steps (and during INIT, the initial-cut install or a checkpoint save), so running the syncer with --island-lease-s expired healthy islands and it stayed off: a dead island was only noticed when its connection closed. Rank 0 now sends HEARTBEAT (local step = the current iteration) from a thread every YETO_DILOCO_HEARTBEAT_S (default 30 s, 0 = off), independent of step ends; the protocol is unchanged (the syncer renews the lease and its step-rate estimate ignores repeats of a step). With YETO_DILOCO_HEARTBEAT_STALL_S > 0 the heartbeats stop once the training loop showed no progress for that long, so a hung island loses its lease. A real-syncer test runs --island-lease-s 2 with an island that stalls 4 s: kept with heartbeats, expired (re-JOIN) without. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ssembly The per-island gathered-anchor guard also refused an asynchronous learner that answers consecutive rounds from the same anchor on one connection, which the max-base-lag carry relies on. Refuse a gathered anchor only when a later connection of the island repeats it; carried anchors stay refused on every connection. Return the plain (type, payload) pair from reassembly when no block digest was computed, as callers outside the live path expect. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…l root The Megatron SFT entry point leaves the root logger with one stdout handler at ERROR. _ensure_visible_logs returned as soon as root had any handler, so every INFO and WARNING of the sync core (resume modes, INIT_PARAMS skipped, anchors released after a reconnect, rebase) and the plugin's LR schedule rebuild warning were dropped on real islands. Decide by whether a handler on the logger's propagation path actually prints INFO; otherwise give the logger its own stderr handler and stop propagating. The callback makes the plugin and sync loggers visible before rebuilding the LR schedule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…the anchor A resumed island shifted every fragment whose global moved past its checkpointed anchor: global_new + (local - anchor_old). When the island had already pushed from that anchor (the record's spent anchor), the delta it pushed may be merged into global_new, and the shift added it a second time. Install such a global as is, the rule the live path already applies to a newer catch-up of a fragment pushed from its current anchor (live_install_mode). A fragment never pushed from its anchor, or a global that went backwards, still shifts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A chunked broadcast of two or more digest blocks reaches _dispatch with its block digest as a fourth argument. SubscriberClient overrides _dispatch with the old three-argument signature, so such a broadcast raised TypeError in its receive thread, which counts as a socket failure: the subscriber redialed and failed again on the next large broadcast. Accept and forward the digest; a test binds every override's signature. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The chunked INIT hands each exported chunk to a sender thread. An export that raised mid-fragment (a collective timeout) left that thread waiting for the rest of the fragment, and the cleanup joined it forever, turning an error into a hung island. End the fragment's queue on the way out and bound the joins on that path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The graceful stop gives gathering rounds --shutdown-timeout-s (600 s by default) to close. A round that holds no delta and none of whose islands is still connected can never close, so a syncer stopped after its learners exited (the usual end of a launcher or a test) waited the whole timeout before checkpointing. Abandon such rounds immediately; rounds with deltas or connected islands keep the timeout. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
michaellchung
force-pushed
the
pr/swift-megatron-resume-rejoin
branch
from
November 11, 2026 00:52
f10af40 to
3b5b378
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 #341 (stacked on
pr/swift-megatron-learner-scale).Summary
Restart and rejoin correctness for ms-swift Megatron islands and the syncer: a resumed island keeps its local progress through the catch-up, a running island survives a syncer restart, and the syncer checkpoints on stop and on start.
Motivation
On a 120B run a restarted island overwrote its local progress with the syncer's catch-up, a syncer restart could stall every island that had pushed shortly before it, a stopped syncer lost every merge since its last periodic checkpoint, and a resumed learner 0 resent ~480 GB of INIT while its catch-up queued in memory.
Changes
Island resume / rejoin
resume_fragment_mode: same version keeps local values bit for bit; a newer one becomesglobal_new + (local - anchor_old)when the old anchors were saved, except after a push from that anchor (installed as is, the live-path rule); no record is the cold-start overwrite.global_new + (local - anchor)) and accepts a rewound one (a syncer resumed from an older checkpoint).YETO_DILOCO_LR_RESUME=offdisables).--island-lease-scan be on (YETO_DILOCO_HEARTBEAT_S, optional stall cut-off).Syncer / protocol
--shutdown-timeout-s, commit computed rounds, write--checkpoint-path, exit 0 (docs/PROTOCOL.md).--checkpoint-on-startwrites the checkpoint before the first round.Testing
cargo testinsyncer/: 325 passed.pytest tests/test_swift_megatron_sync.py tests/test_swift_megatron_plugin.py tests/test_diloco_sync.py tests/test_reconnect.py tests/test_integration.py tests/test_head_mode.py tests/test_run_flags.py tests/test_rl_island_pool_launcher.py tests/test_spot_rollout_syncer.py tests/test_spot_health_fixes.py: 189 passed;tests/test_rl_full_parameter_restart_rejoin.py tests/test_integration.py: 12 passed after the last commit.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
SubscriberClient._dispatchaccepts the block digest (a large chunked BCAST raisedTypeErrorthere and made the spot-rollout subscriber redial in a loop); an INIT export that raises on rank 0 now fails instead of hanging the sender thread.YETO_DILOCO_HEARTBEAT_STALL_Sdefaults to 0, so by default a hung (not dead) island keeps its lease; the block digest is also computed for large non-BCAST chunked frames.🤖 Generated with Claude Code