diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh index 4b8493c0ba..52dcfa32a8 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh @@ -18,6 +18,7 @@ export FP8=${FP8:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} export FAST_CE=${FAST_CE:-1} export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} +export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1} export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1} export FUSED_SILU_W13=${FUSED_SILU_W13:-1} export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh index 71ce310be2..7369f12199 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh @@ -16,7 +16,8 @@ export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1} export FP8=${FP8:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} export FAST_CE=${FAST_CE:-0} -export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} +export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0} +export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0} export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0} export FUSED_SILU_W13=${FUSED_SILU_W13:-0} export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh index f38282a354..3589721923 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh @@ -18,6 +18,7 @@ export FP8=${FP8:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} export FAST_CE=${FAST_CE:-1} export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1} +export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1} export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1} export FUSED_SILU_W13=${FUSED_SILU_W13:-1} export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh index 04b40de827..3703f0ab69 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh @@ -10,9 +10,19 @@ export DEVICE_IN_FUNCTION_BUG=1 export DEBUG=${DEBUG:-0} export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1} export ALL2ALL=${ALL2ALL:-1} -export USE_ATOMICS=${USE_ATOMICS:-0} +export USE_ATOMICS=${USE_ATOMICS:-1} 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 FAST_CE=${FAST_CE:-0} +export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0} +export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0} +export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0} +export FUSED_SILU_W13=${FUSED_SILU_W13:-0} +export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0} +export SPLIT_W13=${SPLIT_W13:-1} export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1} export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh index 21c1ace10c..d8ee6a0538 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh @@ -19,6 +19,7 @@ export FP8=1 export ALLREDUCE_CAST=1 export FAST_CE=1 export FUSED_INPUT_QUANTIZE=1 +export FUSED_GRAD_QUANTIZE=1 export FUSED_ADD_NORM_MUL_QUANTIZE=1 export FUSED_SILU_W13=1 export FUSED_PAD_GRAD_ACCUM=1 diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index 2220f76a12..5486c7598d 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -2713,12 +2713,20 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp): gbase = gradient.base if hasattr(gradient, "base") else gradient mailbox_entry = _grad_fp8_mailbox.pop(gbase, None) or _grad_fp8_mailbox.pop(gradient, None) if mailbox_entry is not None: - g_fp8_u, inv_scale_u, _new_amax_u, store_effect = mailbox_entry + g_fp8_u, inv_scale_u = mailbox_entry g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]] g_scale = Tensor(inv_scale_u, device=a.device) else: assert grad_amax_state is not None, "fp8 matmul bwd needs either a mailbox entry or a grad_amax_state" - g_fp8, g_scale, _, store_effect = quantize_fp8_delayed(g_t, Tensor(grad_amax_state, device=a.device)) + if getenv("FUSED_GRAD_QUANTIZE", 0): + g_fp8, g_scale, _, store_effect = quantize_fp8_delayed(g_t, Tensor(grad_amax_state, device=a.device)) + assert g_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {g_fp8.uop.op}" + g_fp8 = Tensor(g_fp8.uop.replace(src=g_fp8.uop.src + (store_effect,)), device=a.device) + else: + grad_amax_t = Tensor(grad_amax_state, device=a.device) + g_fp8, g_scale, new_grad_amax = quantize_fp8(g_t, amax_state=grad_amax_t) + store_effect = grad_amax_state.store(new_grad_amax.uop) + g_fp8 = Tensor(g_fp8.contiguous().uop.after(store_effect), device=a.device) # dgrad: uses g_scale * x_scale * w_scale grad_a = asm_gemm(g_fp8, b_t, x_scale=g_scale * s_x_t, w_scale=s_w_t) # wgrad: no w_scale @@ -2729,8 +2737,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp): else: g_fp8_T = g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1) grad_b = asm_gemm(g_fp8_T, a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t) - # Attach the delayed-amax store effect (if any) to grad_a so realizing grads commits the amax update. - ret = (None, grad_a.uop.after(store_effect), grad_b.uop, None, None) + ret = (None, grad_a.uop, grad_b.uop, None, None) if len(inputs) == 6: ret = ret + (None,) return ret else: diff --git a/extra/llama_kernels/cast_amax/__init__.py b/extra/llama_kernels/cast_amax/__init__.py index c7d79fc764..45161bc8a1 100644 --- a/extra/llama_kernels/cast_amax/__init__.py +++ b/extra/llama_kernels/cast_amax/__init__.py @@ -53,8 +53,10 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp): inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX new_grad_amax = scalar_amax(grad_amax_buf) store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop) - # Stash fp8 companion + amax store for cdna_asm_gemm's bwd to attach to grad_a. - _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8.uop, inv_scale.uop, new_grad_amax.uop, store_effect) + assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}" + grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,)) + # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a. + _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8_uop, inv_scale.uop) return (None, None, grad_xw13.uop, None, None) def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]: