From 5b2846304cb9e250f55170c4f06ccfd3d495b459 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 06:32:51 +0800
Subject: [PATCH 01/17] swift-megatron: keep a resumed island's progress
 through the rejoin catch-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>
---
 tests/test_diloco_sync.py           |  35 +++++
 tests/test_swift_megatron_plugin.py |  44 +++++-
 tests/test_swift_megatron_sync.py   | 143 ++++++++++++++++++
 yeto/diloco_sync.py                 |  51 +++++++
 yeto/swift_megatron/swift_plugin.py |  64 +++++++-
 yeto/swift_megatron/sync.py         | 223 ++++++++++++++++++++++++++--
 6 files changed, 548 insertions(+), 12 deletions(-)

diff --git a/tests/test_diloco_sync.py b/tests/test_diloco_sync.py
index bc086886..987cb1f8 100644
--- a/tests/test_diloco_sync.py
+++ b/tests/test_diloco_sync.py
@@ -364,3 +364,38 @@ def test_a_resent_pull_of_the_same_attempt_is_answered_once():
     state.pending_pulls = [pull, PullRequest(fragment_id=0, global_step=3, round_attempt=1)]
     ready = take_ready_pulls(state, steps_total=5, units_total=50, anchored=lambda fid: True)
     assert [p for p, _, _ in ready] == [pull] and state.pending_pulls == []
+
+
+def test_resume_record_round_trips_and_picks_the_install_mode():
+    import json
+
+    from yeto.diloco_sync import (
+        RESUME_ANCHOR,
+        RESUME_OVERWRITE,
+        RESUME_SHIFT,
+        resume_fragment_mode,
+        shifted_resume_value,
+    )
+
+    state = DiLoCoSyncState.create(3, track_anchors=False)
+    state.fragment_versions = [4, 7, 0]
+    state.steps_at_reset = [0, 12, 0]
+    state.units_at_reset = [0, 500, 0]
+    state.global_step = 7
+    record = json.loads(json.dumps(state.resume_record(
+        anchored=[True, True, False], units_total=900, spent=[None, (7, 9), None])))
+    assert record["spent"] == [None, [7, 9], None] and record["units_total"] == 900
+    # Same version: the broadcast is the checkpoint's own anchor.
+    assert resume_fragment_mode(record, 0, 4, have_anchor=False) == RESUME_ANCHOR
+    # Newer: move the local progress only when the old anchor is at hand.
+    assert resume_fragment_mode(record, 1, 9, have_anchor=True) == RESUME_SHIFT
+    assert resume_fragment_mode(record, 1, 9, have_anchor=False) == RESUME_OVERWRITE
+    # Never anchored before the checkpoint, or no record at all: a cold start.
+    assert resume_fragment_mode(record, 2, 0, have_anchor=True) == RESUME_OVERWRITE
+    assert resume_fragment_mode(None, 0, 4, have_anchor=True) == RESUME_OVERWRITE
+    fresh = DiLoCoSyncState.create(3, track_anchors=False)
+    fresh.steps_at_reset = [15, 15, 15]
+    fresh.restore_fragment(record, 1)
+    assert fresh.steps_at_reset == [15, 12, 15] and fresh.units_at_reset == [0, 500, 0]
+    new, local, old = torch.tensor([1.0, 2.0]), torch.tensor([0.5, 0.75]), torch.tensor([0.25, 1.0])
+    assert torch.equal(shifted_resume_value(new, local, old), torch.tensor([1.25, 1.75]))
diff --git a/tests/test_swift_megatron_plugin.py b/tests/test_swift_megatron_plugin.py
index 14bfe373..e37aa1de 100644
--- a/tests/test_swift_megatron_plugin.py
+++ b/tests/test_swift_megatron_plugin.py
@@ -19,8 +19,10 @@
 class FakeSync:
     instances = []
 
-    def __init__(self, model_chunks, optimizer, cfg):
+    def __init__(self, model_chunks, optimizer, cfg, *, resume_dir=None, resume_iteration=None):
         self.model_chunks, self.optimizer, self.cfg = model_chunks, optimizer, cfg
+        self.resume_dir, self.resume_iteration = resume_dir, resume_iteration
+        self.saves = []
         self.calls = []
         self.closed = 0
         FakeSync.instances.append(self)
@@ -28,6 +30,10 @@ def __init__(self, model_chunks, optimizer, cfg):
     def after_optimizer_step(self, iteration, target_tokens, *, update_successful=True):
         self.calls.append((iteration, target_tokens))
 
+    def save_state(self, output_dir, *, iteration, save_anchors=False):
+        self.saves.append((output_dir, iteration, save_anchors))
+        return True
+
     def close(self):
         self.closed += 1
 
@@ -129,6 +135,42 @@ def test_resume_continues_iteration(make):
     assert FakeSync.instances[0].calls == [(101, 2), (102, 2)]
 
 
+def test_cold_start_has_no_resume_dir(make):
+    make([([1], True)])
+    assert FakeSync.instances[0].resume_dir is None and FakeSync.instances[0].resume_iteration == 0
+
+
+def test_resume_reads_the_record_of_the_loaded_checkpoint(make, tmp_path):
+    from yeto.swift_megatron.sync import RESUME_STATE_FILE
+
+    ckpt = tmp_path / "checkpoint-15"
+    ckpt.mkdir()
+    trainer = FakeTrainer([([1], True)], start_iteration=15)
+    trainer.args.mcore_model = str(ckpt)
+    cfg = sp.DiLoCoHookConfig.from_env()
+    assert sp.resume_dir_for(cfg, trainer.args) is None  # no record there: cold start
+    (ckpt / RESUME_STATE_FILE).write_text("{}")
+    assert sp.resume_dir_for(cfg, trainer.args) == str(ckpt)
+    explicit = sp.DiLoCoHookConfig(**{**cfg.__dict__, "resume_dir": "/elsewhere"})
+    assert sp.resume_dir_for(explicit, trainer.args) == "/elsewhere"
+    _, cb = make([([1], True)], start_iteration=15)
+    assert FakeSync.instances[-1].resume_iteration == 15
+
+
+def test_on_save_records_the_sync_state_next_to_the_checkpoint(make, monkeypatch):
+    monkeypatch.setenv("YETO_DILOCO_SAVE_ANCHORS", "1")
+    trainer, cb = make([([1], True), ([1], True)])
+    trainer.args.output_dir = "/runs/x"
+    trainer.state.iteration = 2  # the step whose checkpoint was just written
+    cb.on_save()
+    cb.on_save(output_dir="/runs/x/checkpoint-2-explicit")
+    assert FakeSync.instances[0].saves == [("/runs/x/checkpoint-2", 2, True),
+                                           ("/runs/x/checkpoint-2-explicit", 2, True)]
+    cb.on_train_end()
+    cb.on_save()  # closed: nothing more to record
+    assert len(FakeSync.instances[0].saves) == 2
+
+
 def test_close_is_idempotent(make):
     trainer, cb = make([([1], True)])
     trainer.train()
diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py
index 5b136d19..a5c3fa4d 100644
--- a/tests/test_swift_megatron_sync.py
+++ b/tests/test_swift_megatron_sync.py
@@ -498,6 +498,149 @@ def test_initial_cut_is_installed_fragment_by_fragment_before_training(tmp_path)
     _join(procs, 200, tmp_path)
 
 
+def _resume_worker(rank, world, port, out_dir):
+    _guard(out_dir, f"resume-{rank}", lambda: _resume_body(rank, world, port, out_dir))
+
+
+def _resume_body(rank, world, port, out_dir):
+    """An island checkpoints after three steps and restarts from it; the
+    syncer's rejoin catch-up must not undo the island's local progress."""
+
+    from yeto.diloco_sync import shifted_resume_value
+    from yeto.protocol import BcastFragment
+    from yeto.tensor_io import pack_tensor
+
+    dist = _init_dist(rank, world, port)
+    tp, dp = rank % 2, rank // 2
+    tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])]
+    dp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])]
+    singles = [dist.new_group([r]) for r in range(world)]
+    ctx = sm.ParallelContext(rank, world, tp_groups[dp], tp, 2, singles[rank], 0, 1)
+    chunk = _chunk(tp=tp, pp=0, tp_size=2, pp_size=1)
+    trainable = [p for p in chunk.parameters() if p.requires_grad]
+    opt = _DistOpt(trainable, dp_rank=dp, dp_size=2, dp_group=dp_groups[tp])
+    full_shapes = _global_shapes(2, 1)
+    old_g = {n: (_full(n, s) + 0.125).to(torch.bfloat16).float() for n, s in full_shapes.items()}
+    new_g = {n: (_full(n, s) - 0.25).to(torch.bfloat16).float() for n, s in full_shapes.items()}
+    cut = {}  # fid -> (version, values) of the next sync's catch-up
+
+    def factory(addrs, learner_id, layout, wire, streams, **kwargs):
+        bcasts = [BcastFragment(fid, v, pack_tensor(_flat(layout, fid, values), 2))
+                  for fid, (v, values) in sorted(cut.items())]
+        return _ScriptedClient(layout, {("bcast", 1): bcasts})
+
+    def make(resume_dir=None, resume_iteration=None):
+        return sm.MegatronDiLoCoSync([chunk], opt, syncer="tcp://127.0.0.1:1234", learner_id=1, num_learners=2,
+                                     fragments=3, sync_interval_steps=100, chunk_numel=40, parallel_context=ctx,
+                                     client_factory=factory, initial_cut_timeout=60, resume_dir=resume_dir,
+                                     resume_iteration=resume_iteration)
+
+    def shard(full, p):
+        if getattr(p, "tensor_model_parallel", False):
+            size = p._full_shape[p.partition_dim] // 2
+            full = full.narrow(p.partition_dim, tp * size, size)
+        return full
+
+    ckpt, bare = Path(out_dir, "checkpoint-3"), Path(out_dir, "checkpoint-3-bare")
+    # The first run: the initial cut (version 4), three local steps, a checkpoint.
+    cut.update({fid: (4, old_g) for fid in range(3)})
+    first = make()
+    for it in (1, 2, 3):
+        opt.step({id(p): _local_update(p, 1, it, tp, 2) for p in trainable})
+        first.after_optimizer_step(it, 10)
+    assert first.save_state(str(ckpt), iteration=3, save_anchors=True)
+    assert first.save_state(str(bare), iteration=3, save_anchors=False)
+    fid_of = {name: fid for fid, frag in enumerate(first.layout.fragments) for name, _ in frag.tensors}
+    first.close()
+    if rank == 0:
+        record = json.loads((ckpt / sm.RESUME_STATE_FILE).read_text())
+        assert record["versions"] == [4, 4, 4] and record["anchors"] and record["iteration"] == 3
+        assert record["steps_at_reset"] == [0, 0, 0] and record["units_total"] == 30
+        assert not json.loads((bare / sm.RESUME_STATE_FILE).read_text())["anchors"]
+    saved_master = {id(p): opt.master(p).clone() for p in trainable if id(p) in opt.owned}
+    saved_param = {id(p): p.detach().clone() for p in trainable}
+
+    def restart():  # the trainer loads the checkpoint again
+        for p in trainable:
+            if id(p) in opt.owned:
+                opt.master(p).copy_(saved_master[id(p)])
+            p.data.copy_(saved_param[id(p)])
+
+    def check(sync, expect):  # expect(p) -> the full local FP32 value or None (= the checkpoint, bit for bit)
+        assert sync.anchored == [True] * 3
+        for p in trainable:
+            want = expect(p)
+            if want is None:
+                if id(p) in opt.owned:
+                    assert torch.equal(opt.master(p), saved_master[id(p)]), p._full_name
+                assert torch.equal(p.detach(), saved_param[id(p)]), p._full_name
+                continue
+            if id(p) in opt.owned:
+                s, e, _ = opt.owned[id(p)]
+                assert torch.equal(opt.master(p), want.reshape(-1)[s:e]), p._full_name
+                assert torch.equal(p.detach().reshape(-1)[s:e], want.reshape(-1)[s:e].to(torch.bfloat16))
+
+    def check_anchors(sync, values):
+        for p in trainable:
+            assert torch.equal(sync.anchors[id(p)].float(), shard(values[p._full_name], p)), p._full_name
+
+    # (a) The syncer is still at the checkpoint's anchor version: anchor only.
+    restart()
+    cut.update({fid: (4, old_g) for fid in range(3)})
+    same = make(str(ckpt), 3)
+    check(same, lambda p: None)
+    check_anchors(same, old_g)
+    assert same.resumed == {0: "anchor", 1: "anchor", 2: "anchor"}
+    same.after_optimizer_step(4, 10)
+    assert same.state.steps_at_reset == [0, 0, 0], "the H gate restarted with the island"
+    assert same.units_total == 40
+    same.close()
+
+    # (b) The syncer moved on: global_new + (local - anchor_old), exactly.
+    restart()
+    cut.update({fid: (9, new_g) for fid in range(3)})
+    moved = make(str(ckpt), 3)
+    for p in trainable:
+        if id(p) in opt.owned:
+            s, e, _ = opt.owned[id(p)]
+            want = shifted_resume_value(shard(new_g[p._full_name], p).reshape(-1)[s:e], saved_master[id(p)],
+                                        shard(old_g[p._full_name], p).reshape(-1)[s:e])
+            assert torch.equal(opt.master(p), want), p._full_name
+            assert torch.equal(p.detach().reshape(-1)[s:e], want.to(torch.bfloat16)), p._full_name
+            assert not torch.equal(want, shard(new_g[p._full_name], p).reshape(-1)[s:e])
+    check_anchors(moved, new_g)
+    assert set(moved.resumed.values()) == {"shift"} and not moved._resume_anchors
+    moved.after_optimizer_step(4, 10)
+    assert moved.state.steps_at_reset == [0, 0, 0] and moved.state.fragment_versions == [9, 9, 9]
+    moved.close()
+
+    # (c) Without saved anchors a moved fragment takes the global; an unmoved one is kept.
+    restart()
+    cut.update({0: (4, old_g), 1: (9, new_g), 2: (9, new_g)})
+    mixed = make(str(bare), 3)
+    check(mixed, lambda p: None if fid_of[p._full_name] == 0 else shard(new_g[p._full_name], p))
+    assert mixed.resumed == {0: "anchor", 1: "overwrite", 2: "overwrite"}
+    mixed.after_optimizer_step(4, 10)
+    assert mixed.state.steps_at_reset == [0, 3, 3]
+    mixed.close()
+
+    # (d) A cold start (no record, or a record of another iteration) takes the global.
+    for resume in ((None, None), (str(ckpt), 7)):
+        restart()
+        cut.update({fid: (9, new_g) for fid in range(3)})
+        cold = make(*resume)
+        check(cold, lambda p: shard(new_g[p._full_name], p))
+        assert cold.resumed == {} and cold.resume_record is None
+        cold.close()
+    dist.destroy_process_group()
+
+
+@pytest.mark.timeout(240)
+def test_resumed_island_keeps_its_checkpoint_through_the_rejoin_catch_up(tmp_path):
+    procs = _run(_resume_worker, 4, free_port(), str(tmp_path))
+    _join(procs, 200, tmp_path)
+
+
 # ----------------------------------------------------------------------------- two islands, real syncer (TP2 x PP2)
 def build_syncer() -> Path:
     binary = ROOT / "syncer/target/debug/yeto-syncer"
diff --git a/yeto/diloco_sync.py b/yeto/diloco_sync.py
index 7afbb042..3388e76a 100644
--- a/yeto/diloco_sync.py
+++ b/yeto/diloco_sync.py
@@ -148,6 +148,57 @@ def create(
             anchors=[None] * num_fragments if track_anchors else None,
         )
 
+    def resume_record(self, *, anchored: list[bool], units_total: int,
+                      spent: list[tuple[int, int] | None] | None = None) -> dict:
+        """The per-fragment part of an island checkpoint (JSON-safe): which
+        version each fragment's anchor is, and the H-gate / token counters
+        measured from it.  The anchor values themselves are not included."""
+        n = len(self.fragment_versions)
+        spent = spent if spent is not None else [None] * n
+        return {
+            "versions": [int(v) for v in self.fragment_versions],
+            "anchored": [bool(a) for a in anchored],
+            "steps_at_reset": [int(s) for s in self.steps_at_reset],
+            "units_at_reset": [int(u) for u in self.units_at_reset],
+            "spent": [None if s is None else [int(s[0]), int(s[1])] for s in spent],
+            "units_total": int(units_total),
+            "global_step": int(self.global_step),
+        }
+
+    def restore_fragment(self, record: dict, fid: int) -> None:
+        """Take fragment ``fid``'s counters from a ``resume_record`` (a resumed
+        island that kept its local progress keeps its H-gate clock too)."""
+        self.steps_at_reset[fid] = int(record["steps_at_reset"][fid])
+        self.units_at_reset[fid] = int(record["units_at_reset"][fid])
+
+
+# What a resumed island does with the first broadcast of a fragment.
+RESUME_OVERWRITE = "overwrite"  # install the global (no usable record: a cold start)
+RESUME_ANCHOR = "anchor"        # same version as the checkpoint's anchor: anchor only
+RESUME_SHIFT = "shift"          # newer global: global_new + (local - anchor_old)
+
+
+def resume_fragment_mode(record: dict | None, fid: int, version: int, *, have_anchor: bool) -> str:
+    """How a resumed island installs the syncer's first broadcast of ``fid``.
+
+    The checkpoint holds ``local``, trained from the anchor of
+    ``record["versions"][fid]``.  The same version again means the broadcast
+    *is* that anchor, so the local parameters stay as they are.  Any other
+    version moves the island's local progress onto the new global, which
+    needs the old anchor value (``have_anchor``); without it the fragment
+    falls back to the cold-start overwrite."""
+    if record is None or not record["anchored"][fid]:
+        return RESUME_OVERWRITE
+    if int(version) == int(record["versions"][fid]):
+        return RESUME_ANCHOR
+    return RESUME_SHIFT if have_anchor else RESUME_OVERWRITE
+
+
+def shifted_resume_value(global_new: torch.Tensor, local: torch.Tensor, anchor_old: torch.Tensor) -> torch.Tensor:
+    """``global_new + (local - anchor_old)`` in FP32: the island's progress since
+    its checkpointed anchor, moved onto the newer global."""
+    return global_new.to(torch.float32) + (local.to(torch.float32) - anchor_old.to(torch.float32))
+
 
 def sync_diloco_boundary(
     client: SyncerClient | None,
diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py
index 0ac2d695..0f4fb9d5 100644
--- a/yeto/swift_megatron/swift_plugin.py
+++ b/yeto/swift_megatron/swift_plugin.py
@@ -57,6 +57,21 @@
   YETO_DILOCO_INITIAL_CUT_TIMEOUT seconds to install the syncer's initial cut
                                   fragment by fragment before training (0 =
                                   apply it at the boundaries; default 0)
+  YETO_DILOCO_RESUME_DIR          checkpoint directory whose sync record to
+                                  resume from (default: ``args.mcore_model``
+                                  when it holds one)
+  YETO_DILOCO_SAVE_ANCHORS        1 = also save every rank's anchor shards
+                                  with each checkpoint (~2 bytes per local
+                                  parameter per rank for bf16); default 0
+
+Resume: every checkpoint also gets the sync's record (``on_save`` ->
+``MegatronDiLoCoSync.save_state``): each fragment's anchor version and
+H-gate counters.  On a restart from that checkpoint the syncer's catch-up
+no longer overwrites the island: a fragment still at its checkpointed
+version keeps its local values (anchor only), a newer one becomes
+``global_new + (local - anchor_old)`` if the anchors were saved and the
+global otherwise; without a record (a cold start) the island takes the
+syncer's cut as before.
 """
 
 from __future__ import annotations
@@ -109,6 +124,8 @@ class DiLoCoHookConfig:
     final_cut_payload: str = "full"
     transport: str = "cpu"
     initial_cut_timeout: float = 0.0
+    resume_dir: Optional[str] = None
+    save_anchors: bool = False
 
     @classmethod
     def from_env(cls, env=None) -> "DiLoCoHookConfig":
@@ -147,6 +164,8 @@ def get_float(name: str, default: float) -> float:
             final_cut_payload=get("FINAL_CUT_PAYLOAD", default="full"),  # type: ignore[arg-type]
             transport=get("TRANSPORT", default="cpu"),  # type: ignore[arg-type]
             initial_cut_timeout=get_float("INITIAL_CUT_TIMEOUT", 0.0),
+            resume_dir=get("RESUME_DIR"),
+            save_anchors=(get("SAVE_ANCHORS", default="0") or "").strip().lower() in {"1", "true", "yes", "on"},
         )
         if cfg.initial_cut_timeout < 0:
             raise ValueError(f"{ENV_PREFIX}INITIAL_CUT_TIMEOUT must be >= 0")
@@ -223,7 +242,8 @@ def trainable_filter(name: str, param) -> bool:
     return bool(getattr(param, "requires_grad", False))
 
 
-def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig):
+def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig, *,
+                         resume_dir: Optional[str] = None, resume_iteration: Optional[int] = None):
     from yeto.swift_megatron.sync import MegatronDiLoCoSync
 
     return MegatronDiLoCoSync(
@@ -240,9 +260,35 @@ def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig):
         final_cut_payload=cfg.final_cut_payload,
         transport=cfg.transport,
         initial_cut_timeout=cfg.initial_cut_timeout,
+        resume_dir=resume_dir,
+        resume_iteration=resume_iteration,
     )
 
 
+def resume_dir_for(cfg: DiLoCoHookConfig, args) -> Optional[str]:
+    """The checkpoint directory whose sync record a resumed island reads:
+    ``YETO_DILOCO_RESUME_DIR``, else the checkpoint ms-swift loaded
+    (``--mcore_model``) when it holds a record.  None: a cold start."""
+    from yeto.swift_megatron.sync import RESUME_STATE_FILE
+
+    if cfg.resume_dir:
+        return cfg.resume_dir
+    loaded = getattr(args, "mcore_model", None)
+    if isinstance(loaded, str) and os.path.isfile(os.path.join(loaded, RESUME_STATE_FILE)):
+        return loaded
+    return None
+
+
+def save_dir_for(args, iteration: int, *call_args, **call_kwargs) -> Optional[str]:
+    """The directory of the checkpoint ``on_save`` reports: its
+    ``output_dir`` argument, else ``<args.output_dir>/checkpoint-<iteration>``."""
+    for value in (call_kwargs.get("output_dir"), call_kwargs.get("save_dir"), *call_args):
+        if isinstance(value, (str, os.PathLike)):
+            return os.fspath(value)
+    root = getattr(args, "output_dir", None) or getattr(args, "save", None)
+    return os.path.join(root, f"checkpoint-{iteration}") if root else None
+
+
 def _model_chunks(trainer) -> list:
     # Pass the DDP-wrapped chunks: the sync core needs the grad/param buffers
     # and the DistributedOptimizer's mapping, both keyed on the wrapped models.
@@ -290,7 +336,10 @@ def __init__(self, trainer):
         # Build here (not in on_train_begin) so a bad config fails before data
         # loading; the trainer has already loaded the resume checkpoint, which
         # is what learner 0's INIT_PARAMS carries.
-        self.sync = type(self).sync_factory(_model_chunks(trainer), trainer.optimizer, self.config)
+        start = int(self.state.iteration)
+        resume_dir = resume_dir_for(self.config, self.args) if start > 0 else None
+        self.sync = type(self).sync_factory(_model_chunks(trainer), trainer.optimizer, self.config,
+                                            resume_dir=resume_dir, resume_iteration=start)
         atexit.register(self.close)
         logger.info("yeto DiLoCo sync attached: learner %d/%d, interval %d, fragments %d, start iteration %d",
                     self.config.learner_id, self.config.num_learners, self.config.sync_interval_steps,
@@ -366,6 +415,17 @@ def _log_step(self, iteration: int, target_tokens: int, sync_s: float,
         except OSError as error:  # never let bookkeeping kill a training step
             logger.warning("yeto DiLoCo: cannot write step log %s: %s", path, error)
 
+    def on_save(self, *args, **kwargs):
+        """Record the sync state next to the checkpoint just written (every rank)."""
+        if self.sync is None or self._closed or getattr(self.sync, "finished", False):
+            return
+        iteration = int(self.state.iteration)
+        output_dir = save_dir_for(self.args, iteration, *args, **kwargs)
+        if output_dir is None:
+            logger.warning("yeto DiLoCo: on_save without a checkpoint directory; sync state not recorded")
+            return
+        self.sync.save_state(output_dir, iteration=iteration, save_anchors=self.config.save_anchors)
+
     def on_train_end(self):
         self.close()
 
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index d87585e8..db451b87 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -142,7 +142,9 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import logging
+import os
 import re
 import time
 from collections import deque
@@ -158,6 +160,14 @@
 DEFAULT_CHUNK_NUMEL = 1 << 28
 CONTRACT_DOMAIN = b"yeto-swift-megatron-native-v1\0"
 FINAL_CUT_PAYLOADS = ("full", "manifest")
+# Island checkpoint record of the sync (``save_state`` / ``resume_dir``), next
+# to the trainer's checkpoint: the JSON on global rank 0, the optional anchor
+# values one file per rank.
+RESUME_STATE_FILE = "yeto_diloco_state.json"
+RESUME_ANCHOR_DIR = "yeto_diloco_anchors"
+RESUME_FORMAT = 1
+# ``yeto.diloco_sync.resume_fragment_mode`` results (that module imports torch).
+RESUME_OVERWRITE, RESUME_ANCHOR, RESUME_SHIFT = "overwrite", "anchor", "shift"
 
 _LAYER_MODULE_RE = re.compile(r"(^|\.)decoder\.layers\.(\d+)$")
 _GROUPED_EXPERT_RE = re.compile(r"^(.*?)(weight|bias)(\d+)$")
@@ -563,7 +573,8 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
                  final_cut_payload: str = "full", chunk_numel: int = DEFAULT_CHUNK_NUMEL,
                  transport: str = "cpu", parallel_context: ParallelContext | None = None,
                  client_factory=None, observer: Callable[[dict], None] | None = None,
-                 initial_cut_timeout: float = 0.0) -> None:
+                 initial_cut_timeout: float = 0.0, resume_dir: str | None = None,
+                 resume_iteration: int | None = None) -> None:
         import torch
         import torch.distributed as dist
 
@@ -663,6 +674,16 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
         self.sender: _PushSender | None = None  # rank 0: background PUSH streaming
         # (iteration, "apply" | "push", fragment, version | pull step), same on every rank.
         self.events: deque[tuple[int, str, int, int]] = deque(maxlen=4096)
+        # Resume (``resume_dir``): the record ``save_state`` wrote with the
+        # checkpoint the trainer resumed from, the old anchor of every local
+        # shard when it was saved too, and how each fragment's first broadcast
+        # after the restart was installed (fid -> mode, same on every rank).
+        self.resume_record: dict | None = None
+        self._resume_anchors: dict[int, Any] = {}  # id(param) -> old anchor (host)
+        self._resume_has_anchors = False
+        self.resumed: dict[int, str] = {}
+        if resume_dir:
+            self._load_resume(resume_dir, resume_iteration)
 
         if self.rank == 0:
             log.info("swift-megatron sync: %d trainable tensors (%.2f B elements) -> %d fragments",
@@ -765,7 +786,7 @@ def _install_initial_cut(self, timeout: float) -> None:
             updates = self._take_updates(plan) if self.rank == 0 else {}
             for fid, version in plan.applies:
                 data = updates.pop(fid) if self.rank == 0 else None
-                self._apply_fragment(fid, data, self.bulk_torch, blend=True)
+                self._apply_first_or_blend(fid, data, version)
                 del data
                 self.state.fragment_versions[fid] = version
                 self.state.global_step = max(self.state.global_step, version)
@@ -780,6 +801,166 @@ def _install_initial_cut(self, timeout: float) -> None:
             log.info("initial cut: %d of %d fragments installed before training in %.0f s",
                      sum(self.anchored), self.num_fragments, time.monotonic() - started)
 
+    # ------------------------------------------------------------------ resume
+    def _apply_first_or_blend(self, fid: int, data: Any, version: int) -> str | None:
+        """Apply a broadcast; the first one of a fragment after a resume by
+        ``resume_fragment_mode``.  Returns that mode (None: not a resumed
+        first install) and records it in ``self.resumed``."""
+
+        mode = None
+        if self.resume_record is not None and not self.anchored[fid] and fid not in self.resumed:
+            from yeto.diloco_sync import resume_fragment_mode
+
+            mode = resume_fragment_mode(self.resume_record, fid, version, have_anchor=self._resume_has_anchors)
+            self.resumed[fid] = mode
+            if self.rank == 0:
+                log.info("resume: fragment %d version %d (checkpoint anchor %d) -> %s", fid, version,
+                         self.resume_record["versions"][fid], mode)
+        self._apply_fragment(fid, data, self.bulk_torch, blend=True, mode=mode)
+        if mode is not None and len(self.resumed) == self.num_fragments:
+            self._resume_anchors.clear()  # every old anchor is used or no longer needed
+        return mode
+
+    def _restore_counters(self, fid: int, mode: str) -> None:
+        """A fragment whose checkpointed progress was kept (anchor / shift)
+        keeps its H-gate and token counters, and an unchanged anchor its
+        spent-anchor guard; an overwritten one starts afresh."""
+
+        if mode == RESUME_OVERWRITE:
+            return
+        self.state.restore_fragment(self.resume_record, fid)
+        if mode == RESUME_ANCHOR:
+            spent = self.resume_record["spent"][fid]
+            self.spent[fid] = None if spent is None else (int(spent[0]), int(spent[1]))
+
+    def _resume_paths(self, resume_dir: str) -> tuple[str, str]:
+        return (os.path.join(resume_dir, RESUME_STATE_FILE),
+                os.path.join(resume_dir, RESUME_ANCHOR_DIR, f"rank{self.rank:05d}.pt"))
+
+    def _shard_keys(self) -> list[tuple[str, int]]:
+        return [(s.name, s.tp_rank) for s in self.shards]
+
+    def _load_resume(self, resume_dir: str, resume_iteration: int | None) -> None:
+        """Read the record ``save_state`` left in ``resume_dir`` (collective).
+
+        A missing or foreign record (another learner, layout, parallel shape
+        or iteration) is a cold start: every fragment takes the global as
+        before.  Old anchors are used only if every rank has its file."""
+
+        import torch
+
+        record, reason = None, None
+        if self.rank == 0:
+            path = self._resume_paths(resume_dir)[0]
+            try:
+                with open(path, encoding="utf-8") as handle:
+                    record = json.load(handle)
+            except FileNotFoundError:
+                reason = f"no {RESUME_STATE_FILE} in {resume_dir}"
+            except (OSError, ValueError) as exc:
+                reason = f"unreadable {path}: {exc}"
+            if record is not None:
+                checks = [
+                    (record.get("format") == RESUME_FORMAT, "format"),
+                    (record.get("learner_id") == self.learner_id, "learner id"),
+                    (record.get("contract") == self.contract.hex(), "layout / parallel shape"),
+                    (len(record.get("versions", ())) == self.num_fragments, "fragment count"),
+                    (resume_iteration is None or record.get("iteration") == int(resume_iteration), "iteration"),
+                ]
+                bad = [what for ok, what in checks if not ok]
+                if bad:
+                    reason = f"{path} does not match this island ({', '.join(bad)})"
+                    record = None
+            if record is None:
+                log.warning("resume: %s; the syncer's cut replaces every fragment (cold start)", reason)
+        record = self._bcast_object(record)
+        if record is None:
+            return
+        loaded: dict[int, Any] = {}
+        ok = bool(record.get("anchors"))
+        if ok:
+            try:
+                saved = torch.load(self._resume_paths(resume_dir)[1], map_location="cpu", weights_only=True)
+                if saved["keys"] != [list(k) for k in self._shard_keys()]:
+                    raise ValueError("shard list differs")
+                for shard, value in zip(self.shards, saved["values"], strict=True):
+                    if tuple(value.shape) != tuple(shard.local_shape):
+                        raise ValueError(f"{shard.name} has shape {tuple(value.shape)}")
+                    loaded[id(shard.param)] = value
+            except Exception as exc:  # agreed on below; a missing anchor file is not fatal
+                log.warning("resume: rank %d cannot use its saved anchors: %s", self.rank, exc)
+                ok = False
+        ok = all(self._gather_object(ok))
+        self.resume_record = record
+        self._resume_has_anchors = ok
+        self._resume_anchors = loaded if ok else {}
+        self.units_total = int(record["units_total"])
+        self.state.global_step = int(record["global_step"])
+        if self.rank == 0:
+            log.info("resume: island record of iteration %s (anchor versions %s..%s, old anchors %s)",
+                     record.get("iteration"), min(record["versions"]), max(record["versions"]),
+                     "loaded" if ok else "not saved: a newer fragment takes the global")
+
+    def _gather_object(self, value: Any) -> list[Any]:
+        import torch.distributed as dist
+
+        if self.world == 1:
+            return [value]
+        out: list[Any] = [None] * self.world
+        dist.all_gather_object(out, value, group=self.comm.world)
+        return out
+
+    def save_state(self, output_dir: str, *, iteration: int, save_anchors: bool = False) -> bool:
+        """Record the sync's per-fragment state next to the trainer checkpoint
+        ``output_dir`` of ``iteration`` (call it on every rank after the
+        checkpoint is written; collective only with ``save_anchors``).
+
+        Global rank 0 writes ``RESUME_STATE_FILE``: each fragment's anchor
+        version, the H-gate / token counters and the spent-anchor guard.
+        ``save_anchors`` also writes every rank's anchor shards (the
+        broadcast dtype, ~``2 * n_r`` bytes per rank for bf16) to
+        ``RESUME_ANCHOR_DIR``; only with them can a fragment that the syncer
+        moved past the checkpoint's anchor keep the island's local progress
+        (``global_new + (local - anchor_old)``).  Never raises on an I/O
+        error (training goes on); returns whether rank 0 wrote the record
+        (True on the other ranks)."""
+
+        import torch
+
+        anchors_ok = False
+        if save_anchors:
+            path = self._resume_paths(output_dir)[1]
+            try:
+                os.makedirs(os.path.dirname(path), exist_ok=True)
+                values = [self.anchors.get(id(s.param)) for s in self.shards]
+                if any(v is None for v in values):
+                    raise ValueError("a local shard has no anchor yet")
+                tmp = f"{path}.tmp"
+                torch.save({"keys": [list(k) for k in self._shard_keys()], "values": values}, tmp)
+                os.replace(tmp, path)
+                anchors_ok = True
+            except Exception as exc:
+                log.warning("save_state: rank %d cannot write its anchors: %s", self.rank, exc)
+        anchors_ok = all(self._gather_object(anchors_ok)) if save_anchors else False
+        written = True
+        if self.rank == 0:
+            record = {"format": RESUME_FORMAT, "learner_id": self.learner_id, "iteration": int(iteration),
+                      "contract": self.contract.hex(), "anchors": anchors_ok,
+                      **self.state.resume_record(anchored=self.anchored, units_total=self.units_total,
+                                                 spent=self.spent)}
+            path = self._resume_paths(output_dir)[0]
+            try:
+                os.makedirs(output_dir, exist_ok=True)
+                tmp = f"{path}.tmp"
+                with open(tmp, "w", encoding="utf-8") as handle:
+                    json.dump(record, handle)
+                os.replace(tmp, path)
+                log.info("save_state: %s (iteration %d, anchors %s)", path, iteration, anchors_ok)
+            except OSError as exc:
+                log.warning("save_state: cannot write %s: %s", path, exc)
+                written = False
+        return written
+
     # ------------------------------------------------------------------ collectives
     def _bcast_object(self, value: Any) -> Any:
         import torch.distributed as dist
@@ -882,9 +1063,16 @@ def _push_part_dtype(self) -> int:
 
         return DTYPE_F32 if bulk_dtype(self.wire) == DTYPE_F32 else push_wire_dtype(self.wire)
 
-    def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, blend: bool) -> None:
+    def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, blend: bool,
+                     mode: str | None = None) -> None:
         """Install the raw global values of ``names`` (``flat`` on rank 0, in the
-        broadcast dtype) as every rank's anchor, master and model copy."""
+        broadcast dtype) as every rank's anchor, master and model copy.
+
+        A resumed island's first install of a fragment (``mode``,
+        ``yeto.diloco_sync.resume_fragment_mode``): ``"anchor"`` only records
+        the anchor and leaves the master and model copy untouched,
+        ``"shift"`` writes ``global + (master - old anchor)``; ``"overwrite"``
+        (or None) is the ordinary install."""
 
         import torch
         import torch.distributed as dist
@@ -937,7 +1125,16 @@ def piece(buffer, name, part_dim, tp_rank, size):
             raw = whole.to(dtype=torch.float32).contiguous()
             self.anchors[id(shard.param)] = whole.to(device="cpu", dtype=self.anchor_dtype).contiguous().clone()
             targets.append(raw)
-        if blend and self.merge_alpha > 0:
+        if mode == RESUME_ANCHOR:
+            return  # the broadcast is the checkpoint's own anchor: keep the local values
+        if mode == RESUME_SHIFT:
+            from yeto.diloco_sync import shifted_resume_value
+
+            masters = self._masters(local)
+            targets = [shifted_resume_value(t, m, self._resume_anchors.pop(id(s.param)))
+                       for s, m, t in zip(local, masters, targets, strict=True)]
+            del masters
+        elif blend and self.merge_alpha > 0:
             masters = self._masters(local)
             targets = [self.merge_alpha * m.to(torch.float32) + (1.0 - self.merge_alpha) * t
                        for m, t in zip(masters, targets, strict=True)]
@@ -961,8 +1158,10 @@ def _check_payload(self, fid: int, data: Any, dtype: Any) -> None:
         if nbytes != expected:
             raise SwiftMegatronSyncError(f"fragment {fid} payload has {nbytes} bytes, expected {expected}")
 
-    def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool) -> None:
-        """Apply one fragment's payload (``data`` on rank 0) chunk by chunk."""
+    def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool,
+                        mode: str | None = None) -> None:
+        """Apply one fragment's payload (``data`` on rank 0) chunk by chunk
+        (``mode``: a resumed island's first install, ``_apply_chunk``)."""
 
         import torch
 
