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)