Skip to content

Fix group norm backward for inputs with more than three dimensions #1376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
truong-v wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix group norm backward for inputs with more than three dimensions #1376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/liger_kernel/ops/group_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,9 +242,11 @@ def group_norm_forward(X, num_channels, num_groups, W, B, eps):
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)
# Number of elements per channel, so it has to be measured after the flatten:
# the last dimension of an unflattened (N, C, H, W) input is only W
hidden_size = dY.shape[-1] // channels_per_group
DX = torch.empty(
(batch_size, num_groups, hidden_size * channels_per_group),
dtype=X.dtype,
Expand Down
58 changes: 40 additions & 18 deletions test/transformers/test_group_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,26 +7,10 @@
device = infer_device()


@pytest.mark.parametrize(
"batch_size, num_channels, num_groups, hidden_size",
[
(1, 1, 1, 3), # minimal
(1, 32, 32, 4), # group == channel
(16, 32, 1, 4096), # single group
(2, 63, 21, 2163), # non-aligned hidden
(16, 48, 12, 8192), # large hidden
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-4, 1e-4),
],
)
def test_liger_group_norm(batch_size, num_channels, num_groups, hidden_size, dtype, atol, rtol):
def _test_liger_group_norm(shape, num_channels, num_groups, dtype, atol, rtol):
torch.manual_seed(0)

_tensor = torch.randn(batch_size, num_channels, hidden_size, dtype=dtype, device=device)
_tensor = torch.randn(*shape, dtype=dtype, device=device)

liger_x = _tensor.clone().detach().requires_grad_(True)
torch_x = _tensor.clone().detach().requires_grad_(True)
Expand All @@ -50,3 +34,41 @@ def test_liger_group_norm(batch_size, num_channels, num_groups, hidden_size, dty
assert torch.allclose(liger_x.grad, torch_x.grad, atol=atol, rtol=rtol)
assert torch.allclose(liger_ln.bias.grad, torch_ln.bias.grad, atol=atol, rtol=rtol), "Bias grads different"
assert torch.allclose(liger_ln.weight.grad, torch_ln.weight.grad, atol=atol, rtol=rtol), "Weight grads different"


@pytest.mark.parametrize(
"batch_size, num_channels, num_groups, hidden_size",
[
(1, 1, 1, 3), # minimal
(1, 32, 32, 4), # group == channel
(16, 32, 1, 4096), # single group
(2, 63, 21, 2163), # non-aligned hidden
(16, 48, 12, 8192), # large hidden
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-4, 1e-4),
],
)
def test_liger_group_norm(batch_size, num_channels, num_groups, hidden_size, dtype, atol, rtol):
_test_liger_group_norm((batch_size, num_channels, hidden_size), num_channels, num_groups, dtype, atol, rtol)


@pytest.mark.parametrize(
"shape, num_channels, num_groups",
[
((2, 6, 4, 8), 6, 3), # convolutional features
((4, 32, 16, 16), 32, 8), # larger spatial plane
((2, 8, 2, 4, 4), 8, 4), # volumetric features
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-4, 1e-4),
],
)
def test_liger_group_norm_spatial_dims(shape, num_channels, num_groups, dtype, atol, rtol):
_test_liger_group_norm(shape, num_channels, num_groups, dtype, atol, rtol)