@@ -974,7 +1173,7 @@ def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool) -> No
             flat = None
             if view is not None:
                 flat = torch.frombuffer(bytearray(view[offset:offset + nbytes]), dtype=torch.uint8).view(dtype)
-            self._apply_chunk(names, flat, dtype, blend=blend)
+            self._apply_chunk(names, flat, dtype, blend=blend, mode=mode)
             offset += nbytes
 
     # ------------------------------------------------------------------ boundary
@@ -1000,6 +1199,10 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *,
         iteration = int(iteration)
         if not self._counters_started:  # count from the (possibly resumed) first step
             state.steps_at_reset = [iteration - 1] * self.num_fragments
+            # A fragment that kept its checkpointed progress keeps its clock:
+            # the H gate is not restarted by the island's own restart.
+            for fid, mode in self.resumed.items():
+                self._restore_counters(fid, mode)
             self._counters_started = True
         self.iteration = iteration
         if update_successful:
@@ -1017,10 +1220,12 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *,
         for fid, version in plan.applies:
             data = updates.get(fid) if self.rank == 0 else None
             started = time.perf_counter()
-            self._apply_fragment(fid, data, self.bulk_torch, blend=True)
+            mode = self._apply_first_or_blend(fid, data, version)
             self.stats["apply_s"] += time.perf_counter() - started
             state.steps_at_reset[fid] = iteration
             state.units_at_reset[fid] = self.units_total
+            if mode is not None:
+                self._restore_counters(fid, mode)
             state.fragment_versions[fid] = version
             self.anchored[fid] = True
             state.global_step = max(state.global_step, version)

From f3ea120ae7ef6231dbce675aeb2fa507a56daba3 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 06:32:51 +0800
Subject: [PATCH 02/17] syncer: checkpoint and exit cleanly on SIGTERM / SIGINT

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>
---
 docs/PROTOCOL.md          |  11 ++++
 syncer/src/main.rs        |  11 ++++
 syncer/src/server.rs      |  90 ++++++++++++++++++++++++---
 syncer/src/shutdown.rs    | 124 ++++++++++++++++++++++++++++++++++++++
 tests/test_integration.py |  72 ++++++++++++++++++++++
 5 files changed, 300 insertions(+), 8 deletions(-)
 create mode 100644 syncer/src/shutdown.rs

diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md
index f8c52ec2..0dbce353 100644
--- a/docs/PROTOCOL.md
+++ b/docs/PROTOCOL.md
@@ -593,6 +593,17 @@ coordinator always rewrites `--checkpoint-path` at the final quiescent cut
 before terminal delivery. Thus a total-step count not divisible by the
 periodic interval still leaves an atomic final checkpoint.
 
