From e3c2e34c269216483942ebb90b67d910df6f4053 Mon Sep 17 00:00:00 2001 From: Prathamesh Jadhav <55660103+lollinng@users.noreply.github.com> Date: Fri, 5 Jun 2026 02:54:09 +0530 Subject: [PATCH] Cast GroupNorm fp16 backward grad accumulators to fp16, not bfloat16 group_norm_backward set triton_dtype = bfloat16 for every non-fp32 input. For fp16 inputs that means the dW/dB gradient accumulators were rounded to bfloat16 (8 mantissa bits) before being atomic-added into the fp16 DW/DB buffers (10 mantissa bits) -- losing precision for no reason, since the buffers are fp16. Map fp16 -> fp16 so the atomic_add dtype matches the buffer dtype. This is a type-consistency / precision fix, not a crash fix: on the Triton version tested the bf16->fp16 atomic_add is silently coerced rather than erroring. Co-authored-by: Claude Opus 4.8 (1M context) --- src/liger_kernel/ops/group_norm.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/liger_kernel/ops/group_norm.py b/src/liger_kernel/ops/group_norm.py index 865fc337f..fcdacf801 100644 --- a/src/liger_kernel/ops/group_norm.py +++ b/src/liger_kernel/ops/group_norm.py @@ -252,7 +252,17 @@ def group_norm_backward(dY, X, W, B, Mean, RSTD, num_channels, num_groups): ) DW = torch.zeros((num_channels), dtype=W.dtype, device=W.device) DB = torch.zeros((num_channels), dtype=B.dtype, device=B.device) - triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16 + # `triton_dtype` is the dtype the dW/dB accumulators are cast to before being + # atomic-added into the DW/DB buffers, so it should match those buffers' + # dtype (W.dtype / B.dtype). Mapping every non-fp32 dtype to bfloat16 meant + # that for fp16 inputs a bfloat16-rounded gradient was atomic-added into an + # fp16 buffer, which loses precision unnecessarily. Map fp16 -> fp16. + if X.dtype == torch.float32: + triton_dtype = tl.float32 + elif X.dtype == torch.float16: + triton_dtype = tl.float16 + else: + triton_dtype = tl.bfloat16 BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(hidden_size)) _group_norm_backward_kernel[(batch_size, num_groups)](