allreduce in precast dtype (#15689)

2026-04-13 20:24:12 -07:00
parent 5683126844
commit 2b8d303f75
7 changed files with 49 additions and 5 deletions
@@ -14,6 +14,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
@@ -14,6 +14,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
@@ -15,6 +15,7 @@ export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
 export FP8=1
+export ALLREDUCE_CAST=1

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=8 MP=1 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=4
@@ -1,6 +1,7 @@
 import gc, unittest
 from tinygrad import Tensor, GlobalCounters, dtypes
 from tinygrad.engine.jit import TinyJit
+from tinygrad.helpers import Context

 class TestMultiRamUsage(unittest.TestCase):
  def setUp(self):
@@ -139,6 +140,19 @@ class TestMultiRamUsage(unittest.TestCase):
    mem_4 = run_layers(4)
    self.assertEqual(mem_2, mem_4, f"graph memory should not grow with layers: 2 layers={mem_2}, 4 layers={mem_4}")

+  def test_allreduce_cast_dtype_memory(self):
+    N = 32
+    devices_2 = ("NULL:1", "NULL:2")
+    mem = {}
+    for allreduce_cast in (0, 1):
+      GlobalCounters.reset()
+      with Context(ALLREDUCE_CAST=allreduce_cast, SCACHE=0):
+        x = Tensor.empty((N, N), dtype=dtypes.bfloat16, device="NULL:1").shard(devices_2, axis=0)
+        x.sum(0).realize()
+      mem[allreduce_cast] = GlobalCounters.global_mem
+    # with ALLREDUCE_CAST, allreduce copies happen in bf16 (2 bytes) instead of fp32 (4 bytes)
+    self.assertLess(mem[1], mem[0])
+
 class TestMultiAxis(unittest.TestCase):
  def test_reshape_shard_invalid(self):
    devices = ("NULL:0", "NULL:1")
@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor
+from tinygrad import Tensor, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops

@@ -25,5 +25,28 @@ class TestRingAllReduce(unittest.TestCase):
      out = t.sum(0)
      self.assertListEqual(out.tolist(), [4]*N*100)

+class TestAllreduceCast(unittest.TestCase):
+  def _get_copy_dtypes(self, dtype, allreduce_cast):
+    ds = tuple(f"CPU:{i}" for i in range(2))
+    with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0):
+      t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0)
+      schedules = t.sum(0).schedule_with_vars()[0]
+      return {si.bufs[0].dtype.scalar() for si in schedules if si.ast.op is Ops.COPY}
+
+  def test_allreduce_cast_bf16(self):
+    # with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32
+    self.assertNotIn(dtypes.float, self._get_copy_dtypes(dtypes.bfloat16, allreduce_cast=1))
+    self.assertIn(dtypes.float, self._get_copy_dtypes(dtypes.bfloat16, allreduce_cast=0))
+
+  def test_allreduce_cast_half(self):
+    self.assertNotIn(dtypes.float, self._get_copy_dtypes(dtypes.half, allreduce_cast=1))
+    self.assertIn(dtypes.float, self._get_copy_dtypes(dtypes.half, allreduce_cast=0))
+
+  def test_allreduce_cast_float32_noop(self):
+    # float32 should not be affected by ALLREDUCE_CAST (no promotion happens)
+    dtypes_on = self._get_copy_dtypes(dtypes.float, allreduce_cast=1)
+    dtypes_off = self._get_copy_dtypes(dtypes.float, allreduce_cast=0)
+    self.assertEqual(dtypes_on, dtypes_off)
+
 if __name__ == '__main__':
  unittest.main()
@@ -224,7 +224,7 @@ WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1),
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
-RING, ALL2ALL = ContextVar("RING", 1), ContextVar("ALL2ALL", 0)
+RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
@@ -1,4 +1,4 @@
-from tinygrad.helpers import all_same, prod, getenv
+from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, should_resolve_call
 from tinygrad.dtype import dtypes
 from tinygrad.schedule.allreduce import handle_allreduce
@@ -66,8 +66,12 @@ def alu_multi(root:UOp):
 def reduce_multi(root:UOp, multi:UOp):
  op, axis = root.arg
  if multi.axis is not None and multi.axis in axis:
-    # all-reduce on sharded axes
-    return multi.src[0]._rop(op, axis).allreduce(op, multi.device)
+    local = multi.src[0]._rop(op, axis)
+    # allreduce in pre-cast dtype when sum_acc_dtype promoted from bf16/half
+    if ALLREDUCE_CAST and multi.src[0].op is Ops.CAST and multi.src[0].src[0].dtype.scalar() in (dtypes.bfloat16, dtypes.half):
+      orig_dtype = multi.src[0].src[0].dtype
+      return local.cast(orig_dtype).allreduce(op, multi.device).cast(local.dtype)
+    return local.allreduce(op, multi.device)
  # reduce on non sharded axes, piecewise is fine. if axis is None this is also correct
  return multi.src[0]._rop(op, axis).multi(axis=multi.axis)