Fix group norm backward for inputs with more than three dimensions

[truong-v]

Summary

group_norm_backward reads hidden_size = dY.shape[-1] before flattening the gradient while group_norm_forward flattens first and then measures, so for an (N, C, H, W) input the backward pass uses W where it means the per-channel element count H*W: it allocates the input-gradient buffer too small, addresses it with the input's real strides (so the kernel writes past the allocation), reduces dW/dB over one row per channel, and finally raises RuntimeError: shape '[N, C, H, W]' is invalid for input of size ... on the closing view. LigerGroupNorm accepts any input with dim() >= 3 and its forward pass is correct for 4-D and 5-D, so this hits the standard convolutional case at the first backward pass; 3-D inputs are unaffected, since there dY.shape[-1] already equals the per-channel element count.

Fixes #1375

Details

• src/liger_kernel/ops/group_norm.py: move the hidden_size computation below the flatten and divide by channels_per_group, matching the forward pass. For 3-D inputs the kernel receives exactly the same arguments as before, so that path is unchanged; this is host-side scalar arithmetic, no kernel change and nothing to benchmark.
• test/transformers/test_group_norm.py: factor the existing test body into _test_liger_group_norm(shape, ...) and add test_liger_group_norm_spatial_dims covering (2, 6, 4, 8), (4, 32, 16, 16) and (2, 8, 2, 4, 4). The existing cases keep their parameters.

The NPU backend's own group_norm_backward (ops/backends/_ascend/ops/group_norm.py) already flattens first and then derives its per-channel count as hidden_size // channels_per_group, so this brings the default path in line with it.

(2, 6, 4, 8), 3 groups, fp32 on H100, against torch.nn.GroupNorm:

	forward	dX	dW	dB
before	2.4e-07	RuntimeError	RuntimeError	RuntimeError
after	2.4e-07	4.8e-07	1.9e-06	0.0

Under compute-sanitizer --tool memcheck with PYTORCH_NO_CUDA_MEMORY_CACHING=1, the same 4-D backward reports 28 invalid global writes at the DX store (30 errors total) before the fix and 0 after.

Testing Done

• test/transformers/test_group_norm.py: 8 passed (5 pre-existing + 3 new).

• Reverting only the group_norm.py change makes the 3 new cases fail and leaves the 5 pre-existing ones passing.

• test_group_norm.py is the only file in the suite that exercises this op, so nothing else in make test is affected. I ran ruff check . and ruff format --check . over the repo, but not the full suite or the convergence suite locally — boxes below reflect that.

• Hardware Type: H100 NVL

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

Environment

• Liger-Kernel a5d795efd2c1436549e70118ef519134e9c27833 (main) + this branch, editable install
• GPU: NVIDIA H100 NVL
• torch 2.6.0+cu124, triton 3.2.0, CUDA 12.4, transformers 5.14.1, Python 3.10.20