+SIGTERM or SIGINT stops the syncer gracefully: it launches no new round,
+gives the rounds still gathering `--shutdown-timeout-s` (default 600) to
+close, then abandons the rest (they have not touched the state, so a resumed
+run launches those steps again), commits every computed round, writes
+`--checkpoint-path` synchronously (after any background checkpoint) and exits
+0 with `stop: checkpoint written step=<t>`. The checkpoint's global step is
+then the last `outer step` on the event tape, and `--resume` continues the
+tape at `t + 1`. A signal before training starts exits at once (nothing was
+merged), and a second signal exits immediately without a checkpoint. The
+terminal cut is not marked final.
+
 `--checkpoint-at-steps T1,T2,...` (launcher: `--rl-sync-checkpoint-at-steps`,
 asynchronous full-parameter rounds) also writes the cut at each listed global
 step `t` (one fragment merge, the `step=` of the `outer step` log, the unit of
diff --git a/syncer/src/main.rs b/syncer/src/main.rs
index 9bb05cca..073e00c9 100644
--- a/syncer/src/main.rs
+++ b/syncer/src/main.rs
@@ -19,6 +19,7 @@ mod protocol;
 mod sample_relay;
 mod server;
 mod sha256;
+mod shutdown;
 mod state;
 mod subscriber;
 
@@ -167,6 +168,13 @@ struct Args {
     /// sessions are refused.
     #[arg(long)]
     bcast_delta_full_every: Option<u64>,
+    /// On SIGTERM / SIGINT: stop launching rounds, give the rounds still
+    /// gathering this many seconds to close (then drop them; they have not
+    /// touched the state), commit every computed round, write
+    /// --checkpoint-path synchronously and exit 0. A second signal exits at
+    /// once without a checkpoint.
+    #[arg(long, default_value_t = 600)]
+    shutdown_timeout_s: u64,
     /// Resume from --checkpoint-path if it exists.
     #[arg(long, default_value_t = false)]
     resume: bool,
@@ -518,6 +526,9 @@ fn main() -> anyhow::Result<()> {
         island_hmac_key,
         expected_island_contract,
         syncer_epoch: args.syncer_epoch,
+        shutdown: Some(std::sync::Arc::new(shutdown::Shutdown::new(std::time::Duration::from_secs(
+            args.shutdown_timeout_s,
+        )))),
         membership,
         recovery,
         inline_resend: args.inline_resend,
diff --git a/syncer/src/server.rs b/syncer/src/server.rs
index f15d96b1..461d3a1a 100644
--- a/syncer/src/server.rs
+++ b/syncer/src/server.rs
@@ -104,6 +104,8 @@ pub struct Config {
     pub quorum_timeout_s: u64,
     /// Bounded wait for learners to apply and acknowledge the terminal cut.
     pub final_ack_timeout_s: u64,
+    /// Graceful stop (`crate::shutdown`); None: signals keep their default action.
+    pub shutdown: Option<Arc<crate::shutdown::Shutdown>>,
     pub total_steps: u64,
     /// Strict dense-policy mode: one logical local step contributes every
     /// fragment in a serial sweep.  None preserves the legacy independent
@@ -2080,6 +2082,9 @@ fn check_and_record_profile(cfg: &Config, profile: &[u8; 32]) -> Result<()> {
 
 pub async fn run(cfg: Config) -> Result<()> {
     validate_config(&cfg)?;
+    if let Some(stop) = &cfg.shutdown {
+        stop.install_signal_handlers()?;
+    }
     check_resume_checkpoint(&cfg)?;
     let semantic_profile_hash = cfg.semantic_profile_hash();
     check_and_record_profile(&cfg, &semantic_profile_hash)?;
@@ -3848,14 +3853,20 @@ async fn scheduler(
                 break;
             }
         }
-        match recv_init_event(
-            &mut events,
-            &budget_cutoff,
-            budget_reports.len(),
-            cfg.learners,
-        )
-        .await?
-        {
+        let init_event = recv_init_event(&mut events, &budget_cutoff, budget_reports.len(), cfg.learners);
+        let event = match cfg.shutdown.as_deref() {
+            Some(stop) => tokio::select! {
+                biased;
+                _ = stop.requested() => {
+                    // Nothing merged yet: the checkpoint on disk (if any) is current.
+                    info!("stop requested before training started; exiting");
+                    return Ok(());
+                }
+                event = init_event => event?,
+            },
+            None => init_event.await?,
+        };
+        match event {
             Event::Fatal { metric, message } => {
                 append_strict_failure(cfg.event_tape.as_deref(), metric, &message);
                 bail!("RL strict failure {metric}: {message}");
@@ -4134,6 +4145,9 @@ async fn scheduler(
     if dynamic {
         write_pool_status(&cfg, &registry, st.global_step, "running");
     }
+    // A stop signal (`crate::shutdown`): set once every round has drained.
+    let mut stopped = false;
+    let mut stop_logged = false;
     'outer: while next_t <= last_step
         || !inflight.is_empty()
         || !computing.is_empty()
@@ -4144,6 +4158,31 @@ async fn scheduler(
             cutoff_established = true;
             break 'outer;
         }
+        let stop_deadline = cfg.shutdown.as_ref().and_then(|stop| stop.deadline());
+        if let Some(deadline) = stop_deadline {
+            if !stop_logged {
+                stop_logged = true;
+                info!(
+                    step = st.global_step,
+                    gathering = inflight.len(),
+                    computing = computing.len() + ready.len(),
+                    "stopping: no new rounds; waiting for the rounds in flight"
+                );
+            }
+            if Instant::now() >= deadline && !inflight.is_empty() {
+                // A gathering round has not touched the state; after --resume
+                // its step is launched again.
+                let dropped: Vec<u64> = inflight.iter().map(|round| round.t).collect();
+                warn!(steps = ?dropped, "stop timeout: abandoning the rounds still gathering");
+                for round in inflight.drain(..) {
+                    busy_fragments.remove(&round.p);
+                }
+            }
+            if inflight.is_empty() && computing.is_empty() && ready.is_empty() && publishing.is_none() {
+                stopped = true;
+                break 'outer;
+            }
+        }
         let expired = expire_member_leases(&registry, &mut inflight, &mut step_rates).await;
         for member in &expired {
             lease_expired.insert(*member);
@@ -4241,6 +4280,7 @@ async fn scheduler(
         // Keep the pipeline full (throttled by min_round_interval_ms).
         while inflight.len() + computing.len() + ready.len() < depth
             && next_t <= last_step
+            && stop_deadline.is_none()
             && Instant::now() >= next_launch
         {
             if !budget_cutoff.try_linearize_work() {
@@ -4372,6 +4412,13 @@ async fn scheduler(
                         cutoff_established = true;
                         break 'outer;
                     }
+                    if stop_deadline.is_some() {
+                        // Stopping: no new attempt; the round is launched again after --resume.
+                        let round = inflight.remove(i);
+                        busy_fragments.remove(&round.p);
+                        warn!(step = round.t, "stop: quorum timeout, round abandoned");
+                        continue;
+                    }
                     let r = &mut inflight[i];
                     let groups = current_groups(&registry);
                     if fixed_roster {
@@ -4569,6 +4616,7 @@ async fn scheduler(
             .min();
         let next_fragment_available = inflight.len() + computing.len() + ready.len() < depth
             && next_t <= last_step
+            && stop_deadline.is_none()
             && !busy_fragments.contains(&schedule.fragment_for_step(next_t));
         if next_fragment_available {
             earliest = Some(earliest.map_or(next_launch, |d| d.min(next_launch)));
@@ -4582,6 +4630,11 @@ async fn scheduler(
             let check = deadline_now + RESEND_CHECK_INTERVAL;
             earliest = Some(earliest.map_or(check, |d| d.min(check)));
         }
+        if cfg.shutdown.is_some() {
+            // Notice a stop signal (and its timeout) within a fraction of a second.
+            let check = deadline_now + SHUTDOWN_CHECK_INTERVAL;
+            earliest = Some(earliest.map_or(check, |d| d.min(check)));
+        }
         let wake = if let Some(pending) = publishing.as_mut() {
             // Same event-first bias; a durable checkpoint outranks compute
             // so the oldest committed round is exposed first.
@@ -5107,6 +5160,24 @@ async fn scheduler(
         resends.stop(&registry)
     };
 
+    if stopped {
+        drop(forwarders); // a catch-up in progress is resent after the restart
+        st.drain_iso_backend().await?;
+        eval_checkpoints.finish().await?;
+        if let Some(path) = &cfg.checkpoint_path {
+            // Waits for a background checkpoint first, then writes this cut.
+            st.save_checkpoint(path)?;
+            info!(step = st.global_step, path = %path.display(), "stop: checkpoint written");
+        } else {
+            warn!(step = st.global_step, "stop without --checkpoint-path: the merged state is not saved");
+        }
+        if dynamic {
+            write_pool_status(&cfg, &registry, st.global_step, "stopped");
+        }
+        info!(step = st.global_step, "stopped; restart with --resume to continue");
+        return Ok(());
+    }
+
     if cfg.learner_budget_steps.is_some() {
         wait_forwarders(std::mem::take(&mut forwarders)).await;
         cutoff_established |= budget_cutoff.deadline().is_some();
@@ -7696,6 +7767,8 @@ impl<G: Clone> CatchUps<G> {
 
 /// How often the scheduler checks background resends while any is active.
 const RESEND_CHECK_INTERVAL: Duration = Duration::from_millis(100);
+/// How often the scheduler looks for a stop signal while it waits.
+const SHUTDOWN_CHECK_INTERVAL: Duration = Duration::from_millis(200);
 /// At the end of the round loop, wait this long for a forwarder to hand
 /// over the fragment it is sending before it is aborted.
 const RESEND_STOP_WAIT: Duration = Duration::from_secs(30);
@@ -9465,6 +9538,7 @@ mod tests {
             delta_correction: false,
             quorum_timeout_s: 1,
             final_ack_timeout_s: 1,
+            shutdown: None,
             total_steps,
             policy_sweep_fragments: None,
             outer_lr: 1.0,
diff --git a/syncer/src/shutdown.rs b/syncer/src/shutdown.rs
new file mode 100644
index 00000000..9caabd90
--- /dev/null
+++ b/syncer/src/shutdown.rs
@@ -0,0 +1,124 @@
+//! Graceful stop on SIGTERM / SIGINT.
+//!
+//! The first signal asks the scheduler to stop: no new round is launched,
+//! the rounds already gathering get `--shutdown-timeout-s` to close, every
+//! computed round commits, and the consistent cut is written synchronously
+//! to `--checkpoint-path` before the process exits 0, so `--resume` starts
+//! at the last outer step on the event tape. A second signal exits at once.
+
+use std::time::{Duration, Instant};
+
+use tokio::sync::watch;
+use tracing::warn;
+
+/// Exit status of a second signal while the first one is being served.
+pub const FORCED_EXIT_CODE: i32 = 143;
+
+pub struct Shutdown {
+    tx: watch::Sender<Option<Instant>>,
+    timeout: Duration,
+}
+
+impl Shutdown {
+    pub fn new(timeout: Duration) -> Self {
+        Self {
+            tx: watch::channel(None).0,
+            timeout,
+        }
+    }
+
+    /// Ask for a stop; returns false if one was already requested.
+    pub fn request(&self) -> bool {
+        let mut first = false;
+        self.tx.send_if_modified(|at| {
+            first = at.is_none();
+            if first {
+                *at = Some(Instant::now());
+            }
+            first
+        });
+        first
+    }
+
+    pub fn requested_at(&self) -> Option<Instant> {
+        *self.tx.borrow()
+    }
+
+    /// When the rounds still gathering are abandoned.
+    pub fn deadline(&self) -> Option<Instant> {
+        self.requested_at().map(|at| at + self.timeout)
+    }
+
+    /// Resolves once a stop is requested.
+    pub async fn requested(&self) {
+        let mut rx = self.tx.subscribe();
+        let _ = rx.wait_for(Option::is_some).await;
+    }
+
+    /// Route SIGTERM and SIGINT here (inside a tokio runtime).
+    #[cfg(unix)]
+    pub fn install_signal_handlers(self: &std::sync::Arc<Self>) -> anyhow::Result<()> {
+        use tokio::signal::unix::{signal, SignalKind};
+
+        let mut term = signal(SignalKind::terminate())?;
+        let mut int = signal(SignalKind::interrupt())?;
+        let this = self.clone();
+        tokio::spawn(async move {
+            loop {
+                let name = tokio::select! {
+                    _ = term.recv() => "SIGTERM",
+                    _ = int.recv() => "SIGINT",
+                };
+                if this.request() {
+                    warn!(
+                        signal = name,
+                        timeout_s = this.timeout.as_secs(),
+                        "stop requested: no new rounds; draining in-flight rounds, then checkpoint and exit \
+                         (a second signal exits at once)"
+                    );
+                } else {
+                    warn!(signal = name, "second stop signal: exiting without a checkpoint");
+                    std::process::exit(FORCED_EXIT_CODE);
+                }
+            }
+        });
+        Ok(())
+    }
+
+    #[cfg(not(unix))]
+    pub fn install_signal_handlers(self: &std::sync::Arc<Self>) -> anyhow::Result<()> {
+        let this = self.clone();
+        tokio::spawn(async move {
+            while tokio::signal::ctrl_c().await.is_ok() {
+                if !this.request() {
+                    std::process::exit(FORCED_EXIT_CODE);
+                }
+            }
+        });
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn first_request_wins_and_wakes_waiters() {
+        let stop = std::sync::Arc::new(Shutdown::new(Duration::from_secs(5)));
+        assert!(stop.requested_at().is_none() && stop.deadline().is_none());
+        let waiter = {
+            let stop = stop.clone();
+            tokio::spawn(async move { stop.requested().await })
+        };
+        tokio::task::yield_now().await;
+        assert!(stop.request());
+        let at = stop.requested_at().unwrap();
+        assert!(!stop.request());
+        assert_eq!(stop.requested_at(), Some(at));
+        assert_eq!(stop.deadline(), Some(at + Duration::from_secs(5)));
+        tokio::time::timeout(Duration::from_secs(1), waiter).await.unwrap().unwrap();
+        // A waiter that starts after the request returns at once.
+        tokio::time::timeout(Duration::from_secs(1), stop.requested()).await.unwrap();
+    }
+}
diff --git a/tests/test_integration.py b/tests/test_integration.py
index a5aa75a5..f582925e 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -757,3 +757,75 @@ def test_single_learner_roundtrip():
         if proc.poll() is None:
             proc.kill()
         print((proc.stdout.read() if proc.stdout else "")[-3000:])
+
+
+@pytest.mark.timeout(180)
+def test_sigterm_checkpoints_the_last_outer_step_and_resume_continues_the_tape(tmp_path):
+    """SIGTERM stops a mid-run syncer cleanly: no new rounds, the round in
+    flight closes, the cut is written to --checkpoint-path (although no
+    --checkpoint-every boundary was reached) and the process exits 0; the
+    checkpoint's global step is the last outer step on the tape, and a
+    --resume run continues the tape from the next step."""
+    import signal
+
+    binary = build_syncer()
+    named = [("model.embed.weight", DIM // 4), ("model.body.weight", DIM)]
+    layout = build_layout(named, 2)
+    checkpoint = tmp_path / "state.ckpt"
+    tape = tmp_path / "tape.jsonl"
+
+    def syncer(port, total_steps, *extra):
+        return group_popen(
+            [str(binary), "--port", str(port), "--learners", "1", "--quorum", "1",
+             "--grace-ms", "20", "--sync-interval-steps", "0", "--min-round-interval-ms", "50",
+             "--total-steps", str(total_steps), "--checkpoint-path", str(checkpoint),
+             "--checkpoint-every", "1000", "--event-tape", str(tape),
+             "--shutdown-timeout-s", "5", *extra],
+            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
+        )
+
+    def steps():
+        if not tape.exists():
+            return []
+        return [json.loads(line)["step"] for line in tape.read_text().splitlines()
+                if line.strip() and "step" in json.loads(line)]
+
+    port = free_port()
+    first = syncer(port, 1000)
+    outputs = []
+    try:
+        learner = ToyLearner(0, port, torch.ones(DIM + DIM // 4), layout, dtype=DTYPE_F32)
+        learner.start()
+        deadline = time.monotonic() + 60
+        while len(steps()) < 4:
+            assert time.monotonic() < deadline, "no outer steps before the stop"
+            assert first.poll() is None, first.communicate()[0]
+            time.sleep(0.05)
+        first.send_signal(signal.SIGTERM)
+        output, _ = first.communicate(timeout=60)
+        outputs.append(output)
+        assert first.returncode == 0, output
+        assert "stop: checkpoint written" in output
+        stopped = steps()
+        assert stopped == list(range(1, len(stopped) + 1))
+        assert read_checkpoint_global_step(checkpoint) == stopped[-1]
+        learner.join(timeout=30)
+
+        port = free_port()
+        total = stopped[-1] + 4
+        second = syncer(port, total, "--resume")
+        learner = ToyLearner(0, port, torch.ones(DIM + DIM // 4), layout, dtype=DTYPE_F32)
+        learner.start()
+        output, _ = second.communicate(timeout=90)
+        outputs.append(output)
+        assert second.returncode == 0, output
+        assert "resumed from checkpoint" in output
+        learner.join(timeout=30)
+        assert learner.final_manifest is not None and learner.final_manifest.global_step == total
+        assert steps() == list(range(1, total + 1))
+        assert read_checkpoint_global_step(checkpoint) == total
+    finally:
+        for process in (first, locals().get("second")):
+            if process is not None and process.poll() is None:
+                process.kill()
+        print("\n".join(outputs)[-4000:])

From 764ae9065007eabc6158c80780e78c46b164bb5b Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 06:33:24 +0800
Subject: [PATCH 03/17] swift-megatron: accept a resume record without the
 layout contract

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>
---
 tests/test_swift_megatron_sync.py | 7 ++++++-
 yeto/swift_megatron/sync.py       | 5 ++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py
index a5c3fa4d..ad464af0 100644
--- a/tests/test_swift_megatron_sync.py
+++ b/tests/test_swift_megatron_sync.py
@@ -614,7 +614,12 @@ def check_anchors(sync, values):
     assert moved.state.steps_at_reset == [0, 0, 0] and moved.state.fragment_versions == [9, 9, 9]
     moved.close()
 
-    # (c) Without saved anchors a moved fragment takes the global; an unmoved one is kept.
+    # (c) Without saved anchors a moved fragment takes the global; an unmoved one is kept
+    # (a record without the contract, as written by hand for an older checkpoint, is accepted).
+    if rank == 0:
+        record = json.loads((bare / sm.RESUME_STATE_FILE).read_text())
+        del record["contract"]
+        (bare / sm.RESUME_STATE_FILE).write_text(json.dumps(record))
     restart()
     cut.update({0: (4, old_g), 1: (9, new_g), 2: (9, new_g)})
     mixed = make(str(bare), 3)
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index db451b87..a26dbfe0 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -863,7 +863,10 @@ def _load_resume(self, resume_dir: str, resume_iteration: int | None) -> None:
                 checks = [
                     (record.get("format") == RESUME_FORMAT, "format"),
                     (record.get("learner_id") == self.learner_id, "learner id"),
-                    (record.get("contract") == self.contract.hex(), "layout / parallel shape"),
+                    # A hand-written bridge record (an island checkpoint from before
+                    # these records) may leave the contract out.
+                    (record.get("contract", self.contract.hex()) == self.contract.hex(),
+                     "layout / parallel shape"),
                     (len(record.get("versions", ())) == self.num_fragments, "fragment count"),
                     (resume_iteration is None or record.get("iteration") == int(resume_iteration), "iteration"),
                 ]

From ac8d4a97815c8c8536bb90d55b9faea36aa159fd Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 14:20:42 +0800
Subject: [PATCH 04/17] swift-megatron: release spent anchors of undelivered
 pushes and after 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>
---
 syncer/src/server.rs              |  78 ++++++++++----
 tests/test_swift_megatron_sync.py | 166 ++++++++++++++++++++++++++++++
 yeto/swift_megatron/sync.py       |  77 ++++++++++++--
 3 files changed, 291 insertions(+), 30 deletions(-)

diff --git a/syncer/src/server.rs b/syncer/src/server.rs
index 461d3a1a..de24ed33 100644
--- a/syncer/src/server.rs
+++ b/syncer/src/server.rs
@@ -5949,21 +5949,24 @@ fn route_push(
 /// same fragment against its own base, weighted by its counters (and the
 /// carry discount), as long as its base lag stays within `--max-base-lag`.
 ///
-/// A carried delta counts once: after the syncer carries a delta from an
-/// anchor (fragment version), a later push of that island connection from
-/// the same or an older anchor of that fragment is refused
-/// (`anchor_reused`), and a late push from an anchor whose delta was
-/// already merged is not carried. Pushes that never involve a carry keep
-/// the existing stale-base rules.
+/// A delta counts once: after the syncer carries or gathers a delta from an
+/// anchor (fragment version), a later push of that island from the same or
+/// an older anchor of that fragment is refused (`anchor_reused`), on any of
+/// its connections. An island that redials (a WAN drop, a syncer it lost
+/// for a while) releases the anchors its lost pushes held, so it may answer
+/// the next pull from an anchor whose delta did reach this syncer; that
+/// second delta repeats the first one's local steps and must not be merged
+/// again. A round restart frees the anchors of the pushes it dropped
+/// (`release_anchor`). Pushes from newer anchors keep the stale-base rules.
 struct LateCarry {
     limit: u64,
     /// Late pushes waiting for their fragment's next round.
     queued: HashMap<usize, Vec<(Member, Push)>>,
-    /// Newest base version of a carried delta per island connection and
-    /// fragment.
-    anchors: HashMap<(Member, usize), u64>,
-    /// Newest base version of a delta gathered by its own round.
-    gathered: HashMap<(Member, usize), u64>,
+    /// Newest base version of a carried delta per island and fragment.
+    anchors: HashMap<(u32, usize), u64>,
+    /// Newest base version of a delta gathered by its own round, per island
+    /// and fragment.
+    gathered: HashMap<(u32, usize), u64>,
 }
 
 #[derive(Debug, Eq, PartialEq)]
@@ -5991,27 +5994,30 @@ impl LateCarry {
         }
     }
 
-    /// True when `push` reuses an anchor this connection already spent.
+    /// True when `push` reuses an anchor whose delta this island already
+    /// had carried or gathered (by any of its connections).
     fn reuses_anchor(&self, member: Member, push: &Push) -> bool {
-        self.anchors
-            .get(&(member, push.fragment_id as usize))
-            .is_some_and(|&spent| push.base_version <= spent)
+        let key = (member.learner_id, push.fragment_id as usize);
+        [self.anchors.get(&key), self.gathered.get(&key)]
+            .into_iter()
+            .flatten()
+            .any(|&taken| push.base_version <= taken)
     }
 
     fn spend_anchor(&mut self, member: Member, fragment: usize, base_version: u64) {
-        let spent = self.anchors.entry((member, fragment)).or_insert(0);
+        let spent = self.anchors.entry((member.learner_id, fragment)).or_insert(0);
         *spent = (*spent).max(base_version);
     }
 
     fn note_gathered(&mut self, member: Member, fragment: usize, base_version: u64) {
-        let taken = self.gathered.entry((member, fragment)).or_insert(0);
+        let taken = self.gathered.entry((member.learner_id, fragment)).or_insert(0);
         *taken = (*taken).max(base_version);
     }
 
     /// A restarted round drops its gathered pushes; their islands answer the
     /// new attempt from the same anchor, so the anchor is free again.
     fn release_anchor(&mut self, member: Member, push: &Push) {
-        let key = (member, push.fragment_id as usize);
+        let key = (member.learner_id, push.fragment_id as usize);
         if self.gathered.get(&key) == Some(&push.base_version) {
             if push.base_version == 0 {
                 self.gathered.remove(&key);
@@ -6045,7 +6051,7 @@ impl LateCarry {
         }
         if self
             .gathered
-            .get(&(member, p))
+            .get(&(member.learner_id, p))
             .is_some_and(|&taken| push.base_version <= taken)
         {
             return CarryOutcome::Rejected(PushDisposition::Duplicate);
@@ -12531,6 +12537,30 @@ mod tests {
         assert!(carry.queued.is_empty());
     }
 
+    #[test]
+    fn a_gathered_anchor_yields_one_delta_across_reconnects() {
+        // Island 1 answered round t=7 of fragment 1 from anchor 5; the delta
+        // is gathered. The link drops, the learner redials (generation 11),
+        // releases its spent anchor and answers the next pull from the same
+        // anchor 5 before the round's broadcast reached it: refused.
+        let mut carry = LateCarry::new(Some(4), false).unwrap();
+        let (old, new) = (member(1, 10), member(1, 11));
+        carry.note_gathered(old, 1, 5);
+        let again = test_push_for(1, 5);
+        assert!(carry.reuses_anchor(new, &again));
+        assert!(carry.reuses_anchor(old, &again));
+        // From the broadcast it then gets (version 7) it pushes again.
+        assert!(!carry.reuses_anchor(new, &test_push_for(1, 7)));
+        // A restart of that round drops the gathered delta: the anchor is
+        // free for the new attempt, from either connection.
+        carry.release_anchor(old, &test_push_for(1, 5));
+        assert!(!carry.reuses_anchor(new, &again));
+        // A restarted syncer starts with nothing taken: the delta it lost is
+        // merged once when the island answers from that anchor again.
+        let fresh = LateCarry::new(Some(4), false).unwrap();
+        assert!(!fresh.reuses_anchor(new, &again));
+    }
+
     #[test]
     fn a_carried_anchor_yields_one_delta() {
         let mut carry = LateCarry::new(Some(2), false).unwrap();
@@ -12541,8 +12571,14 @@ mod tests {
         assert!(carry.reuses_anchor(island, &push));
         assert!(carry.reuses_anchor(island, &test_push_for(1, 3)));
         assert!(!carry.reuses_anchor(island, &test_push_for(1, 7)));
-        // Another connection or fragment is unaffected.
-        assert!(!carry.reuses_anchor(member(1, 11), &push));
+        // The island's next connection may not spend it again (a redialed
+        // learner answers from the same anchor); another island or
+        // fragment is unaffected.
+        assert!(carry.reuses_anchor(member(1, 11), &push));
+        assert!(!carry.reuses_anchor(member(3, 10), &test_push_for(1, 5)));
+        let mut elsewhere = test_push_for(1, 5);
+        elsewhere.fragment_id += 1;
+        assert!(!carry.reuses_anchor(island, &elsewhere));
         // A late push from an anchor already merged by its own round is not
         // carried; a restart clears the gathered push and frees the anchor.
         let other = member(2, 10);
diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py
index ad464af0..1720e9c5 100644
--- a/tests/test_swift_megatron_sync.py
+++ b/tests/test_swift_megatron_sync.py
@@ -74,6 +74,98 @@ def test_chunks_are_bounded_and_keep_order():
 
 
 # ----------------------------------------------------------------------------- fake Megatron island
+class _LinkClient:
+    """The client surface ``_PushSender`` and ``_ready_pushes`` use: a link
+    generation the test bumps, and a push result it chooses."""
+
+    def __init__(self, result=True):
+        self.connection_generation, self.result, self.pushes, self.pulls = 1, result, [], []
+
+    def push_fragment_parts(self, *head_and_parts):
+        if isinstance(self.result, BaseException):
+            raise self.result
+        self.pushes.append(head_and_parts[:-1])
+        return self.result
+
+    def drain_pulls(self):
+        pulls, self.pulls = self.pulls, []
+        return pulls
+
+
+@pytest.mark.parametrize("outcome", ["not-queued", "torso-lost", "socket-error"])
+def test_push_sender_treats_an_undelivered_push_as_a_drop(outcome):
+    """``push_fragment_parts`` returns False when the link is down (nothing
+    queued) and raises PartialMessageGenerationLost when the connection died
+    mid-message: neither reaches the syncer, so the anchor is released
+    instead of blocking that fragment's later pulls (or failing the island)."""
+
+    from yeto.protocol import PartialMessageGenerationLost
+
+    result = {"not-queued": False,
+              "torso-lost": PartialMessageGenerationLost(connection_generation=1, connection_epoch=1,
+                                                         queued_chunks=1, operation="PUSH_FRAGMENT"),
+              "socket-error": ConnectionResetError("reset")}[outcome]
+    stats = {"pushed": 0, "push_bytes": 0, "dropped_pushes": 0}
+    dropped, sent = [], []
+    sender = sm._PushSender(_LinkClient(result), stats, lambda *a: dropped.append(a), sent.append)
+    sender.submit((2, 9, 1, 7, 40, 3, 300), parts=[b"\0" * 8])
+    assert sender.wait_idle(5)
+    sender.stop()
+    assert dropped == [(2, 7, 9)] and sent == [2]
+    assert stats["dropped_pushes"] == 1 and stats["pushed"] == 0
+    sender.raise_error()  # not an island failure
+
+
+def _bare_sync(client, fragments=2):
+    from yeto.diloco_sync import DiLoCoSyncState
+
+    sync = object.__new__(sm.MegatronDiLoCoSync)
+    sync.client, sync.sender, sync.interval = client, None, 1
+    sync.state = DiLoCoSyncState.create(fragments, track_anchors=False)
+    sync.state.fragment_versions = [5] * fragments
+    sync.anchored = [True] * fragments
+    sync.spent, sync.spent_link = [None] * fragments, [None] * fragments
+    sync.iteration, sync.units_total, sync.stats = 20, 100, {"deferred_pushes": 0}
+    return sync
+
+
+def test_a_spent_anchor_is_released_once_the_island_redials():
+    """A PUSH queued on a connection that later died may never have reached
+    the syncer, and a restarted syncer has lost the round it was carried
+    into; after the redial the next pull of that fragment is answered from
+    the same anchor (the syncer refuses it if it does hold that delta)."""
+
+    from yeto.protocol import PullRequest
+
+    client = _LinkClient()
+    sync = _bare_sync(client)
+    sync.spent[0] = (5, 7)  # pushed from anchor 5 for round 7, on link 1
+    sync._push_sent(0)
+    assert sync.spent_link[0] == 1
+    client.pulls = [PullRequest(0, 9, 1)]  # another round of the same anchor
+    assert sync._ready_pushes() == []  # same link: the syncer has the delta
+    assert sync.state.pending_pulls == [PullRequest(0, 9, 1)]
+    client.connection_generation = 2  # the link was replaced
+    assert [p[:3] for p in sync._ready_pushes()] == [(0, 9, 1)]
+    assert sync.spent == [None, None] and sync.stats["released_anchors"] == 1
+    # A record from the process the island resumed from: released likewise.
+    sync.spent[1], sync.spent_link[1] = (5, 4), sm._PREVIOUS_PROCESS
+    client.pulls = [PullRequest(1, 10, 1)]
+    assert [p[:3] for p in sync._ready_pushes()] == [(1, 10, 1)]
+
+
+def test_a_dropped_push_frees_its_anchor_on_the_same_link():
+    from yeto.protocol import PullRequest
+
+    client = _LinkClient()
+    sync = _bare_sync(client)
+    sync.spent[0] = (5, 7)
+    sync._push_sent(0)
+    sync._push_dropped(0, 5, 7)
+    client.pulls = [PullRequest(0, 9, 1)]
+    assert [p[:3] for p in sync._ready_pushes()] == [(0, 9, 1)]
+
+
 def _full(name: str, shape) -> torch.Tensor:
     """The initial full value of a global tensor (bf16-exact, so INIT is lossless)."""
 
@@ -859,3 +951,77 @@ def test_decoupled_quorum_one_keeps_the_fast_island_merging_while_the_other_stal
     # A late-carried delta counts towards its round's quorum, so the fast
     # island's rounds keep closing while the other island stalls.
     assert merged_during >= 3, merged_during
+
+
+@pytest.mark.timeout(300)
+def test_islands_keep_merging_across_a_syncer_restart_from_its_shutdown_checkpoint(tmp_path):
+    """The syncer is stopped (SIGTERM: it writes its checkpoint and exits)
+    while both islands keep training and answering the pulls they already
+    hold, then comes back with --resume on the same port.  During the outage
+    a PUSH is not queued at all (``push_fragment_parts`` returns False) and is
+    a drop; a PUSH that went out on the old connection holds an anchor whose
+    round the restarted syncer no longer has, and is released once the
+    island redials.  Both islands must reach the restarted syncer's final cut
+    instead of refusing its pulls forever."""
+
+    import signal
+
+    binary = build_syncer()
+    from process_groups import group_popen
+
+    port = free_port()
+    final_state = tmp_path / "final.bin"
+    ckpt = tmp_path / "syncer.ckpt"
+    flags = [str(binary), "--port", str(port), "--learners", "2", "--grace-ms", "200",
+             "--terminal-grace-ms", "2000", "--total-steps", "30", "--final-state", str(final_state),
+             "--checkpoint-path", str(ckpt), "--checkpoint-every", "4", "--shutdown-timeout-s", "5",
+             *DECOUPLED_SYNCER_FLAGS]
+    logs: list[str] = []
+
+    def start(extra):
+        proc = group_popen(flags + extra, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+        logs.append("")
+        return proc
+
+    def read(i, proc):  # the log file written so far
+        logs[i] += proc.stdout.read()
+        return logs[i]
+
+    first = start([])
+    second = None
+    procs = []
+    try:
+        procs += _run(_island_worker, 4, 0, free_port(), port, str(tmp_path), None, {"step_s": 0.05})
+        procs += _run(_island_worker, 4, 1, free_port(), port, str(tmp_path), None, {"step_s": 0.05})
+        deadline = time.monotonic() + 120
+        while read(0, first).count("outer step") < 8:
+            assert time.monotonic() < deadline, "the first syncer never merged 8 rounds"
+            assert first.poll() is None, "the first syncer exited early"
+            time.sleep(0.05)
+        first.send_signal(signal.SIGTERM)
+        assert first.wait(timeout=60) == 0, "SIGTERM did not end in a clean checkpoint"
+        assert ckpt.exists()
+        time.sleep(3.0)  # the outage: both islands keep stepping
+        second = start(["--resume"])
+        _join(procs, 240, tmp_path)
+        assert second.wait(timeout=60) == 0, "the restarted syncer exited nonzero"
+    finally:
+        for p in procs:
+            if p.is_alive():
+                p.kill()
+        for i, proc in enumerate((first, second)):
+            if proc is not None:
+                if proc.poll() is None:
+                    proc.kill()
+                read(i, proc)
+                (tmp_path / f"syncer{i}.log").write_text(logs[i])
+        print(logs[-1][-4000:])
+
+    islands = [torch.load(tmp_path / f"island{i}.pt") for i in range(2)]
+    for name in islands[0]["names"]:
+        assert torch.equal(islands[0]["values"][name], islands[1]["values"][name]), name
+    final = read_final_state(final_state)
+    for fid, frag_names in enumerate(islands[0]["fragments"]):
+        assert torch.equal(torch.cat([islands[0]["values"][n].reshape(-1) for n in frag_names]), final[fid])
+    assert "outer step" in logs[1], "the restarted syncer merged nothing"
+    print(json.dumps([i["stats"] for i in islands]))
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index a26dbfe0..d2bdeea7 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -169,6 +169,9 @@
 # ``yeto.diloco_sync.resume_fragment_mode`` results (that module imports torch).
 RESUME_OVERWRITE, RESUME_ANCHOR, RESUME_SHIFT = "overwrite", "anchor", "shift"
 
+# ``spent_link`` of an anchor spent by the process the island resumed from.
+_PREVIOUS_PROCESS = "previous-process"
+
 _LAYER_MODULE_RE = re.compile(r"(^|\.)decoder\.layers\.(\d+)$")
 _GROUPED_EXPERT_RE = re.compile(r"^(.*?)(weight|bias)(\d+)$")
 
@@ -656,6 +659,11 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
         # for a newer broadcast instead of a gather that would be refused.  A
         # restarted attempt of the same round may reuse it.
         self.spent: list[tuple[int, int] | None] = [None] * self.num_fragments
+        # The syncer link (connection generation) each spent PUSH went out on:
+        # once the link is replaced, an undelivered delta is gone with the old
+        # connection and a restarted syncer lost the rounds it was carried to,
+        # so the guard is released (``_release_spent_after_reconnect``).
+        self.spent_link: list[Any] = [None] * self.num_fragments
         # The torch learner's per-fragment counters and versions, kept alike on
         # every rank (anchors live in ``self.anchors``, sharded, not here).
         self.state = DiLoCoSyncState.create(self.num_fragments, track_anchors=False)
@@ -692,7 +700,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
         if self.rank == 0:
             try:
                 self.client = self._connect(syncer, hmac_key_path, wan_streams, client_factory)
-                self.sender = _PushSender(self.client, self.stats, self._push_dropped)
+                self.sender = _PushSender(self.client, self.stats, self._push_dropped, self._push_sent)
             except Exception as exc:  # raised on every rank below
                 error = f"{type(exc).__name__}: {exc}"
         self._raise_if(self._bcast_object(error))
@@ -832,6 +840,9 @@ def _restore_counters(self, fid: int, mode: str) -> None:
         if mode == RESUME_ANCHOR:
             spent = self.resume_record["spent"][fid]
             self.spent[fid] = None if spent is None else (int(spent[0]), int(spent[1]))
+            # Pushed by the previous process, on a link that is gone: released
+            # at the first boundary like any spent anchor after a reconnect.
+            self.spent_link[fid] = None if spent is None else _PREVIOUS_PROCESS
 
     def _resume_paths(self, resume_dir: str) -> tuple[str, str]:
         return (os.path.join(resume_dir, RESUME_STATE_FILE),
@@ -1310,6 +1321,7 @@ def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str:
             if state.pending_pulls:
                 self.stats["deferred_pushes"] += 1
             return []
+        self._release_spent_after_reconnect()
         ready = take_ready_pulls(state, steps_total=self.iteration, units_total=self.units_total,
                                  anchored=self.anchored.__getitem__, min_steps=self.interval)
         # One fragment per boundary: its PUSH streams in the background while
@@ -1323,6 +1335,38 @@ def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str:
             out.append((int(pull.fragment_id), int(pull.global_step), int(pull.round_attempt), c_steps, c_units))
         return out
 
+    def link_of(self, fid: int) -> Any:
+        """Rank 0: the connection generation fragment ``fid`` is pushed on
+        (None for a client that does not expose it)."""
+        client = self.client
+        of = getattr(client, "connection_generation_of", None)
+        if of is not None:
+            return of(fid)
+        return getattr(client, "connection_generation", None)
+
+    def _release_spent_after_reconnect(self) -> None:
+        """Rank 0, sender idle: free every spent anchor whose PUSH went out
+        on a connection that has since been replaced.
+
+        Without this an island that pushed just before a syncer restart (or
+        whose PUSH was carried into a round the restarted syncer no longer
+        has) refuses every pull of that fragment until a newer broadcast,
+        which no syncer sends once every island waits like this.  Answering
+        from the same anchor again cannot double-count it: a syncer that
+        still holds the delta refuses the second one per island
+        (``anchor_reused``, syncer ``LateCarry``), and one that lost it
+        merges it once."""
+        for fid, spent in enumerate(self.spent):
+            if spent is None or self.spent_link[fid] is None:
+                continue
+            link = self.link_of(fid)
+            if link is not None and link != self.spent_link[fid]:
+                log.info("fragment %d: anchor %d released after a reconnect (push of step %d)",
+                         fid, spent[0], spent[1])
+                self.spent[fid] = None
+                self.spent_link[fid] = None
+                self.stats["released_anchors"] = self.stats.get("released_anchors", 0) + 1
+
     def _pushable(self, pull: Any) -> bool:
         fid = int(pull.fragment_id)
         spent = self.spent[fid]
@@ -1332,6 +1376,11 @@ def _push_dropped(self, fid: int, base: int, step: int) -> None:
         """Sender thread: a dropped PUSH never reached the syncer, so its anchor is not spent."""
         if self.spent[fid] == (base, step):
             self.spent[fid] = None
+            self.spent_link[fid] = None
+
+    def _push_sent(self, fid: int) -> None:
+        """Sender thread, before the PUSH is queued: the link it goes out on."""
+        self.spent_link[fid] = self.link_of(fid)
 
     def _take_updates(self, plan: _Plan) -> dict[int, Any]:
         updates, self._updates = self._updates, {}
@@ -1351,6 +1400,7 @@ def _push(self, fid: int, step: int, attempt: int, c_steps: int, c_units: int,
             return f"PUSH fragment {fid}: {self._export_error}"
         head = (fid, step, attempt, self.state.fragment_versions[fid], self.iteration, c_steps, c_units)
         self.spent[fid] = (self.state.fragment_versions[fid], step)
+        self.spent_link[fid] = None  # set by the sender once it goes out
         if self.wire == DTYPE_Q4:
             # q4 blocks span tensors: the shared encoder takes the whole delta.
             self.sender.submit(head, payload=encode_push_delta(self.client, torch.cat(parts)))
@@ -1456,11 +1506,12 @@ class _PushSender:
     push (the syncer re-requests it after the redial); any other failure is
     re-raised by the next boundary on every rank."""
 
-    def __init__(self, client: Any, stats: dict, on_drop: Callable[[int, int, int], None]) -> None:
+    def __init__(self, client: Any, stats: dict, on_drop: Callable[[int, int, int], None],
+                 on_sent: Callable[[int], None] | None = None) -> None:
         import queue
         import threading
 
-        self.client, self.stats, self.on_drop = client, stats, on_drop
+        self.client, self.stats, self.on_drop, self.on_sent = client, stats, on_drop, on_sent
         self.jobs: queue.Queue = queue.Queue()
         self.idle = threading.Event()
         self.idle.set()
@@ -1488,6 +1539,11 @@ def stop(self) -> None:
         self.thread.join(timeout=5.0)
 
     def _run(self) -> None:
+        from yeto.protocol import PartialMessageGenerationLost
+
+        # A torso on a dead connection generation is dropped by the syncer,
+        # so all three mean the delta was never delivered.
+        undelivered = (ConnectionError, OSError, PartialMessageGenerationLost)
         while True:
             job = self.jobs.get()
             if job is None:
@@ -1496,14 +1552,17 @@ def _run(self) -> None:
             fid = head[0]
             try:
                 if payload is not None:
-                    self.client.push_fragment(*head, payload)
-                    nbytes = len(payload)
-                else:
-                    self.client.push_fragment_parts(*head, parts)
-                    nbytes = sum(len(part) for part in parts)
+                    parts = [payload]  # q4: already wire bytes, sent unchanged
+                nbytes = sum(len(part) for part in parts)
+                if self.on_sent is not None:
+                    self.on_sent(fid)  # before the send: the link it is queued on
+                if self.client.push_fragment_parts(*head, parts) is False:
+                    # Nothing was queued: the link is down (an outage, a syncer
+                    # restart).  The syncer never sees this delta.
+                    raise ConnectionError("syncer link down, nothing queued")
                 self.stats["pushed"] += 1
                 self.stats["push_bytes"] += nbytes
-            except (ConnectionError, OSError) as exc:
+            except undelivered as exc:  # the anchor is not spent
                 self.stats["dropped_pushes"] += 1
                 self.on_drop(fid, head[3], head[1])
                 log.warning("push of fragment %d dropped (%s); the syncer re-requests it", fid, exc)

From 91f002c09416727250cf78396eb33f1260d32234 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 14:21:34 +0800
Subject: [PATCH 05/17] swift-megatron: rebuild the LR schedule of a resume
 that did not load it

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>
---
 tests/test_swift_megatron_plugin.py | 76 +++++++++++++++++++++++++++++
 yeto/swift_megatron/swift_plugin.py | 35 +++++++++++++
 2 files changed, 111 insertions(+)

diff --git a/tests/test_swift_megatron_plugin.py b/tests/test_swift_megatron_plugin.py
index e37aa1de..9ae733af 100644
--- a/tests/test_swift_megatron_plugin.py
+++ b/tests/test_swift_megatron_plugin.py
@@ -297,3 +297,79 @@ def test_step_log_only_on_sync_rank_zero(make, monkeypatch, tmp_path):
     FakeSync.instances[0].rank = 5
     trainer.train()
     assert not path.exists()
+
+
+class _Scheduler:
+    """Megatron-Core ``OptimizerParamScheduler`` in miniature: sample-based
+    ``num_steps``, linear warmup then cosine decay, ``step(increment)`` sets
+    every parameter group's lr (the constructor steps by 0)."""
+
+    def __init__(self, groups, *, max_lr=3e-4, min_lr=3e-5, warmup=5 * 8, decay=140 * 8):
+        self.groups, self.max_lr, self.min_lr, self.warmup, self.decay = groups, max_lr, min_lr, warmup, decay
+        self.num_steps = 0
+        self.step(0)
+
+    def get_lr(self):
+        import math
+
+        if self.warmup > 0 and self.num_steps <= self.warmup:
+            return self.max_lr * self.num_steps / self.warmup
+        ratio = min(1.0, (self.num_steps - self.warmup) / (self.decay - self.warmup))
+        return self.min_lr + (self.max_lr - self.min_lr) * 0.5 * (math.cos(math.pi * ratio) + 1.0)
+
+    def step(self, increment):
+        self.num_steps += increment
+        for group in self.groups:
+            group["lr"] = self.get_lr()
+
+
+def _lr_trainer(iteration, gbs=8, scheduler=None):
+    return types.SimpleNamespace(args=types.SimpleNamespace(global_batch_size=gbs),
+                                 state=types.SimpleNamespace(iteration=iteration),
+                                 opt_param_scheduler=scheduler or _Scheduler([{}]))
+
+
+def test_a_resume_without_scheduler_state_continues_the_lr_schedule_in_place(monkeypatch):
+    """A params-only checkpoint (--no_save_optim) restores the iteration but
+    no scheduler: without the rebuild the resumed lr at iteration i is the
+    uninterrupted one at i - 5 (warmup again)."""
+
+    monkeypatch.delenv("YETO_DILOCO_LR_RESUME", raising=False)
+    reference = _Scheduler([{}])
+    want = {}
+    for it in range(1, 141):
+        reference.step(8)
+        want[it] = reference.groups[0]["lr"]
+    trainer = _lr_trainer(5)
+    assert sp.restore_lr_schedule(trainer) == 40
+    resumed = trainer.opt_param_scheduler
+    assert resumed.groups[0]["lr"] == want[5]
+    for it in range(6, 141):
+        resumed.step(8)
+        assert resumed.groups[0]["lr"] == want[it], it
+    # Without it: shifted by the resumed iteration.
+    shifted = _Scheduler([{}])
+    for it in range(6, 141):
+        shifted.step(8)
+        assert shifted.groups[0]["lr"] == want[it - 5]
+
+
+def test_lr_schedule_is_left_alone_when_loaded_cold_or_off(monkeypatch):
+    loaded = _Scheduler([{}])
+    loaded.step(5 * 8)  # restored from a checkpoint with optimizer state
+    assert sp.restore_lr_schedule(_lr_trainer(5, scheduler=loaded)) is None and loaded.num_steps == 40
+    cold = _lr_trainer(0)
+    assert sp.restore_lr_schedule(cold) is None and cold.opt_param_scheduler.num_steps == 0
+    assert sp.restore_lr_schedule(types.SimpleNamespace(state=types.SimpleNamespace(iteration=5))) is None
+    monkeypatch.setenv("YETO_DILOCO_LR_RESUME", "off")
+    off = _lr_trainer(5)
+    assert sp.restore_lr_schedule(off) is None and off.opt_param_scheduler.num_steps == 0
+
+
+def test_callback_rebuilds_the_lr_schedule_on_attach(make):
+    trainer = FakeTrainer([([1], True)], start_iteration=15)
+    trainer.args.global_batch_size = 384
+    trainer.opt_param_scheduler = _Scheduler([{}], warmup=10 * 384, decay=1000 * 384)
+    cb = sp.YetoDiLoCoCallback(trainer)
+    assert trainer.opt_param_scheduler.num_steps == 15 * 384
+    cb.close()
diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py
index 0f4fb9d5..4c651ffd 100644
--- a/yeto/swift_megatron/swift_plugin.py
+++ b/yeto/swift_megatron/swift_plugin.py
@@ -63,6 +63,9 @@
   YETO_DILOCO_SAVE_ANCHORS        1 = also save every rank's anchor shards
                                   with each checkpoint (~2 bytes per local
                                   parameter per rank for bf16); default 0
+  YETO_DILOCO_LR_RESUME           auto (default) | off: rebuild the learning
+                                  rate schedule of a resume that did not load
+                                  it (``restore_lr_schedule``)
 
 Resume: every checkpoint also gets the sync's record (``on_save`` ->
 ``MegatronDiLoCoSync.save_state``): each fragment's anchor version and
@@ -289,6 +292,37 @@ def save_dir_for(args, iteration: int, *call_args, **call_kwargs) -> Optional[st
     return os.path.join(root, f"checkpoint-{iteration}") if root else None
 
 
+def restore_lr_schedule(trainer) -> Optional[int]:
+    """Put a resumed run's LR schedule where an uninterrupted run would be.
+
+    ms-swift loads the ``OptimizerParamScheduler`` state only together with
+    the optimizer (``load_mcore_checkpoint``: ``not finetune and not
+    no_load_optim``), and a checkpoint saved with ``--no_save_optim`` holds
+    no scheduler state at all.  So an island that saves parameters only and
+    resumes with ``--no_load_optim true --finetune false`` gets back its
+    iteration and data position but restarts the schedule at step 0: warmup
+    again, the whole decay shifted by the resumed iteration, silently.
+
+    The scheduler counts samples (ms-swift steps it by ``global_batch_size``
+    per successful update), so a fresh scheduler (``num_steps == 0``) at a
+    resumed iteration ``i > 0`` is stepped by ``i * global_batch_size`` here,
+    which sets every parameter group's lr and weight decay the way ``i``
+    uninterrupted updates would.  A scheduler restored from the checkpoint
+    (``num_steps > 0``) and a cold start are left alone.  Returns the samples
+    stepped, or None.  Runs on every rank (each holds its own optimizer)."""
+    if (os.environ.get(ENV_PREFIX + "LR_RESUME", "auto").strip().lower() or "auto") == "off":
+        return None
+    scheduler = getattr(trainer, "opt_param_scheduler", None)
+    iteration = int(getattr(trainer.state, "iteration", 0) or 0)
+    if scheduler is None or iteration <= 0 or int(getattr(scheduler, "num_steps", 0) or 0) != 0:
+        return None
+    samples = iteration * int(trainer.args.global_batch_size)
+    scheduler.step(increment=samples)
+    logger.warning("yeto DiLoCo: the checkpoint of iteration %d restored no LR schedule; stepped it by %d "
+                   "samples to resume it in place", iteration, samples)
+    return samples
+
+
 def _model_chunks(trainer) -> list:
     # Pass the DDP-wrapped chunks: the sync core needs the grad/param buffers
     # and the DistributedOptimizer's mapping, both keyed on the wrapped models.
@@ -332,6 +366,7 @@ def __init__(self, trainer):
             logger.warning("yeto DiLoCo callback disabled via %sDISABLE", ENV_PREFIX)
             return
         self.config = DiLoCoHookConfig.from_env()
+        restore_lr_schedule(trainer)
         self._wrap_train_step()
         # Build here (not in on_train_begin) so a bad config fails before data
         # loading; the trainer has already loaded the resume checkpoint, which

From f1066d0a21457bc1decd6d40b0bf703f2cff4001 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 14:29:34 +0800
Subject: [PATCH 06/17] swift-megatron: rebase a running island onto a redial
 catch-up, take 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>
---
 tests/test_diloco_sync.py         |  17 ++++
 tests/test_reconnect.py           |  53 ++++++++++
 tests/test_swift_megatron_sync.py | 154 +++++++++++++++++++++++++++---
 yeto/diloco_sync.py               |  30 ++++++
 yeto/protocol.py                  |  98 ++++++++++++-------
 yeto/swift_megatron/sync.py       |  97 ++++++++++++++++---
 6 files changed, 390 insertions(+), 59 deletions(-)

diff --git a/tests/test_diloco_sync.py b/tests/test_diloco_sync.py
index 987cb1f8..684416a9 100644
--- a/tests/test_diloco_sync.py
+++ b/tests/test_diloco_sync.py
@@ -399,3 +399,20 @@ def test_resume_record_round_trips_and_picks_the_install_mode():
     assert fresh.steps_at_reset == [15, 12, 15] and fresh.units_at_reset == [0, 500, 0]
     new, local, old = torch.tensor([1.0, 2.0]), torch.tensor([0.5, 0.75]), torch.tensor([0.25, 1.0])
     assert torch.equal(shifted_resume_value(new, local, old), torch.tensor([1.25, 1.75]))
+
+
+def test_live_install_mode_keeps_a_running_islands_progress_through_a_catch_up():
+    from yeto.diloco_sync import LIVE_CONFIRM, LIVE_REBASE, live_install_mode
+
+    def mode(version, catch_up, pushed_from=None, anchored=True):
+        return live_install_mode(anchored=anchored, current_version=5, version=version,
+                                 catch_up=catch_up, pushed_from=pushed_from)
+
+    assert mode(9, False, anchored=False) is None and mode(9, True, anchored=False) is None
+    assert mode(5, False) == LIVE_CONFIRM  # a replay of the anchor
+    assert mode(9, False) is None  # an ordinary merge broadcast
+    assert mode(9, True) == LIVE_REBASE  # a merge missed while away, not pushed from this anchor
+    assert mode(9, True, pushed_from=5) is None  # its delta may be in it
+    assert mode(9, True, pushed_from=3) == LIVE_REBASE
+    assert mode(3, True, pushed_from=5) == LIVE_REBASE  # a rewound syncer
+    assert mode(5, True) == LIVE_REBASE  # same number, other values (the client compared the bytes)
diff --git a/tests/test_reconnect.py b/tests/test_reconnect.py
index ce9710d2..5c36d091 100644
--- a/tests/test_reconnect.py
+++ b/tests/test_reconnect.py
@@ -440,3 +440,56 @@ def test_lease_expired_error_redials_with_new_generation_even_without_reconnect_
     finally:
         client.close()
         listener.close()
+
+
+def _redial(client):
+    """What _connect_group does for a redial: every fragment awaits its catch-up."""
+    with client._bcast_lock:
+        client._catch_up_pending = [True] * client.layout.num_fragments
+
+
+def test_first_broadcast_after_a_redial_is_flagged_and_a_rewound_one_is_taken_on_request():
+    layout = make_layout()
+    v2 = struct.pack("<4f", 2.0, 2.0, 2.0, 2.0)
+    v1 = struct.pack("<4f", 1.0, 1.0, 1.0, 1.0)
+    for accept in (False, True):
+        client = SyncerClient(("127.0.0.1", 1), LEARNER_ID, layout, dtype=DTYPE_F32, num_streams=0,
+                              accept_rewound_catch_up=accept)
+        client._gen = 1
+        client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 2) + v2)
+        assert [(u.version, u.catch_up) for u in client.drain_updates()] == [(2, False)]
+        _redial(client)
+        # The syncer resumed from an older checkpoint: its catch-up is version 1.
+        client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 1) + v1)
+        got = [(u.version, u.data, u.catch_up) for u in client.drain_updates()]
+        assert got == ([(1, v1, True)] if accept else [])
+        # Only the first broadcast after the redial may go backwards.
+        client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 0) + v1)
+        assert client.drain_updates() == []
+        # A replay of the version held is dropped (the island keeps its values).
+        _redial(client)
+        held = (1, v1) if accept else (2, v2)
+        client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, held[0]) + held[1])
+        assert client.drain_updates() == []
+        # A newer one is the catch-up of a merge the island missed.
+        _redial(client)
+        client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 5) + v2)
+        assert [(u.version, u.catch_up) for u in client.drain_updates()] == [(5, True)]
+        client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 6) + v1)
+        assert [(u.version, u.catch_up) for u in client.drain_updates()] == [(6, False)]
+        assert client.rewound_catch_ups == int(accept)
+
+
+def test_a_rewound_catch_up_with_the_same_version_and_other_bytes_is_taken_on_request():
+    layout = make_layout()
+    a, b = struct.pack("<4f", 2.0, 2.0, 2.0, 2.0), struct.pack("<4f", 3.0, 3.0, 3.0, 3.0)
+    client = SyncerClient(("127.0.0.1", 1), LEARNER_ID, layout, dtype=DTYPE_F32, num_streams=0,
+                          accept_rewound_catch_up=True)
+    client._gen = 1
+    client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 3) + a)
+    client.drain_updates()
+    with pytest.raises(ProtocolError, match="conflicting BCAST_FRAGMENT"):  # not after a redial
+        client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 3) + b)
+    _redial(client)
+    client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 3) + b)
+    assert [(u.version, u.data, u.catch_up) for u in client.drain_updates()] == [(3, b, True)]
diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py
index 1720e9c5..8bfdc1e0 100644
--- a/tests/test_swift_megatron_sync.py
+++ b/tests/test_swift_megatron_sync.py
@@ -124,7 +124,7 @@ def _bare_sync(client, fragments=2):
     sync.state = DiLoCoSyncState.create(fragments, track_anchors=False)
     sync.state.fragment_versions = [5] * fragments
     sync.anchored = [True] * fragments
-    sync.spent, sync.spent_link = [None] * fragments, [None] * fragments
+    sync.spent, sync.spent_link, sync.pushed_from = [None] * fragments, [None] * fragments, [None] * fragments
     sync.iteration, sync.units_total, sync.stats = 20, 100, {"deferred_pushes": 0}
     return sync
 
@@ -738,6 +738,121 @@ def test_resumed_island_keeps_its_checkpoint_through_the_rejoin_catch_up(tmp_pat
     _join(procs, 200, tmp_path)
 
 
+def _catch_up_worker(rank, world, port, out_dir):
+    _guard(out_dir, f"catch-up-{rank}", lambda: _catch_up_body(rank, world, port, out_dir))
+
+
+def _catch_up_body(rank, world, port, out_dir):
+    """A running island redials and gets the syncer's catch-up: a replay of
+    its anchor changes nothing (bit for bit), a rewound or newer catch-up of
+    a fragment it did not push from keeps its local progress (global_new +
+    (local - anchor)), a newer one of a fragment it pushed from its anchor is
+    installed as an ordinary merge."""
+
+    from yeto.diloco_sync import shifted_resume_value
+    from yeto.protocol import BcastFragment, PullRequest
+    from yeto.tensor_io import pack_tensor
+
+    dist = _init_dist(rank, world, port)
+    tp, dp = rank % 2, rank // 2
+    tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])]
+    dp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])]
+    singles = [dist.new_group([r]) for r in range(world)]
+    ctx = sm.ParallelContext(rank, world, tp_groups[dp], tp, 2, singles[rank], 0, 1)
+    chunk = _chunk(tp=tp, pp=0, tp_size=2, pp_size=1)
+    trainable = [p for p in chunk.parameters() if p.requires_grad]
+    opt = _DistOpt(trainable, dp_rank=dp, dp_size=2, dp_group=dp_groups[tp])
+    full_shapes = _global_shapes(2, 1)
+    old_g = {n: (_full(n, s) + 0.125).to(torch.bfloat16).float() for n, s in full_shapes.items()}
+    new_g = {n: (_full(n, s) - 0.25).to(torch.bfloat16).float() for n, s in full_shapes.items()}
+
+    def factory(addrs, learner_id, layout, wire, streams, **kwargs):
+        assert kwargs["accept_rewound_catch_up"] is True
+        client = _ScriptedClient(layout, script)
+        script[("layout",)] = layout
+        return client
+
+    def bcast(fid, version, values, catch_up):
+        layout = script[("layout",)]
+        return BcastFragment(fid, version, pack_tensor(_flat(layout, fid, values), 2), catch_up=catch_up)
+
+    plan = {1: [(fid, 4, old_g, False) for fid in range(3)],  # the initial cut
+            3: "pull",                                         # iteration 2: a pull of fragment 2
+            5: [(0, 4, old_g, False), (1, 2, new_g, True), (2, 9, new_g, True)],
+            6: [(0, 9, new_g, True)]}
+
+    class _Script(dict):  # call -> scripted pulls / broadcasts (packed once the layout is known)
+        def get(self, key, default=None):
+            kind, call = key
+            entry = plan.get(call)
+            if kind == "pull":
+                return [PullRequest(2, 7, 1)] if entry == "pull" else []
+            if isinstance(entry, list):
+                return [bcast(*row) for row in entry]
+            return []
+
+    script = _Script()
+    sync = sm.MegatronDiLoCoSync([chunk], opt, syncer="tcp://127.0.0.1:1234", learner_id=1, num_learners=2,
+                                 fragments=3, sync_interval_steps=1, chunk_numel=40, parallel_context=ctx,
+                                 client_factory=factory, initial_cut_timeout=60)
+    fid_of = {name: fid for fid, frag in enumerate(sync.layout.fragments) for name, _ in frag.tensors}
+    for it in (1, 2, 3):
+        opt.step({id(p): _local_update(p, 1, it, tp, 2) for p in trainable})
+        sync.after_optimizer_step(it, 10)
+    if rank == 0:
+        assert sync.sender.wait_idle(10) and sync.stats["pushed"] == 1 and sync.pushed_from == [None, None, 4]
+    master = {id(p): opt.master(p).clone() for p in trainable if id(p) in opt.owned}
+    param = {id(p): p.detach().clone() for p in trainable}
+
+    def shard(full, p):
+        if getattr(p, "tensor_model_parallel", False):
+            size = p._full_shape[p.partition_dim] // 2
+            full = full.narrow(p.partition_dim, tp * size, size)
+        return full
+
+    def owned(p, full):
+        s, e, _ = opt.owned[id(p)]
+        return shard(full, p).reshape(-1)[s:e]
+
+    sync.after_optimizer_step(4, 10)  # call 5: no local step since the snapshot
+    for p in trainable:
+        fid = fid_of[p._full_name]
+        if fid == 0:  # a replay of its anchor: untouched
+            assert torch.equal(p.detach(), param[id(p)]), p._full_name
+            if id(p) in opt.owned:
+                assert torch.equal(opt.master(p), master[id(p)])
+            assert torch.equal(sync.anchors[id(p)].float(), shard(old_g[p._full_name], p))
+            continue
+        assert torch.equal(sync.anchors[id(p)].float(), shard(new_g[p._full_name], p)), p._full_name
+        if id(p) not in opt.owned:
+            continue
+        if fid == 1:  # rewound to version 2: rebased exactly
+            want = shifted_resume_value(owned(p, new_g[p._full_name]), master[id(p)], owned(p, old_g[p._full_name]))
+            assert not torch.equal(want, owned(p, new_g[p._full_name]))
+        else:  # newer, pushed from its anchor: the ordinary install
+            want = owned(p, new_g[p._full_name])
+        assert torch.equal(opt.master(p), want), (fid, p._full_name)
+        assert torch.equal(p.detach().reshape(-1)[opt.owned[id(p)][0]:opt.owned[id(p)][1]], want.to(torch.bfloat16))
+    assert sync.state.fragment_versions == [4, 2, 9]
+    assert sync.state.steps_at_reset[1] == 0 and sync.state.steps_at_reset[2] == 4, sync.state.steps_at_reset
+    before = {id(p): opt.master(p).clone() for p in trainable if id(p) in opt.owned and fid_of[p._full_name] == 0}
+    sync.after_optimizer_step(5, 10)  # call 6: a newer catch-up of fragment 0, never pushed: rebased
+    for p in trainable:
+        if fid_of[p._full_name] == 0 and id(p) in opt.owned:
+            want = shifted_resume_value(owned(p, new_g[p._full_name]), before[id(p)], owned(p, old_g[p._full_name]))
+            assert torch.equal(opt.master(p), want), p._full_name
+    if rank == 0:
+        assert sync.stats["confirmed"] == 1 and sync.stats["rebased"] == 2, sync.stats
+    sync.close()
+    dist.destroy_process_group()
+
+
+@pytest.mark.timeout(240)
+def test_running_island_keeps_its_progress_through_a_catch_up(tmp_path):
+    procs = _run(_catch_up_worker, 4, free_port(), str(tmp_path))
+    _join(procs, 200, tmp_path)
+
+
 # ----------------------------------------------------------------------------- two islands, real syncer (TP2 x PP2)
 def build_syncer() -> Path:
     binary = ROOT / "syncer/target/debug/yeto-syncer"
@@ -954,15 +1069,19 @@ def test_decoupled_quorum_one_keeps_the_fast_island_merging_while_the_other_stal
 
 
 @pytest.mark.timeout(300)
-def test_islands_keep_merging_across_a_syncer_restart_from_its_shutdown_checkpoint(tmp_path):
-    """The syncer is stopped (SIGTERM: it writes its checkpoint and exits)
-    while both islands keep training and answering the pulls they already
-    hold, then comes back with --resume on the same port.  During the outage
-    a PUSH is not queued at all (``push_fragment_parts`` returns False) and is
-    a drop; a PUSH that went out on the old connection holds an anchor whose
-    round the restarted syncer no longer has, and is released once the
-    island redials.  Both islands must reach the restarted syncer's final cut
-    instead of refusing its pulls forever."""
+@pytest.mark.parametrize("stop", ["sigterm", "sigkill"])
+def test_islands_keep_merging_across_a_syncer_restart(tmp_path, stop):
+    """The syncer is stopped while both islands keep training and answering
+    the pulls they already hold, then comes back with --resume on the same
+    port.  ``sigterm``: it writes its checkpoint and exits.  During the
+    outage a PUSH is not queued at all (``push_fragment_parts`` returns
+    False) and is a drop; a PUSH that went out on the old connection holds an
+    anchor whose round the restarted syncer no longer has, and is released
+    once the island redials.  ``sigkill``: a crash; the syncer resumes from
+    its last periodic checkpoint, older than what the islands hold, and its
+    catch-up rewinds them (rebased onto it) instead of being dropped while
+    every push from the lost versions is refused as a future base.  Both
+    islands must reach the restarted syncer's final cut."""
 
     import signal
 
@@ -994,12 +1113,17 @@ def read(i, proc):  # the log file written so far
         procs += _run(_island_worker, 4, 0, free_port(), port, str(tmp_path), None, {"step_s": 0.05})
         procs += _run(_island_worker, 4, 1, free_port(), port, str(tmp_path), None, {"step_s": 0.05})
         deadline = time.monotonic() + 120
-        while read(0, first).count("outer step") < 8:
-            assert time.monotonic() < deadline, "the first syncer never merged 8 rounds"
+        # sigkill: past a periodic checkpoint (every 4 rounds), not on one.
+        while read(0, first).count("outer step") < (8 if stop == "sigterm" else 10):
+            assert time.monotonic() < deadline, "the first syncer never merged enough rounds"
             assert first.poll() is None, "the first syncer exited early"
             time.sleep(0.05)
-        first.send_signal(signal.SIGTERM)
-        assert first.wait(timeout=60) == 0, "SIGTERM did not end in a clean checkpoint"
+        if stop == "sigterm":
+            first.send_signal(signal.SIGTERM)
+            assert first.wait(timeout=60) == 0, "SIGTERM did not end in a clean checkpoint"
+        else:
+            first.kill()
+            first.wait(timeout=60)
         assert ckpt.exists()
         time.sleep(3.0)  # the outage: both islands keep stepping
         second = start(["--resume"])
@@ -1024,4 +1148,6 @@ def read(i, proc):  # the log file written so far
     for fid, frag_names in enumerate(islands[0]["fragments"]):
         assert torch.equal(torch.cat([islands[0]["values"][n].reshape(-1) for n in frag_names]), final[fid])
     assert "outer step" in logs[1], "the restarted syncer merged nothing"
+    if stop == "sigkill":
+        assert sum(i["stats"].get("rebased", 0) for i in islands) > 0, "no island was rewound"
     print(json.dumps([i["stats"] for i in islands]))
diff --git a/yeto/diloco_sync.py b/yeto/diloco_sync.py
index 3388e76a..db8cb509 100644
--- a/yeto/diloco_sync.py
+++ b/yeto/diloco_sync.py
@@ -194,6 +194,36 @@ def resume_fragment_mode(record: dict | None, fid: int, version: int, *, have_an
     return RESUME_SHIFT if have_anchor else RESUME_OVERWRITE
 
 
+# What a running island does with a broadcast of a fragment it already
+# holds an anchor of (``live_install_mode``); None = the ordinary install.
+LIVE_CONFIRM = "confirm"  # a replay of its anchor: nothing to install
+LIVE_REBASE = "rebase"    # a catch-up: global_new + (local - anchor)
+
+
+def live_install_mode(*, anchored: bool, current_version: int, version: int, catch_up: bool,
+                      pushed_from: int | None) -> str | None:
+    """How a running island installs a broadcast of a fragment.
+
+    A broadcast of the version the island is anchored at is a replay of
+    that anchor (``LIVE_CONFIRM``): the local parameters stay bit for bit.
+    The first broadcast after a redial (``catch_up``) is the syncer's
+    current value, not the merge of a round the island answered: the
+    island's progress since its anchor is moved onto it (``LIVE_REBASE``),
+    also when the version went backwards (a syncer resumed from an older
+    checkpoint) or kept its number with other values.  The exception is a
+    newer catch-up of a fragment the island pushed from its current anchor
+    (``pushed_from``): that delta may already be merged into it, so it is
+    installed like any merge broadcast.  None: the ordinary install."""
+    if not anchored:
+        return None
+    if not catch_up:
+        return LIVE_CONFIRM if int(version) == int(current_version) else None
+    if int(version) > int(current_version) and pushed_from is not None \
+            and int(pushed_from) == int(current_version):
+        return None
+    return LIVE_REBASE
+
+
 def shifted_resume_value(global_new: torch.Tensor, local: torch.Tensor, anchor_old: torch.Tensor) -> torch.Tensor:
     """``global_new + (local - anchor_old)`` in FP32: the island's progress since
     its checkpointed anchor, moved onto the newer global."""
diff --git a/yeto/protocol.py b/yeto/protocol.py
index 87d17989..34edcef2 100644
--- a/yeto/protocol.py
+++ b/yeto/protocol.py
@@ -1010,6 +1010,10 @@ class BcastFragment:
     payload_hash: bytes | None = field(default=None, compare=False)
     stored: bool = field(default=False, compare=False)
     discard: Callable[[], None] | None = field(default=None, repr=False, compare=False)
+    # The first broadcast of this fragment after a redial: the syncer's
+    # current value sent as the re-HELLO catch-up (or a newer one), not
+    # necessarily a merge of a round this island answered.
+    catch_up: bool = field(default=False, compare=False)
 
 
 @dataclass
@@ -1172,8 +1176,18 @@ def __init__(
         recover: bool = False,
         owned_fragments: range | None = None,
         expert_tokens: bool = False,
+        accept_rewound_catch_up: bool = False,
     ):
-        """``syncer_epoch``: the syncer incarnation this island was launched
+        """``accept_rewound_catch_up``: after a redial, take the first
+        broadcast of a fragment even when its version is older than (or equal
+        to, with other bytes) one already received: a syncer that crashed and
+        resumed from an older checkpoint has rewound, and its catch-up is the
+        only global it will merge pushes against (the island rebases onto it;
+        otherwise every push from the lost version is refused as a future
+        base).  Off: such a broadcast is dropped (older) or a protocol error
+        (same version), the historical behaviour.
+
+        ``syncer_epoch``: the syncer incarnation this island was launched
         for; a syncer with a newer epoch refuses it (not retried). ``recover``:
         this island lost its trainer state; every HELLO asks for peer recovery
         until the syncer sends RECOVERY_DONE (see ``yeto.rl.peer_recovery``).
@@ -1305,6 +1319,10 @@ def __init__(
             None
         ] * layout.num_fragments
         self._bcast_lock = threading.Lock()
+        # Set for every fragment on each redial; cleared by its first BCAST.
+        self._catch_up_pending = [False] * layout.num_fragments
+        self.accept_rewound_catch_up = bool(accept_rewound_catch_up)
+        self.rewound_catch_ups = 0
         # Raw globals are retained independently of the drain queue so a
         # final fragment that outruns its control-stream manifest can be
         # re-applied exactly after the manifest arrives.
@@ -1438,6 +1456,9 @@ def _connect_group(self, patient: bool) -> None:
                 _close_socket(s)
             raise
         with self._lock:
+            if self.connection_generation != 0:  # a redial: the syncer catches us up
+                with self._bcast_lock:
+                    self._catch_up_pending = [True] * self.layout.num_fragments
             self._gen += 1
             gen = self._gen
             self.connection_generation = connection_generation
@@ -2792,6 +2813,33 @@ def _dispatch(
             # then publish a stale terminal manifest or shutdown afterward.
             self._dispatch_live(gen, msg_type, payload)
 
+    def _admit_bcast_locked(self, fid: int, version: int, digest: bytes) -> bool | None:
+        """Under ``_bcast_lock``: whether a BCAST of ``fid`` is queued (its
+        ``catch_up`` flag) or dropped (None).  Versions only move forward;
+        a replay of the version already held (same bytes) is dropped, other
+        bytes under that version are a protocol error, except for the first
+        BCAST after a redial with ``accept_rewound_catch_up`` (a rewound
+        syncer's catch-up, taken as the fragment's new global)."""
+        first = self._catch_up_pending[fid]
+        self._catch_up_pending[fid] = False
+        seen = self._bcast_seen[fid]
+        if seen is not None:
+            seen_version, seen_digest = seen
+            if version == seen_version and digest == seen_digest:
+                return None
+            if version <= seen_version:
+                if not (first and self.accept_rewound_catch_up and self._downlink is None):
+                    if version < seen_version:
+                        return None
+                    raise ProtocolError(
+                        f"conflicting BCAST_FRAGMENT payloads for fragment {fid} version {version}"
+                    )
+                self.rewound_catch_ups += 1
+                _connect_log(f"fragment {fid}: the syncer rewound from version {seen_version} to "
+                             f"{version} (resumed from an older checkpoint); taking its catch-up")
+        self._bcast_seen[fid] = (version, digest)
+        return first
+
     def _dispatch_streamed(
         self,
         gen: int,
@@ -2811,25 +2859,16 @@ def _dispatch_streamed_live(self, inbound: _StreamedInboundFragment) -> None:
         digest = committed.payload_hash
         if inbound.msg_type == MSG_BCAST_FRAGMENT:
             with self._bcast_lock:
-                seen = self._bcast_seen[fragment_id]
-                if seen is not None:
-                    seen_version, seen_digest = seen
-                    if version < seen_version:
-                        if committed.discard is not None:
-                            committed.discard()
-                        return
-                    if version == seen_version:
-                        if digest != seen_digest:
-                            if committed.discard is not None:
-                                committed.discard()
-                            raise ProtocolError(
-                                "conflicting BCAST_FRAGMENT payloads for fragment "
-                                f"{fragment_id} version {version}"
-                            )
-                        if committed.discard is not None:
-                            committed.discard()
-                        return
-                self._bcast_seen[fragment_id] = (version, digest)
+                try:
+                    catch_up = self._admit_bcast_locked(fragment_id, version, digest)
+                except ProtocolError:
+                    if committed.discard is not None:
+                        committed.discard()
+                    raise
+                if catch_up is None:
+                    if committed.discard is not None:
+                        committed.discard()
+                    return
                 self._bcasts.put(
                     BcastFragment(
                         fragment_id,
@@ -2838,6 +2877,7 @@ def _dispatch_streamed_live(self, inbound: _StreamedInboundFragment) -> None:
                         payload_hash=digest,
                         stored=True,
                         discard=committed.discard,
+                        catch_up=catch_up,
                     )
                 )
             return
@@ -2965,24 +3005,14 @@ def _dispatch_live(
                 data = payload[12:]
             digest = hashlib.sha256(data).digest()
             with self._bcast_lock:
-                seen = self._bcast_seen[fid]
-                if seen is not None:
-                    seen_version, seen_digest = seen
-                    if version < seen_version:
-                        return
-                    if version == seen_version:
-                        if digest != seen_digest:
-                            raise ProtocolError(
-                                f"conflicting BCAST_FRAGMENT payloads for fragment "
-                                f"{fid} version {version}"
-                            )
-                        return
-                self._bcast_seen[fid] = (version, digest)
+                catch_up = self._admit_bcast_locked(fid, version, digest)
+                if catch_up is None:
+                    return
                 # Receiver threads run once per striped socket. Queue the
                 # accepted update under the same lock as the monotonicity
                 # check so two versions cannot validate in order but enqueue
                 # in reverse order after a thread reschedule.
-                self._bcasts.put(BcastFragment(fid, version, data))
+                self._bcasts.put(BcastFragment(fid, version, data, catch_up=catch_up))
         elif msg_type == MSG_FINAL_FRAGMENT:
             if len(payload) < 12:
                 self._set_protocol_error("truncated FINAL_FRAGMENT")
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index d2bdeea7..98554693 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -76,6 +76,19 @@
     param all-gather reads).  The Adam moments are kept, exactly as the torch
     learner keeps its AdamW state when ``apply_fragment`` overwrites a fragment.
 
+Redial (catch-up)
+    After a redial the syncer sends every fragment's current value.  The
+    client drops a replay of a version already held (same bytes), so an
+    unchanged fragment stays bit for bit; a broadcast of the anchored
+    version that does reach the sync is only confirmed.  The first broadcast
+    of a fragment after the redial (``BcastFragment.catch_up``) moves the
+    island's progress since its anchor onto it, ``global_new + (local -
+    anchor)``, unless the island pushed from that anchor (the delta may be in
+    the new version: the ordinary install).  A syncer that crashed and
+    resumed from an older checkpoint rewinds; the client takes its catch-up
+    (``accept_rewound_catch_up``) and the island rebases onto it, instead of
+    keeping the lost versions as bases the syncer refuses forever.
+
 Only global rank 0 talks to the syncer.  Every boundary starts with one
 ``broadcast_object_list`` of rank 0's plan (fragments to apply, fragments to
 push, finalization, shutdown, error), so every rank runs the same collectives
@@ -168,6 +181,8 @@
 RESUME_FORMAT = 1
 # ``yeto.diloco_sync.resume_fragment_mode`` results (that module imports torch).
 RESUME_OVERWRITE, RESUME_ANCHOR, RESUME_SHIFT = "overwrite", "anchor", "shift"
+# ``yeto.diloco_sync.live_install_mode`` results.
+LIVE_CONFIRM, LIVE_REBASE = "confirm", "rebase"
 
 # ``spent_link`` of an anchor spent by the process the island resumed from.
 _PREVIOUS_PROCESS = "previous-process"
@@ -535,6 +550,7 @@ class _Plan:
     """Rank 0's decision for one boundary, broadcast to every rank."""
 
     applies: list[tuple[int, int]] = field(default_factory=list)  # (fragment, version)
+    modes: dict[int, str] = field(default_factory=dict)  # fragment -> LIVE_* (else ordinary)
     finalizing: bool = False
     shutdown: bool = False
     error: str | None = None
@@ -664,6 +680,10 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
         # connection and a restarted syncer lost the rounds it was carried to,
         # so the guard is released (``_release_spent_after_reconnect``).
         self.spent_link: list[Any] = [None] * self.num_fragments
+        # The anchor version of each fragment's last PUSH that may have been
+        # delivered (a drop clears it; a redial does not): a newer catch-up
+        # of that fragment may already hold the delta (``live_install_mode``).
+        self.pushed_from: list[int | None] = [None] * self.num_fragments
         # The torch learner's per-fragment counters and versions, kept alike on
         # every rank (anchors live in ``self.anchors``, sharded, not here).
         self.state = DiLoCoSyncState.create(self.num_fragments, track_anchors=False)
@@ -713,7 +733,9 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
     def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, client_factory):
         from yeto.diloco_sync import connect_syncer
 
-        kwargs = {"session_contract_hash": self.contract}
+        # A syncer resumed from an older checkpoint rewinds; its catch-up is
+        # taken (and the island rebased onto it) instead of being dropped.
+        kwargs = {"session_contract_hash": self.contract, "accept_rewound_catch_up": True}
         key = _read_hmac_key(hmac_key_path)
         if key is not None:
             kwargs["hmac_key"] = key
@@ -794,7 +816,10 @@ def _install_initial_cut(self, timeout: float) -> None:
             updates = self._take_updates(plan) if self.rank == 0 else {}
             for fid, version in plan.applies:
                 data = updates.pop(fid) if self.rank == 0 else None
-                self._apply_first_or_blend(fid, data, version)
+                live = plan.modes.get(fid)
+                if live == LIVE_CONFIRM:
+                    continue
+                self._apply_first_or_blend(fid, data, version, live=live)
                 del data
                 self.state.fragment_versions[fid] = version
                 self.state.global_step = max(self.state.global_step, version)
@@ -810,12 +835,21 @@ def _install_initial_cut(self, timeout: float) -> None:
                      sum(self.anchored), self.num_fragments, time.monotonic() - started)
 
     # ------------------------------------------------------------------ resume
-    def _apply_first_or_blend(self, fid: int, data: Any, version: int) -> str | None:
+    def _apply_first_or_blend(self, fid: int, data: Any, version: int, *, live: str | None = None) -> str | None:
         """Apply a broadcast; the first one of a fragment after a resume by
-        ``resume_fragment_mode``.  Returns that mode (None: not a resumed
-        first install) and records it in ``self.resumed``."""
+        ``resume_fragment_mode``, an anchored one by ``live`` (``_plan``:
+        ``LIVE_REBASE`` or the ordinary install).  Returns the resume mode
+        (None: not a resumed first install) and records it in
+        ``self.resumed``."""
 
         mode = None
+        if live == LIVE_REBASE and self.anchored[fid]:
+            if self.rank == 0:
+                log.info("fragment %d: catch-up version %d onto anchor %d -> rebase (local progress kept)",
+                         fid, version, self.state.fragment_versions[fid])
+            self._apply_fragment(fid, data, self.bulk_torch, blend=False, mode=LIVE_REBASE)
+            self.stats["rebased"] = self.stats.get("rebased", 0) + 1
+            return None
         if self.resume_record is not None and not self.anchored[fid] and fid not in self.resumed:
             from yeto.diloco_sync import resume_fragment_mode
 
@@ -843,6 +877,7 @@ def _restore_counters(self, fid: int, mode: str) -> None:
             # Pushed by the previous process, on a link that is gone: released
             # at the first boundary like any spent anchor after a reconnect.
             self.spent_link[fid] = None if spent is None else _PREVIOUS_PROCESS
+            self.pushed_from[fid] = None if spent is None else int(spent[0])
 
     def _resume_paths(self, resume_dir: str) -> tuple[str, str]:
         return (os.path.join(resume_dir, RESUME_STATE_FILE),
@@ -1134,14 +1169,29 @@ def piece(buffer, name, part_dim, tp_rank, size):
                 n = int(shard.param.numel())
                 pieces.append(incoming[at:at + n].view(shard.local_shape))
                 at += n
-        targets = []
+        targets, previous = [], []
         for shard, whole in zip(local, pieces, strict=True):
             raw = whole.to(dtype=torch.float32).contiguous()
+            previous.append(self.anchors.get(id(shard.param)))
             self.anchors[id(shard.param)] = whole.to(device="cpu", dtype=self.anchor_dtype).contiguous().clone()
             targets.append(raw)
         if mode == RESUME_ANCHOR:
             return  # the broadcast is the checkpoint's own anchor: keep the local values
-        if mode == RESUME_SHIFT:
+        if mode == LIVE_REBASE:
+            # The island's progress since its anchor moves onto the catch-up;
+            # a shard whose global did not change keeps its values bit for bit.
+            from yeto.diloco_sync import shifted_resume_value
+
+            moved = [i for i, (s, old) in enumerate(zip(local, previous, strict=True))
+                     if old is None or not torch.equal(old, self.anchors[id(s.param)])]
+            if not moved:
+                return
+            local = [local[i] for i in moved]
+            masters = self._masters(local)
+            targets = [targets[i] if previous[i] is None else shifted_resume_value(targets[i], m, previous[i])
+                       for i, m in zip(moved, masters, strict=True)]
+            del masters
+        elif mode == RESUME_SHIFT:
             from yeto.diloco_sync import shifted_resume_value
 
             masters = self._masters(local)
@@ -1233,11 +1283,17 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *,
         updates = self._take_updates(plan) if self.rank == 0 else {}
         for fid, version in plan.applies:
             data = updates.get(fid) if self.rank == 0 else None
+            live = plan.modes.get(fid)
+            if live == LIVE_CONFIRM:  # a replay of the anchor: nothing changes
+                self.stats["confirmed"] = self.stats.get("confirmed", 0) + 1
+                self.events.append((iteration, "confirm", fid, version))
+                continue
             started = time.perf_counter()
-            mode = self._apply_first_or_blend(fid, data, version)
+            mode = self._apply_first_or_blend(fid, data, version, live=live)
             self.stats["apply_s"] += time.perf_counter() - started
-            state.steps_at_reset[fid] = iteration
-            state.units_at_reset[fid] = self.units_total
+            if live != LIVE_REBASE:  # a rebased fragment keeps its local progress and clock
+                state.steps_at_reset[fid] = iteration
+                state.units_at_reset[fid] = self.units_total
             if mode is not None:
                 self._restore_counters(fid, mode)
             state.fragment_versions[fid] = version
@@ -1290,19 +1346,35 @@ def _plan(self) -> _Plan:
             if plan.finalizing:
                 return plan
             self._updates = {}
+            catch_up: set[int] = set()
             for update in client.drain_updates():
-                # A later broadcast of the same fragment supersedes an earlier one.
+                # A later broadcast of the same fragment supersedes an earlier
+                # one; a catch-up among them is still the syncer's current value.
                 self._updates[update.fragment_id] = update
+                if getattr(update, "catch_up", False):
+                    catch_up.add(update.fragment_id)
             for fid, update in self._updates.items():
                 # Checked here, before the plan goes out: a bad payload found
                 # while applying would fail rank 0 inside the chunk broadcasts.
                 self._check_payload(fid, update.data, self.bulk_torch)
             plan.applies = sorted((fid, int(u.version)) for fid, u in self._updates.items())
+            plan.modes = self._live_modes(plan.applies, catch_up)
             plan.shutdown = bool(client.shutdown.is_set())
         except Exception as exc:
             plan.error = f"{type(exc).__name__}: {exc}"
         return plan
 
+    def _live_modes(self, applies: Sequence[tuple[int, int]], catch_up: set[int]) -> dict[int, str]:
+        from yeto.diloco_sync import live_install_mode
+
+        modes = {}
+        for fid, version in applies:
+            mode = live_install_mode(anchored=self.anchored[fid], current_version=self.state.fragment_versions[fid],
+                                     version=version, catch_up=fid in catch_up, pushed_from=self.pushed_from[fid])
+            if mode is not None:
+                modes[fid] = mode
+        return modes
+
     def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str:
         """Rank 0: the pulls answered now, by the torch learner's rule
         (``take_ready_pulls``) with the H-step gate; one push per fragment
@@ -1377,6 +1449,8 @@ def _push_dropped(self, fid: int, base: int, step: int) -> None:
         if self.spent[fid] == (base, step):
             self.spent[fid] = None
             self.spent_link[fid] = None
+        if self.pushed_from[fid] == base:
+            self.pushed_from[fid] = None
 
     def _push_sent(self, fid: int) -> None:
         """Sender thread, before the PUSH is queued: the link it goes out on."""
@@ -1401,6 +1475,7 @@ def _push(self, fid: int, step: int, attempt: int, c_steps: int, c_units: int,
         head = (fid, step, attempt, self.state.fragment_versions[fid], self.iteration, c_steps, c_units)
         self.spent[fid] = (self.state.fragment_versions[fid], step)
         self.spent_link[fid] = None  # set by the sender once it goes out
+        self.pushed_from[fid] = self.state.fragment_versions[fid]
         if self.wire == DTYPE_Q4:
             # q4 blocks span tensors: the shared encoder takes the whole delta.
             self.sender.submit(head, payload=encode_push_delta(self.client, torch.cat(parts)))

From 4d3692d47628bd803377c4a4b48f848a8f565340 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 14:43:32 +0800
Subject: [PATCH 07/17] protocol: hash a large BCAST in parallel blocks while
 its chunks arrive

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>
---
 tests/test_reconnect.py |  57 +++++++++++
 yeto/protocol.py        | 217 +++++++++++++++++++++++++++++++++-------
 2 files changed, 238 insertions(+), 36 deletions(-)

diff --git a/tests/test_reconnect.py b/tests/test_reconnect.py
index 5c36d091..a097b8c4 100644
--- a/tests/test_reconnect.py
+++ b/tests/test_reconnect.py
@@ -493,3 +493,60 @@ def test_a_rewound_catch_up_with_the_same_version_and_other_bytes_is_taken_on_re
     _redial(client)
     client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 3) + b)
     assert [(u.version, u.data, u.catch_up) for u in client.drain_updates()] == [(3, b, True)]
+
+
+def _chunk_frames(fid, version, data, msg_id=7):
+    """The CHUNK payloads of one BCAST_FRAGMENT, as the syncer stripes them."""
+    from yeto.protocol import _CHUNK_HEAD, _HEADER, CHUNK_SIZE, MAGIC
+
+    inner = _HEADER.pack(MAGIC, MSG_BCAST_FRAGMENT, 12 + len(data)) + struct.pack("<IQ", fid, version) + data
+    return [bytearray(_CHUNK_HEAD.pack(msg_id, len(inner), off) + inner[off:off + CHUNK_SIZE])
+            for off in range(0, len(inner), CHUNK_SIZE)]
+
+
+def test_block_digest_is_computed_while_chunks_arrive_in_any_order():
+    import random
+
+    from yeto.protocol import BCAST_DIGEST_BLOCK, bcast_payload_digest
+
+    numel = (2 * BCAST_DIGEST_BLOCK + 3 * 1024 * 1024 + 4) // 4  # not a whole number of blocks
+    layout = FragmentLayout([Fragment(MERGE_RDA, [("w", numel)])])
+    data = bytes(random.Random(0).randbytes(4 * numel))
+    want = bcast_payload_digest(data)
+    assert want != bcast_payload_digest(data[:-1] + b"\x01")
+    for seed in (1, 2):
+        client = SyncerClient(("127.0.0.1", 1), LEARNER_ID, layout, dtype=DTYPE_F32, num_streams=0)
+        client._gen = 1
+        seen = []
+        client._dispatch = lambda gen, msg_type, payload, digest=None: seen.append((bytes(payload[12:]), digest))
+        frames = _chunk_frames(0, 3, data)
+        random.Random(seed).shuffle(frames)
+        out = [client._reassemble(1, frame) for frame in frames]
+        assert all(o is None for o in out[:-1])
+        client._dispatch(1, *out[-1])
+        assert seen == [(data, want)]
+        client.close()
+
+
+def test_chunked_and_direct_broadcasts_share_the_replay_and_conflict_checks():
+    from yeto.protocol import BCAST_DIGEST_BLOCK
+
+    numel = 2 * BCAST_DIGEST_BLOCK // 4
+    layout = FragmentLayout([Fragment(MERGE_RDA, [("w", numel)])])
+    data = bytes(4 * numel)
+    other = b"\x01" + data[1:]
+    client = SyncerClient(("127.0.0.1", 1), LEARNER_ID, layout, dtype=DTYPE_F32, num_streams=0)
+    client._gen = 1
+    for frame in _chunk_frames(0, 3, data):
+        inner = client._reassemble(1, frame)
+    client._dispatch(1, *inner)
+    assert [u.version for u in client.drain_updates()] == [3]
+    # The same version again, unchunked: a replay, dropped.
+    client._dispatch(1, MSG_BCAST_FRAGMENT, struct.pack("<IQ", 0, 3) + data)
+    assert client.drain_updates() == []
+    # Other bytes under it, chunked: a conflict.
+    for frame in _chunk_frames(0, 3, other, msg_id=8):
+        inner = client._reassemble(1, frame)
+    with pytest.raises(ProtocolError, match="conflicting BCAST_FRAGMENT"):
+        client._dispatch(1, *inner)
+    client.close()
diff --git a/yeto/protocol.py b/yeto/protocol.py
index 34edcef2..08832b08 100644
--- a/yeto/protocol.py
+++ b/yeto/protocol.py
@@ -575,7 +575,7 @@ def write_frame(sock: socket.socket, msg_type: int, payload: bytes) -> None:
     sock.sendall(_HEADER.pack(MAGIC, msg_type, len(payload)) + payload)
 
 
-def _read_exact(sock: socket.socket, n: int) -> bytes:
+def _read_exact(sock: socket.socket, n: int, *, copy: bool = True) -> bytes | bytearray:
     buf = bytearray(n)
     view = memoryview(buf)
     got = 0
@@ -584,14 +584,18 @@ def _read_exact(sock: socket.socket, n: int) -> bytes:
         if r == 0:
             raise ConnectionError("syncer closed connection")
         got += r
-    return bytes(buf)
+    return bytes(buf) if copy else buf
 
 
 def read_frame(
     sock: socket.socket,
     max_payload: int | Callable[[int], int] = MAX_HANDSHAKE_FRAME,
+    *,
+    raw_types: tuple[int, ...] = (),
 ) -> tuple[int, bytes]:
-    """Read one frame after checking its type-aware limit before allocation."""
+    """Read one frame after checking its type-aware limit before allocation.
+    The payload of a type in ``raw_types`` is the receive buffer itself (a
+    bytearray, one copy less), every other payload is bytes."""
     magic, msg_type, length = _HEADER.unpack(_read_exact(sock, _HEADER.size))
     if magic != MAGIC:
         raise ValueError(f"bad magic {magic:#x}")
@@ -600,7 +604,7 @@ def read_frame(
         raise ValueError(
             f"frame type {msg_type} has length {length}, exceeds limit {limit}"
         )
-    return msg_type, _read_exact(sock, length)
+    return msg_type, _read_exact(sock, length, copy=msg_type not in raw_types)
 
 
 HELLO_MAC_DOMAIN = b"yeto-hello-mac-v1\0"
@@ -1052,6 +1056,113 @@ class _StagedInboundChunk:
     token: object | None
 
 
+# Digest of a BCAST_FRAGMENT payload on the in-memory receive path: SHA-256
+# over the SHA-256 of every BCAST_DIGEST_BLOCK bytes, so the blocks of a
+# large payload are hashed in parallel (hashlib releases the GIL) while its
+# chunks still arrive.  It only ever meets digests of this client's own
+# receive path (``_bcast_seen``): a replay of the version held is one with
+# the same digest, another payload under that version a conflict, as with a
+# flat SHA-256.
+BCAST_DIGEST_BLOCK = 32 << 20
+BCAST_DIGEST_THREADS = 4
+_BCAST_DIGEST_DOMAIN = b"yeto-bcast-blocks-sha256-v1\0"
+
+
+def _block_digest_join(nbytes: int, blocks: Sequence[bytes]) -> bytes:
+    return hashlib.sha256(_BCAST_DIGEST_DOMAIN + struct.pack("<Q", nbytes) + b"".join(blocks)).digest()
+
+
+def bcast_payload_digest(data: bytes | bytearray | memoryview) -> bytes:
+    """``_BlockDigest``'s digest of ``data``, computed in one call."""
+    view = memoryview(data).cast("B")
+    blocks = [hashlib.sha256(view[i:i + BCAST_DIGEST_BLOCK]).digest()
+              for i in range(0, view.nbytes, BCAST_DIGEST_BLOCK)]
+    return _block_digest_join(view.nbytes, blocks)
+
+
+class _BlockDigest:
+    """``bcast_payload_digest(buf[start:])``, computed while a chunked
+    message arrives: each block is hashed on the client's digest pool as
+    soon as all its bytes are in."""
+
+    def __init__(self, buf: bytearray, start: int, submit: Callable[[Callable[[], None]], None]) -> None:
+        self.view = memoryview(buf)[start:]
+        self.start = start
+        self.nbytes = self.view.nbytes
+        count = -(-self.nbytes // BCAST_DIGEST_BLOCK)
+        self.missing = [min(BCAST_DIGEST_BLOCK, self.nbytes - i * BCAST_DIGEST_BLOCK) for i in range(count)]
+        self.blocks: list[bytes | None] = [None] * count
+        self.submit = submit
+        self.lock = threading.Lock()
+        self.done = threading.Condition(self.lock)
+        self.pending = 0  # submitted, not yet hashed
+        self.cancelled = False
+
+    def arrived(self, offset: int, end: int) -> None:
+        """``buf[offset:end]`` is in (called once per chunk, under the
+        reassembly lock): queue every block it completes."""
+        lo, hi = max(offset - self.start, 0), max(end - self.start, 0)
+        ready = []
+        with self.lock:
+            for block in range(lo // BCAST_DIGEST_BLOCK, -(-hi // BCAST_DIGEST_BLOCK)):
+                first = block * BCAST_DIGEST_BLOCK
+                self.missing[block] -= min(hi, first + BCAST_DIGEST_BLOCK) - max(lo, first)
+                if self.missing[block] == 0:
+                    ready.append(block)
+            self.pending += len(ready)
+        for block in ready:
+            self.submit(lambda block=block: self._hash(block))
+
+    def _hash(self, block: int) -> None:
+        try:
+            if not self.cancelled:
+                at = block * BCAST_DIGEST_BLOCK
+                self.blocks[block] = hashlib.sha256(self.view[at:at + BCAST_DIGEST_BLOCK]).digest()
+        finally:
+            with self.lock:
+                self.pending -= 1
+                self.done.notify_all()
+
+    def finish(self) -> bytes:
+        """Every byte arrived: wait for the last blocks, return the digest."""
+        with self.lock:
+            while self.pending:
+                self.done.wait()
+        for block, digest in enumerate(self.blocks):  # a block the pool never got (e.g. closed)
+            if digest is None:
+                at = block * BCAST_DIGEST_BLOCK
+                self.blocks[block] = hashlib.sha256(self.view[at:at + BCAST_DIGEST_BLOCK]).digest()
+        self.view.release()
+        return _block_digest_join(self.nbytes, self.blocks)  # type: ignore[arg-type]
+
+    def cancel(self) -> None:
+        self.cancelled = True
+
+
+class _Reassembly:
+    """One chunked inbound message being put together in memory."""
+
+    __slots__ = ("buf", "filled", "ranges", "digest")
+
+    def __init__(self, total: int, digest_start: int | None, submit=None) -> None:
+        self.buf = bytearray(total)
+        self.filled = 0
+        self.ranges: list[tuple[int, int]] = []  # sorted, disjoint
+        self.digest = None if digest_start is None else _BlockDigest(self.buf, digest_start, submit)
+
+    def overlaps(self, offset: int, end: int) -> bool:
+        import bisect
+
+        i = bisect.bisect_left(self.ranges, (offset, end))
+        return (i > 0 and self.ranges[i - 1][1] > offset) or (i < len(self.ranges) and self.ranges[i][0] < end)
+
+    def add(self, offset: int, end: int) -> None:
+        import bisect
+
+        bisect.insort(self.ranges, (offset, end))
+        self.filled += end - offset
+
+
 @dataclass
 class _StreamedReassembly:
     total: int
@@ -1315,7 +1426,9 @@ def __init__(
         self._threads: list[threading.Thread] = []
         self._pulls: queue.Queue[PullRequest] = queue.Queue()
         self._bcasts: queue.Queue[BcastFragment] = queue.Queue()
-        self._bcast_seen: list[tuple[int, bytes] | None] = [
+        # (version, digest, digest kind) of each fragment's newest BCAST: the
+        # in-memory path's block digest or a streamed sink's payload hash.
+        self._bcast_seen: list[tuple[int, bytes, str] | None] = [
             None
         ] * layout.num_fragments
         self._bcast_lock = threading.Lock()
@@ -1332,11 +1445,9 @@ def __init__(
         self.finalizing = threading.Event()
         self.finalized = threading.Event()
         self._final_cond = threading.Condition()
-        self._reasm: dict[
-            int,
-            tuple[bytearray, list[int], list[tuple[int, int]]]
-            | _StreamedReassembly,
-        ] = {}
+        self._reasm: dict[int, _Reassembly | _StreamedReassembly] = {}
+        # Hashes the blocks of large inbound BCAST payloads (_BlockDigest).
+        self._digest_pool: Any = None
         self._reasm_lock = threading.Lock()
         self._reasm_touched: dict[int, float] = {}  # msg_id -> last chunk (monotonic)
         self._inbound_chunk_sink: object | None = None
@@ -1559,6 +1670,10 @@ def _supervise(self) -> None:
     def close(self) -> None:
         self._closed.set()
         self._failure.set()  # unblock the supervisor so it exits instead of redialing
+        with self._reasm_lock:
+            pool, self._digest_pool = self._digest_pool, None
+        if pool is not None:
+            pool.shutdown(wait=False)  # queued block hashes still run: a finish() may wait on them
         with self._lock:
             queues, socks = list(self._queues), list(self._socks)
         for q in queues:
@@ -2518,7 +2633,7 @@ def _recv_loop(self, gen: int, sock: socket.socket) -> None:
             sock = _RateLimitedSocket(sock, test_rate)
         try:
             while True:
-                msg_type, payload = read_frame(sock, self._incoming_frame_limit)
+                msg_type, payload = read_frame(sock, self._incoming_frame_limit, raw_types=(MSG_CHUNK,))
                 if msg_type == MSG_CHUNK:
                     inner = self._reassemble(gen, payload)
                     if inner is not None:
@@ -2526,6 +2641,7 @@ def _recv_loop(self, gen: int, sock: socket.socket) -> None:
                             self._dispatch_streamed(gen, inner)
                         else:
                             self._dispatch(gen, *inner)
+                    del inner, payload
                 else:
                     self._dispatch(gen, msg_type, payload)
         except (ValueError, ProtocolError) as e:
@@ -2554,13 +2670,25 @@ def _incoming_frame_limit(self, msg_type: int) -> int:
         except KeyError:
             raise ValueError(f"unexpected syncer message type {msg_type}") from None
 
+    def _submit_digest(self, job: Callable[[], None]) -> None:
+        """Run one block hash of a ``_BlockDigest`` on the digest pool
+        (called under ``_reasm_lock``, which also guards the lazy start)."""
+        if self._digest_pool is None:
+            from concurrent.futures import ThreadPoolExecutor
+
+            self._digest_pool = ThreadPoolExecutor(BCAST_DIGEST_THREADS, thread_name_prefix="yeto-digest")
+        try:
+            self._digest_pool.submit(job)
+        except RuntimeError:  # the client closed its pool
+            job()
+
     def _reassemble(
-        self, gen: int, payload: bytes
-    ) -> tuple[int, memoryview] | _StreamedInboundFragment | None:
+        self, gen: int, payload: bytes | bytearray
+    ) -> tuple[int, memoryview, bytes | None] | _StreamedInboundFragment | None:
         if len(payload) < _CHUNK_HEAD.size:
             raise ValueError("chunk header truncated")
         msg_id, total, offset = _CHUNK_HEAD.unpack_from(payload)
-        data = payload[_CHUNK_HEAD.size :]
+        data = memoryview(payload)[_CHUNK_HEAD.size :]
         if not data:
             raise ValueError("empty chunk")
         if total < _HEADER.size:
@@ -2580,7 +2708,11 @@ def _reassemble(
                 if len(self._reasm) >= MAX_PARTIAL_MESSAGES:
                     raise ValueError("too many partial chunked messages")
                 if self._inbound_chunk_sink is None:
-                    self._reasm[msg_id] = (bytearray(total), [0], [])
+                    # The digest of a large BCAST / FINAL payload (the bytes
+                    # after its 12-byte head) runs while the chunks arrive.
+                    start = _HEADER.size + 12
+                    big = total - start >= 2 * BCAST_DIGEST_BLOCK
+                    self._reasm[msg_id] = _Reassembly(total, start if big else None, self._submit_digest)
                 else:
                     sink = self._inbound_chunk_sink
                     self._reasm[msg_id] = _StreamedReassembly(
@@ -2610,24 +2742,30 @@ def _reassemble(
                     self._reasm.pop(msg_id, None)
                     self._reasm_touched.pop(msg_id, None)
                 return completed
-            buf, filled, ranges = partial
+            buf = partial.buf
             if len(buf) != total:
                 raise ValueError("chunk total changed within one message")
-            if any(offset < old_end and old_start < end for old_start, old_end in ranges):
+            if partial.overlaps(offset, end):
                 raise ValueError("overlapping chunk")
             buf[offset:end] = data
-            filled[0] += len(data)
-            ranges.append((offset, end))
-            if filled[0] < total:
-                return None
-            if filled[0] != total:
-                raise ValueError("chunk byte count exceeds frame length")
-            del self._reasm[msg_id]
-            self._reasm_touched.pop(msg_id, None)
+            partial.add(offset, end)
+            digest = partial.digest
+            if digest is not None:
+                digest.arrived(offset, end)
+            complete = partial.filled >= total
+            if complete:
+                if partial.filled != total:
+                    raise ValueError("chunk byte count exceeds frame length")
+                del self._reasm[msg_id]
+                self._reasm_touched.pop(msg_id, None)
+        if not complete:
+            return None
         magic, msg_type, length = _HEADER.unpack_from(buf)
         if magic != MAGIC or length != total - _HEADER.size:
+            if digest is not None:
+                digest.cancel()
             raise ValueError("corrupt reassembled frame")
-        return msg_type, memoryview(buf)[_HEADER.size :]
+        return msg_type, memoryview(buf)[_HEADER.size :], None if digest is None else digest.finish()
 
     def _stream_chunk_locked(
         self,
@@ -2793,6 +2931,8 @@ def _clear_partial_reassemblies(self) -> None:
             self._reasm.clear()
             self._reasm_touched.clear()
         for partial in partials:
+            if isinstance(partial, _Reassembly) and partial.digest is not None:
+                partial.digest.cancel()
             if isinstance(partial, _StreamedReassembly):
                 try:
                     self._abort_streamed_partial(partial)
@@ -2802,8 +2942,10 @@ def _clear_partial_reassemblies(self) -> None:
                     self._last_err = exc
 
     def _dispatch(
-        self, gen: int, msg_type: int, payload: bytes | memoryview
+        self, gen: int, msg_type: int, payload: bytes | memoryview, digest: bytes | None = None
     ) -> None:
+        """``digest``: SHA-256 of ``payload[12:]`` already computed while a
+        chunked BCAST payload arrived (``_BlockDigest``)."""
         with self._lock:
             if gen != self._gen:
                 return  # late message from a dead group
@@ -2811,21 +2953,23 @@ def _dispatch(
             # atomic with respect to teardown/redial. Otherwise an obsolete
             # receiver could pass the check, lose the race to a reconnect,
             # then publish a stale terminal manifest or shutdown afterward.
-            self._dispatch_live(gen, msg_type, payload)
+            self._dispatch_live(gen, msg_type, payload, digest)
 
-    def _admit_bcast_locked(self, fid: int, version: int, digest: bytes) -> bool | None:
+    def _admit_bcast_locked(self, fid: int, version: int, digest: bytes, kind: str = "blocks") -> bool | None:
         """Under ``_bcast_lock``: whether a BCAST of ``fid`` is queued (its
         ``catch_up`` flag) or dropped (None).  Versions only move forward;
         a replay of the version already held (same bytes) is dropped, other
         bytes under that version are a protocol error, except for the first
         BCAST after a redial with ``accept_rewound_catch_up`` (a rewound
-        syncer's catch-up, taken as the fragment's new global)."""
+        syncer's catch-up, taken as the fragment's new global).  Digests of
+        two kinds (a sink installed in between) cannot be compared: the same
+        version is then taken as a replay."""
         first = self._catch_up_pending[fid]
         self._catch_up_pending[fid] = False
         seen = self._bcast_seen[fid]
         if seen is not None:
-            seen_version, seen_digest = seen
-            if version == seen_version and digest == seen_digest:
+            seen_version, seen_digest, seen_kind = seen
+            if version == seen_version and (digest == seen_digest or kind != seen_kind):
                 return None
             if version <= seen_version:
                 if not (first and self.accept_rewound_catch_up and self._downlink is None):
@@ -2837,7 +2981,7 @@ def _admit_bcast_locked(self, fid: int, version: int, digest: bytes) -> bool | N
                 self.rewound_catch_ups += 1
                 _connect_log(f"fragment {fid}: the syncer rewound from version {seen_version} to "
                              f"{version} (resumed from an older checkpoint); taking its catch-up")
-        self._bcast_seen[fid] = (version, digest)
+        self._bcast_seen[fid] = (version, digest, kind)
         return first
 
     def _dispatch_streamed(
@@ -2860,7 +3004,7 @@ def _dispatch_streamed_live(self, inbound: _StreamedInboundFragment) -> None:
         if inbound.msg_type == MSG_BCAST_FRAGMENT:
             with self._bcast_lock:
                 try:
-                    catch_up = self._admit_bcast_locked(fragment_id, version, digest)
+                    catch_up = self._admit_bcast_locked(fragment_id, version, digest, "sink")
                 except ProtocolError:
                     if committed.discard is not None:
                         committed.discard()
@@ -2924,7 +3068,7 @@ def _dispatch_streamed_live(self, inbound: _StreamedInboundFragment) -> None:
         )
 
     def _dispatch_live(
-        self, gen: int, msg_type: int, payload: bytes | memoryview
+        self, gen: int, msg_type: int, payload: bytes | memoryview, raw_digest: bytes | None = None
     ) -> None:
         if msg_type == MSG_ERROR:
             message = bytes(payload).decode("utf-8", errors="replace")
@@ -2991,6 +3135,7 @@ def _dispatch_live(
             if fid >= self.layout.num_fragments:
                 raise ValueError(f"BCAST_FRAGMENT for unknown fragment {fid}")
             if self._downlink is not None:
+                raw_digest = None  # the digest is of the decoded values
                 data = self._downlink.decode_whole(
                     fid, version, payload[12:], self.layout.fragments[fid].numel)
             else:
@@ -3003,7 +3148,7 @@ def _dispatch_live(
                         f"expected {expected}"
                     )
                 data = payload[12:]
-            digest = hashlib.sha256(data).digest()
+            digest = raw_digest if raw_digest is not None else bcast_payload_digest(data)
             with self._bcast_lock:
                 catch_up = self._admit_bcast_locked(fid, version, digest)
                 if catch_up is None:

From 3961ce6c392a5da06c28877b2d6c55825dc84c4f Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 15:01:01 +0800
Subject: [PATCH 08/17] swift-megatron: skip INIT for a resumed learner 0,
 stream INIT by chunk, 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>
---
 tests/test_swift_megatron_sync.py   |  66 ++++++++++++++
 yeto/swift_megatron/swift_plugin.py |   7 ++
 yeto/swift_megatron/sync.py         | 130 ++++++++++++++++++++++------
 3 files changed, 177 insertions(+), 26 deletions(-)

diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py
index 8bfdc1e0..00f1e278 100644
--- a/tests/test_swift_megatron_sync.py
+++ b/tests/test_swift_megatron_sync.py
@@ -853,6 +853,72 @@ def test_running_island_keeps_its_progress_through_a_catch_up(tmp_path):
     _join(procs, 200, tmp_path)
 
 
+def _resumed_init_worker(rank, world, port, out_dir):
+    _guard(out_dir, f"resumed-init-{rank}", lambda: _resumed_init_body(rank, world, port, out_dir))
+
+
+def _resumed_init_body(rank, world, port, out_dir):
+    """Learner 0 resumed from a checkpoint with a record: an initialized
+    syncer answers its HELLO with the catch-up, so no INIT_PARAMS is sent;
+    a syncer that sends nothing within the probe gets INIT_PARAMS, streamed
+    chunk by chunk, byte-identical to a cold start's."""
+
+    from yeto.protocol import BcastFragment
+    from yeto.tensor_io import pack_tensor
+
+    dist = _init_dist(rank, world, port)
+    tp, dp = rank % 2, rank // 2
+    tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])]
+    dp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])]
+    singles = [dist.new_group([r]) for r in range(world)]
+    ctx = sm.ParallelContext(rank, world, tp_groups[dp], tp, 2, singles[rank], 0, 1)
+    chunk = _chunk(tp=tp, pp=0, tp_size=2, pp_size=1)
+    trainable = [p for p in chunk.parameters() if p.requires_grad]
+    opt = _DistOpt(trainable, dp_rank=dp, dp_size=2, dp_group=dp_groups[tp])
+    full = {n: (_full(n, s) + 0.125).to(torch.bfloat16).float() for n, s in _global_shapes(2, 1).items()}
+    clients, cut = [], {"calls": ()}
+
+    def factory(addrs, learner_id, layout, wire, streams, **kwargs):
+        script = {("bcast", c): [BcastFragment(fid, 4, pack_tensor(_flat(layout, fid, full), 2)) for fid in range(3)]
+                  for c in cut["calls"]}
+        clients.append(_ScriptedClient(layout, script))
+        return clients[-1]
+
+    def make(resume_dir=None, **kw):
+        return sm.MegatronDiLoCoSync([chunk], opt, syncer="tcp://127.0.0.1:1234", learner_id=0, num_learners=2,
+                                     fragments=3, sync_interval_steps=100, chunk_numel=40, parallel_context=ctx,
+                                     client_factory=factory, resume_dir=resume_dir, resume_iteration=3, **kw)
+
+    cut["calls"] = (1,)
+    first = make(initial_cut_timeout=60)  # a cold start: INIT, then the initial cut
+    ckpt = Path(out_dir, "checkpoint-3")
+    first.save_state(str(ckpt), iteration=3)
+    first.close()
+    cold_init = dict(clients[-1].inits) if rank == 0 else {}
+    if rank == 0:
+        assert sorted(cold_init) == [0, 1, 2]
+    resumed = make(str(ckpt), initial_cut_timeout=60, init_probe_s=30)  # the catch-up arrives: no INIT
+    assert resumed.anchored == [True] * 3 and resumed.resumed == {0: "anchor", 1: "anchor", 2: "anchor"}
+    if rank == 0:
+        assert clients[-1].inits == {}
+    resumed.close()
+    cut["calls"] = ()
+    started = time.monotonic()
+    fresh = make(str(ckpt), initial_cut_timeout=1, init_probe_s=0.5)  # nothing arrives: INIT after the probe
+    assert time.monotonic() - started >= 0.5
+    if rank == 0:  # the island's current masters (the installed cut), streamed by chunk
+        layout = clients[-1].layout
+        assert clients[-1].inits == {fid: bytes(pack_tensor(_flat(layout, fid, full), fresh.bulk)) for fid in range(3)}
+    fresh.close()
+    dist.destroy_process_group()
+
+
+@pytest.mark.timeout(240)
+def test_a_resumed_learner_zero_sends_init_only_when_the_syncer_needs_it(tmp_path):
+    procs = _run(_resumed_init_worker, 4, free_port(), str(tmp_path))
+    _join(procs, 200, tmp_path)
+
+
 # ----------------------------------------------------------------------------- two islands, real syncer (TP2 x PP2)
 def build_syncer() -> Path:
     binary = ROOT / "syncer/target/debug/yeto-syncer"
diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py
index 4c651ffd..900df664 100644
--- a/yeto/swift_megatron/swift_plugin.py
+++ b/yeto/swift_megatron/swift_plugin.py
@@ -63,6 +63,10 @@
   YETO_DILOCO_SAVE_ANCHORS        1 = also save every rank's anchor shards
                                   with each checkpoint (~2 bytes per local
                                   parameter per rank for bf16); default 0
+  YETO_DILOCO_INIT_PROBE_S        resumed learner 0 with an initial-cut wait:
+                                  send INIT_PARAMS only if the syncer sends no
+                                  catch-up within this many seconds (an
+                                  initialized syncer ignores INIT); default 1800
   YETO_DILOCO_LR_RESUME           auto (default) | off: rebuild the learning
                                   rate schedule of a resume that did not load
                                   it (``restore_lr_schedule``)
@@ -129,6 +133,7 @@ class DiLoCoHookConfig:
     initial_cut_timeout: float = 0.0
     resume_dir: Optional[str] = None
     save_anchors: bool = False
+    init_probe_s: float = 1800.0
 
     @classmethod
     def from_env(cls, env=None) -> "DiLoCoHookConfig":
@@ -169,6 +174,7 @@ def get_float(name: str, default: float) -> float:
             initial_cut_timeout=get_float("INITIAL_CUT_TIMEOUT", 0.0),
             resume_dir=get("RESUME_DIR"),
             save_anchors=(get("SAVE_ANCHORS", default="0") or "").strip().lower() in {"1", "true", "yes", "on"},
+            init_probe_s=get_float("INIT_PROBE_S", 1800.0),
         )
         if cfg.initial_cut_timeout < 0:
             raise ValueError(f"{ENV_PREFIX}INITIAL_CUT_TIMEOUT must be >= 0")
@@ -265,6 +271,7 @@ def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig, *,
         initial_cut_timeout=cfg.initial_cut_timeout,
         resume_dir=resume_dir,
         resume_iteration=resume_iteration,
+        init_probe_s=cfg.init_probe_s,
     )
 
 
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index 98554693..36e26cca 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -167,6 +167,21 @@
 
 log = logging.getLogger("swift-megatron-sync")
 
+
+def _ensure_visible_logs() -> None:
+    """ms-swift configures only its own ``swift`` logger, so with an
+    unconfigured root this module's INFO lines (``resume: fragment ...``,
+    per-fragment installs) were dropped and only WARNING reached stderr
+    through logging's last resort.  Give the module a stderr handler at INFO
+    unless the process configured logging itself."""
+    if log.handlers or logging.getLogger().handlers:
+        return
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter("[swift-megatron-sync] %(asctime)s %(levelname)s %(message)s"))
+    log.addHandler(handler)
+    if log.level == logging.NOTSET:
+        log.setLevel(logging.INFO)
+
 # Megatron's tied/untied output projection is an embedding-like tensor for the
 # direct-averaged fragment, as ``lm_head`` is for the torch learner.
 OUTPUT_LAYER_AVG_REGEX = r"(^|\.)output_layer\.weight$"
@@ -553,6 +568,7 @@ class _Plan:
     modes: dict[int, str] = field(default_factory=dict)  # fragment -> LIVE_* (else ordinary)
     finalizing: bool = False
     shutdown: bool = False
+    send_init: bool = False  # initial-cut wait of a resumed learner 0: the syncer needs INIT
     error: str | None = None
 
 
@@ -593,7 +609,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
                  transport: str = "cpu", parallel_context: ParallelContext | None = None,
                  client_factory=None, observer: Callable[[dict], None] | None = None,
                  initial_cut_timeout: float = 0.0, resume_dir: str | None = None,
-                 resume_iteration: int | None = None) -> None:
+                 resume_iteration: int | None = None, init_probe_s: float = 1800.0) -> None:
         import torch
         import torch.distributed as dist
 
@@ -615,6 +631,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
             raise ValueError("the ms-swift Megatron sync talks to one syncer; sharded syncers are not supported")
         if not dist.is_initialized():
             raise SwiftMegatronSyncError("torch.distributed must be initialized (Megatron does this)")
+        _ensure_visible_logs()
         self.ctx = parallel_context or ParallelContext.from_megatron()
         self.rank, self.world = self.ctx.rank, self.ctx.world
         self.chunks = list(model_chunks)
@@ -724,10 +741,17 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
             except Exception as exc:  # raised on every rank below
                 error = f"{type(exc).__name__}: {exc}"
         self._raise_if(self._bcast_object(error))
-        if self.learner_id == 0:
+        # A resumed learner 0 (it holds a record, so the syncer of this session
+        # was initialized) first waits for the syncer's catch-up and sends
+        # INIT_PARAMS (~4 bytes per parameter, 480 GB for 120B) only if nothing
+        # arrives within ``init_probe_s``: an initialized syncer ignores INIT,
+        # and receiving the catch-up while gathering INIT doubled rank 0's RSS.
+        defer_init = self.learner_id == 0 and self.resume_record is not None and initial_cut_timeout > 0
+        if self.learner_id == 0 and not defer_init:
             self._send_init()
         if initial_cut_timeout > 0:
-            self._install_initial_cut(float(initial_cut_timeout))
+            self._install_initial_cut(float(initial_cut_timeout),
+                                      init_probe=float(init_probe_s) if defer_init else None)
 
     # ------------------------------------------------------------------ setup
     def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, client_factory):
@@ -747,42 +771,63 @@ def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, cli
     def _send_init(self) -> None:
         """INIT_PARAMS of every fragment from this island's current masters.
 
-        Rank 0 streams fragment k on a background thread while every rank
-        gathers fragment k + 1, so the island's gather and the link overlap;
-        at most two fragments' wire parts are on rank 0's host."""
+        Streamed chunk by chunk: rank 0 hands each exported chunk to a
+        background sender, which writes it into the fragment's INIT_PARAMS
+        frame and drops it, while every rank gathers the next chunk.  At most
+        a few chunks (``chunk_numel`` elements each) are on rank 0's host
+        instead of whole fragments."""
 
         import queue
         import threading
 
+        end = object()
         jobs: queue.Queue = queue.Queue(maxsize=1)
         failure: list[str] = []
         sender = None
         if self.rank == 0:
+            def chunks_of(pending: queue.Queue, ended: list):
+                while (item := pending.get()) is not end:
+                    if isinstance(item, BaseException):
+                        raise item
+                    yield from item
+                ended.append(True)
+
             def send() -> None:
                 while (job := jobs.get()) is not None:
-                    fid, parts = job
-                    if failure:
-                        continue  # drain: the gathers go on until the error is broadcast
+                    fid, pending = job
+                    ended: list = []
                     try:
-                        self.client.send_init_parts(fid, parts)
+                        if not failure and self.client.send_init_parts(fid, chunks_of(pending, ended)) is False:
+                            failure.append(f"INIT fragment {fid}: syncer link down")
                     except Exception as exc:
                         failure.append(f"INIT fragment {fid}: {type(exc).__name__}: {exc}")
-                    del parts, job
+                    # Drain what was not sent: the gathers go on until the error is broadcast.
+                    while not ended and pending.get() is not end:
+                        pass
+                    del job, pending
 
             sender = threading.Thread(target=send, name="yeto-swift-init", daemon=True)
             sender.start()
         error = None
         try:
             for fid in range(self.num_fragments):
-                parts = self._export_fragment(fid, delta=False)
-                if self.rank == 0 and error is None:
-                    if self._export_error is not None:
+                pending = queue.Queue(maxsize=2) if self.rank == 0 else None
+                if pending is not None:
+                    jobs.put((fid, pending))
+                for parts in self._export_fragment_chunks(fid, delta=False):
+                    if pending is None:
+                        continue
+                    if error is None and self._export_error is not None:
                         error = f"INIT fragment {fid}: {self._export_error}"
-                    elif failure:
+                        failure.append(error)  # the sender skips the fragments after it
+                        pending.put(SwiftMegatronSyncError(error))
+                    elif error is None:
+                        pending.put(parts)
+                    del parts
+                if pending is not None:
+                    pending.put(end)
+                    if error is None and failure:
                         error = failure[0]
-                    else:
-                        jobs.put((fid, parts))
-                del parts
         finally:
             if sender is not None:
                 jobs.put(None)
@@ -793,7 +838,7 @@ def send() -> None:
         if self.rank == 0:
             log.info("sent INIT_PARAMS for %d fragments", self.num_fragments)
 
-    def _install_initial_cut(self, timeout: float) -> None:
+    def _install_initial_cut(self, timeout: float, *, init_probe: float | None = None) -> None:
         """Apply the syncer's first broadcast of every fragment as it arrives,
         before training starts (collective: every rank loops together).
 
@@ -807,12 +852,26 @@ def _install_initial_cut(self, timeout: float) -> None:
         import time
 
         started = time.monotonic()
+        init_pending = init_probe is not None  # resumed learner 0: INIT only if the syncer needs it
         while not all(self.anchored):
             plan = self._plan() if self.rank == 0 else None
-            if self.rank == 0 and time.monotonic() - started > timeout and not plan.applies:
-                plan.shutdown = True  # only ends the wait; the client's flag is unchanged
+            if self.rank == 0:
+                waited = time.monotonic() - started
+                if init_pending and not plan.applies and waited > init_probe:
+                    plan.send_init = True
+                elif not init_pending and waited > timeout and not plan.applies:
+                    plan.shutdown = True  # only ends the wait; the client's flag is unchanged
             plan = self._bcast_object(plan)
             self._raise_if(plan.error)
+            if init_pending and (plan.applies or plan.send_init):
+                init_pending = False
+                if plan.send_init:
+                    if self.rank == 0:
+                        log.warning("no catch-up from the syncer within %.0f s: sending INIT_PARAMS", init_probe)
+                    self._send_init()
+                elif self.rank == 0:
+                    log.warning("the syncer is initialized (catch-up arriving): INIT_PARAMS skipped")
+                started = time.monotonic()
             updates = self._take_updates(plan) if self.rank == 0 else {}
             for fid, version in plan.applies:
                 data = updates.pop(fid) if self.rank == 0 else None
@@ -830,9 +889,11 @@ def _install_initial_cut(self, timeout: float) -> None:
                 break
             if not plan.applies:
                 time.sleep(0.2)
+        if init_pending:  # the wait ended before the probe did: keep the old order's guarantee
+            self._send_init()
         if self.rank == 0:
-            log.info("initial cut: %d of %d fragments installed before training in %.0f s",
-                     sum(self.anchored), self.num_fragments, time.monotonic() - started)
+            log.warning("initial cut: %d of %d fragments installed before training in %.0f s",
+                        sum(self.anchored), self.num_fragments, time.monotonic() - started)
 
     # ------------------------------------------------------------------ resume
     def _apply_first_or_blend(self, fid: int, data: Any, version: int, *, live: str | None = None) -> str | None:
@@ -861,6 +922,10 @@ def _apply_first_or_blend(self, fid: int, data: Any, version: int, *, live: str
         self._apply_fragment(fid, data, self.bulk_torch, blend=True, mode=mode)
         if mode is not None and len(self.resumed) == self.num_fragments:
             self._resume_anchors.clear()  # every old anchor is used or no longer needed
+            if self.rank == 0:
+                counts = {m: list(self.resumed.values()).count(m) for m in (RESUME_ANCHOR, RESUME_SHIFT, RESUME_OVERWRITE)}
+                log.warning("resume: all %d fragments installed: anchor %d, shift %d, overwrite %d",
+                            self.num_fragments, counts[RESUME_ANCHOR], counts[RESUME_SHIFT], counts[RESUME_OVERWRITE])
         return mode
 
     def _restore_counters(self, fid: int, mode: str) -> None:
@@ -1085,23 +1150,36 @@ def _export_fragment(self, fid: int, *, delta: bool) -> list[Any] | None:
         """The fragment's values on rank 0 as wire byte parts (INIT, or PUSH
         for bf16/f32 sessions) or, for a q4 push, one f32 tensor."""
 
+        parts: list[Any] | None = [] if self.rank == 0 else None
+        for chunk in self._export_fragment_chunks(fid, delta=delta):
+            if parts is not None:
+                parts.extend(chunk)
+        return parts
+
+    def _export_fragment_chunks(self, fid: int, *, delta: bool):
+        """``_export_fragment`` one chunk at a time (collective): yields each
+        chunk's parts on rank 0, None elsewhere; ``_export_error`` is set once
+        a non-finite value was seen."""
+
         from yeto.protocol import DTYPE_Q4
         from yeto.tensor_io import pack_tensor
 
-        parts: list[Any] | None = [] if self.rank == 0 else None
         self._export_error = None
         q4 = delta and self.wire == DTYPE_Q4
         for names in self.fragment_chunks[fid]:
             tensors = self._export_chunk(names, delta=delta)
-            if parts is None:
+            if tensors is None:
+                yield None
                 continue
+            parts = []
             for tensor in tensors:
                 if self._export_error is None and not bool(tensor.isfinite().all()):
                     # Raised on every rank by the caller's error broadcast.
                     self._export_error = f"fragment {fid} holds a non-finite value"
                 parts.append(tensor.reshape(-1) if q4 else
                              pack_tensor(tensor.reshape(-1), self.bulk if not delta else self._push_part_dtype()))
-        return parts
+            del tensors
+            yield parts
 
     def _push_part_dtype(self) -> int:
         """The dtype of the PUSH parts handed to the client: f32 on the

From 357f7e134dd210fb5323d1c28b6adc90cec2a904 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 15:02:08 +0800
Subject: [PATCH 09/17] syncer: --checkpoint-on-start writes the checkpoint
 before the first 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>
---
 syncer/src/main.rs   |  8 ++++++
 syncer/src/server.rs | 66 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/syncer/src/main.rs b/syncer/src/main.rs
index 073e00c9..8b2d1417 100644
--- a/syncer/src/main.rs
+++ b/syncer/src/main.rs
@@ -159,6 +159,13 @@ struct Args {
     /// during the write) host memory peaks at about twice params + momentum.
     #[arg(long)]
     background_checkpoint: bool,
+    /// Write --checkpoint-path once before the first round (after INIT or
+    /// --resume), so the rounds before the first --checkpoint-every one are
+    /// covered too. Skipped when --resume found that checkpoint (it is the
+    /// state just loaded). On a background thread with
+    /// --background-checkpoint.
+    #[arg(long)]
+    checkpoint_on_start: bool,
     /// Policy sweep: write per-round checkpoints off the scheduler; each
     /// round is broadcast only after its checkpoint is durable.
     #[arg(long, default_value_t = false)]
@@ -513,6 +520,7 @@ fn main() -> anyhow::Result<()> {
         checkpoint_every: args.checkpoint_every,
         checkpoint_at_steps: args.checkpoint_at_steps,
         background_checkpoint: args.background_checkpoint,
+        checkpoint_on_start: args.checkpoint_on_start,
         async_checkpoint: args.async_checkpoint,
         bcast_delta_full_every: args.bcast_delta_full_every,
         resume: args.resume,
diff --git a/syncer/src/server.rs b/syncer/src/server.rs
index de24ed33..f5e24d74 100644
--- a/syncer/src/server.rs
+++ b/syncer/src/server.rs
@@ -127,6 +127,9 @@ pub struct Config {
     /// rounds keep merging (`GlobalState::save_checkpoint_in_background`)
     /// instead of on the scheduler. Every other checkpoint waits for it.
     pub background_checkpoint: bool,
+    /// Write the checkpoint once before the first round, unless `resume`
+    /// loaded it (`--checkpoint-on-start`).
+    pub checkpoint_on_start: bool,
     /// Global steps whose cut is also written to `<checkpoint_path>.step-<t>`
     /// off the scheduler (`crate::eval_checkpoint`); strictly increasing,
     /// below `total_steps`. Empty: none.
@@ -1500,6 +1503,28 @@ fn round_grace(
     (grace, floor)
 }
 
+/// `--checkpoint-on-start`: write the checkpoint before the first round.
+/// Without it the first rounds (16 at ~860 s per write on the 120B run) have
+/// no checkpoint, and a restart then loses the INIT and every merge so far.
+/// A resume already has one (the state it loaded); a finished run needs none.
+/// Returns whether a write was started.
+fn save_start_checkpoint(cfg: &Config, st: &GlobalState, resumed: bool, last_step: u64) -> Result<bool> {
+    let Some(path) = cfg.checkpoint_path.as_ref().filter(|_| cfg.checkpoint_on_start) else {
+        return Ok(false);
+    };
+    if resumed || st.global_step >= last_step {
+        return Ok(false);
+    }
+    if cfg.background_checkpoint {
+        st.save_checkpoint_in_background(path)?;
+        info!(step = st.global_step, path = %path.display(), "start checkpoint started in the background");
+    } else {
+        st.save_checkpoint(path)?;
+        info!(step = st.global_step, path = %path.display(), "start checkpoint written");
+    }
+    Ok(true)
+}
+
 /// Whether step `t` is the last round of its fragment within `total_steps`.
 fn is_terminal_round(schedule: &crate::shard::ShardSchedule, t: u64, total_steps: u64) -> bool {
     schedule.next_merge_of_fragment(t) > total_steps
@@ -1831,6 +1856,9 @@ fn validate_config(cfg: &Config) -> Result<()> {
              without --policy-sweep-fragments (per-round checkpoints persist before their broadcast)"
         );
     }
+    if cfg.checkpoint_on_start && cfg.checkpoint_path.is_none() {
+        bail!("--checkpoint-on-start needs --checkpoint-path");
+    }
     if cfg.async_checkpoint {
         if cfg.policy_sweep_fragments.is_none() {
             bail!("--async-checkpoint requires --policy-sweep-fragments");
@@ -4072,6 +4100,8 @@ async fn scheduler(
                 info!(step = st.global_step, path = %path.display(), "checkpoint committed");
             }
         }
+    } else {
+        save_start_checkpoint(&cfg, &st, resumed, last_step)?;
     }
     // Send everyone the initial (or resumed) global parameters so all
     // learners start bit-identical (also serves recovery for late joiners).
@@ -9556,6 +9586,7 @@ mod tests {
             checkpoint_every: 0,
             checkpoint_at_steps: Vec::new(),
             background_checkpoint: false,
+            checkpoint_on_start: false,
             async_checkpoint: false,
             bcast_delta_full_every: None,
             resume: false,
@@ -11223,6 +11254,41 @@ mod tests {
         std::fs::remove_dir_all(&directory).ok();
     }
 
+    #[test]
+    fn checkpoint_on_start_covers_the_rounds_before_the_first_periodic_one() {
+        let directory = unique_test_directory("start-checkpoint");
+        for background in [false, true] {
+            let checkpoint = directory.join(format!("start-{background}.ckpt"));
+            let mut config = round_test_config(8);
+            config.checkpoint_path = Some(checkpoint.clone());
+            config.checkpoint_every = 4;
+            config.background_checkpoint = background;
+            let mut state = two_fragment_sweep_state();
+            state.policy_sweep_fragments = None;
+            // Off: nothing before the first periodic checkpoint.
+            assert!(!save_start_checkpoint(&config, &state, false, 8).unwrap());
+            assert!(!checkpoint.exists());
+            config.checkpoint_on_start = true;
+            validate_config(&config).unwrap();
+            // A resume already holds the state on disk; a finished run needs none.
+            assert!(!save_start_checkpoint(&config, &state, true, 8).unwrap());
+            state.global_step = 8;
+            assert!(!save_start_checkpoint(&config, &state, false, 8).unwrap());
+            state.global_step = 0;
+            assert!(!checkpoint.exists());
+            assert!(save_start_checkpoint(&config, &state, false, 8).unwrap());
+            state.wait_for_background_checkpoint().unwrap();
+            let mut on_disk = two_fragment_sweep_state();
+            on_disk.policy_sweep_fragments = None;
+            on_disk.load_checkpoint(&checkpoint).unwrap();
+            assert_eq!((on_disk.global_step, on_disk.versions.clone()), (0, state.versions.clone()));
+        }
+        let mut config = round_test_config(8);
+        config.checkpoint_on_start = true;
+        assert!(validate_config(&config).is_err(), "needs a checkpoint path");
+        std::fs::remove_dir_all(&directory).ok();
+    }
+
     #[tokio::test]
     async fn background_checkpoint_keeps_rounds_going_and_writes_the_same_cuts() {
         let directory = unique_test_directory("background-checkpoint");

From ae3eadae1dd59d48a194885abd9952a1c2f707a2 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 15:05:32 +0800
Subject: [PATCH 10/17] sft presets: initial-cut wait and background / start
 checkpoints on by 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>
---
 tests/test_head_mode.py               |  2 ++
 tests/test_rl_island_pool_launcher.py |  7 +++++--
 tests/test_run_flags.py               |  3 ++-
 tests/test_swift_megatron_plugin.py   |  9 +++++++++
 yeto/launch/syncer.py                 |  3 +++
 yeto/swift_megatron/swift_plugin.py   | 12 ++++++++----
 6 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/tests/test_head_mode.py b/tests/test_head_mode.py
index 91f74e34..9392b08e 100644
--- a/tests/test_head_mode.py
+++ b/tests/test_head_mode.py
@@ -838,6 +838,8 @@ def test_local_syncer_command_matches_cluster_syncer_flags():
     assert "--resume" in cmd
     assert "--mark-final-checkpoint" in cmd
     assert f"--port {launcher.SYNCER_PORT}" in cmd
+    # The SFT preset: periodic checkpoints off the scheduler, one before the first round.
+    assert "--background-checkpoint" in cmd and "--checkpoint-on-start" in cmd
 
 
 def test_rl_launch_records_the_run_under_its_own_name(fake_sky, monkeypatch, tmp_path):
diff --git a/tests/test_rl_island_pool_launcher.py b/tests/test_rl_island_pool_launcher.py
index 60475005..ddb964da 100644
--- a/tests/test_rl_island_pool_launcher.py
+++ b/tests/test_rl_island_pool_launcher.py
@@ -48,7 +48,8 @@ def _sft():
 def test_fixed_roster_syncer_command_is_byte_identical(head_launcher):
     """secret-handling-hardening: the only change to the fixed-roster RL line is the
     island HMAC key check in front (the key itself is a secret env); SFT only
-    gains the explicit --allow-unauthenticated-islands."""
+    gains the explicit --allow-unauthenticated-islands and its background and
+    start checkpoints."""
     args = _args()
     # The default data split (the dispenser) adds --dispense flags on purpose;
     # the row slots keep the line this test compares
@@ -61,8 +62,10 @@ def test_fixed_roster_syncer_command_is_byte_identical(head_launcher):
         commands = launcher.role_syncer_commands(args, n)
         assert list(commands) == ["actor"]
         assert launcher.syncer_command(args, n).endswith(HMAC_CHECK + commands["actor"])
+    checkpoints = " --background-checkpoint --checkpoint-on-start"
     assert launcher.syncer_command(_sft(), 2) == (
-        head_launcher.syncer_command(_sft(), 2) + " --allow-unauthenticated-islands")
+        head_launcher.syncer_command(_sft(), 2).replace(
+            " $(test -f", checkpoints + " $(test -f", 1) + " --allow-unauthenticated-islands")
 
 
 def test_fixed_roster_learner_flags_unchanged():
diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py
index e709dd47..aebdcfd8 100644
--- a/tests/test_run_flags.py
+++ b/tests/test_run_flags.py
@@ -259,7 +259,8 @@ def test_forward_tables_name_flags_of_both_layers():
 _NON_RL_LITERALS = frozenset({
     "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter",
     "--allow-unauthenticated-islands", "--require-profile-record", "--async-checkpoint", "--inline-resend", "--inline-broadcast", "--bcast-delta-full-every", "--bind", "--cap-add", "--carry-gamma",
-    "--checkpoint-at-steps", "--checkpoint-every", "--checkpoint-path", "--config-json",
+    "--background-checkpoint", "--checkpoint-at-steps", "--checkpoint-every", "--checkpoint-on-start",
+    "--checkpoint-path", "--config-json",
     "--create", "--dashboard-host", "--max-learners", "--dearmor", "--depth", "--detach", "--dtype",
     # task dispenser (syncer flags)
     "--dispense", "--dispense-lease-s", "--dispense-log", "--dispense-max-redo", "--dispense-passes",
diff --git a/tests/test_swift_megatron_plugin.py b/tests/test_swift_megatron_plugin.py
index 9ae733af..c402d5e6 100644
--- a/tests/test_swift_megatron_plugin.py
+++ b/tests/test_swift_megatron_plugin.py
@@ -373,3 +373,12 @@ def test_callback_rebuilds_the_lr_schedule_on_attach(make):
     cb = sp.YetoDiLoCoCallback(trainer)
     assert trainer.opt_param_scheduler.num_steps == 15 * 384
     cb.close()
+
+
+def test_initial_cut_wait_is_on_by_default(monkeypatch):
+    for k, v in ENV.items():
+        monkeypatch.setenv(k, v)
+    monkeypatch.delenv("YETO_DILOCO_INITIAL_CUT_TIMEOUT", raising=False)
+    assert sp.DiLoCoHookConfig.from_env().initial_cut_timeout == 3600.0
+    monkeypatch.setenv("YETO_DILOCO_INITIAL_CUT_TIMEOUT", "0")
+    assert sp.DiLoCoHookConfig.from_env().initial_cut_timeout == 0.0
diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py
index 99a2b344..77bbce57 100644
--- a/yeto/launch/syncer.py
+++ b/yeto/launch/syncer.py
@@ -880,6 +880,9 @@ def _sft_syncer_command(args, num_learners: int, binary: str,
         f" --total-steps {args.total_steps}"
         f"{_syncer_outer_optimizer(args)}"
         f" --checkpoint-path {checkpoint}"
+        # The SFT preset writes its periodic checkpoints (every 8 rounds by
+        # default) off the scheduler, and one before the first round.
+        f" --background-checkpoint --checkpoint-on-start"
         f" {_resume_if_exists(checkpoint)}"
         f" --mark-final-checkpoint"
         f" --event-tape {_shard_path(SYNCER_EVENT_TAPE, index)}"
diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py
index 900df664..d9fd3799 100644
--- a/yeto/swift_megatron/swift_plugin.py
+++ b/yeto/swift_megatron/swift_plugin.py
@@ -55,8 +55,12 @@
   YETO_DILOCO_HMAC_KEY_PATH       path to the HMAC key file (optional)
   YETO_DILOCO_DISABLE             set to 1 to make the callback a no-op
   YETO_DILOCO_INITIAL_CUT_TIMEOUT seconds to install the syncer's initial cut
-                                  fragment by fragment before training (0 =
-                                  apply it at the boundaries; default 0)
+                                  (or a resumed island's catch-up) fragment by
+                                  fragment before training; default 3600.  0 =
+                                  apply it at the first boundaries instead: the
+                                  whole cut (4 bytes per parameter) then queues
+                                  on rank 0 first, and a resumed island takes it
+                                  only after one more step
   YETO_DILOCO_RESUME_DIR          checkpoint directory whose sync record to
                                   resume from (default: ``args.mcore_model``
                                   when it holds one)
@@ -130,7 +134,7 @@ class DiLoCoHookConfig:
     hmac_key_path: Optional[str] = None
     final_cut_payload: str = "full"
     transport: str = "cpu"
-    initial_cut_timeout: float = 0.0
+    initial_cut_timeout: float = 3600.0
     resume_dir: Optional[str] = None
     save_anchors: bool = False
     init_probe_s: float = 1800.0
@@ -171,7 +175,7 @@ def get_float(name: str, default: float) -> float:
             hmac_key_path=get("HMAC_KEY_PATH"),
             final_cut_payload=get("FINAL_CUT_PAYLOAD", default="full"),  # type: ignore[arg-type]
             transport=get("TRANSPORT", default="cpu"),  # type: ignore[arg-type]
-            initial_cut_timeout=get_float("INITIAL_CUT_TIMEOUT", 0.0),
+            initial_cut_timeout=get_float("INITIAL_CUT_TIMEOUT", 3600.0),
             resume_dir=get("RESUME_DIR"),
             save_anchors=(get("SAVE_ANCHORS", default="0") or "").strip().lower() in {"1", "true", "yes", "on"},
             init_probe_s=get_float("INIT_PROBE_S", 1800.0),

From 6b17b8c87c6e0549d4f027225725c9b3e47c984f Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 15:09:25 +0800
Subject: [PATCH 11/17] swift-megatron: heartbeat thread so the syncer can
 lease the island

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>
---
 tests/test_swift_megatron_sync.py   | 71 ++++++++++++++++++++++++++++-
 yeto/swift_megatron/swift_plugin.py | 13 ++++++
 yeto/swift_megatron/sync.py         | 52 ++++++++++++++++++++-
 3 files changed, 133 insertions(+), 3 deletions(-)

diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py
index 00f1e278..90fad13d 100644
--- a/tests/test_swift_megatron_sync.py
+++ b/tests/test_swift_megatron_sync.py
@@ -74,6 +74,33 @@ def test_chunks_are_bounded_and_keep_order():
 
 
 # ----------------------------------------------------------------------------- fake Megatron island
+def test_heartbeat_thread_beats_between_steps_and_stops_on_a_stall():
+    class Beats:
+        def __init__(self):
+            self.steps = []
+
+        def heartbeat(self, step):
+            self.steps.append(step)
+
+    client, progress = Beats(), [time.monotonic()]
+    beat = sm._Heartbeat(client, 0.02, 0.3, lambda: 7, lambda: progress[0])
+    deadline = time.monotonic() + 5
+    while len(client.steps) < 3:
+        assert time.monotonic() < deadline
+        time.sleep(0.01)
+    assert set(client.steps) == {7}
+    time.sleep(0.5)  # no progress for longer than the stall limit: silence
+    quiet = len(client.steps)
+    time.sleep(0.2)
+    assert len(client.steps) == quiet
+    progress[0] = time.monotonic()
+    while len(client.steps) == quiet:
+        assert time.monotonic() < deadline
+        time.sleep(0.01)
+    beat.stop()
+    assert not beat.thread.is_alive()
+
+
 class _LinkClient:
     """The client surface ``_PushSender`` and ``_ready_pushes`` use: a link
     generation the test bumps, and a push result it chooses."""
@@ -959,7 +986,8 @@ def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at,
     syncer = sm.MegatronDiLoCoSync([chunk], opt, syncer=f"127.0.0.1:{syncer_port}", learner_id=island,
                                    num_learners=2, fragments=3, sync_interval_steps=2,
                                    wire_dtype=opts.get("wire", "bf16"),
-                                   chunk_numel=64, parallel_context=ctx, final_cut_payload="full")
+                                   chunk_numel=64, parallel_context=ctx, final_cut_payload="full",
+                                   heartbeat_s=opts.get("heartbeat_s", 30.0))
     it = 0
     dropped = False
     started = time.monotonic()
@@ -992,7 +1020,7 @@ def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at,
     if rank == 0:
         torch.save({"names": names, "values": dict(zip(names, values, strict=True)),
                     "events": list(syncer.events), "stats": dict(syncer.stats), "dropped": dropped,
-                    "timeline": timeline,
+                    "timeline": timeline, "lease_expiries": syncer.client.lease_expiries,
                     "fragments": [[n for n, _ in f.tensors] for f in syncer.layout.fragments]},
                    Path(out_dir, f"island{island}.pt"))
     syncer.close()
@@ -1217,3 +1245,42 @@ def read(i, proc):  # the log file written so far
     if stop == "sigkill":
         assert sum(i["stats"].get("rebased", 0) for i in islands) > 0, "no island was rewound"
     print(json.dumps([i["stats"] for i in islands]))
+
+
+@pytest.mark.timeout(300)
+@pytest.mark.parametrize("heartbeat_s", [0.3, 0.0])
+def test_heartbeats_keep_a_stalled_island_inside_its_lease(tmp_path, heartbeat_s):
+    """--island-lease-s 2 with an island that makes no boundary for 4 s:
+    rank 0's heartbeat thread keeps its lease; without heartbeats the syncer
+    expires it (the island redials as a re-JOIN)."""
+
+    binary = build_syncer()
+    from process_groups import group_popen
+
+    port = free_port()
+    final_state = tmp_path / "final.bin"
+    proc = group_popen([str(binary), "--port", str(port), "--learners", "2", "--grace-ms", "200",
+                        "--terminal-grace-ms", "2000", "--total-steps", "30", "--final-state", str(final_state),
+                        "--island-lease-s", "2", *DECOUPLED_SYNCER_FLAGS],
+                       stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    procs = []
+    try:
+        procs += _run(_island_worker, 4, 0, free_port(), port, str(tmp_path), None,
+                      {"step_s": 0.05, "heartbeat_s": heartbeat_s})
+        procs += _run(_island_worker, 4, 1, free_port(), port, str(tmp_path), None,
+                      {"step_s": 0.05, "stall_at": 6, "stall_s": 4.0, "heartbeat_s": heartbeat_s})
+        _join(procs, 240, tmp_path)
+        assert proc.wait(timeout=60) == 0, "syncer exited nonzero"
+    finally:
+        for p in procs:
+            if p.is_alive():
+                p.kill()
+        if proc.poll() is None:
+            proc.kill()
+        log = proc.stdout.read() if proc.stdout else ""
+        (tmp_path / "syncer.log").write_text(log)
+    islands = [torch.load(tmp_path / f"island{i}.pt") for i in range(2)]
+    for name in islands[0]["names"]:
+        assert torch.equal(islands[0]["values"][name], islands[1]["values"][name]), name
+    expired = '"reason":"lease_expired"' in log or islands[1]["lease_expiries"] > 0
+    assert expired == (heartbeat_s == 0), (islands[1]["lease_expiries"], heartbeat_s)
diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py
index d9fd3799..cfdb33c3 100644
--- a/yeto/swift_megatron/swift_plugin.py
+++ b/yeto/swift_megatron/swift_plugin.py
@@ -71,6 +71,13 @@
                                   send INIT_PARAMS only if the syncer sends no
                                   catch-up within this many seconds (an
                                   initialized syncer ignores INIT); default 1800
+  YETO_DILOCO_HEARTBEAT_S         seconds between rank 0's HEARTBEATs to the
+                                  syncer, sent from a thread (independent of
+                                  step ends) so its --island-lease-s can be on;
+                                  default 30, 0 = none
+  YETO_DILOCO_HEARTBEAT_STALL_S   stop heartbeating once the training loop
+                                  made no progress for this long, so a hung
+                                  island loses its lease; default 0 (never)
   YETO_DILOCO_LR_RESUME           auto (default) | off: rebuild the learning
                                   rate schedule of a resume that did not load
                                   it (``restore_lr_schedule``)
@@ -138,6 +145,8 @@ class DiLoCoHookConfig:
     resume_dir: Optional[str] = None
     save_anchors: bool = False
     init_probe_s: float = 1800.0
+    heartbeat_s: float = 30.0
+    heartbeat_stall_s: float = 0.0
 
     @classmethod
     def from_env(cls, env=None) -> "DiLoCoHookConfig":
@@ -179,6 +188,8 @@ def get_float(name: str, default: float) -> float:
             resume_dir=get("RESUME_DIR"),
             save_anchors=(get("SAVE_ANCHORS", default="0") or "").strip().lower() in {"1", "true", "yes", "on"},
             init_probe_s=get_float("INIT_PROBE_S", 1800.0),
+            heartbeat_s=get_float("HEARTBEAT_S", 30.0),
+            heartbeat_stall_s=get_float("HEARTBEAT_STALL_S", 0.0),
         )
         if cfg.initial_cut_timeout < 0:
             raise ValueError(f"{ENV_PREFIX}INITIAL_CUT_TIMEOUT must be >= 0")
@@ -276,6 +287,8 @@ def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig, *,
         resume_dir=resume_dir,
         resume_iteration=resume_iteration,
         init_probe_s=cfg.init_probe_s,
+        heartbeat_s=cfg.heartbeat_s,
+        heartbeat_stall_s=cfg.heartbeat_stall_s,
     )
 
 
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index 36e26cca..9dbeb21c 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -609,7 +609,8 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
                  transport: str = "cpu", parallel_context: ParallelContext | None = None,
                  client_factory=None, observer: Callable[[dict], None] | None = None,
                  initial_cut_timeout: float = 0.0, resume_dir: str | None = None,
-                 resume_iteration: int | None = None, init_probe_s: float = 1800.0) -> None:
+                 resume_iteration: int | None = None, init_probe_s: float = 1800.0,
+                 heartbeat_s: float = 30.0, heartbeat_stall_s: float = 0.0) -> None:
         import torch
         import torch.distributed as dist
 
@@ -717,6 +718,13 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
                       "deferred_pushes": 0, "apply_s": 0.0, "export_s": 0.0}
         self._observed_push_bytes = 0  # rank 0: push_bytes at the last observer report
         self.sender: _PushSender | None = None  # rank 0: background PUSH streaming
+        # Rank 0: HEARTBEAT every ``heartbeat_s`` from a thread, so the syncer
+        # can lease the island (``--island-lease-s``) across ~10 min steps,
+        # INIT and checkpoint saves; ``_progress_at`` is the last sign of life
+        # of the training loop (``heartbeat_stall_s``: none for that long and
+        # the heartbeats stop, so a hung island loses its lease).
+        self.heartbeat: _Heartbeat | None = None
+        self._progress_at = time.monotonic()
         # (iteration, "apply" | "push", fragment, version | pull step), same on every rank.
         self.events: deque[tuple[int, str, int, int]] = deque(maxlen=4096)
         # Resume (``resume_dir``): the record ``save_state`` wrote with the
@@ -738,6 +746,9 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int,
             try:
                 self.client = self._connect(syncer, hmac_key_path, wan_streams, client_factory)
                 self.sender = _PushSender(self.client, self.stats, self._push_dropped, self._push_sent)
+                if heartbeat_s > 0:
+                    self.heartbeat = _Heartbeat(self.client, float(heartbeat_s), float(heartbeat_stall_s),
+                                                lambda: self.iteration, lambda: self._progress_at)
             except Exception as exc:  # raised on every rank below
                 error = f"{type(exc).__name__}: {exc}"
         self._raise_if(self._bcast_object(error))
@@ -815,6 +826,7 @@ def send() -> None:
                 if pending is not None:
                     jobs.put((fid, pending))
                 for parts in self._export_fragment_chunks(fid, delta=False):
+                    self._progress_at = time.monotonic()
                     if pending is None:
                         continue
                     if error is None and self._export_error is not None:
@@ -854,6 +866,7 @@ def _install_initial_cut(self, timeout: float, *, init_probe: float | None = Non
         started = time.monotonic()
         init_pending = init_probe is not None  # resumed learner 0: INIT only if the syncer needs it
         while not all(self.anchored):
+            self._progress_at = time.monotonic()
             plan = self._plan() if self.rank == 0 else None
             if self.rank == 0:
                 waited = time.monotonic() - started
@@ -1337,6 +1350,7 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *,
 
         if self.finished or self._closed:
             return
+        self._progress_at = time.monotonic()
         state = self.state
         iteration = int(iteration)
         if not self._counters_started:  # count from the (possibly resumed) first step
@@ -1622,6 +1636,8 @@ def close(self) -> None:
         if self._closed:
             return
         self._closed = True
+        if self.heartbeat is not None:
+            self.heartbeat.stop()
         if self.sender is not None:
             self.sender.stop()
         client, self.client = self.client, None
@@ -1652,6 +1668,40 @@ def session_dtype(name: str) -> int:
     return code | DTYPE_FLAG_TENSOR_BLOCKS if key != "bf16" else code
 
 
+class _Heartbeat:
+    """Rank 0's HEARTBEAT thread (``client.heartbeat(local_step)`` every
+    ``interval`` seconds): the syncer renews the island's lease on it, so a
+    lease can be on although a step (or an INIT, a checkpoint save) takes far
+    longer.  ``stall`` > 0: no heartbeat once the training loop showed no
+    progress (``progress_at``) for that long.  A send during an outage is
+    dropped by the client; the next one goes out after the redial."""
+
+    def __init__(self, client: Any, interval: float, stall: float, step: Callable[[], int],
+                 progress_at: Callable[[], float]) -> None:
+        import threading
+
+        self.client, self.interval, self.stall = client, interval, stall
+        self.step, self.progress_at = step, progress_at
+        self.sent = 0
+        self._stop = threading.Event()
+        self.thread = threading.Thread(target=self._run, name="yeto-swift-heartbeat", daemon=True)
+        self.thread.start()
+
+    def _run(self) -> None:
+        while not self._stop.wait(self.interval):
+            if self.stall > 0 and time.monotonic() - self.progress_at() > self.stall:
+                continue
+            try:
+                self.client.heartbeat(int(self.step()))
+                self.sent += 1
+            except Exception as exc:  # never take the island down for a heartbeat
+                log.debug("heartbeat not sent: %s", exc)
+
+    def stop(self) -> None:
+        self._stop.set()
+        self.thread.join(timeout=5.0)
+
+
 class _PushSender:
     """Rank 0's background PUSH streaming: one PUSH in flight at a time, so a
     slow link or syncer never holds a training step (decoupled DiLoCo, as the

From 51a4f57ecc44774056544b682ae23383c0b314e9 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 15:55:24 +0800
Subject: [PATCH 12/17] syncer, protocol: keep same-connection anchor reuse and
 two-tuple reassembly

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>
---
 syncer/src/server.rs | 46 +++++++++++++++++++++++++++-----------------
 yeto/protocol.py     |  6 ++++--
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/syncer/src/server.rs b/syncer/src/server.rs
index f5e24d74..94d35897 100644
--- a/syncer/src/server.rs
+++ b/syncer/src/server.rs
@@ -5979,10 +5979,11 @@ fn route_push(
 /// same fragment against its own base, weighted by its counters (and the
 /// carry discount), as long as its base lag stays within `--max-base-lag`.
 ///
-/// A delta counts once: after the syncer carries or gathers a delta from an
-/// anchor (fragment version), a later push of that island from the same or
-/// an older anchor of that fragment is refused (`anchor_reused`), on any of
-/// its connections. An island that redials (a WAN drop, a syncer it lost
+/// A delta counts once: after the syncer carries a delta from an anchor
+/// (fragment version), a later push of that island from the same or an
+/// older anchor of that fragment is refused (`anchor_reused`) on any of its
+/// connections; after it gathers one, the same holds for the island's later
+/// connections. An island that redials (a WAN drop, a syncer it lost
 /// for a while) releases the anchors its lost pushes held, so it may answer
 /// the next pull from an anchor whose delta did reach this syncer; that
 /// second delta repeats the first one's local steps and must not be merged
@@ -5995,8 +5996,8 @@ struct LateCarry {
     /// Newest base version of a carried delta per island and fragment.
     anchors: HashMap<(u32, usize), u64>,
     /// Newest base version of a delta gathered by its own round, per island
-    /// and fragment.
-    gathered: HashMap<(u32, usize), u64>,
+    /// and fragment, with the connection generation that pushed it.
+    gathered: HashMap<(u32, usize), (u64, u64)>,
 }
 
 #[derive(Debug, Eq, PartialEq)]
@@ -6024,14 +6025,19 @@ impl LateCarry {
         }
     }
 
-    /// True when `push` reuses an anchor whose delta this island already
-    /// had carried or gathered (by any of its connections).
+    /// True when `push` reuses an anchor whose delta this island already had
+    /// carried (on any connection) or gathered on an earlier connection. One
+    /// connection may answer consecutive rounds from the same anchor (an
+    /// asynchronous learner that has not installed the newer version yet);
+    /// a redialed one repeating the anchor resends a delta already taken.
     fn reuses_anchor(&self, member: Member, push: &Push) -> bool {
         let key = (member.learner_id, push.fragment_id as usize);
-        [self.anchors.get(&key), self.gathered.get(&key)]
-            .into_iter()
-            .flatten()
-            .any(|&taken| push.base_version <= taken)
+        let carried = self.anchors.get(&key).is_some_and(|&taken| push.base_version <= taken);
+        let gathered = self
+            .gathered
+            .get(&key)
+            .is_some_and(|&(taken, generation)| push.base_version <= taken && member.generation != generation);
+        carried || gathered
     }
 
     fn spend_anchor(&mut self, member: Member, fragment: usize, base_version: u64) {
@@ -6040,19 +6046,21 @@ impl LateCarry {
     }
 
     fn note_gathered(&mut self, member: Member, fragment: usize, base_version: u64) {
-        let taken = self.gathered.entry((member.learner_id, fragment)).or_insert(0);
-        *taken = (*taken).max(base_version);
+        let taken = self.gathered.entry((member.learner_id, fragment)).or_insert((0, member.generation));
+        if base_version >= taken.0 {
+            *taken = (base_version, member.generation);
+        }
     }
 
     /// A restarted round drops its gathered pushes; their islands answer the
     /// new attempt from the same anchor, so the anchor is free again.
     fn release_anchor(&mut self, member: Member, push: &Push) {
         let key = (member.learner_id, push.fragment_id as usize);
-        if self.gathered.get(&key) == Some(&push.base_version) {
+        if let Some(taken) = self.gathered.get_mut(&key).filter(|taken| taken.0 == push.base_version) {
             if push.base_version == 0 {
                 self.gathered.remove(&key);
             } else {
-                self.gathered.insert(key, push.base_version - 1);
+                taken.0 = push.base_version - 1;
             }
         }
     }
@@ -6082,7 +6090,7 @@ impl LateCarry {
         if self
             .gathered
             .get(&(member.learner_id, p))
-            .is_some_and(|&taken| push.base_version <= taken)
+            .is_some_and(|&(taken, _)| push.base_version <= taken)
         {
             return CarryOutcome::Rejected(PushDisposition::Duplicate);
         }
@@ -12614,7 +12622,9 @@ mod tests {
         carry.note_gathered(old, 1, 5);
         let again = test_push_for(1, 5);
         assert!(carry.reuses_anchor(new, &again));
-        assert!(carry.reuses_anchor(old, &again));
+        // The connection that pushed it may answer the next round from the
+        // same anchor (it has not installed the newer version yet).
+        assert!(!carry.reuses_anchor(old, &again));
         // From the broadcast it then gets (version 7) it pushes again.
         assert!(!carry.reuses_anchor(new, &test_push_for(1, 7)));
         // A restart of that round drops the gathered delta: the anchor is
diff --git a/yeto/protocol.py b/yeto/protocol.py
index 08832b08..e79bce91 100644
--- a/yeto/protocol.py
+++ b/yeto/protocol.py
@@ -2684,7 +2684,7 @@ def _submit_digest(self, job: Callable[[], None]) -> None:
 
     def _reassemble(
         self, gen: int, payload: bytes | bytearray
-    ) -> tuple[int, memoryview, bytes | None] | _StreamedInboundFragment | None:
+    ) -> tuple[int, memoryview] | tuple[int, memoryview, bytes] | _StreamedInboundFragment | None:
         if len(payload) < _CHUNK_HEAD.size:
             raise ValueError("chunk header truncated")
         msg_id, total, offset = _CHUNK_HEAD.unpack_from(payload)
@@ -2765,7 +2765,9 @@ def _reassemble(
             if digest is not None:
                 digest.cancel()
             raise ValueError("corrupt reassembled frame")
-        return msg_type, memoryview(buf)[_HEADER.size :], None if digest is None else digest.finish()
+        if digest is None:
+            return msg_type, memoryview(buf)[_HEADER.size :]
+        return msg_type, memoryview(buf)[_HEADER.size :], digest.finish()
 
     def _stream_chunk_locked(
         self,

From 43f6bd1dac96f0f553c7d80bc2c9785f13f8c596 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Fri, 6 Nov 2026 17:44:32 +0800
Subject: [PATCH 13/17] swift-megatron: keep sync and plugin logs visible under
 an ERROR-level 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>
---
 tests/test_swift_megatron_plugin.py | 64 +++++++++++++++++++++++++++++
 yeto/swift_megatron/swift_plugin.py | 12 ++++++
 yeto/swift_megatron/sync.py         | 51 +++++++++++++++++------
 3 files changed, 114 insertions(+), 13 deletions(-)

diff --git a/tests/test_swift_megatron_plugin.py b/tests/test_swift_megatron_plugin.py
index c402d5e6..3df7842d 100644
--- a/tests/test_swift_megatron_plugin.py
+++ b/tests/test_swift_megatron_plugin.py
@@ -382,3 +382,67 @@ def test_initial_cut_wait_is_on_by_default(monkeypatch):
     assert sp.DiLoCoHookConfig.from_env().initial_cut_timeout == 3600.0
     monkeypatch.setenv("YETO_DILOCO_INITIAL_CUT_TIMEOUT", "0")
     assert sp.DiLoCoHookConfig.from_env().initial_cut_timeout == 0.0
+
+
+def _logger_under(root_handler_level):
+    """A logger whose root mimics the Megatron SFT entry point's: one stdout
+    handler at ``root_handler_level`` (ERROR in the real process)."""
+    import io
+    import logging
+
+    root = logging.RootLogger(logging.INFO)
+    sink = io.StringIO()
+    handler = logging.StreamHandler(sink)
+    handler.setLevel(root_handler_level)
+    root.addHandler(handler)
+    logger = logging.Logger("yeto-test-visible")
+    logger.parent = root
+    return logger, sink
+
+
+def test_logs_are_made_visible_when_the_root_handler_drops_them(capsys):
+    import logging
+
+    from yeto.swift_megatron.sync import ensure_visible_logs
+
+    logger, root_sink = _logger_under(logging.ERROR)
+    ensure_visible_logs(logger)
+    logger.info("info line")
+    logger.warning("warning line")
+    logger.error("error line")
+    err = capsys.readouterr().err
+    assert "info line" in err and "warning line" in err and err.count("error line") == 1
+    assert root_sink.getvalue() == ""  # no longer propagated (no double ERROR)
+    ensure_visible_logs(logger)  # idempotent
+    assert len(logger.handlers) == 1
+
+
+def test_logs_are_left_alone_when_the_process_prints_info(capsys):
+    import logging
+
+    from yeto.swift_megatron.sync import ensure_visible_logs
+
+    logger, root_sink = _logger_under(logging.INFO)
+    ensure_visible_logs(logger)
+    assert logger.handlers == [] and logger.propagate
+    logger.info("via root")
+    assert "via root" in root_sink.getvalue() and capsys.readouterr().err == ""
+
+
+def test_the_lr_rebuild_warning_survives_an_error_level_root(make, monkeypatch, capsys):
+    import logging
+
+    import yeto.swift_megatron.sync as sync_mod
+
+    plugin_logger, _ = _logger_under(logging.ERROR)
+    sync_logger, _ = _logger_under(logging.ERROR)
+    monkeypatch.setattr(sp, "logger", plugin_logger)
+    monkeypatch.setattr(sync_mod, "log", sync_logger)
+    trainer = FakeTrainer([([1], True)], start_iteration=15)
+    trainer.args.global_batch_size = 384
+    trainer.opt_param_scheduler = _Scheduler([{}], warmup=10 * 384, decay=1000 * 384)
+    cb = sp.YetoDiLoCoCallback(trainer)
+    err = capsys.readouterr().err
+    assert "restored no LR schedule; stepped it by 5760" in err
+    assert sync_logger.handlers and not sync_logger.propagate
+    cb.close()
diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py
index cfdb33c3..0f4839e6 100644
--- a/yeto/swift_megatron/swift_plugin.py
+++ b/yeto/swift_megatron/swift_plugin.py
@@ -347,6 +347,17 @@ def restore_lr_schedule(trainer) -> Optional[int]:
     return samples
 
 
+def _visible_logs() -> None:
+    """This module's and the sync core's WARNING/INFO lines would otherwise
+    go to the root handler the Megatron entry point leaves at ERROR."""
+    _ensure_repo_on_path()
+    try:
+        from yeto.swift_megatron.sync import ensure_visible_logs, log as sync_log
+    except ImportError:  # pragma: no cover - the sync core needs torch
+        return
+    ensure_visible_logs(logger, sync_log)
+
+
 def _model_chunks(trainer) -> list:
     # Pass the DDP-wrapped chunks: the sync core needs the grad/param buffers
     # and the DistributedOptimizer's mapping, both keyed on the wrapped models.
@@ -390,6 +401,7 @@ def __init__(self, trainer):
             logger.warning("yeto DiLoCo callback disabled via %sDISABLE", ENV_PREFIX)
             return
         self.config = DiLoCoHookConfig.from_env()
+        _visible_logs()
         restore_lr_schedule(trainer)
         self._wrap_train_step()
         # Build here (not in on_train_begin) so a bad config fails before data
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index 9dbeb21c..50fa0db9 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -168,19 +168,44 @@
 log = logging.getLogger("swift-megatron-sync")
 
 
-def _ensure_visible_logs() -> None:
-    """ms-swift configures only its own ``swift`` logger, so with an
-    unconfigured root this module's INFO lines (``resume: fragment ...``,
-    per-fragment installs) were dropped and only WARNING reached stderr
-    through logging's last resort.  Give the module a stderr handler at INFO
-    unless the process configured logging itself."""
-    if log.handlers or logging.getLogger().handlers:
-        return
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter("[swift-megatron-sync] %(asctime)s %(levelname)s %(message)s"))
-    log.addHandler(handler)
-    if log.level == logging.NOTSET:
-        log.setLevel(logging.INFO)
+def _emits_info(logger: logging.Logger) -> bool:
+    """Whether an INFO record of ``logger`` reaches some handler that prints it."""
+    if not logger.isEnabledFor(logging.INFO):
+        return False
+    node: logging.Logger | None = logger
+    while node is not None:
+        if any(handler.level <= logging.INFO for handler in node.handlers):
+            return True
+        if not node.propagate:
+            return False
+        node = node.parent  # type: ignore[assignment]
+    return False
+
+
+def ensure_visible_logs(*loggers: logging.Logger) -> None:
+    """Make these loggers' INFO lines reach stderr.
+
+    ms-swift configures only its own ``swift`` logger, and the Megatron SFT
+    entry point leaves the root logger with a single stdout handler at ERROR:
+    every INFO *and* WARNING of a yeto module (``resume: fragment ...``,
+    ``INIT_PARAMS skipped``, ``anchor ... released after a reconnect``, the LR
+    schedule rebuild) propagated into that handler and was dropped, and
+    logging's last resort never fires once root has a handler.  A logger whose
+    records some handler on its propagation path already prints at INFO is
+    left alone; any other gets its own stderr handler at INFO and stops
+    propagating (so an ERROR is not printed twice)."""
+    for logger in loggers or (log,):
+        if logger.handlers or _emits_info(logger):
+            continue
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter(f"[{logger.name}] %(asctime)s %(levelname)s %(message)s"))
+        logger.addHandler(handler)
+        if logger.level == logging.NOTSET or logger.level > logging.INFO:
+            logger.setLevel(logging.INFO)
+        logger.propagate = False
+
+
+_ensure_visible_logs = ensure_visible_logs
 
 # Megatron's tied/untied output projection is an embedding-like tensor for the
 # direct-averaged fragment, as ``lm_head`` is for the torch learner.

From bcd1f893eb87039729c9f79d0a1993c227976cbc Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Sat, 7 Nov 2026 20:08:42 +0800
Subject: [PATCH 14/17] swift-megatron: install, not shift, a newer global
 after a push from 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>
---
 tests/test_diloco_sync.py | 13 ++++++++++---
 yeto/diloco_sync.py       | 11 +++++++++--
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/tests/test_diloco_sync.py b/tests/test_diloco_sync.py
index 684416a9..2ad902a4 100644
--- a/tests/test_diloco_sync.py
+++ b/tests/test_diloco_sync.py
@@ -387,9 +387,16 @@ def test_resume_record_round_trips_and_picks_the_install_mode():
     assert record["spent"] == [None, [7, 9], None] and record["units_total"] == 900
     # Same version: the broadcast is the checkpoint's own anchor.
     assert resume_fragment_mode(record, 0, 4, have_anchor=False) == RESUME_ANCHOR
-    # Newer: move the local progress only when the old anchor is at hand.
-    assert resume_fragment_mode(record, 1, 9, have_anchor=True) == RESUME_SHIFT
-    assert resume_fragment_mode(record, 1, 9, have_anchor=False) == RESUME_OVERWRITE
+    # Newer, after a push from that anchor: the delta may be merged into it,
+    # so it is installed as is rather than counted twice.
+    assert resume_fragment_mode(record, 1, 9, have_anchor=True) == RESUME_OVERWRITE
+    # Newer, nothing pushed from the anchor: move the local progress, only
+    # when the old anchor is at hand.
+    unspent = dict(record, spent=[None, None, None])
+    assert resume_fragment_mode(unspent, 1, 9, have_anchor=True) == RESUME_SHIFT
+    assert resume_fragment_mode(unspent, 1, 9, have_anchor=False) == RESUME_OVERWRITE
+    # A rewound global (a syncer resumed from an older checkpoint) still shifts.
+    assert resume_fragment_mode(record, 1, 5, have_anchor=True) == RESUME_SHIFT
     # Never anchored before the checkpoint, or no record at all: a cold start.
     assert resume_fragment_mode(record, 2, 0, have_anchor=True) == RESUME_OVERWRITE
     assert resume_fragment_mode(None, 0, 4, have_anchor=True) == RESUME_OVERWRITE
diff --git a/yeto/diloco_sync.py b/yeto/diloco_sync.py
index db8cb509..035801b2 100644
--- a/yeto/diloco_sync.py
+++ b/yeto/diloco_sync.py
@@ -186,11 +186,18 @@ def resume_fragment_mode(record: dict | None, fid: int, version: int, *, have_an
     *is* that anchor, so the local parameters stay as they are.  Any other
     version moves the island's local progress onto the new global, which
     needs the old anchor value (``have_anchor``); without it the fragment
-    falls back to the cold-start overwrite."""
+    falls back to the cold-start overwrite.  As in ``live_install_mode``, a
+    newer global of a fragment the island had pushed from that anchor
+    (``record["spent"]``) is installed as is: the pushed delta may already be
+    merged into it, and shifting would count it twice."""
     if record is None or not record["anchored"][fid]:
         return RESUME_OVERWRITE
-    if int(version) == int(record["versions"][fid]):
+    anchor = int(record["versions"][fid])
+    if int(version) == anchor:
         return RESUME_ANCHOR
+    spent = (record.get("spent") or [None] * len(record["versions"]))[fid]
+    if int(version) > anchor and spent is not None and int(spent[0]) == anchor:
+        return RESUME_OVERWRITE
     return RESUME_SHIFT if have_anchor else RESUME_OVERWRITE
 
 

From e9180370e22641704a01674e1bd61da68561713f Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Sat, 7 Nov 2026 20:08:42 +0800
Subject: [PATCH 15/17] protocol: pass the block digest through the
 spot-rollout subscriber

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>
---
 tests/test_reconnect.py            | 12 ++++++++++++
 yeto/rl/spot_rollout/subscriber.py |  4 ++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/tests/test_reconnect.py b/tests/test_reconnect.py
index a097b8c4..5eb224a6 100644
--- a/tests/test_reconnect.py
+++ b/tests/test_reconnect.py
@@ -550,3 +550,15 @@ def test_chunked_and_direct_broadcasts_share_the_replay_and_conflict_checks():
     with pytest.raises(ProtocolError, match="conflicting BCAST_FRAGMENT"):
         client._dispatch(1, *inner)
     client.close()
+
+
+def test_subclass_dispatch_overrides_take_the_block_digest():
+    """A reassembled BCAST of two or more digest blocks reaches ``_dispatch``
+    with its block digest; every override must accept it."""
+    import inspect
+
+    from yeto.protocol import SyncerClient
+    from yeto.rl.spot_rollout.subscriber import SubscriberClient
+
+    for cls in (SyncerClient, SubscriberClient):
+        inspect.signature(cls._dispatch).bind(None, 0, 0, b"", b"\0" * 32)
diff --git a/yeto/rl/spot_rollout/subscriber.py b/yeto/rl/spot_rollout/subscriber.py
index c792a113..d1331d6a 100644
--- a/yeto/rl/spot_rollout/subscriber.py
+++ b/yeto/rl/spot_rollout/subscriber.py
@@ -106,7 +106,7 @@ def __init__(self, addr: tuple[str, int], subscriber_id: int, layout: FragmentLa
         self._beats = threading.Event()
         self._run_finished = threading.Event()
 
-    def _dispatch(self, gen: int, msg_type: int, payload) -> None:
+    def _dispatch(self, gen: int, msg_type: int, payload, digest: bytes | None = None) -> None:
         # The end-of-run notice counts on any connection generation. The
         # syncer closes the data streams right after it; when their end reaches
         # the client first, the client starts a new generation and the base class
@@ -116,7 +116,7 @@ def _dispatch(self, gen: int, msg_type: int, payload) -> None:
             self._run_finished.set()
             self._closed.set()  # no redial: the run is over
             self._failure.set()
-        super()._dispatch(gen, msg_type, payload)
+        super()._dispatch(gen, msg_type, payload, digest)
 
     def _hello_body_suffix(self) -> bytes:
         return self._trailer

From ec6d6d39a84922f2a974889a66db347d93fed430 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Sat, 7 Nov 2026 20:08:42 +0800
Subject: [PATCH 16/17] swift-megatron: fail, not hang, when an INIT export
 raises on rank 0

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>
---
 tests/test_swift_megatron_sync.py | 40 +++++++++++++++++++++++++++++++
 yeto/swift_megatron/sync.py       | 23 ++++++++++++++++--
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py
index 90fad13d..274ec284 100644
--- a/tests/test_swift_megatron_sync.py
+++ b/tests/test_swift_megatron_sync.py
@@ -946,6 +946,46 @@ def test_a_resumed_learner_zero_sends_init_only_when_the_syncer_needs_it(tmp_pat
     _join(procs, 200, tmp_path)
 
 
+def _init_abort_worker(rank, world, port, out_dir):
+    _guard(out_dir, f"init-abort-{rank}", lambda: _init_abort_body(rank, world, port))
+
+
+def _init_abort_body(rank, world, port):
+    dist = _init_dist(rank, world, port)
+    single = dist.new_group([0])
+    ctx = sm.ParallelContext(0, 1, single, 0, 1, single, 0, 1)
+    chunk = _chunk(tp=0, pp=0, tp_size=1, pp_size=1)
+    opt = _DistOpt([p for p in chunk.parameters() if p.requires_grad], dp_rank=0, dp_size=1, dp_group=single)
+    holder = {}
+
+    def factory(addrs, learner_id, layout, wire, streams, **kwargs):
+        holder["client"] = _ScriptedClient(layout, {})
+        return holder["client"]
+
+    # learner 1: no INIT from the constructor; the test sends it by hand
+    sync = sm.MegatronDiLoCoSync([chunk], opt, syncer="127.0.0.1:1", learner_id=1, num_learners=2,
+                                 fragments=2, sync_interval_steps=1, chunk_numel=16,
+                                 parallel_context=ctx, client_factory=factory)
+    original = sync._export_fragment_chunks
+
+    def failing(fid, delta):
+        for i, parts in enumerate(original(fid, delta=delta)):
+            if i == 1:
+                raise RuntimeError("collective timed out")
+            yield parts
+
+    sync._export_fragment_chunks = failing
+    started = time.monotonic()
+    with pytest.raises(RuntimeError, match="collective timed out"):
+        sync._send_init()
+    assert time.monotonic() - started < 30
+
+
+def test_an_init_export_that_raises_fails_instead_of_hanging(tmp_path):
+    procs = _run(_init_abort_worker, 1, free_port(), str(tmp_path))
+    _join(procs, 120, tmp_path)
+
+
 # ----------------------------------------------------------------------------- two islands, real syncer (TP2 x PP2)
 def build_syncer() -> Path:
     binary = ROOT / "syncer/target/debug/yeto-syncer"
diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py
index 50fa0db9..3ed9d44e 100644
--- a/yeto/swift_megatron/sync.py
+++ b/yeto/swift_megatron/sync.py
@@ -845,9 +845,12 @@ def send() -> None:
             sender = threading.Thread(target=send, name="yeto-swift-init", daemon=True)
             sender.start()
         error = None
+        current = None  # rank 0: the queue of the fragment being exported
+        aborted = False
         try:
             for fid in range(self.num_fragments):
                 pending = queue.Queue(maxsize=2) if self.rank == 0 else None
+                current = pending
                 if pending is not None:
                     jobs.put((fid, pending))
                 for parts in self._export_fragment_chunks(fid, delta=False):
@@ -863,12 +866,28 @@ def send() -> None:
                     del parts
                 if pending is not None:
                     pending.put(end)
+                    current = None
                     if error is None and failure:
                         error = failure[0]
+        except BaseException:
+            # An export that raised (a collective timeout) leaves the sender
+            # waiting for the rest of the fragment: end it so the error, not
+            # a hang, reaches the caller.
+            aborted = True
+            if current is not None:
+                failure.append("INIT aborted")
+                try:
+                    current.put(end, timeout=60)
+                except queue.Full:
+                    pass
+            raise
         finally:
             if sender is not None:
-                jobs.put(None)
-                sender.join()
+                try:
+                    jobs.put(None, timeout=60 if aborted else None)
+                except queue.Full:
+                    pass
+                sender.join(timeout=60 if aborted else None)
         if self.rank == 0 and error is None and failure:
             error = failure[0]
         self._raise_if(self._bcast_object(error))

From 3b5b378269378fa4651dd928a7ad67588d86f3a7 Mon Sep 17 00:00:00 2001
From: MichaelChung <michaelchung@MichaelChungdeMac-Studio-3.local>
Date: Sat, 7 Nov 2026 20:55:48 +0800
Subject: [PATCH 17/17] syncer: on a stop, abandon rounds whose islands are all
 gone at once

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>
---
 docs/PROTOCOL.md     |  3 ++-
 syncer/src/server.rs | 24 ++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md
index 0dbce353..bcec53f7 100644
--- a/docs/PROTOCOL.md
+++ b/docs/PROTOCOL.md
@@ -595,7 +595,8 @@ periodic interval still leaves an atomic final checkpoint.
 
 SIGTERM or SIGINT stops the syncer gracefully: it launches no new round,
 gives the rounds still gathering `--shutdown-timeout-s` (default 600) to
-close, then abandons the rest (they have not touched the state, so a resumed
+close (a round with no delta whose islands are all disconnected is abandoned
+at once), then abandons the rest (they have not touched the state, so a resumed
 run launches those steps again), commits every computed round, writes
 `--checkpoint-path` synchronously (after any background checkpoint) and exits
 0 with `stop: checkpoint written step=<t>`. The checkpoint's global step is
diff --git a/syncer/src/server.rs b/syncer/src/server.rs
index 94d35897..d944c457 100644
--- a/syncer/src/server.rs
+++ b/syncer/src/server.rs
@@ -4207,6 +4207,30 @@ async fn scheduler(
                 for round in inflight.drain(..) {
                     busy_fragments.remove(&round.p);
                 }
+            } else if !inflight.is_empty() {
+                // A round that holds no delta and none of whose islands is
+                // still connected cannot close: do not wait the timeout out.
+                let orphaned: Vec<usize> = {
+                    let registry = registry.lock().unwrap();
+                    inflight
+                        .iter()
+                        .enumerate()
+                        .filter(|(_, round)| {
+                            round.pushes.is_empty()
+                                && round.carried.is_empty()
+                                && !round.expected_members.iter().any(|m| registry.groups.contains_key(m))
+                        })
+                        .map(|(i, _)| i)
+                        .collect()
+                };
+                if !orphaned.is_empty() {
+                    let steps: Vec<u64> = orphaned.iter().map(|&i| inflight[i].t).collect();
+                    warn!(?steps, "stopping: abandoning gathering rounds whose islands are all gone");
+                    for &i in orphaned.iter().rev() {
+                        let round = inflight.remove(i);
+                        busy_fragments.remove(&round.p);
+                    }
+                }
             }
             if inflight.is_empty() && computing.is_empty() && ready.is_empty() && publishing.is_none() {
                 stopped = true;