[bhargavchintam]

Fixes #1606

Summary

_group_norm_backward_kernel runs its tl.sum reductions in the input dtype, and group_norm_forward stores Mean and RSTD in the input dtype. In fp16 a channel sum can overflow, which turns the whole input gradient into NaN. In bf16 the sums keep only 8 bits, and the backward reads rounded statistics. luca-888 pointed at the Mean/RSTD part in the review of #1538 as a possible follow-up.

Source changes, all in src/liger_kernel/ops/group_norm.py (the new test is under Testing Done):

• W, X and the upstream gradient are upcast to fp32 when they are loaded, in both backward loops.
• Mean and RSTD are allocated as fp32.
• A new MAX_FUSED_SIZE_BACKWARD = 16384 limits the backward block. The casts double the vector width, and at the old limit of 65536 the backward gets about ten times slower (table below). My guess is register spills, but I did not measure register use.

triton_dtype and dW.to(dtype) are unchanged because #1253 works on those lines.

Details

• [NPU] Optimize GroupNorm: fused 2D fast path and launch overhead reduction #1542 has the same casts and fp32 statistics inside a larger NPU rework of this file. Its author put it on hold on 10-30 and says it is kept for reference, so I kept this change small and separate. If you would rather take it through [NPU] Optimize GroupNorm: fused 2D fast path and launch overhead reduction #1542, I can close this one.
• Max error against an fp64 reference on a B200 (my own script: it runs Liger and F.group_norm in fp64 on the same input and upstream gradient, then prints the largest absolute difference):

case	output	main	this PR	torch, same dtype
fp16 (1, 16, 256, 256), 4 groups, upstream ones	x.grad	all 1048576 non-finite	6.0e-8	1.1e-6
fp16 same	weight.grad	5.99	0.460	0.254
bf16 (1, 512, 64, 64), 32 groups, offset 1, upstream ones	x.grad	0.0176	3.6e-7	0.0176
bf16 same	weight.grad	16.93	0.488	16.08
bf16 (1, 16, 256, 256), 4 groups, upstream ones	weight.grad	34.29	0.957	0.957
bf16 (2, 64, 32, 32), 8 groups, offset 16	weight.grad	1.52	0.454	1.77

• Backward time on a B200 (my own script: group_norm_backward launched 100 times after a warmup, with torch.cuda.synchronize() around the loop, mean in µs). "Casts only" is this PR without the block limit:

shape, groups	dtype	main	casts only	this PR
(1, 16, 256, 256), 4	fp32	881	881	68
	bf16	70	705	58
	fp16	67	797	58
(16, 32, 256, 256), 8	fp32	2068	2070	134
	bf16	139	1518	124
	fp16	136	1645	123
(16, 512, 64, 64), 32	fp32	149	148	149
	bf16	175	142	142
	fp16	177	142	142
(64, 64, 8192), 8	fp32	116	116	116
	bf16	208	212	212
	fp16	203	212	212

I timed block limits of 65536, 16384, 8192 and 4096. 16384 was the fastest at 256×256, and at the other two shapes it matched 65536. The forward kernel is not changed here.

Testing Done

• Hardware Type: NVIDIA B200 (NGC PyTorch 26.10 container: torch 2.15.0a0, Triton 3.8.0, CUDA 13.4)
• run make test to ensure correctness
• run make checkstyle to ensure code style (I ran its two check-only steps, shown at the end, because the other two steps rewrite files)
• run make test-convergence to ensure convergence

I did not run make test or make test-convergence. The B200 was rented for this run, and on it I ran only the changed kernel's test file, test/transformers/test_group_norm.py, as below.

New test test_liger_group_norm_low_precision_backward has three cases (fp16 and bf16 at (1, 16, 256, 256) with an upstream of ones, and bf16 at (2, 8, 16) with an offset), and compares with fp32 autograd. The bias gradient is only compared where the true value fits in the dtype, since 65536 does not fit in fp16.

$ python3 -m pytest test/transformers/test_group_norm.py -k low_precision_backward --disable-warnings -q --tb=line

# main (e44de705) with only the new test applied
FAILED test/transformers/test_group_norm.py::test_liger_group_norm_low_precision_backward[shape0-4-0.0-ones-dtype0-0.001-0.01]
FAILED test/transformers/test_group_norm.py::test_liger_group_norm_low_precision_backward[shape1-4-0.0-ones-dtype1-0.01-0.01]
FAILED test/transformers/test_group_norm.py::test_liger_group_norm_low_precision_backward[shape2-4-16.0-randn-dtype2-0.01-0.01]
================= 3 failed, 11 deselected, 1 warning in 30.15s =================

# this PR
================= 3 passed, 11 deselected, 1 warning in 26.52s =================

$ python3 -m pytest test/transformers/test_group_norm.py --disable-warnings -q --tb=line
======================== 14 passed, 1 warning in 31.98s ========================

The two snippets from the issue, run on this branch (on main they print finite x.grad: 0 / 1048576 and max |x.grad| (expected 0): 466.0):

finite x.grad: 1048576 / 1048576
max |x.grad| (expected 0): 0.0

$ ruff check --output-format=concise .
All checks passed!

$ ruff format --check --diff .
500 files already formatted

[bhargavchintam]

I ran make test on a B200 with this PR and #1607 applied together on main (e44de705; torch 2.15.0a0, Triton 3.8.0, transformers 5.22.2):

= 22 failed, 5702 passed, 1054 skipped, 21 xfailed, 318 warnings, 6 rerun in 2820.14s (0:47:00) =

None of the 22 is in test_group_norm.py or test_dyt.py. I re-ran those 22 on plain e44de705 and on this tree. The same 21 fail on both, and the other one, a DPO loss test, passes on both when run that way. I did not look into why those 21 fail, and I did not run make test-convergence.