swift_megatron: multi-node island support (cross-node gloo, per-rank scatter, INIT overlap, step log) (2/3)

tests/test_cross_node_gloo.py

@@ -0,0 +1,90 @@
+"""cross_node_gloo: groups spanning "nodes" go to gloo; staged collectives and
+P2P keep their semantics. CPU-only: 4 gloo ranks, 2 ranks per simulated node."""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+torch = pytest.importorskip("torch")
+import torch.distributed as dist  # noqa: E402
+import torch.multiprocessing as mp  # noqa: E402
+
+from tests.free_ports import free_port  # noqa: E402
+from yeto.swift_megatron import cross_node_gloo as X  # noqa: E402
+
+WORLD, LOCAL = 4, 2
+
+
+def test_node_mapping():
+    assert X.node_of(5, 4) == 1
+    assert X.spans_nodes([0, 1], 2) is False
+    assert X.spans_nodes([1, 2], 2) is True
+
+
+def _worker(rank, port, q):
+    try:
+        os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port), YETO_CROSS_NODE_LOCAL_WORLD=str(LOCAL))
+        dist.init_process_group("gloo", rank=rank, world_size=WORLD)
+        X.install()
+        intra = [dist.new_group([n * LOCAL + i for i in range(LOCAL)]) for n in range(WORLD // LOCAL)]
+        pps = [dist.new_group([i + n * LOCAL for n in range(WORLD // LOCAL)]) for i in range(LOCAL)]
+        mine_intra, pp = intra[rank // LOCAL], pps[rank % LOCAL]
+        assert id(pp) in X._cross and id(mine_intra) not in X._cross
+
+        t = torch.full((3,), float(rank))
+        dist.all_reduce(t)  # WORLD is redirected to the gloo world group
+        assert t.tolist() == [6.0] * 3
+
+        t = torch.full((2,), float(rank))
+        dist.all_reduce(t, op=dist.ReduceOp.AVG, group=pp)
+        peer = (rank + LOCAL) % WORLD
+        assert t.tolist() == [(rank + peer) / 2] * 2
+
+        out = torch.empty(4)
+        dist.all_gather_into_tensor(out, torch.full((2,), float(rank)), group=pp)
+        lo, hi = sorted((rank, peer))
+        assert out.tolist() == [lo, lo, hi, hi]
+
+        rs = torch.empty(2)
+        dist.reduce_scatter_tensor(rs, torch.arange(4, dtype=torch.float32) + rank, group=pp)
+        r = pp.rank()
+        assert rs.tolist() == [2 * (2 * r) + rank + peer, 2 * (2 * r + 1) + rank + peer]
+
+        objs = [None] * WORLD
+        dist.all_gather_object(objs, rank)
+        assert objs == list(range(WORLD))
+
+        x, y = torch.full((5,), float(rank)), torch.empty(5)
+        ops = [dist.P2POp(dist.isend, x, peer, pp), dist.P2POp(dist.irecv, y, peer, pp)]
+        for w in dist.batch_isend_irecv(ops):
+            w.wait()
+        assert y.tolist() == [float(peer)] * 5
+
+        y2 = torch.empty(5)
+        if rank < LOCAL:
+            a, b = dist.isend(x, peer, group=pp), dist.irecv(y2, peer, group=dist.group.WORLD)
+        else:
+            b, a = dist.irecv(y2, peer, group=pp), dist.isend(x, peer, group=dist.group.WORLD)
+        a.wait()
+        b.wait()
+        assert y2.tolist() == [float(peer)] * 5
+        dist.barrier()
+        dist.destroy_process_group()
+        q.put((rank, "ok"))
+    except Exception as e:  # pragma: no cover - reported to the parent
+        q.put((rank, repr(e)))
+
+
+def test_cross_node_collectives_and_p2p():
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+    port = free_port()
+    procs = [ctx.Process(target=_worker, args=(r, port, q)) for r in range(WORLD)]
+    for p in procs:
+        p.start()
+    res = dict(q.get(timeout=120) for _ in range(WORLD))
+    for p in procs:
+        p.join(30)
+    assert res == {r: "ok" for r in range(WORLD)}

tests/test_diloco_sync.py

@@ -329,3 +329,38 @@ def test_replicas_never_double_report():
         observer=seen.append,
     )
     assert seen == []
+
+
+def test_newer_attempt_of_a_round_drops_the_queued_old_pull():
+    from yeto.diloco_sync import take_ready_pulls
+
+    state = DiLoCoSyncState.create(2, track_anchors=False)
+    state.pending_pulls = [
+        PullRequest(fragment_id=0, global_step=1, round_attempt=0),
+        PullRequest(fragment_id=1, global_step=1, round_attempt=0),
+        PullRequest(fragment_id=0, global_step=1, round_attempt=1),
+    ]
+    ready = take_ready_pulls(state, steps_total=5, units_total=50, anchored=lambda fid: True)
+    assert [(p.fragment_id, p.global_step, p.round_attempt) for p, _, _ in ready] == [(1, 1, 0), (0, 1, 1)]
+    assert state.pending_pulls == []
+
+
+def test_superseded_pull_is_dropped_while_its_fragment_waits():
+    from yeto.diloco_sync import take_ready_pulls
+
+    state = DiLoCoSyncState.create(1, track_anchors=False)
+    state.pending_pulls = [PullRequest(fragment_id=0, global_step=3, round_attempt=0)]
+    assert take_ready_pulls(state, steps_total=0, units_total=0, anchored=lambda fid: True) == []
+    state.pending_pulls.append(PullRequest(fragment_id=0, global_step=3, round_attempt=1))
+    assert take_ready_pulls(state, steps_total=0, units_total=0, anchored=lambda fid: True) == []
+    assert state.pending_pulls == [PullRequest(fragment_id=0, global_step=3, round_attempt=1)]
+
+
+def test_a_resent_pull_of_the_same_attempt_is_answered_once():
+    from yeto.diloco_sync import take_ready_pulls
+
+    state = DiLoCoSyncState.create(1, track_anchors=False)
+    pull = PullRequest(fragment_id=0, global_step=3, round_attempt=1)
+    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 == []

tests/test_swift_megatron_plugin.py

@@ -219,3 +219,39 @@ def test_repo_root_on_sys_path():
     root = os.path.dirname(os.path.dirname(os.path.abspath(sp.__file__)))
     assert os.path.dirname(root) in sys.path or any(
         os.path.abspath(p) == os.path.dirname(root) for p in sys.path)
+
+
+def test_step_log_records_sync_time_and_counters(make, monkeypatch, tmp_path):
+    import json
+
+    path = tmp_path / "steps.jsonl"
+    monkeypatch.setenv("YETO_DILOCO_STEP_LOG", str(path))
+    trainer, cb = make([([3], True), ([2], True)])
+    sync = FakeSync.instances[0]
+    sync.rank = 0
+    sync.stats = {"applied": 0, "pushed": 0, "apply_s": 0.0, "export_s": 0.0}
+    original = sync.after_optimizer_step
+
+    def step(iteration, target_tokens, *, update_successful=True):
+        original(iteration, target_tokens, update_successful=update_successful)
+        if iteration == 2:
+            sync.stats["applied"] += 1
+            sync.stats["apply_s"] += 1.5
+            sync.stats["export_s"] += 0.25
+
+    sync.after_optimizer_step = step
+    trainer.train()
+    rows = [json.loads(line) for line in path.read_text().splitlines()]
+    assert [r["iteration"] for r in rows] == [1, 2]
+    assert rows[0]["applies"] == 0 and rows[0]["apply_s"] == 0
+    assert rows[1]["applies"] == 1 and rows[1]["apply_s"] == 1.5 and rows[1]["export_s"] == 0.25
+    assert rows[1]["target_tokens"] == 4 and rows[1]["sync_s"] >= 0
+
+
+def test_step_log_only_on_sync_rank_zero(make, monkeypatch, tmp_path):
+    path = tmp_path / "steps.jsonl"
+    monkeypatch.setenv("YETO_DILOCO_STEP_LOG", str(path))
+    trainer, cb = make([([3], True)])
+    FakeSync.instances[0].rank = 5
+    trainer.train()
+    assert not path.exists()

tests/test_swift_megatron_sync.py

@@ -439,6 +439,65 @@ def test_tp2_dp2_island_reads_masters_writes_masters_and_pushes_raw_anchor_delta
     _join(procs, 200, tmp_path)
 
 
+def _initial_cut_worker(rank, world, port, out_dir):
+    _guard(out_dir, f"initial-cut-{rank}", lambda: _initial_cut_body(rank, world, port, out_dir))
+
+
+def _initial_cut_body(rank, world, port, out_dir):
+    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)
+    globals_g = {name: (_full(name, shape) + 0.125).to(torch.bfloat16).float()
+                 for name, shape in full_shapes.items()}
+    holder = {}
+
+    def factory(addrs, learner_id, layout, wire, streams, **kwargs):
+        # The initial cut trickles in: fragment 0 at the first poll, the rest later.
+        cut = [BcastFragment(fid, 0, pack_tensor(_flat(layout, fid, globals_g), 2))
+               for fid in range(layout.num_fragments)]
+        holder["client"] = client = _ScriptedClient(layout, {("bcast", 1): cut[:1], ("bcast", 4): cut[1:]})
+        return client
+
+    syncer = 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)
+    # Installed before the first step, fragment by fragment, as raw globals.
+    assert syncer.anchored == [True] * syncer.num_fragments
+    assert [(e[0], e[1], e[3]) for e in syncer.events] == [(0, "apply", 0)] * syncer.num_fragments
+    for p in trainable:
+        full = globals_g[p._full_name]
+        if getattr(p, "tensor_model_parallel", False):
+            size = p._full_shape[p.partition_dim] // 2
+            full = full.narrow(p.partition_dim, tp * size, size)
+        if id(p) in opt.owned:
+            s, e, _ = opt.owned[id(p)]
+            assert torch.equal(opt.master(p), full.reshape(-1)[s:e]), p._full_name
+        assert torch.equal(p.detach(), full.to(torch.bfloat16)), p._full_name
+    if rank == 0:
+        assert holder["client"].calls >= 4 and not holder["client"].inits  # learner 1 sends no INIT
+    # The boundary finds nothing left to apply.
+    syncer.after_optimizer_step(1, 10)
+    assert syncer.stats["applied"] == syncer.num_fragments
+    syncer.close()
+    dist.destroy_process_group()
+
+
+@pytest.mark.timeout(240)
+def test_initial_cut_is_installed_fragment_by_fragment_before_training(tmp_path):
+    procs = _run(_initial_cut_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"

yeto/diloco_sync.py

@@ -49,6 +49,29 @@ def connect_syncer(spec: str, learner_id: int, layout: FragmentLayout, wire_dtyp
     return client
 
 
+def supersede_pull_attempts(pulls: list[PullRequest]) -> list[PullRequest]:
+    """``pulls`` without any pull a newer attempt of the same round replaced.
+
+    A round the syncer reopens after a quorum timeout sends a new pull with
+    a higher ``round_attempt``; a push answering the old attempt would be
+    rejected as out of round, so only the newest attempt is kept, in order.
+    A repeat of the same attempt (resent after a redial) is kept once, so the
+    fragment is not pushed twice for it.
+    """
+    newest: dict[tuple[int, int], int] = {}
+    for pull in pulls:
+        key = (int(pull.fragment_id), int(pull.global_step))
+        newest[key] = max(newest.get(key, -1), int(pull.round_attempt))
+    kept, seen = [], set()
+    for pull in pulls:
+        key = (int(pull.fragment_id), int(pull.global_step))
+        attempt = int(pull.round_attempt)
+        if attempt == newest[key] and (key, attempt) not in seen:
+            seen.add((key, attempt))
+            kept.append(pull)
+    return kept
+
+
 def take_ready_pulls(
     state: DiLoCoSyncState,
     *,
@@ -61,7 +84,7 @@ def take_ready_pulls(
     now (the fragment trained ``min_steps`` steps since its last broadcast and
     has an anchor); the others stay pending, in order."""
     ready, still = [], []
-    for pull in state.pending_pulls:
+    for pull in supersede_pull_attempts(state.pending_pulls):
         fid = pull.fragment_id
         c_steps = steps_total - state.steps_at_reset[fid]
         if c_steps < max(1, min_steps) or not anchored(fid):