From b5477fbe35c21aed5a08b8cfcf6a8159cf3292ca Mon Sep 17 00:00:00 2001 From: Alexander <208520214+AlexanderSoftCode@users.noreply.github.com> Date: Sat, 31 Oct 2026 20:46:36 -0500 Subject: [PATCH] fix(conv): scale gradients before the fp16 cast so small values survive The WMMA backward casts dY to fp16, which rounds anything below about 3e-8 to exactly 0. Per-element gradients are often smaller than that, and shrink as image size grows, so training silently lost accuracy at higher resolution. Multiply dY by a fixed 2^10 in the cast (kernel and astype fallback) and divide it back out of dweights, dbiases and dinputs in the existing kernel epilogues. Gradients stay in true units outside the layer. This lowers the effective per-element dY clip from +-65504 to about +-64. --- aether/custom_kernels/conv_kernel/__init__.py | 2 + aether/custom_kernels/conv_kernel/cast.py | 15 ++++--- aether/custom_kernels/conv_kernel/common.py | 7 +++ aether/custom_kernels/conv_kernel/dinput.py | 9 +++- aether/custom_kernels/conv_kernel/dweight.py | 10 ++++- aether/layers/conv.py | 7 ++- tests/layers/conv2d/test_gpu_parity.py | 43 ++++++++++++++++--- tests/ops/test_conv_cast_kernel.py | 19 ++++---- 8 files changed, 87 insertions(+), 25 deletions(-) diff --git a/aether/custom_kernels/conv_kernel/__init__.py b/aether/custom_kernels/conv_kernel/__init__.py index c1b2094..83b7636 100644 --- a/aether/custom_kernels/conv_kernel/__init__.py +++ b/aether/custom_kernels/conv_kernel/__init__.py @@ -12,6 +12,7 @@ get_compiled_backward_cast_kernel, ) from aether.custom_kernels.conv_kernel.common import ( + GRAD_SCALE, ConvLaunchPlan, ConvShapeMeta, choose_block_tile_n, @@ -37,6 +38,7 @@ ) __all__ = [ + "GRAD_SCALE", "_CONV_BACKWARD_CAST_CACHE", "_CONV_BACKWARD_DINPUT_CACHE", "_CONV_BACKWARD_DWEIGHT_CACHE", diff --git a/aether/custom_kernels/conv_kernel/cast.py b/aether/custom_kernels/conv_kernel/cast.py index 5b5da36..5fb2f86 100644 --- a/aether/custom_kernels/conv_kernel/cast.py +++ b/aether/custom_kernels/conv_kernel/cast.py @@ -1,15 +1,16 @@ import warnings import aether.config as config -from aether.custom_kernels.conv_kernel.common import _COMPILE_OPTIONS +from aether.custom_kernels.conv_kernel.common import _COMPILE_OPTIONS, GRAD_SCALE # One launch casts both fp16 operands of the conv backward kernels. dvalues is -# clipped to the fp16 range first; inputs is not, so dweight sees exactly the -# fp16 inputs the forward pass multiplied. +# scaled by GRAD_SCALE and clipped to the fp16 range first; inputs is neither, so +# dweight sees exactly the fp16 inputs the forward pass multiplied. CONV_BACKWARD_CAST_SOURCE = r''' $fp16_include #define HALF_MAX 65504.0f +#define GRAD_SCALE ${grad_scale}f extern "C" __global__ void conv_backward_cast_fp16( const float* __restrict__ dvalues, // [n_dvalues] fp32, contiguous @@ -21,7 +22,7 @@ ) { const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; if (i < n_dvalues) { - const float v = dvalues[i]; + const float v = dvalues[i] * GRAD_SCALE; // Comparisons rather than fminf/fmaxf, which would turn NaN into // +-HALF_MAX: like xp.clip, NaN passes through and +-inf clamp. dvalues_fp16[i] = __float2half(v > HALF_MAX ? HALF_MAX : (v < -HALF_MAX ? -HALF_MAX : v)); @@ -40,7 +41,9 @@ def get_compiled_backward_cast_kernel(is_hip: bool): return _CONV_BACKWARD_CAST_CACHE[is_hip] fp16_include = "#include " if is_hip else "#include " - source = CONV_BACKWARD_CAST_SOURCE.replace("$fp16_include", fp16_include) + source = CONV_BACKWARD_CAST_SOURCE.replace("$fp16_include", fp16_include).replace( + "${grad_scale}", str(GRAD_SCALE) + ) try: module = config.cp.RawModule(code=source, options=_COMPILE_OPTIONS) @@ -55,7 +58,7 @@ def get_compiled_backward_cast_kernel(is_hip: bool): def cast_backward_operands_fp16(dvalues, inputs): - """(dvalues clipped to +-65504, inputs) as fp16, in one launch. + """(dvalues * GRAD_SCALE clipped to +-65504, inputs) as fp16, in one launch. Both must be contiguous fp32. Returns None when the kernel is unavailable, so the caller can fall back to astype/clip. diff --git a/aether/custom_kernels/conv_kernel/common.py b/aether/custom_kernels/conv_kernel/common.py index 3dc4b9f..43fb5a7 100644 --- a/aether/custom_kernels/conv_kernel/common.py +++ b/aether/custom_kernels/conv_kernel/common.py @@ -105,6 +105,13 @@ def threads_per_block(block_tile_m, block_tile_n): # ROCM_PATH would be searched first, so a stale ROCM_PATH could shadow it. _COMPILE_OPTIONS = ('-std=c++17',) +# dY is multiplied by this before its fp16 cast and the backward kernels divide +# it back out of what they write, because per-element gradients can be far below +# fp16's smallest step (about 6e-8) and would round to 0. A power of two, so +# neither direction rounds in fp32. It lowers the effective dY clip from +-65504 +# to about +-64. +GRAD_SCALE = 1024.0 + # --- Vendor substitution dictionaries -------------------------------------- diff --git a/aether/custom_kernels/conv_kernel/dinput.py b/aether/custom_kernels/conv_kernel/dinput.py index 5aa3d35..ee39832 100644 --- a/aether/custom_kernels/conv_kernel/dinput.py +++ b/aether/custom_kernels/conv_kernel/dinput.py @@ -8,6 +8,7 @@ _CONV_CUDA, _CONV_HIP, _WMMA_DIM, + GRAD_SCALE, ConvLaunchPlan, ConvShapeMeta, _m_total, @@ -31,6 +32,8 @@ #define H_OUT $h_out #define W_OUT $w_out +#define INV_GRAD_SCALE (1.0f / ${grad_scale}f) + #define TAPS_H ((F_H + STR_H - 1) / STR_H) #define TAPS_W ((F_W + STR_W - 1) / STR_W) #define K_GATHER (TAPS_H * TAPS_W * C_OUT) @@ -312,7 +315,7 @@ const long long in_idx = (((long long)s * H_IN + h_in) * W_IN + w_in) * C_IN + c_in; - dinputs[in_idx] = Cs[local_m][local_n]; + dinputs[in_idx] = Cs[local_m][local_n] * INV_GRAD_SCALE; } } } @@ -340,6 +343,7 @@ def render_dinput_source(shape_meta: ConvShapeMeta, vectorize=True): wmma_dim=_WMMA_DIM, block_tile_m=block_tile_m, block_tile_n=block_tile_n, vectorize=int(bool(vectorize)), + grad_scale=GRAD_SCALE, **vendor, ) @@ -373,6 +377,9 @@ def get_compiled_backward_dinput_kernel(shape_meta: ConvShapeMeta): def plan_dinput_launch(shape_meta: ConvShapeMeta, S: int): """conv_backward_dinput_wmma(dvalues_fp16, weight_fp16, dinputs). + dvalues_fp16 must already be scaled by GRAD_SCALE; dinputs come out in true + units. + gridDim.z carries the residue class (h_in % sH, w_in % sW). The class has to be block-uniform because the B operand depends on it, and classes differ in size by at most one row/column, so the grid is sized for the largest and diff --git a/aether/custom_kernels/conv_kernel/dweight.py b/aether/custom_kernels/conv_kernel/dweight.py index 289784c..8849ae4 100644 --- a/aether/custom_kernels/conv_kernel/dweight.py +++ b/aether/custom_kernels/conv_kernel/dweight.py @@ -8,6 +8,7 @@ _CONV_CUDA, _CONV_HIP, _WMMA_DIM, + GRAD_SCALE, ConvLaunchPlan, ConvShapeMeta, _m_total, @@ -51,6 +52,7 @@ def choose_dweight_split_k(num_m_tiles, grid_x, grid_y, threads_per_block): #define W_IN $w_in #define H_OUT $h_out #define W_OUT $w_out +#define INV_GRAD_SCALE (1.0f / ${grad_scale}f) #define K_TOTAL (F_H * F_W * C_IN) #define HW_OUT (H_OUT * W_OUT) #define C_IN_FW (C_IN * F_W) @@ -311,7 +313,7 @@ def choose_dweight_split_k(num_m_tiles, grid_x, grid_y, threads_per_block): const int n_global = n_tile_start + local_n; if (k_global < K_TOTAL && n_global < C_OUT) { - const float val = Cs[local_k][local_n]; + const float val = Cs[local_k][local_n] * INV_GRAD_SCALE; if (gridDim.z > 1) { atomicAdd(&dweights[k_global * C_OUT + n_global], val); } else { @@ -347,7 +349,7 @@ def choose_dweight_split_k(num_m_tiles, grid_x, grid_y, threads_per_block): if (blockIdx.x == 0 && threadIdx.x < BLOCK_TILE_N) { const int n_global = n_tile_start + threadIdx.x; if (n_global < C_OUT) { - atomicAdd(&dbiases[n_global], s_bias[threadIdx.x]); + atomicAdd(&dbiases[n_global], s_bias[threadIdx.x] * INV_GRAD_SCALE); } } } @@ -378,6 +380,7 @@ def render_dweight_source(shape_meta: ConvShapeMeta, vectorize=True): block_tile_m=block_tile_m, block_tile_n=block_tile_n, a_pad=_DWEIGHT_A_PAD, b_pad=_DWEIGHT_B_PAD, vectorize=int(bool(vectorize)), + grad_scale=GRAD_SCALE, **vendor, ) @@ -411,6 +414,9 @@ def get_compiled_backward_dweight_kernel(shape_meta: ConvShapeMeta): def plan_dweight_launch(shape_meta: ConvShapeMeta, S: int): """conv_backward_dweight_wmma(input_fp16, dvalues_fp16, dweights, dbiases). + dvalues_fp16 must already be scaled by GRAD_SCALE; dweights/dbiases come out + in true units. + Note the 3D grid: gridDim.z carries the split-K slices that make this kernel fast. dweights/dbiases must be zeroed by the caller whenever gridDim.z > 1, since the epilogue accumulates with atomicAdd. diff --git a/aether/layers/conv.py b/aether/layers/conv.py index 238250d..933c881 100644 --- a/aether/layers/conv.py +++ b/aether/layers/conv.py @@ -44,6 +44,11 @@ class Conv2d(Layer): Unless a precision is set, on a GPU the layer uses fp16 WMMA kernels (float32 accumulate). WMMA can't take >fp16 operands, so an explicit ``float32`` request uses the slower fallback vectorized path instead. + + On that path backward scales incoming gradients by 2^10 before the fp16 cast + and undoes it on what it outputs, so values far below fp16's smallest step + survive. Gradients stay in true units outside the layer, but each element of + the incoming gradient is silently clipped at about +-64. """ def __init__(self, in_channels, out_channels = 1, filter_size = (3, 3), stride = (1, 1), padding = "same", l1=(), l2=()): @@ -310,7 +315,7 @@ def _backward_gpu(self, dvalues): if cast is not None: dvalues_fp16, inputs_fp16 = cast else: - dvalues_bounded = xp.clip(dvalues_fp32, -65504.0, 65504.0) + dvalues_bounded = xp.clip(dvalues_fp32 * conv_kernel.GRAD_SCALE, -65504.0, 65504.0) dvalues_fp16 = dvalues_bounded.astype(xp.float16, copy=False) inputs_fp16 = inputs_fp32.astype(xp.float16, copy=False) diff --git a/tests/layers/conv2d/test_gpu_parity.py b/tests/layers/conv2d/test_gpu_parity.py index a0fa7c1..9062b99 100644 --- a/tests/layers/conv2d/test_gpu_parity.py +++ b/tests/layers/conv2d/test_gpu_parity.py @@ -40,7 +40,7 @@ def test_forward_gpu_fallback_numeric_parity(self): self.xp.testing.assert_allclose(gpu_out, fallback_out, rtol=1e-2, atol=1e-2) def _run_dweight_kernel_parity(self, S, spatial, in_channels, out_channels, - stride=None, padding=None, filter_size=None): + stride=None, padding=None, filter_size=None, dy_std=1.0): """Drive the fused dweight kernel once and check both of its outputs. conv_backward_dweight_wmma emits dweights AND dbiases from a single @@ -71,9 +71,9 @@ def _run_dweight_kernel_parity(self, S, spatial, in_channels, out_channels, fixed_input = self.xp.random.randn(S, H, W, in_channels).astype(self.xp.float32) meta = layer._get_shape_meta((H, W, in_channels)) - dvalues = self.xp.random.randn( + dvalues = (self.xp.random.randn( S, meta.H_out, meta.W_out, out_channels - ).astype(self.xp.float32) + ) * dy_std).astype(self.xp.float32) layer._forward_fallback(fixed_input, training=True) layer._backward_fallback(dvalues) @@ -119,7 +119,7 @@ def test_conv_gpu_backward_dweight_kernel_parity_valid_strided(self): ) def _run_dinput_kernel_parity(self, S, spatial, in_channels, out_channels, - stride=None, padding=None, filter_size=None): + stride=None, padding=None, filter_size=None, dy_std=1.0): """Drive the gather dinput kernel and check it against the fallback.""" if self.backend_name != 'cupy': self.skipTest("CuPy backend required for GPU kernel tests") @@ -145,9 +145,9 @@ def _run_dinput_kernel_parity(self, S, spatial, in_channels, out_channels, fixed_input = self.xp.random.randn(S, H, W, in_channels).astype(self.xp.float32) meta = layer._get_shape_meta((H, W, in_channels)) - dvalues = self.xp.random.randn( + dvalues = (self.xp.random.randn( S, meta.H_out, meta.W_out, out_channels - ).astype(self.xp.float32) + ) * dy_std).astype(self.xp.float32) layer._forward_fallback(fixed_input, training=True) layer._backward_fallback(dvalues) @@ -290,6 +290,37 @@ def test_conv_gpu_backward_parity_without_cast_kernel(self): if self.backend_name == 'cupy' and conv_kernel.get_is_conv_gpu_available(): self.assertGreater(cast.call_count, 0) + def test_conv_gpu_backward_tiny_gradients_survive_fp16_cast(self): + """Regression: dY near 1e-8 is real at high resolution, and fp16 rounds it to 0. + + Random dY of that size must match the fp32 fallback, and a constant one + must sum to a hand-computed bias gradient, with and without the cast kernel. + C_in = 3 and C_out = 6 take the scalar staging, 8 and 48 the vector one. + """ + if self.backend_name != 'cupy': + self.skipTest("CuPy backend required for GPU kernel tests") + from aether.custom_kernels import conv_kernel + if not conv_kernel.get_is_conv_gpu_available(): + self.skipTest("Matrix-core GPU kernels not available on this device") + + S, H, W = 4, 16, 16 + real_cast = conv_kernel.cast_backward_operands_fp16 + for name, cast in (("cast kernel", real_cast), ("astype fallback", lambda *_: None)): + for C_in, C_out in ((3, 6), (8, 48)): + with self.subTest(path=name, C_in=C_in, C_out=C_out), \ + mock.patch.object(conv_kernel, "cast_backward_operands_fp16", cast): + self._run_dweight_kernel_parity(S, (H, W), C_in, C_out, dy_std=1e-8) + self._run_dinput_kernel_parity(S, (H, W), C_in, C_out, dy_std=1e-8) + + layer = self._make_same_layer((H, W), (3, 3), (1, 1), + in_channels=C_in, out_channels=C_out) + x = self.xp.random.randn(S, H, W, C_in).astype(self.xp.float32) + layer._forward_gpu(x, training=True) + layer._backward_gpu(self.xp.full((S, H, W, C_out), 1e-8, dtype=self.xp.float32)) + self._assert_gemm_close( + layer.dbiases, self.xp.full((C_out,), S * H * W * 1e-8), "constant dY dbiases" + ) + def test_conv_gpu_partial_second_column_tile_parity(self): """84 = 64 + 16 + 4 channels on the GEMM N axis (C_out, or C_in for dinput). diff --git a/tests/ops/test_conv_cast_kernel.py b/tests/ops/test_conv_cast_kernel.py index dd9ff22..7769563 100644 --- a/tests/ops/test_conv_cast_kernel.py +++ b/tests/ops/test_conv_cast_kernel.py @@ -6,7 +6,7 @@ class TestConvBackwardCastKernel(base_case.AetherBaseTestCase): - """conv_backward_cast_fp16: dvalues clipped to the fp16 range, inputs not.""" + """conv_backward_cast_fp16: dvalues scaled by 2^10 then clipped to the fp16 range, inputs neither.""" def setUp(self): super().setUp() @@ -25,34 +25,35 @@ def _cast(self, dvalues, inputs): def test_special_values(self): inf, nan = np.inf, np.nan - dvalues = [-inf, -70000.0, -65504.0, -1.5, 0.1, 0.0, 65504.0, 70000.0, inf, nan] + dvalues = [-inf, -70.0, -64.0, -1.5, 1e-8, 0.0, 63.96875, 64.0, 70000.0, inf, nan] inputs = [70000.0, -70000.0, 0.1, nan, 2.0] dv16, in16 = self._cast(dvalues, inputs) self.assertEqual(dv16.dtype, np.float16) self.assertEqual(in16.dtype, np.float16) + # 1e-8 * 2^10 is 171.8 subnormal steps of 2^-24, which rounds to 172. np.testing.assert_array_equal( dv16, - np.array([-65504, -65504, -65504, -1.5, np.float16(0.1), 0, 65504, 65504, 65504, nan], + np.array([-65504, -65504, -65504, -1536, 172 * 2.0 ** -24, 0, 65504, 65504, 65504, 65504, nan], dtype=np.float16), ) - # Unclipped: out-of-range inputs overflow to inf exactly as forward's astype does. + # Unscaled and unclipped: out-of-range inputs overflow to inf exactly as forward's astype does. np.testing.assert_array_equal( in16, np.array([inf, -inf, np.float16(0.1), nan, 2.0], dtype=np.float16) ) - def test_matches_clip_and_astype(self): - """Both ranges, with lengths that end mid-block, against NumPy's clip + astype.""" + def test_matches_scale_clip_and_astype(self): + """Both ranges, with lengths that end mid-block, against NumPy's scale + clip + astype.""" rng = np.random.default_rng(0) - dvalues = (rng.standard_normal((3, 7, 5, 19)) * 3e4).astype(np.float32) + dvalues = (rng.standard_normal((3, 7, 5, 19)) * 30).astype(np.float32) inputs = (rng.standard_normal((3, 9, 9, 7)) * 3e4).astype(np.float32) dv16, in16 = self._cast(dvalues, inputs) self.assertEqual(dv16.shape, dvalues.shape) self.assertEqual(in16.shape, inputs.shape) - self.assertTrue((np.abs(dvalues) > 65504).any(), "test data must exercise the clip") + self.assertTrue((np.abs(dvalues) * 2 ** 10 > 65504).any(), "test data must exercise the clip") self.assertTrue((np.abs(inputs) > 65504).any(), "test data must exercise the overflow") - np.testing.assert_array_equal(dv16, np.clip(dvalues, -65504, 65504).astype(np.float16)) + np.testing.assert_array_equal(dv16, np.clip(dvalues * 2 ** 10, -65504, 65504).astype(np.float16)) # The overflow to inf is the behavior under test, so NumPy's warning is expected. with np.errstate(over="ignore"): expected_in16 = inputs.astype(np.float16)