diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index ac4bd1f675..d87214934d 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -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):