diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh index 23971cd2e4..f7b9ecdc2a 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh @@ -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} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh index e13b070ac1..d8e0af1996 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh @@ -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} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh index 5f8e71f880..7923437982 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh @@ -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 diff --git a/test/null/test_multitensor.py b/test/null/test_multitensor.py index 2b748386ce..a846de8abd 100644 --- a/test/null/test_multitensor.py +++ b/test/null/test_multitensor.py @@ -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") diff --git a/test/unit/test_allreduce.py b/test/unit/test_allreduce.py index a30309660a..e894268643 100644 --- a/test/unit/test_allreduce.py +++ b/test/unit/test_allreduce.py @@ -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() diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index a8820f1f56..419a3d4ff7 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -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) diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 129d4d9d7d..00e502ea76 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -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)