diff --git a/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass.h b/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass.h index c85181a52..dc54570f4 100644 --- a/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass.h +++ b/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 == 8 || 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; diff --git a/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass_kernel.cuh b/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass_kernel.cuh index 008520414..5b2ddf698 100644 --- a/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass_kernel.cuh +++ b/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 == 8 || __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(&x_ptr[hwj * params.c]); + dy_extra = *reinterpret_cast(&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 == 8 || __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(&dx_ptr[hwj * params.c]) = IOTraits::pack(dx); } } + +#if __CUDA_ARCH__ / 100 == 8 || __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(&x_ptr[hwj * params.c])); + float2 dy_f2 = IOTraits::unpack(*reinterpret_cast(&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(&dx_ptr[hwj * params.c]) = IOTraits::pack(dx); + } + } +#endif // __CUDA_ARCH__ / 100 == 8 || __CUDA_ARCH__ / 100 == 11 || __CUDA_ARCH__ / 100 == 12 } // The completion barrier. diff --git a/apex/contrib/test/group_norm/test_group_norm.py b/apex/contrib/test/group_norm/test_group_norm.py index 99bc7072b..5e82e4d79 100644 --- a/apex/contrib/test/group_norm/test_group_norm.py +++ b/apex/contrib/test/group_norm/test_group_norm.py @@ -107,10 +107,6 @@ def _has_sufficient_cuda_memory(required_bytes: int, *, safety_factor: float = 0 return required_bytes <= int(free_bytes * safety_factor) -@unittest.skipIf( - torch.cuda.get_device_properties().multi_processor_count < 16, - "GroupNorm is unsupported on low SM count devices", -) @unittest.skipIf(SKIP_TEST, f"{SKIP_TEST}") class GroupNormTest(unittest.TestCase): def setUp(self, seed=0):