LigerGroupNorm backward writes out of bounds and then raises for any input with more than 3 dimensions

[truong-v]

🐛 Describe the bug

LigerGroupNorm accepts any input with at least 3 dimensions (the layer asserts exactly that), and the forward pass matches torch.nn.GroupNorm for 4-D and 5-D inputs. The backward pass does not: it measures hidden_size as dY.shape[-1] before flattening the gradient, while the forward pass flattens first. For an (N, C, H, W) input that gives W instead of the per-channel element count H*W.

The consequences run in order: the input-gradient buffer is allocated N*C*W floats but the kernel addresses it with the input's real strides, so it writes past the allocation; the dW/dB sums and the normalising divisor cover only the first row of every channel; and the host finally raises RuntimeError: shape '[N, C, H, W]' is invalid for input of size ... when it views the undersized buffer back. So a training run using LigerGroupNorm on convolutional features — the usual case for group norm — corrupts GPU memory and then dies at the first backward pass. 3-D inputs are unaffected, since there dY.shape[-1] happens to equal the per-channel element count, which is why test/transformers/test_group_norm.py (3-D only) passes. Issue #883 pointed at the same function's addressing without a reproducer; this is one.

Liger-Kernel/src/liger_kernel/ops/group_norm.py

Lines 242 to 252 in a5d795e

	def group_norm_backward(dY, X, W, B, Mean, RSTD, num_channels, num_groups):
	shape = dY.shape
	batch_size = shape[0]
	hidden_size = dY.shape[-1]
	channels_per_group = num_channels // num_groups
	dY = dY.view(batch_size, num_groups, -1)
	DX = torch.empty(
	(batch_size, num_groups, hidden_size * channels_per_group),
	dtype=X.dtype,
	device=X.device,
	)

The op was added in #353.

Reproduce

import torch

from liger_kernel.transformers.group_norm import LigerGroupNorm

x = torch.randn(2, 6, 4, 8, device="cuda", requires_grad=True)  # (N, C, H, W)
gn = LigerGroupNorm(num_channels=6, num_groups=3).cuda()

y = gn(x)  # forward matches torch.nn.GroupNorm
y.backward(torch.randn_like(y))  # raises

  File "/.../liger_kernel/ops/group_norm.py", line 310, in backward
    DX, DW, DB = group_norm_backward(dY, X, W, B, Mean, RSTD, ctx.num_channels, ctx.num_groups)
  File "/.../liger_kernel/ops/group_norm.py", line 278, in group_norm_backward
    return DX.view(*shape), DW, DB
RuntimeError: shape '[2, 6, 4, 8]' is invalid for input of size 96

The out-of-bounds writes happen before that, in the kernel: under compute-sanitizer with the caching allocator disabled (PYTORCH_NO_CUDA_MEMORY_CACHING=1) the same script reports invalid 4-byte global writes at _group_norm_backward_kernel, group_norm.py:203, landing 233 bytes after a 24-byte allocation (ERROR SUMMARY: 30 errors); 5-D inputs fail the same way, and a 3-D input of the same element count is clean.

Versions

• Liger-Kernel commit a5d795efd2c1436549e70118ef519134e9c27833 (main), editable install
• GPU: NVIDIA H100 NVL
• Liger Kernel version: 0.8.1
• PyTorch version: 2.6.0+cu124, CUDA 12.4
• Triton version: 3.2.0
• Transformers version: 5.14.1
• Python 3.10.20, Linux 5.15.0