linkedin / linkedin/Liger-Kernel

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

Open
#1,375 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.