fix(conv): scale gradients before the fp16 cast so small values survive

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",

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 <hip/hip_fp16.h>" if is_hip else "#include <cuda_fp16.h>"
-    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.

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 --------------------------------------
 

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

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.

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)
 

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).
 

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)