nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

test for extra copy in allreduce_cast (#16740)

Browse Source 
* test for extra copy in allreduce_cast

* simpler

* cleanup
This commit is contained in:
qazal
2026-07-28 19:02:54 +09:00
committed by GitHub
parent 1380d6cc5a
commit e3b3eea1e2
1 changed files with 22 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_multitensor.py
+22
View File
@@ -127,6 +127,28 @@ class TestMultiTensor(unittest.TestCase):
with Context(RING=use_ring):
np.testing.assert_equal(t.shard(devices_2, axis=axis).sum().item(), 10)
def test_allreduce_cast_half(self, assign=False, kernel_count=8):
devices = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
a_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
b_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
a = a_src.shard(devices, axis=0).realize()
b = b_src.shard(devices, axis=0).realize()
# assigning creates a copy of the output before allreduce
if assign:
tst = Tensor.empty_like(b)
tst.assign(a + b)
else:
tst = a + b
tst = tst.float().sum(0)
GlobalCounters.reset()
with Context(ALLREDUCE_CAST=1, RING=0, ALL2ALL=0):
tst.realize()
kernel_count = GlobalCounters.kernel_count
np.testing.assert_allclose(tst.numpy(), (a_src.numpy()+b_src.numpy()).sum(0))
self.assertEqual(kernel_count, kernel_count)
def test_allreduce_cast_half_assign(self): self.test_allreduce_cast_half(assign=True, kernel_count=10)
def test_multiple_to_single_device(self):
kernel_counts = {}
for ring in (0, 2):