From 079685384594fbf404308b2f6b286e35b8480765 Mon Sep 17 00:00:00 2001 From: b1tg <33436708+b1tg@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:53:12 +0800 Subject: [PATCH] support symbolic shapes in allreduce (#17364) --- test/unit/test_allreduce.py | 9 ++++++++- tinygrad/schedule/allreduce.py | 11 +++++++---- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/test/unit/test_allreduce.py b/test/unit/test_allreduce.py index 29aec4591b..0997e55979 100644 --- a/test/unit/test_allreduce.py +++ b/test/unit/test_allreduce.py @@ -1,5 +1,5 @@ import unittest -from tinygrad import Tensor, dtypes +from tinygrad import Tensor, UOp, dtypes from tinygrad.helpers import Context from tinygrad.uop.ops import Ops @@ -43,6 +43,13 @@ class TestRingAllReduce(unittest.TestCase): self.assertEqual(len(sinks), 2) self.assertTrue(all(dst != src for dst, src in pairs)) + def test_symbolic_shape(self): + rows = UOp.variable("rows", 1, 4).bind(3) + t = Tensor.ones(4, 4).shard(("CPU:0", "CPU:1"), axis=1).realize() + out = t[:rows].sum(1).realize() + self.assertEqual(out.shape, (rows,)) + self.assertTrue((out == 4).all().item()) + def test_correct_ring(self): with Context(RING=2): N = 4 diff --git a/tinygrad/schedule/allreduce.py b/tinygrad/schedule/allreduce.py index 866b4d3189..f5cc4c8e95 100644 --- a/tinygrad/schedule/allreduce.py +++ b/tinygrad/schedule/allreduce.py @@ -5,24 +5,27 @@ from tinygrad.uop.ops import UOp # *** allreduce implementation *** def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: if not isinstance(buf.device, tuple): return None - assert all_int(buf.shape), f"does not support symbolic shape {buf.shape}" ndev, shape, numel = len(buf.device), buf.shape, prod(buf.shape) op, device = red.arg # ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically) # fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks. - use_all2all = (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1)) - use_ring = not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1)) + concrete = all_int(shape) + use_all2all = concrete and (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1)) + use_ring = concrete and not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1)) if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}") + if not concrete: buf = buf.pad_to(buf.max_shape) # contiguous before we copy it buf = buf.contiguous() # naive: copy to all devices. if you shrink later, that'll be handled if not use_ring and not use_all2all: - return functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)]) + out = functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)]) + return out if concrete else out.shrink_to(shape) # chunk data into ndev pieces + assert isinstance(numel, int) factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1) base, left = divmod(numel // factor, ndev) chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (ndev - left), initial=0)))