Fix GroupNorm NHWC one-pass backward execution

apex/__init__.py

@@ -11,11 +11,9 @@
 # so they expect those backends to be available, but for some reason they actually aren't
 # available (for example because they built improperly in a way that isn't revealed until
 # load time) the error message is timely and visible.
-from . import optimizers
-from . import normalization
+from . import normalization, optimizers
 
-
-__all__ = ["optimizers", "normalization"]
+__all__ = ["normalization", "optimizers"]
 
 
 def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int) -> bool:

apex/_autocast_utils.py

@@ -1,8 +1,8 @@
-from typing import Optional, Sequence
+from collections.abc import Sequence
+from typing import Optional
 
 import torch
 
-
 __all__ = ["_cast_if_autocast_enabled"]
 
 
@@ -12,7 +12,7 @@ def _get_autocast_dtypes() -> Sequence[torch.dtype]:
     return [torch.half]
 
 
-def _get_current_dtype(dtype: Optional[torch.dtype] = None) -> torch.dtype:
+def _get_current_dtype(dtype: torch.dtype | None = None) -> torch.dtype:
     if not torch.is_autocast_enabled():
         return torch.float or dtype
     else:

apex/contrib/bottleneck/__init__.py

@@ -1,7 +1,7 @@
 from .bottleneck import Bottleneck, SpatialBottleneck
 from .halo_exchangers import (
-    HaloExchangerNoComm,
     HaloExchangerAllGather,
-    HaloExchangerSendRecv,
+    HaloExchangerNoComm,
     HaloExchangerPeer,
+    HaloExchangerSendRecv,
 )

apex/contrib/bottleneck/bottleneck.py

@@ -1,12 +1,11 @@
 import functools as func
 
+import fast_bottleneck
+import nccl_p2p_cuda as inc
 import torch
 from torch import nn
 
 from apex import check_cudnn_version_and_warn
-import fast_bottleneck
-import nccl_p2p_cuda as inc
-
 
 assert check_cudnn_version_and_warn(__name__, 8400)
 
