linkedin / linkedin/Liger-Kernel
LigerGroupNorm backward writes out of bounds and then raises for any input with more than 3 dimensions
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
## 🐛 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.
https://github.com/linkedin/Liger-Kernel/blob/a5d795efd2c1436549e70118ef519134e9c27833/src/liger_kernel/ops/group_norm.py#L242-L252
The op was added in #353.
## Reproduce
```python
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
Contributor guide
Research direction
Start with the reproducer in the issue, then read src/liger_kernel/ops/group_norm.py around lines 242-252 and the backward path near line 310. Compare the existing 3-D coverage in test/transformers/test_group_norm.py with equivalent 4-D and 5-D cases. Done means multidimensional backward passes no longer raises or writes out of bounds, while results remain consistent with torch.nn.GroupNorm.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100