fix groupnorm int32 index overflow

[tlogn]

This PR fix groupnorm int32 index calculate overflow when hwc is large, as hwc is of int data type. The problem could be reproduced by code below.
@crcrpar please review, thanks !

from apex.contrib.group_norm import GroupNorm as ApexGroupNorm
import torch
layer = ApexGroupNorm(32, 128, dtype=torch.bfloat16, device='cuda', act='silu')

x = torch.randn(1, 128, 16128, 1200, dtype=torch.bfloat16, device='cuda').to(memory_format=torch.channels_last)

o = layer(x)
print(o[0][0][0][0])

[crcrpar]

sorry for having let this one slip off, could you add a test case which this PR allows?

[tlogn]

@crcrpar thanks for your advise. I've added a test case and tested with H100.

with the main branch, error occurs:

Mismatched elements: 1866930499 / 2477260800 (75.4%)
Greatest absolute difference: nan at index (0, 0, 2146, 1184) (up to 0.04 allowed)
Greatest relative difference: nan at index (0, 0, 2146, 1184) (up to 0 allowed)

----------------------------------------------------------------------
Ran 16 tests in 8.897s

FAILED (failures=1)

with the fixed branch, the atol should be adjusted to 7e-2 to account for potential increased reduction accuracy error

................
----------------------------------------------------------------------
Ran 16 tests in 8.079s

OK

[crcrpar]

is this for the large tensor?

[tlogn]

yes, float reduction accuracy error would grow up with larger tensor

[crcrpar]

this seems to require about 50 GB, how about checking the memory size like

diff --git a/apex/contrib/test/group_norm/test_group_norm.py b/apex/contrib/test/group_norm/test_group_norm.py
index 5675749..687c2e1 100644
--- a/apex/contrib/test/group_norm/test_group_norm.py
+++ b/apex/contrib/test/group_norm/test_group_norm.py
@@ -177,8 +177,9 @@ class GroupNormTest(unittest.TestCase):
             [8, 1920, 32, 32],
             [8, 1920, 16, 16],
             [8, 2560, 8, 8],
-            [1, 128, 16128, 1200],
         ]
+        if torch.cuda.get_device_properties().total_memory > 50_000_000_000:
+            sizes.append([1, 128, 16128, 1200])
         for sz in sizes:
             n, c, h, w = sz
             self.verify_group_norm(GroupNorm,

[tlogn]

I've checked seriously, this tensor gots 128 * 16128 * 1200=2,477,260,800 elements, about 5GB. This kernel will take about 10GB, so this check seems not so necessary ?

[tlogn]

@crcrpar
Hi there, is there any problem remained ?