diff --git a/src/tensorfold/cuda/comm.py b/src/tensorfold/cuda/comm.py index 9e9b16a1d6..bc322252d1 100644 --- a/src/tensorfold/cuda/comm.py +++ b/src/tensorfold/cuda/comm.py @@ -72,6 +72,8 @@ def __init__(self, rank: int, world: int, master: str, port: int, *, timeout_s: ctypes.c_void_p] lib.ncclAllReduce.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p] + lib.ncclReduceScatter.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, + ctypes.c_void_p, ctypes.c_void_p] self.store = TCPStore(master, port, world, rank == 0, timeout=timedelta(seconds=timeout_s)) uid = _UniqueId() if rank == 0: @@ -108,6 +110,15 @@ def all_reduce(self, value: torch.Tensor) -> None: self._check(self.lib.ncclAllReduce(value.data_ptr(), value.data_ptr(), value.numel(), _DTYPES[value.dtype], NCCL_SUM, self.comm, stream)) + def reduce_scatter(self, send: torch.Tensor, recv: torch.Tensor) -> None: + """recv [n] <- ncclSum over the ranks of their send's rank-th block (send [world * n]).""" + + if send.numel() != recv.numel() * self.world or send.dtype != recv.dtype: + raise ValueError("reduce_scatter: send must hold world x recv of the same dtype") + stream = torch.cuda.current_stream().cuda_stream + self._check(self.lib.ncclReduceScatter(send.data_ptr(), recv.data_ptr(), recv.numel(), _DTYPES[send.dtype], + NCCL_SUM, self.comm, stream)) + def ready(self, label: str, *, every: float = 60.0, timeout: float = 3600.0) -> None: """Every rank finishes ``label`` before any goes on; a rank missing after ``timeout`` s is named.""" diff --git a/src/tensorfold/families/deepseek_v41/cuda/model.py b/src/tensorfold/families/deepseek_v41/cuda/model.py index c33ee49303..df2cc88fe3 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/model.py +++ b/src/tensorfold/families/deepseek_v41/cuda/model.py @@ -21,6 +21,9 @@ RAW = 64 # per-position compressor inputs kept (a verify window rolls back by length alone) RING_EXTRA = 16 # window ring slots beyond the 128-token window (verify windows never clobber) KV_QUANT = os.environ.get("TF_DS_KV", "native") != "bf16" +# TF_DS_RS=1 (two ranks): a prompt chunk's attention and MoE partials reach the mHC post as bf16(p0 + p1), from an +# fp32 reduce-scatter and a bf16 all-gather (Comm.reduced). The same bits as the fp32 gather, 3/4 of its bytes. +RS = os.environ.get("TF_DS_RS", "0") == "1" KERNELS = os.environ.get("TF_DS_KERNELS", "1") != "0" # 0: the plain-torch phase-1 path (A/B and debugging) TIMING = os.environ.get("TF_DS_TIMING", "0") == "1" # per-section wall times (synchronizing; profiling only) # TF_DS_REPLAY_FLOOR=1: a prompt chunk of RING_EXTRA rows or fewer whose window reaches below the bounded replay's cut @@ -115,6 +118,21 @@ def gather(self, x: torch.Tensor) -> torch.Tensor: self.nccl.all_gather(x.view(-1), out.view(-1)) return out + def reduced(self, x: torch.Tensor) -> torch.Tensor: + """The partials [n, d] the mHC post takes: ``gather(x)`` [world, n, d], or TF_DS_RS on two ranks [1, n, d] + bf16 = bf16(p0 + p1) from an fp32 reduce-scatter and a bf16 all-gather (one fp32 add of two operands, the same + bits in either order: the same y from 3/4 of the bytes).""" + + rdma = getattr(self.nccl, "rdma", None) + if not (RS and self.world == 2 and x.shape[0] % 2 == 0 and x.dtype == F32) or (rdma and rdma.fits(x)): + return self.gather(x) + x = x.contiguous() + half = torch.empty((x.shape[0] // 2, *x.shape[1:]), dtype=F32, device=x.device) + self.nccl.reduce_scatter(x.view(-1), half.view(-1)) + y = torch.empty((1, *x.shape), dtype=BF16, device=x.device) + self.nccl.all_gather(half.to(BF16).view(-1), y.view(-1)) + return y + def sum(self, x: torch.Tensor) -> torch.Tensor: """Sum rank partials; NCCL can modify x, so callers must treat x as consumed.""" @@ -1377,7 +1395,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r pa = self.attention_k(lay, x, sc, start, shared, pos, floor=floor, kv_done=kv_done and lay.idx == c.n_layers // 2) with _T("gather"): - g = self.comm.gather(pa) + g = self.comm.reduced(pa) with _T("hc"): fn, scale, base = lay.hc_ffn if hc_pf: @@ -1394,7 +1412,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r with _T("moe"): pm = self.moe(lay, x, img=img) with _T("gather"): - g = self.comm.gather(pm) + g = self.comm.reduced(pm) with _T("hc"): K.hc_post(g, h, post, comb, h) pre, pre_f = pre_f, pre diff --git a/tests/test_dsv41_reduced.py b/tests/test_dsv41_reduced.py new file mode 100644 index 0000000000..ce712d8907 --- /dev/null +++ b/tests/test_dsv41_reduced.py @@ -0,0 +1,101 @@ +"""Comm.reduced (TF_DS_RS): two ranks get bf16(p0 + p1), the y the mHC post kernels form from the fp32 gather, and +the gather itself where the reduce-scatter does not apply.""" + +from __future__ import annotations + +import threading + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") + +from tensorfold.families.deepseek_v41.cuda import model as M + + +class Link: + """Two ranks' collectives on the host: each call waits for both ranks' buffers, as NCCL does.""" + + def __init__(self): + self.barrier, self.sends, self.calls = threading.Barrier(2), [None, None], [] + + def swap(self, rank, send): + self.sends[rank] = send.clone() + self.barrier.wait() + both = list(self.sends) + self.barrier.wait() + return both + + +class Rank: + def __init__(self, link, rank, rdma=None): + self.link, self.rank, self.world = link, rank, 2 + if rdma is not None: + self.rdma = rdma + + def reduce_scatter(self, send, recv): + a, b = self.link.swap(self.rank, send) + self.link.calls.append(("reduce_scatter", self.rank)) + recv.copy_((a + b).view(2, -1)[self.rank]) # ncclSum of two operands + + def all_gather(self, send, recv): + a, b = self.link.swap(self.rank, send) + self.link.calls.append(("all_gather", send.dtype)) + recv.copy_(torch.cat([a, b])) + + +class Fits: + def fits(self, x): + return True + + +def both_ranks(p0, p1, rdma=None): + """Comm.reduced on two threads, one a rank: (rank 0's result, rank 1's result, the calls).""" + + link, out = Link(), [None, None] + + def run(rank, mine): + c = M.Comm.__new__(M.Comm) + c.nccl, c.world, c._buf = Rank(link, rank, rdma), 2, {} + out[rank] = c.reduced(mine) + + ts = [threading.Thread(target=run, args=(r, p)) for r, p in ((0, p0), (1, p1))] + for t in ts: + t.start() + for t in ts: + t.join() + return out[0], out[1], link.calls + + +def parts(n=8, d=64, seed=0): + g = torch.Generator().manual_seed(seed) + return torch.randn((n, d), generator=g) * 3, torch.randn((n, d), generator=g) * 1e-3 + + +def test_both_ranks_get_the_posts_y(monkeypatch): + monkeypatch.setattr(M, "RS", True) + p0, p1 = parts() + y0, y1, calls = both_ranks(p0, p1) + want = (p0 + p1).to(torch.bfloat16)[None] # what _hc_post forms from the gather, before the mix + for y in (y0, y1): + assert y.dtype == torch.bfloat16 and torch.equal(y.view(torch.int16), want.view(torch.int16)) + assert sorted(c[0] for c in calls) == ["all_gather", "all_gather", "reduce_scatter", "reduce_scatter"] + assert all(c[1] == torch.bfloat16 for c in calls if c[0] == "all_gather") + + +@pytest.mark.parametrize("case", ["off", "odd rows", "rdma holds it"]) +def test_falls_back_to_the_gather(monkeypatch, case): + monkeypatch.setattr(M, "RS", case != "off") + p0, p1 = parts(n=7 if case == "odd rows" else 8) + y0, y1, calls = both_ranks(p0, p1, Fits() if case == "rdma holds it" else None) + for y in (y0, y1): + assert y.dtype == torch.float32 and torch.equal(y, torch.stack([p0, p1])) + assert [c[0] for c in calls] == ["all_gather", "all_gather"] + + +def test_one_rank_is_the_gather(monkeypatch): + monkeypatch.setattr(M, "RS", True) + c = M.Comm.__new__(M.Comm) + c.nccl, c.world, c._buf = None, 1, {} + p0, _ = parts() + assert torch.equal(c.reduced(p0), p0[None])