@@ -35,7 +34,7 @@ class FrozenBatchNorm2d(torch.jit.ScriptModule):
     """
 
     def __init__(self, n):
-        super(FrozenBatchNorm2d, self).__init__()
+        super().__init__()
         self.register_buffer("weight", torch.ones(n))
         self.register_buffer("bias", torch.zeros(n))
         self.register_buffer("running_mean", torch.zeros(n))
@@ -171,7 +170,7 @@ def __init__(
         use_cudnn=False,
         explicit_nhwc=False,
     ):
-        super(Bottleneck, self).__init__()
+        super().__init__()
         if groups != 1:
             raise RuntimeError("Only support groups == 1")
         if dilation != 1:
@@ -224,8 +223,6 @@ def __init__(
                 with torch.no_grad():
                     p.data = p.data.permute(0, 2, 3, 1).contiguous()
 
-        return
-
     # Returns single callable that recomputes scale and bias for all frozen batch-norms.
     # This method must be called before cuda graphing.
     # The callable it returns can be called anytime.
@@ -851,7 +848,7 @@ def __init__(
         explicit_nhwc=False,
         spatial_parallel_args=None,
     ):
-        super(SpatialBottleneck, self).__init__()
+        super().__init__()
         if groups != 1:
             raise RuntimeError("Only support groups == 1")
         if dilation != 1:
@@ -911,7 +908,6 @@ def __init__(
             self.spatial_parallel_args = (1, 0, None, None, 0, False)
         else:
             self.spatial_parallel_args = spatial_parallel_args
-        return
 
     # Returns single callable that recomputes scale and bias for all frozen batch-norms.
     # This method must be called before cuda graphing.

apex/contrib/bottleneck/halo_exchangers.py

@@ -1,13 +1,13 @@
-import torch
 import nccl_p2p_cuda as inc
 import peer_memory_cuda as pm
+import torch
 
 
 # Communication free halo exchanger.
 # NB! This halo exchanger does not exchange halos with neighbors as it should, it merely swaps the inputs
 # NB! This is only useful for performance testing.
 # NB! Do not use for actual production runs
-class HaloExchanger(object):
+class HaloExchanger:
     def __init__(self, ranks, rank_in_group):
         self.stream1 = torch.cuda.Stream()
         self.stream2 = torch.cuda.Stream()
@@ -27,7 +27,7 @@ def __init__(self, ranks, rank_in_group):
 
 class HaloExchangerNoComm(HaloExchanger):
     def __init__(self, ranks, rank_in_group):
-        super(HaloExchangerNoComm, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
 
     def left_right_halo_exchange(
         self,
@@ -45,7 +45,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerAllGather(HaloExchanger):
     def __init__(self, ranks, rank_in_group, comm):
-        super(HaloExchangerAllGather, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         # self.comm must be NCCL process_group created with torch.distributed.new_group(ranks=ranks)
         self.comm = comm
 
@@ -94,7 +94,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerSendRecv(HaloExchanger):
     def __init__(self, ranks, rank_in_group):
-        super(HaloExchangerSendRecv, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         nccl_id = inc.get_unique_nccl_id(1).cuda()
         torch.distributed.broadcast(nccl_id, 0)
         nccl_id = nccl_id.cpu()
@@ -145,7 +145,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerPeer(HaloExchanger):
     def __init__(self, ranks, rank_in_group, peer_pool, explicit_nhwc, numSM=0):
-        super(HaloExchangerPeer, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         self.diagnostics = False
         self.explicit_nhwc = explicit_nhwc
         self.numSM = numSM

apex/contrib/bottleneck/test.py

@@ -29,7 +29,7 @@
         print("[DEBUG] ref dx :", d_grad.sum().item())
         # print wgrad. we don't need to reset since later cpp print before accumulation
         for i, w in enumerate(model.w_conv):
-            print("[DEBUG] ref wgrad{} :".format(i + 1), w.grad.sum().item())
+            print(f"[DEBUG] ref wgrad{i + 1} :", w.grad.sum().item())
 
     wgrads = []
     for w in model.w_conv:
@@ -56,7 +56,7 @@
     )
     for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):
         print(
-            "max error wgrad{}:".format(i + 1),
+            f"max error wgrad{i + 1}:",
             (wgrad - w.grad.float()).abs().max().item(),
             "max elem:",
             wgrad.abs().max().item(),
@@ -104,7 +104,7 @@
     )
     for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):
         print(
-            "max error wgrad{}:".format(i + 1),
+            f"max error wgrad{i + 1}:",
             (wgrad - w.grad.float()).abs().max().item(),
             "max elem:",
             wgrad.abs().max().item(),

apex/contrib/clip_grad/clip_grad.py

@@ -1,10 +1,12 @@
-from typing import Union, Iterable
+from collections.abc import Iterable
+from typing import Union
 
 import torch
 
 _kernel_import_succeeded = False
 try:
     import amp_C
+
     from apex.multi_tensor_apply import multi_tensor_applier
 
     _kernel_import_succeeded = True

apex/contrib/conv_bias_relu/__init__.py

@@ -1,6 +1,6 @@
 from .conv_bias_relu import (
-    ConvBiasReLU,
     ConvBias,
     ConvBiasMaskReLU,
+    ConvBiasReLU,
     ConvFrozenScaleBiasReLU,
 )

apex/contrib/conv_bias_relu/conv_bias_relu.py

@@ -1,7 +1,7 @@
+import fused_conv_bias_relu
 import torch
 
 from apex import check_cudnn_version_and_warn
-import fused_conv_bias_relu
 
 check_cudnn_version_and_warn(__name__, 8400)
 

apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass.h

@@ -128,8 +128,14 @@ void group_norm_nhwc_bwd_one_pass_setup(Group_norm_nhwc_bwd_params& params, size
   // The number of blocks per grid.
   int max_blocks_per_grid = blocks_per_sm * props.multiProcessorCount;
 
-  // Make sure we are safe to run that many blocks
-  assert(blocks_per_slice <= max_blocks_per_grid);
+  if (props.major == 11 || props.major == 12) {
+    // Cooperative kernels require all blocks to be resident concurrently. Blocks process
+    // additional activation tiles in a grid-stride loop when the full grid does not fit.
+    blocks_per_slice = std::min(blocks_per_slice, max_blocks_per_grid);
+  } else {
+    // Make sure we are safe to run that many blocks
+    assert(blocks_per_slice <= max_blocks_per_grid);
+  }
 
   // The number of blocks per slice is the X dimension of the grid.
   grid.x = blocks_per_slice;

apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass_kernel.cuh

@@ -173,6 +173,49 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK_) void group_norm_nhwc_bwd_one_pa
       mean_2 += dx_norm_y;
     }
 
+#if __CUDA_ARCH__ / 100 == 11 || __CUDA_ARCH__ / 100 == 12
+    // A cooperative launch may use fewer blocks than activation tiles. Accumulate any
+    // additional tiles assigned to this block without increasing its register footprint.
+    for (int extra_hwi = hwi + gridDim.x * params.acts_per_block; extra_hwi < params.hw;
+         extra_hwi += gridDim.x * params.acts_per_block) {
+#pragma unroll
+      for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {
+        int hwj = extra_hwi + ii * ACTS_PER_LOOP;
+        IOType2 x_extra = IOTraits::zero();
+        IOType2 dy_extra = IOTraits::zero();
+        if (is_active && hwj < params.hw) {
+          x_extra = *reinterpret_cast<const IOType2*>(&x_ptr[hwj * params.c]);
+          dy_extra = *reinterpret_cast<const IOType2*>(&dy_ptr[hwj * params.c]);
+        }
+
+        float2 x_f2 = IOTraits::unpack(x_extra);
+        float2 dy_f2 = IOTraits::unpack(dy_extra);
+
+        float x_norm_x = (x_f2.x - x_mean) * rcp_x_stddev;
+        float x_norm_y = (x_f2.y - x_mean) * rcp_x_stddev;
+
+        if (params.with_swish) {
+          float x_gn_x = x_norm_x * gamma_f2.x + beta_f2.x;
+          float x_gn_y = x_norm_y * gamma_f2.y + beta_f2.y;
+          float s_x = sigmoid(x_gn_x);
+          float s_y = sigmoid(x_gn_y);
+          dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));
+          dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));
+        }
+
+        dgamma_dbeta.x += dy_f2.x * x_norm_x;
+        dgamma_dbeta.y += dy_f2.y * x_norm_y;
+        dgamma_dbeta.z += dy_f2.x;
+        dgamma_dbeta.w += dy_f2.y;
+
+        float dx_norm_x = dy_f2.x * gamma_f2.x;
+        float dx_norm_y = dy_f2.y * gamma_f2.y;
+        mean_1 += dx_norm_x * x_norm_x + dx_norm_y * x_norm_y;
+        mean_2 += dx_norm_x + dx_norm_y;
+      }
+    }
+#endif  // __CUDA_ARCH__ / 100 == 11 || __CUDA_ARCH__ / 100 == 12
+
     // Pack valid gradients.
     float2 sums = make_float2(0.f, 0.f);
     if (ACTIVE_THREADS == THREADS_PER_BLOCK || is_active) {
@@ -308,6 +351,46 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK_) void group_norm_nhwc_bwd_one_pa
         *reinterpret_cast<IOType2*>(&dx_ptr[hwj * params.c]) = IOTraits::pack(dx);
       }
     }
+
+#if __CUDA_ARCH__ / 100 == 11 || __CUDA_ARCH__ / 100 == 12
+    // Store gradients for any additional activation tiles assigned to this block.
+    for (int extra_hwi = hwi + gridDim.x * params.acts_per_block; extra_hwi < params.hw;
+         extra_hwi += gridDim.x * params.acts_per_block) {
+#pragma unroll
+      for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {
+        int hwj = extra_hwi + ii * ACTS_PER_LOOP;
+        if (!is_active || hwj >= params.hw) {
+          continue;
+        }
+
+        float2 x_f2 = IOTraits::unpack(*reinterpret_cast<const IOType2*>(&x_ptr[hwj * params.c]));
+        float2 dy_f2 = IOTraits::unpack(*reinterpret_cast<const IOType2*>(&dy_ptr[hwj * params.c]));
+
+        float2 x_norm;
+        x_norm.x = (x_f2.x - x_mean) * rcp_x_stddev;
+        x_norm.y = (x_f2.y - x_mean) * rcp_x_stddev;
+
+        if (params.with_swish) {
+          float x_gn_x = x_norm.x * gamma_f2.x + beta_f2.x;
+          float x_gn_y = x_norm.y * gamma_f2.y + beta_f2.y;
+          float s_x = sigmoid(x_gn_x);
+          float s_y = sigmoid(x_gn_y);
+          dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));
+          dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));
+        }
+
+        float2 dx_norm;
+        dx_norm.x = dy_f2.x * gamma_f2.x;
+        dx_norm.y = dy_f2.y * gamma_f2.y;
+
+        float2 dx;
+        dx.x = (dx_norm.x - (x_norm.x * mean_1 + mean_2)) * rcp_x_stddev;
+        dx.y = (dx_norm.y - (x_norm.y * mean_1 + mean_2)) * rcp_x_stddev;
+
+        *reinterpret_cast<IOType2*>(&dx_ptr[hwj * params.c]) = IOTraits::pack(dx);
+      }
+    }
+#endif  // __CUDA_ARCH__ / 100 == 11 || __CUDA_ARCH__ / 100 == 12
   }
 
   // The completion barrier.

apex/contrib/csrc/group_norm_v2/generate_gn_cuda_inst.py

@@ -1,6 +1,5 @@
 import pathlib
 
-
 hw_c_list = [
     (8 * 8, 1280),
     (8 * 8, 2560),

apex/contrib/cudnn_gbn/batch_norm.py

@@ -1,10 +1,10 @@
+import cudnn_gbn_lib
+import peer_memory_cuda as pm
 import torch
-from torch.nn.modules.batchnorm import _BatchNorm
-from torch.nn import functional as F
 from torch import Tensor
-import peer_memory_cuda as pm
-import cudnn_gbn_lib
-from torch.cuda.amp import custom_fwd, custom_bwd
+from torch.cuda.amp import custom_bwd, custom_fwd
+from torch.nn import functional as F
+from torch.nn.modules.batchnorm import _BatchNorm
 
 
 class _GroupBatchNorm2d(torch.autograd.Function):
@@ -128,7 +128,7 @@ def __init__(
         affine=True,
         track_running_stats=True,
     ):
-        super(GroupBatchNorm2d, self).__init__(
+        super().__init__(
             num_features,
             eps=eps,
             momentum=momentum,
@@ -165,7 +165,7 @@ def get_peer_buffers(self, num_features):
 
     def _check_input_dim(self, input):
         if input.dim() != 4:
-            raise ValueError("expected 4D input (got {}D input)".format(input.dim()))
+            raise ValueError(f"expected 4D input (got {input.dim()}D input)")
 
     def _check_input_channels(self, input):
         if input.size(1) % 8 != 0:

apex/contrib/examples/gpu_direct_storage/benchmark_load.py

@@ -1,6 +1,8 @@
 import timeit
-import torch
+
 import apex.contrib.gpu_direct_storage as gds
+import torch
+
 
 def run_benchmark_torch_load():
     sizes = [2 ** i for i in range(16, 28)]

apex/contrib/examples/gpu_direct_storage/benchmark_save.py

@@ -1,7 +1,9 @@
 import os
 import timeit
-import torch
+
 import apex.contrib.gpu_direct_storage as gds
+import torch
+
 
 def run_benchmark(func):
     sizes = [2 ** i for i in range(16, 28)]

apex/contrib/examples/gpu_direct_storage/example_load.py

@@ -1,5 +1,5 @@
-import torch
 import apex.contrib.gpu_direct_storage as gds
+import torch
 
 for size in [128, 1024, 8192]:
     x = torch.empty(size, device = "cuda")