diff --git a/tests/test_cross_node_gloo.py b/tests/test_cross_node_gloo.py new file mode 100644 index 00000000..9cafcb7d --- /dev/null +++ b/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)} diff --git a/tests/test_diloco_sync.py b/tests/test_diloco_sync.py index c16f028c..bc086886 100644 --- a/tests/test_diloco_sync.py +++ b/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 == [] diff --git a/tests/test_swift_megatron_plugin.py b/tests/test_swift_megatron_plugin.py index bd6ae6b8..14bfe373 100644 --- a/tests/test_swift_megatron_plugin.py +++ b/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() diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py index 7d58723a..5b136d19 100644 --- a/tests/test_swift_megatron_sync.py +++ b/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" diff --git a/yeto/diloco_sync.py b/yeto/diloco_sync.py index cdf089c4..7afbb042 100644 --- a/yeto/diloco_sync.py +++ b/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): diff --git a/yeto/swift_megatron/cross_node_gloo.py b/yeto/swift_megatron/cross_node_gloo.py new file mode 100644 index 00000000..0c6944d9 --- /dev/null +++ b/yeto/swift_megatron/cross_node_gloo.py @@ -0,0 +1,341 @@ +"""Route cross-node process groups through gloo with host staging. + +Use case: a pipeline-parallel island built from two (or more) Ascend nodes that +have no RoCE between them. HCCL cannot form a communicator across such nodes +(EI0015 on the first collective), but PP only needs one activation P2P per +micro-batch per direction at the stage boundary, which plain TCP handles. + +With this module installed (``YETO_CROSS_NODE_GLOO=1``): + +* ``torch.distributed.new_group`` creates every group whose ranks span more + than one node with the gloo backend; groups inside one node keep the default + (HCCL) backend, so TP/EP/DP traffic stays on HCCS. +* Collectives and P2P ops issued on such a group -- or on the default/WORLD + group when the job spans nodes -- are redirected to the gloo group and run on + host copies of the device tensors; results are copied back to the device. + +Rank-to-node mapping assumes torchrun's contiguous layout +(``rank // LOCAL_WORLD_SIZE``; ``YETO_CROSS_NODE_LOCAL_WORLD`` overrides the +ranks per node). A group whose backend string names gloo (a mixed one too) +is created as requested, without host staging. Load as an ms-swift external +plugin before Megatron builds its parallel state. +""" + +from __future__ import annotations + +import functools +import inspect +import logging +import os + +import torch +import torch.distributed as dist +import torch.distributed.distributed_c10d as c10d + +logger = logging.getLogger(__name__) + +_ENV = "YETO_CROSS_NODE_GLOO" +_installed = False +_orig: dict = {} +_cross: dict = {} # id(group) -> group, gloo groups spanning nodes +_world_gloo = None +stats = {"p2p_bytes": 0, "coll_bytes": 0} + + +def _local_world() -> int: + return int(os.environ.get("YETO_CROSS_NODE_LOCAL_WORLD") or os.environ.get("LOCAL_WORLD_SIZE") or 8) + + +def node_of(rank: int, local_world: int | None = None) -> int: + return rank // (local_world or _local_world()) + + +def spans_nodes(ranks, local_world: int | None = None) -> bool: + return len({node_of(r, local_world) for r in ranks}) > 1 + + +def _job_spans_nodes() -> bool: + return dist.is_initialized() and spans_nodes(range(dist.get_world_size())) + + +def _get_world_gloo(): + global _world_gloo + if _world_gloo is None: + _world_gloo = _orig["new_group"](backend="gloo") + _cross[id(_world_gloo)] = _world_gloo + return _world_gloo + + +def _redirect(group): + """Return the gloo group to use instead of ``group``, or None to keep it.""" + if group is None or group is c10d.GroupMember.WORLD: + return _get_world_gloo() if _job_spans_nodes() else None + return group if id(group) in _cross else None + + +def _new_group(ranks=None, timeout=None, backend=None, pg_options=None, **kwargs): + if _job_spans_nodes(): + _get_world_gloo() # all ranks reach the first new_group together + rr = range(dist.get_world_size()) if ranks is None else ranks + wants_device = backend is None or "gloo" not in str(backend).lower() + if spans_nodes(list(rr)) and wants_device: + kwargs.pop("device_id", None) + g = _orig["new_group"](ranks, timeout=timeout, backend="gloo", **kwargs) + if g is not None and g is not c10d.GroupMember.NON_GROUP_MEMBER: + _cross[id(g)] = g + return g + return _orig["new_group"](ranks, timeout=timeout, backend=backend, pg_options=pg_options, **kwargs) + + +class _Done: + """Work handle for an op that already completed synchronously.""" + + def __init__(self, finish=None): + self._finish = finish + + def wait(self, timeout=None): + if self._finish is not None: + f, self._finish = self._finish, None + with torch.no_grad(): # recv buffers may be leaves that require grad + f() + return True + + def is_completed(self): + return self._finish is None + + def is_success(self): + return True + + def exception(self): + return None + + def synchronize(self): + self.wait() + + def get_future(self): + fut = torch.futures.Future() + self.wait() + fut.set_result(None) + return fut + + +class _Pending(_Done): + """Wraps an async gloo work; on wait, copies host results back to device.""" + + def __init__(self, work, finish=None): + super().__init__(finish) + self._work = work + + def wait(self, timeout=None): + if self._work is not None: + w, self._work = self._work, None + w.wait() + return super().wait() + + def is_completed(self): + return self._work is None and self._finish is None + + +def _host(t): + if t is None or t.device.type == "cpu": + return t + return t.detach().to("cpu") + + +def _bytes(ts): + return sum(t.numel() * t.element_size() for t in ts if t is not None) + + +def _staged_collective(name): + orig = _orig[name] + sig = inspect.signature(orig) + + @functools.wraps(orig) + def wrapper(*args, **kwargs): + try: + ba = sig.bind(*args, **kwargs) + except TypeError: + return orig(*args, **kwargs) + g = _redirect(ba.arguments.get("group")) + if g is None: + return orig(*args, **kwargs) + a = dict(ba.arguments) + a["group"] = g + async_op = bool(a.pop("async_op", False)) + a.pop("device_ids", None) + _run_on_host(name, a) + return _Done() if async_op else None + + return wrapper + + +def _avg_fix(a): + op = a.get("op", dist.ReduceOp.SUM) + if op == dist.ReduceOp.AVG: + a["op"] = dist.ReduceOp.SUM + return True + return False + + +@torch.no_grad() +def _run_on_host(name, a): + o = _orig + g = a["group"] + if name == "barrier": + o["barrier"](group=g) + elif name in ("all_reduce", "reduce", "broadcast"): + t = a["tensor"] + h = _host(t) + stats["coll_bytes"] += _bytes([h]) + avg = _avg_fix(a) if name != "broadcast" else False + a["tensor"] = h + o[name](**a) + if avg: + h.div_(g.size()) + if h is not t: + t.copy_(h) + elif name in ("all_gather_into_tensor", "_all_gather_base"): + out, inp = a["output_tensor"], a["input_tensor"] + hin = _host(inp).contiguous() + parts = [torch.empty_like(hin) for _ in range(g.size())] + stats["coll_bytes"] += _bytes([hin]) * g.size() + o["all_gather"](parts, hin, group=g) + out.copy_(torch.cat([p.reshape(-1) for p in parts]).reshape(out.shape)) + elif name in ("reduce_scatter_tensor", "_reduce_scatter_base"): + out, inp = a["output"], a["input"] + h = _host(inp).clone() + avg = _avg_fix(a) + stats["coll_bytes"] += _bytes([h]) + o["all_reduce"](h, op=a.get("op", dist.ReduceOp.SUM), group=g) + if avg: + h.div_(g.size()) + r = g.rank() + flat = h.reshape(-1) + n = out.numel() + out.copy_(flat[r * n:(r + 1) * n].reshape(out.shape)) + elif name == "all_gather": + outs, t = a["tensor_list"], a["tensor"] + ht = _host(t) + houts = [torch.empty_like(ht) for _ in outs] + stats["coll_bytes"] += _bytes([ht]) * len(outs) + o["all_gather"](houts, ht, group=g) + for d, s in zip(outs, houts): + d.copy_(s) + elif name == "all_to_all_single": + out, inp = a["output"], a["input"] + hout = torch.empty(out.shape, dtype=out.dtype) + o["all_to_all_single"](hout, _host(inp), a.get("output_split_sizes"), a.get("input_split_sizes"), group=g) + out.copy_(hout) + else: # pragma: no cover - guarded by _COLLECTIVES + raise NotImplementedError(name) + + +def _isend(tensor, dst=None, group=None, tag=0, group_dst=None): + g = _redirect(group) + if g is None: + return _orig["isend"](tensor, dst, group=group, tag=tag, group_dst=group_dst) + h = _host(tensor).contiguous() + stats["p2p_bytes"] += _bytes([h]) + w = _orig["isend"](h, dst, group=g, tag=tag, group_dst=group_dst) + return _Pending(w, finish=lambda: h) # keeps h alive until wait + + +def _irecv(tensor, src=None, group=None, tag=0, group_src=None): + g = _redirect(group) + if g is None: + return _orig["irecv"](tensor, src, group=group, tag=tag, group_src=group_src) + h = torch.empty(tensor.shape, dtype=tensor.dtype) + stats["p2p_bytes"] += _bytes([h]) + w = _orig["irecv"](h, src, group=g, tag=tag, group_src=group_src) + return _Pending(w, finish=lambda: tensor.copy_(h)) + + +def _send(tensor, dst=None, group=None, tag=0, group_dst=None): + if _redirect(group) is None: + return _orig["send"](tensor, dst, group=group, tag=tag, group_dst=group_dst) + _isend(tensor, dst, group=group, tag=tag, group_dst=group_dst).wait() + + +def _recv(tensor, src=None, group=None, tag=0, group_src=None): + if _redirect(group) is None: + return _orig["recv"](tensor, src, group=group, tag=tag, group_src=group_src) + w = _irecv(tensor, src, group=group, tag=tag, group_src=group_src) + w.wait() + return src + + +def _batch_isend_irecv(p2p_op_list): + if not any(_redirect(op.group) is not None for op in p2p_op_list): + return _orig["batch_isend_irecv"](p2p_op_list) + reqs = [] + for op in p2p_op_list: + peer = op.peer + if op.op in (_isend, _orig["isend"]): + reqs.append(_isend(op.tensor, peer, group=op.group, tag=op.tag)) + else: + reqs.append(_irecv(op.tensor, peer, group=op.group, tag=op.tag)) + return reqs + + +def _object_coll(name): + orig = _orig[name] + sig = inspect.signature(orig) + + @functools.wraps(orig) + def wrapper(*args, **kwargs): + ba = sig.bind(*args, **kwargs) + g = _redirect(ba.arguments.get("group")) + if g is not None: + ba.arguments["group"] = g + ba.arguments.pop("device", None) + return orig(*ba.args, **ba.kwargs) + + return wrapper + + +_COLLECTIVES = ( + "all_reduce", "reduce", "broadcast", "barrier", "all_gather", "all_gather_into_tensor", + "_all_gather_base", "reduce_scatter_tensor", "_reduce_scatter_base", "all_to_all_single", +) +_OBJECT = ("all_gather_object", "broadcast_object_list", "gather_object", "scatter_object_list") +_P2P = {"isend": _isend, "irecv": _irecv, "send": _send, "recv": _recv, "batch_isend_irecv": _batch_isend_irecv} + + +def _set(name, fn): + setattr(c10d, name, fn) + if hasattr(dist, name): + setattr(dist, name, fn) + + +def install() -> bool: + """Patch torch.distributed in place. Idempotent.""" + global _installed + if _installed: + return False + names = ("new_group",) + _COLLECTIVES + _OBJECT + tuple(_P2P) + for n in names: + if hasattr(c10d, n): + _orig[n] = getattr(c10d, n) + _set("new_group", _new_group) + for n in _COLLECTIVES: + if n in _orig: + _set(n, _staged_collective(n)) + for n in _OBJECT: + if n in _orig: + _set(n, _object_coll(n)) + for n, fn in _P2P.items(): + _set(n, fn) + _installed = True + if _job_spans_nodes(): + _get_world_gloo() + logger.warning("[cross_node_gloo] installed (local_world=%d)", _local_world()) + print(f"[cross_node_gloo] installed local_world={_local_world()}", flush=True) + return True + + +def cross_groups(): + return list(_cross.values()) + + +if os.environ.get(_ENV) == "1": + install() diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py index 1b036e99..0ac2d695 100644 --- a/yeto/swift_megatron/swift_plugin.py +++ b/yeto/swift_megatron/swift_plugin.py @@ -54,13 +54,18 @@ run (default cpu, gloo mirrors) 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) """ from __future__ import annotations import atexit +import json import logging import os +import time from dataclasses import dataclass from typing import Any, Callable, Optional @@ -103,6 +108,7 @@ class DiLoCoHookConfig: hmac_key_path: Optional[str] = None final_cut_payload: str = "full" transport: str = "cpu" + initial_cut_timeout: float = 0.0 @classmethod def from_env(cls, env=None) -> "DiLoCoHookConfig": @@ -123,6 +129,13 @@ def get_int(name: str, default: Optional[int] = None, required: bool = False) -> except (TypeError, ValueError): raise ValueError(f"{ENV_PREFIX}{name} must be an integer, got {raw!r}") from None + def get_float(name: str, default: float) -> float: + raw = get(name, str(default)) + try: + return float(raw) # type: ignore[arg-type] + except (TypeError, ValueError): + raise ValueError(f"{ENV_PREFIX}{name} must be a number, got {raw!r}") from None + cfg = cls( syncer=get("SYNCER", required=True), # type: ignore[arg-type] learner_id=get_int("LEARNER_ID", required=True), @@ -133,7 +146,10 @@ def get_int(name: str, default: Optional[int] = None, required: bool = False) -> 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), ) + if cfg.initial_cut_timeout < 0: + raise ValueError(f"{ENV_PREFIX}INITIAL_CUT_TIMEOUT must be >= 0") if cfg.num_learners < 1 or not 0 <= cfg.learner_id < cfg.num_learners: raise ValueError(f"invalid learner_id={cfg.learner_id} for num_learners={cfg.num_learners}") if cfg.fragments < 1 or cfg.sync_interval_steps < 1: @@ -223,6 +239,7 @@ def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig): trainable_filter=trainable_filter, final_cut_payload=cfg.final_cut_payload, transport=cfg.transport, + initial_cut_timeout=cfg.initial_cut_timeout, ) @@ -314,8 +331,12 @@ def on_step_end(self): # step was skipped; the skipped step then contributes zero tokens. global_tokens = int(type(self).token_reducer(local_tokens)) target_tokens = global_tokens if update_successful else 0 + stats = getattr(self.sync, "stats", None) or {} + before = {k: stats.get(k, 0) for k in ("applied", "pushed", "apply_s", "export_s")} + started = time.perf_counter() self.sync.after_optimizer_step(iteration, target_tokens, update_successful=update_successful) + self._log_step(iteration, target_tokens, time.perf_counter() - started, before, stats) if getattr(self.sync, "finished", False) and iteration < int(self.args.train_iters): # The final cut is installed: stop here (ms-swift's default_flow # saves at train_iters) instead of training past the global result. @@ -323,6 +344,28 @@ def on_step_end(self): iteration, int(self.args.train_iters)) self.args.train_iters = iteration + def _log_step(self, iteration: int, target_tokens: int, sync_s: float, + before: dict, stats: dict) -> None: + """Append one JSON line per step (sync rank 0 only) to ``STEP_LOG``. + + ``sync_s`` is the wall time the training loop spent inside the sync + boundary; ``export_s`` / ``apply_s`` split out the fp32 master gather + for a push and the write-back of a received broadcast.""" + path = os.environ.get(ENV_PREFIX + "STEP_LOG") + if not path or getattr(self.sync, "rank", 0) != 0: + return + record = {"iteration": iteration, "time": time.time(), "target_tokens": target_tokens, + "sync_s": round(sync_s, 4), + "export_s": round(stats.get("export_s", 0) - before["export_s"], 4), + "apply_s": round(stats.get("apply_s", 0) - before["apply_s"], 4), + "pushes_sent": stats.get("pushed", 0) - before["pushed"], + "applies": stats.get("applied", 0) - before["applied"]} + try: + with open(path, "a", encoding="utf-8") as handle: + handle.write(json.dumps(record) + "\n") + 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_train_end(self): self.close() diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py index 82780699..d87585e8 100644 --- a/yeto/swift_megatron/sync.py +++ b/yeto/swift_megatron/sync.py @@ -144,6 +144,7 @@ import hashlib import logging import re +import time from collections import deque from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass, field @@ -561,7 +562,8 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, fragment_pattern: str = "binpack", wan_streams: int = 4, 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) -> None: + client_factory=None, observer: Callable[[dict], None] | None = None, + initial_cut_timeout: float = 0.0) -> None: import torch import torch.distributed as dist @@ -609,6 +611,17 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self.device = _current_accelerator() self.comm = Transport(transport, self.ctx, optimizer, self.device) self.table = GlobalTable.build(self.shards, self.comm.world) + # Rank 0 scatters each applied chunk as per-rank slices (see _apply_chunk), + # so it needs every rank's (name, TP partition) list, in that rank's order. + self.peer_shards: list[list[tuple[str, int, int, int]]] = [[] for _ in range(self.world)] + mine = [(s.name, s.part_dim, s.tp_rank, s.local_shape[s.part_dim] if s.part_dim >= 0 else 0) + for s in self.shards] + if self.world > 1: + gathered: list[Any] = [None] * self.world + dist.all_gather_object(gathered, mine, group=self.comm.world) + self.peer_shards = [list(rows) for rows in gathered] + else: + self.peer_shards = [mine] names = sorted(self.table.shapes) if not names: raise SwiftMegatronSyncError("no trainable parameter to synchronize") @@ -645,7 +658,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self._closed = False self.client = None self.stats = {"applied": 0, "pushed": 0, "push_bytes": 0, "dropped_pushes": 0, - "deferred_pushes": 0} + "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 # (iteration, "apply" | "push", fragment, version | pull step), same on every rank. @@ -664,6 +677,8 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self._raise_if(self._bcast_object(error)) if self.learner_id == 0: self._send_init() + if initial_cut_timeout > 0: + self._install_initial_cut(float(initial_cut_timeout)) # ------------------------------------------------------------------ setup def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, client_factory): @@ -679,22 +694,92 @@ def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, cli return client def _send_init(self) -> None: - """INIT_PARAMS of every fragment from this island's current masters.""" + """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.""" + + import queue + import threading + jobs: queue.Queue = queue.Queue(maxsize=1) + failure: list[str] = [] + sender = None + if self.rank == 0: + 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 + try: + self.client.send_init_parts(fid, parts) + except Exception as exc: + failure.append(f"INIT fragment {fid}: {type(exc).__name__}: {exc}") + del parts, job + + sender = threading.Thread(target=send, name="yeto-swift-init", daemon=True) + sender.start() error = None - for fid in range(self.num_fragments): - parts = self._export_fragment(fid, delta=False) - if self.rank == 0 and error is None: - try: + 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: - raise SwiftMegatronSyncError(self._export_error) - self.client.send_init_parts(fid, parts) - except Exception as exc: - error = f"INIT fragment {fid}: {type(exc).__name__}: {exc}" + error = f"INIT fragment {fid}: {self._export_error}" + elif failure: + error = failure[0] + else: + jobs.put((fid, parts)) + del parts + finally: + if sender is not None: + jobs.put(None) + sender.join() + if self.rank == 0 and error is None and failure: + error = failure[0] self._raise_if(self._bcast_object(error)) if self.rank == 0: log.info("sent INIT_PARAMS for %d fragments", self.num_fragments) + def _install_initial_cut(self, timeout: float) -> None: + """Apply the syncer's first broadcast of every fragment as it arrives, + before training starts (collective: every rank loops together). + + Without this the whole initial cut (``4 * N`` bytes on an f32-bulk + session, ~480 GB for 120B) queues in rank 0's client until the + first boundary drains it, one optimizer step after INIT. Here each + fragment is installed and freed as it lands. Gives up after + ``timeout`` seconds (the rest is applied at the boundaries as + before), and stops early at finalization or shutdown.""" + + import time + + started = time.monotonic() + 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 + plan = self._bcast_object(plan) + self._raise_if(plan.error) + 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) + del data + self.state.fragment_versions[fid] = version + self.state.global_step = max(self.state.global_step, version) + self.anchored[fid] = True + self.stats["applied"] += 1 + self.events.append((0, "apply", fid, version)) + if plan.finalizing or plan.shutdown: + break + if not plan.applies: + time.sleep(0.2) + 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) + # ------------------------------------------------------------------ collectives def _bcast_object(self, value: Any) -> Any: import torch.distributed as dist @@ -806,26 +891,49 @@ def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, bl from yeto.rl.megatron_masters import has_complete_master, write_masters - total = sum(self.table.numel(name) for name in names) - if self.rank == 0: - buffer = flat.to(device=self.comm.buffer_device) - else: - buffer = torch.empty(total, dtype=dtype, device=self.comm.buffer_device) - if self.world > 1: - dist.broadcast(buffer, src=0, group=self.comm.world) + wanted = set(names) offsets, offset = {}, 0 for name in names: offsets[name] = offset offset += self.table.numel(name) - wanted = set(names) + + def piece(buffer, name, part_dim, tp_rank, size): + whole = buffer[offsets[name]:offsets[name] + self.table.numel(name)].view(self.table.shapes[name]) + if part_dim >= 0: + whole = whole.narrow(part_dim, tp_rank * size, size) + return whole + + # Point-to-point scatter instead of a world broadcast: every rank gets + # only its own TP slices of this chunk (about DP copies of the chunk + # cross the wire in total, not world copies). local = [s for s in self.shards if s.name in wanted] + pieces = [] + if self.rank == 0: + buffer = flat.to(device=self.comm.buffer_device) + pending = [] + for peer in range(1, self.world): + rows = [row for row in self.peer_shards[peer] if row[0] in wanted] + if not rows: + continue + out = torch.cat([piece(buffer, *row).reshape(-1) for row in rows]) + pending.append((dist.isend(out, dst=peer, group=self.comm.world), out)) + for shard in local: + pieces.append(piece(buffer, shard.name, shard.part_dim, shard.tp_rank, + shard.local_shape[shard.part_dim] if shard.part_dim >= 0 else 0)) + for work, _ in pending: + work.wait() + del pending + elif local: + total = sum(int(s.param.numel()) for s in local) + incoming = torch.empty(total, dtype=dtype, device=self.comm.buffer_device) + dist.recv(incoming, src=0, group=self.comm.world) + at = 0 + for shard in local: + n = int(shard.param.numel()) + pieces.append(incoming[at:at + n].view(shard.local_shape)) + at += n targets = [] - for shard in local: - start = offsets[shard.name] - whole = buffer[start:start + self.table.numel(shard.name)].view(self.table.shapes[shard.name]) - if shard.part_dim >= 0: - size = shard.local_shape[shard.part_dim] - whole = whole.narrow(shard.part_dim, shard.tp_rank * size, size) + for shard, whole in zip(local, pieces, strict=True): 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) @@ -908,7 +1016,9 @@ 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 + started = time.perf_counter() self._apply_fragment(fid, data, self.bulk_torch, blend=True) + self.stats["apply_s"] += time.perf_counter() - started state.steps_at_reset[fid] = iteration state.units_at_reset[fid] = self.units_total state.fragment_versions[fid] = version @@ -922,7 +1032,9 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *, self._raise_if(pushes) error = None for fid, step, attempt, c_steps, c_units in pushes: + started = time.perf_counter() parts = self._export_fragment(fid, delta=True) + self.stats["export_s"] += time.perf_counter() - started self.events.append((iteration, "push", fid, step)) if self.rank == 0 and error is None: error = self._push(fid, step, attempt, c_steps, c_units, parts)