From 6cb9daf7a87099e68fef16b30e6073e78dbef1d8 Mon Sep 17 00:00:00 2001 From: Alexander <208520214+AlexanderSoftCode@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:17:50 -0500 Subject: [PATCH 1/2] perf(batchnorm): apply a leaky slope inside the BatchNorm kernels Every BatchNorm in the CIFAR-10 model feeds a ReLU, which costs its own read + write pass in each direction (0.28 + 0.44 ms per step on the RX 9070 XT). The BatchNorm kernels now take a slope alpha and apply z = fmaf(x - mean, scale, beta) out = z > 0 ? z : alpha * z so alpha = 1 is no activation (the default, and the behavior until now), 0 is ReLU and anything else LeakyReLU. One alpha covers every case, so there is no separate fused code path. The backward recomputes z from x, which it already reads, and uses g = z > 0 ? dy : (alpha != 0 ? alpha * dy : 0) in place of dy in both grad_sums and the elementwise pass. No extra memory and no extra bandwidth. z, the activation and its gradient are one source string (_ACTIVATION), included in the RawKernel prelude and passed to the ElementwiseKernels as their preamble, so the backward's mask is bit-identical to the forward's. As in ReLU's GPU backward, a NaN gradient at an inactive ReLU unit gives 0 rather than NaN. BatchNorm._fuse_activation(alpha) sets the slope and returns True only when the GPU path is bound; _compile_for_device resets it to 1. The slope is fp32 (BatchNorm is precision-exempt), so a fused LeakyReLU under fp16 is slightly more accurate than the unfused one, whose alpha is fp16. With alpha = 1 the CIFAR-10 step time is unchanged (5.913 vs 5.889 ms). Tests: fused ReLU and LeakyReLU against the fallback BatchNorm followed by the activation's own fallback (out, dx, dgamma, dbeta, inference), on inputs that keep every z away from 0 so rounding can't flip a mask; and the NaN case for both slopes. --- aether/custom_kernels/batchnorm_kernel.py | 46 ++++++++++++-- aether/layers/normalization.py | 27 ++++++-- tests/layers/test_batch_normalization.py | 75 +++++++++++++++++++++++ 3 files changed, 136 insertions(+), 12 deletions(-) diff --git a/aether/custom_kernels/batchnorm_kernel.py b/aether/custom_kernels/batchnorm_kernel.py index 0313b06..fa6477e 100644 --- a/aether/custom_kernels/batchnorm_kernel.py +++ b/aether/custom_kernels/batchnorm_kernel.py @@ -5,6 +5,25 @@ BLOCK_Y = 8 +# The activation every BatchNorm kernel applies to z = (x - mean) * scale + beta: +# out = z > 0 ? z : alpha * z, so alpha = 1 is none, 0 is ReLU, else LeakyReLU. +# Forward and backward compute z from this one source, so the backward's mask is +# bit-identical to the forward's. As in ReLU's GPU backward, a NaN/inf gradient +# at an inactive ReLU unit gives 0, not 0 * NaN. +_ACTIVATION = r''' +__device__ __forceinline__ float bn_z(float x, float mean, float scale, float beta) { + return fmaf(x - mean, scale, beta); +} + +__device__ __forceinline__ float bn_act(float z, float alpha) { + return z > 0.0f ? z : alpha * z; +} + +__device__ __forceinline__ float bn_act_grad(float z, float dy, float alpha) { + return z > 0.0f ? dy : (alpha != 0.0f ? alpha * dy : 0.0f); +} +''' + # Shared by every kernel below. block_sum adds each thread's (a, b) pair over the # block's BLOCK_Y rows and returns true on the row-0 threads that hold a channel's # total; sum_partials does the same over the G rows of a partials[2, G, C] buffer. @@ -13,6 +32,8 @@ #define WAVEFRONT 32 #define BLOCK_Y $block_y +$activation + __device__ __forceinline__ bool block_sum(float& a, float& b, int c, int C) { __shared__ float sh_a[BLOCK_Y][WAVEFRONT]; __shared__ float sh_b[BLOCK_Y][WAVEFRONT]; @@ -54,6 +75,9 @@ const float* __restrict__ p, const float* __restrict__ q, // unused (nullptr) ops still take the param const float* __restrict__ mean, // channel mean [C], unused ops pass nullptr + const float* __restrict__ scale, // grad_sums only, like beta and alpha + const float* __restrict__ beta, + const float alpha, float* __restrict__ partials, // [2, gridDim.y, C] const int N, const int C @@ -156,8 +180,11 @@ _GRAD_SUMS_OP = { "name": "grad_sums", "template": _STATS_REDUCE_TEMPLATE, - "prologue": "float mu = mean[c];", - "accumulate": "float g = p[idx]; acc_a += g; acc_b += g * (q[idx] - mu);", + "prologue": "float mu = mean[c]; float sc = scale[c]; float bt = beta[c];", + "accumulate": ( + "float xv = q[idx]; float g = bn_act_grad(bn_z(xv, mu, sc, bt), p[idx], alpha);" + " acc_a += g; acc_b += g * (xv - mu);" + ), } _FORWARD_FINALIZE_OP = { @@ -190,6 +217,7 @@ def _get_compiled_kernel(op_dict: dict, variant: str, block_y: int = BLOCK_Y): block_y=block_y, prologue=op_dict.get("prologue", ""), accumulate=op_dict.get("accumulate", ""), + activation=_ACTIVATION, **vendor, ) @@ -226,20 +254,26 @@ def get_bn_backward_finalize_kernel(variant: str): # folding (mean, scale, beta) into one bias out of the output. _bn_apply_ew = config.build_kernel( lambda: config.cp.ElementwiseKernel( - 'float32 x, float32 mean, float32 scale, float32 beta', + 'float32 x, float32 mean, float32 scale, float32 beta, float32 alpha', 'float32 out', - 'out = fmaf(x - mean, scale, beta);', + 'out = bn_act(bn_z(x, mean, scale, beta), alpha);', 'bn_apply_ew', + preamble=_ACTIVATION, ), name='_bn_apply_ew', ) _bn_backward_ew = config.build_kernel( lambda: config.cp.ElementwiseKernel( - 'float32 dvalues, float32 x, float32 mean, float32 scale, float32 k1, float32 k2', + 'float32 dvalues, float32 x, float32 mean, float32 scale, float32 beta, ' + 'float32 k1, float32 k2, float32 alpha', 'float32 dinputs', - 'dinputs = dvalues * scale + (x - mean) * k1 + k2;', + ''' + float g = bn_act_grad(bn_z(x, mean, scale, beta), dvalues, alpha); + dinputs = g * scale + (x - mean) * k1 + k2; + ''', 'bn_backward_ew', + preamble=_ACTIVATION, ), name='_bn_backward_ew', ) diff --git a/aether/layers/normalization.py b/aether/layers/normalization.py index f67b3fb..ec2aace 100644 --- a/aether/layers/normalization.py +++ b/aether/layers/normalization.py @@ -71,6 +71,9 @@ def __init__(self, epsilon=1e-5, momentum=0.9): self._forward_finalize_kernel = None self._grad_sums_kernel = None self._backward_finalize_kernel = None + # Leaky slope of an activation the GPU kernels apply after normalizing: + # 1 is none (see _fuse_activation). + self._alpha = np.float32(1.0) self.forward = self._forward_fallback self.backward = self._backward_fallback @@ -133,6 +136,8 @@ def build(self, input_shape: tuple[int, ...], seed: int | None = None) -> tuple[ def _compile_for_device(self, device): """Triggered by Model.to(device) to map low-level hardware paths.""" + # Model re-plans activation fusion after every device sync. + self._alpha = np.float32(1.0) if device == "cupy" and config.HAS_CUPY and gpu_bn.is_gpu_bn_available(): variant, _ = config.resolve_gpu_launch_geometry() @@ -148,6 +153,15 @@ def _compile_for_device(self, device): self.forward = self._forward_fallback self.backward = self._backward_fallback + def _fuse_activation(self, alpha) -> bool: + """Apply `out = z > 0 ? z : alpha * z` to the output inside the GPU kernels, + so the activation layer that follows can be skipped: 0 is ReLU, else + LeakyReLU. Returns False, changing nothing, when the GPU path isn't bound.""" + if self.forward != self._forward_gpu: + return False + self._alpha = np.float32(alpha) + return True + def _forward_gpu(self, inputs, training): xp = config.cp # The kernels read raw float* buffers; a precision policy can hand us @@ -160,7 +174,7 @@ def _forward_gpu(self, inputs, training): if not training: scale = self.gamma / xp.sqrt(self.running_var + self.epsilon) - gpu_bn._bn_apply_ew(inputs, self.running_mean, scale, self.beta, output) + gpu_bn._bn_apply_ew(inputs, self.running_mean, scale, self.beta, self._alpha, output) return output meta = self._get_shape_meta(inputs.shape) @@ -170,7 +184,7 @@ def _forward_gpu(self, inputs, training): partials = xp.empty(meta["partials_shape"], dtype=xp.float32) self._stats_kernel( meta["grid_dim"], meta["block_dim"], - (inputs_flat, None, None, partials) + meta["static_args"], + (inputs_flat, None, None, None, None, np.float32(1.0), partials) + meta["static_args"], ) # Rows: mean, var, inv_std, scale. The running statistics update in place. @@ -186,7 +200,7 @@ def _forward_gpu(self, inputs, training): self.inputs = inputs self.batch_mean, self.batch_var, self.inv_std, self.scale = stats - gpu_bn._bn_apply_ew(inputs, self.batch_mean, self.scale, self.beta, output) + gpu_bn._bn_apply_ew(inputs, self.batch_mean, self.scale, self.beta, self._alpha, output) return output def _backward_gpu(self, dvalues): @@ -200,8 +214,8 @@ def _backward_gpu(self, dvalues): partials = xp.empty(meta["partials_shape"], dtype=xp.float32) self._grad_sums_kernel( meta["grid_dim"], meta["block_dim"], - (dvalues.reshape(N, C), self.inputs.reshape(N, C), self.batch_mean, partials) - + meta["static_args"], + (dvalues.reshape(N, C), self.inputs.reshape(N, C), self.batch_mean, + self.scale, self.beta, self._alpha, partials) + meta["static_args"], ) # Rows: dgamma, dbeta, then the dx coefficients k1 and k2. @@ -214,7 +228,8 @@ def _backward_gpu(self, dvalues): self.dinputs = xp.empty_like(dvalues) gpu_bn._bn_backward_ew( - dvalues, self.inputs, self.batch_mean, self.scale, k1, k2, self.dinputs + dvalues, self.inputs, self.batch_mean, self.scale, self.beta, k1, k2, + self._alpha, self.dinputs, ) return self.dinputs diff --git a/tests/layers/test_batch_normalization.py b/tests/layers/test_batch_normalization.py index 66e138f..b0a226c 100644 --- a/tests/layers/test_batch_normalization.py +++ b/tests/layers/test_batch_normalization.py @@ -1,5 +1,7 @@ import aether.config as config +import aether.custom_kernels.relu_kernel as gpu_relu import tests.base_case as base_case +from aether.layers.activations import LeakyReLU, ReLU from aether.layers.normalization import BatchNorm @@ -312,6 +314,79 @@ def test_gpu_float64_running_stats(self): expected = (self.dense_inputs - mean64) / self.xp.sqrt(var64 + self.EPSILON) self.assertTrue(self.xp.allclose(out, expected, atol=1e-5, rtol=1e-5)) + def test_gpu_fused_activation_matches_unfused_chain(self): + """BatchNorm with a fused ReLU (alpha 0) or LeakyReLU (alpha 0.1) on GPU against + the fallback BatchNorm followed by the activation layer's own fallback: output, + dinputs, dweights and dbiases in training, and the output in inference.""" + if self.backend_name != "cupy": + self.skipTest("GPU fusion test is only applicable on CuPy backend.") + + C = self.RAGGED_SHAPE[-1] + # Every value is -3, -1, 1 or 3, so |x_hat| >= ~0.4 and z = gamma * x_hat + beta + # stays >= ~0.1 from 0: the two paths' rounding can't flip an element's mask. + levels = self.xp.array([-3.0, -1.0, 1.0, 3.0], dtype=self.xp.float32) + inputs = levels[self.xp.random.randint(0, 4, self.RAGGED_SHAPE)] + dvalues = self.xp.random.randn(*self.RAGGED_SHAPE).astype(self.xp.float32) + + for alpha, activation in [(0.0, ReLU()), (0.1, LeakyReLU(alpha=0.1))]: + with self.subTest(activation=type(activation).__name__): + layer = self.make_built_layer( + BatchNorm, input_shape=self.RAGGED_SHAPE[1:], epsilon=self.EPSILON, momentum=self.MOMENTUM + ) + layer.gamma = self.xp.random.uniform(0.5, 2.0, C).astype(self.xp.float32) + layer.beta = self.xp.random.uniform(-0.1, 0.1, C).astype(self.xp.float32) + + expected_out = activation._forward_fallback(layer._forward_fallback(inputs, training=True), True) + expected_dinputs = layer._backward_fallback(activation._backward_fallback(dvalues)) + expected_dweights, expected_dbiases = layer.dweights.copy(), layer.dbiases.copy() + + self.assertTrue(layer._fuse_activation(alpha)) + out = layer._forward_gpu(inputs, training=True) + dinputs = layer._backward_gpu(dvalues) + + self.assertTrue(self.xp.allclose(out, expected_out, atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(dinputs, expected_dinputs, atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.dweights, expected_dweights, atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.dbiases, expected_dbiases, atol=1e-4, rtol=1e-4)) + + # Running statistics (0, 5) put x_hat at +-0.45 and +-1.34 again. + layer.running_mean = self.xp.zeros(C, dtype=self.xp.float32) + layer.running_var = self.xp.full(C, 5.0, dtype=self.xp.float32) + expected = activation._forward_fallback(layer._forward_fallback(inputs, training=False), False) + out = layer._forward_gpu(inputs, training=False) + self.assertTrue(self.xp.allclose(out, expected, atol=1e-5, rtol=1e-5)) + + def test_gpu_fused_relu_nan_gradient_at_inactive_unit(self): + """As in ReLU's GPU backward, a NaN gradient at a unit the fused ReLU zeroed + contributes 0 to dinputs, dweights and dbiases, matching the unfused GPU + chain. With a LeakyReLU slope it propagates, as unfused LeakyReLU does.""" + if self.backend_name != "cupy": + self.skipTest("GPU fusion test is only applicable on CuPy backend.") + + layer, inputs = self._make_ragged() + C = self.RAGGED_SHAPE[-1] + z = layer._forward_gpu(inputs, training=True) + row, channel = (int(i) for i in self.xp.argwhere(z.reshape(-1, C) < 0)[0]) + dvalues = self.xp.random.randn(*self.RAGGED_SHAPE).astype(self.xp.float32) + dvalues.reshape(-1, C)[row, channel] = self.xp.nan + + # Unfused GPU chain: BatchNorm, then ReLU's select backward. + relu_dvalues = gpu_relu.relu_backward(dvalues, self.xp.maximum(0, z)) + expected = [layer._backward_gpu(relu_dvalues).copy(), layer.dweights.copy(), layer.dbiases.copy()] + + layer._fuse_activation(0.0) + layer._forward_gpu(inputs, training=True) + fused = [layer._backward_gpu(dvalues), layer.dweights, layer.dbiases] + for name, got, want in zip(("dinputs", "dweights", "dbiases"), fused, expected): + with self.subTest(name=name): + self.assertFalse(self.xp.any(self.xp.isnan(got))) + self.assertTrue(self.xp.allclose(got, want, atol=1e-4, rtol=1e-4)) + + layer._fuse_activation(0.1) + layer._forward_gpu(inputs, training=True) + layer._backward_gpu(dvalues) + self.assertTrue(self.xp.isnan(layer.dbiases[channel])) + def test_gpu_non_float32_inputs_match_float32(self): """Regression: under set_precision() a Dense before BatchNorm emits fp16/fp64, which the float* reduction kernels misread and the fp32-only elementwise From 025747df3921533fd956a2f562a6becfea0ccec6 Mon Sep 17 00:00:00 2001 From: Alexander <208520214+AlexanderSoftCode@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:18:12 -0500 Subject: [PATCH 2/2] perf(model): fuse BatchNorm + activation pairs at finalize Model.finalize(input_shape, fuse_kernels=True) now looks for a BatchNorm directly followed by a ReLU (slope 0) or LeakyReLU (its alpha), matched on exact types so subclasses never fuse. On CuPy the BatchNorm applies the activation inside its kernels and the activation layer is skipped in both directions: Model.forward runs the new _forward_layers and backward runs _backward_layers, both without the fused activations. Fusion is planned in _bind_fusion, which runs at the end of finalize() and after every to() on a finalized model. Every _compile_for_device resets BatchNorm to no activation, so a model fused on CuPy and moved with to('numpy') runs every layer again (otherwise it would silently skip its ReLUs), and moving it back re-fuses. fuse_kernels=False keeps every layer on both devices. self.layers is unchanged, so save(), the manifest, the weight keys and validate_graph are untouched. fuse_kernels is a runtime choice and is not saved: load() calls finalize() with the default and fuses again. A fused activation's `output` attribute is not updated; nothing in the package reads it (only the loss-fused SoftMax's output is read, and SoftMax doesn't fuse). CIFAR-10 on the RX 9070 XT: 5.889 -> 5.119 ms per step (-13.1 %). The ReLU rows (0.72 ms GPU) drop out of the breakdown and BatchNorm's time doesn't grow. Tests, on a Conv -> BatchNorm -> ReLU -> Conv -> BatchNorm -> LeakyReLU model in strict fp32: fused and unfused give the same outputs, gradients and Adam step (except the conv biases in front of BatchNorm, whose true gradient is 0, so Adam scales rounding noise to ~lr); a fused model moved to NumPy runs every layer and matches the GPU predictions, then re-fuses on the way back; and save + load keeps the activations in the manifest and re-fuses. --- aether/model.py | 53 ++++++++-- tests/integration/test_device_migration.py | 113 +++++++++++++++++++++ 2 files changed, 160 insertions(+), 6 deletions(-) diff --git a/aether/model.py b/aether/model.py index 22dfb75..061ae8c 100644 --- a/aether/model.py +++ b/aether/model.py @@ -58,6 +58,15 @@ "weights.safetensors": (zipfile.ZIP_STORED,), } +# (normalization, activation) pairs whose activation BatchNorm can apply inside its +# own GPU kernels, mapped to the leaky slope it takes. Exact types, so subclasses +# that change the activation never fuse. +_FUSIBLE_PAIRS = { + (layer_module.BatchNorm, layer_module.ReLU): lambda activation: 0.0, + (layer_module.BatchNorm, layer_module.LeakyReLU): lambda activation: activation.alpha, +} + + def _serialize_component(obj): if obj is None or isinstance(obj, (NullOptimizer, NullAccuracy, NullPreprocessor)): return None @@ -284,6 +293,9 @@ def __init__(self): self.device = None self.precision_policy = None self.trainable_layers = [] + self._fuse_kernels = True + self._forward_layers = () + self._backward_layers = () self._seed = None self._predict_activation = None self._rng_clock = None @@ -422,6 +434,7 @@ def to(self, device): self._sync_device(target_device=target_device) if self.is_finalized: + self._bind_fusion() for layer in self.trainable_layers: param_dict = layer.get_parameters() migrated_dict = {} @@ -464,10 +477,19 @@ def set_precision(self, compute_dtype): ): self.preprocessor._apply_precision(self.precision_policy) - def finalize(self, input_shape: tuple[int, ...]): + def finalize(self, input_shape: tuple[int, ...], fuse_kernels: bool = True): """ Finalizes the model architecture, propagating tensor dimensions and allocating weights/buffers across all registered layers. + + Args: + input_shape (tuple[int, ...]): Shape of one sample, without the batch axis. + fuse_kernels (bool): On CuPy, a BatchNorm directly followed by ReLU or + LeakyReLU applies the activation inside its own kernels and the + activation layer is skipped, saving a read and a write of the + activations in both directions. Fusion is re-planned on every `to()`, + and never applies on NumPy. Not saved: `load()` fuses by default. + Defaults to True. """ if self.is_finalized: raise RuntimeError("Cannot modify model after finalize() has been called.") @@ -504,12 +526,11 @@ def finalize(self, input_shape: tuple[int, ...]): if getattr(layer, "weights", None) is not None or getattr(layer, "biases", None) is not None ] - # Nothing below the first trainable layer consumes a gradient, so backward stops - # there and that layer skips its dX. - first = self.layers.index(self.trainable_layers[0]) if self.trainable_layers else 0 - self._backward_layers = tuple(reversed(self.layers[first:])) + # Nothing below the first trainable layer consumes a gradient, so that layer + # skips its dX (and backward stops there, see _bind_fusion). if self.trainable_layers: self.trainable_layers[0].needs_input_grad = False + self._fuse_kernels = fuse_kernels if self.loss is not None: self.loss.validate_graph(layers=self.layers) @@ -545,8 +566,28 @@ def finalize(self, input_shape: tuple[int, ...]): self.device = "cupy" if config.xp.__name__ == "cupy" else "numpy" self._sync_device() + self._bind_fusion() self.is_finalized = True + def _bind_fusion(self): + """Fuses BatchNorm + activation pairs for the device the layers were just + compiled for, and rebuilds the forward/backward layer sequences without the + fused activations. Every _compile_for_device resets BatchNorm to no + activation, so a pair fused on CuPy runs unfused after `to('numpy')`.""" + fused = set() + if self._fuse_kernels: + for norm, activation in zip(self.layers, self.layers[1:]): + alpha_of = _FUSIBLE_PAIRS.get((type(norm), type(activation))) + if alpha_of is not None and norm._fuse_activation(alpha_of(activation)): + fused.add(id(activation)) + + self._forward_layers = tuple(layer for layer in self.layers if id(layer) not in fused) + # Backward stops at the first trainable layer. + first = self.layers.index(self.trainable_layers[0]) if self.trainable_layers else 0 + self._backward_layers = tuple( + layer for layer in reversed(self.layers[first:]) if id(layer) not in fused + ) + def _check_ready(self, op_name): """Guard shared by train/evaluate/predict: the model must be finalized.""" if not self.is_finalized: @@ -658,7 +699,7 @@ def forward(self, X, training=True): Sequentially propagate input batch through all layers. """ output = X - for layer in self.layers: + for layer in self._forward_layers: output = layer.forward(output, training=training) return output diff --git a/tests/integration/test_device_migration.py b/tests/integration/test_device_migration.py index 6970dcd..dac05bd 100644 --- a/tests/integration/test_device_migration.py +++ b/tests/integration/test_device_migration.py @@ -1,7 +1,11 @@ import functools +import os +import tempfile +import warnings import numpy as np +import aether as ae import aether.config as config from aether.base import Layer from tests.base_case import register_test_suites @@ -117,4 +121,113 @@ def test_training_resumes_after_migration_both_directions(self): for layer in model.trainable_layers: self.assertIsInstance(layer.weight_momentums, self.xp.ndarray) + def _build_bn_model(self, *, fuse_kernels=True): + """Conv -> BatchNorm -> ReLU -> Conv -> BatchNorm -> LeakyReLU -> GAP -> Dense, in + strict fp32: Conv2d's fp16 path would put its own rounding between the CPU and + GPU results, and Adam's first step (about lr * sign(g)) would turn it into + lr-sized parameter differences.""" + self.enterContext(warnings.catch_warnings()) + warnings.filterwarnings("ignore", message=r"\[aether\] Conv2d: precision policy requests 'float32'") + model = ae.Model() + model.add(ae.Conv2d(3, 8, (3, 3), (1, 1), padding="same")) + model.add(ae.BatchNorm()) + model.add(ae.ReLU()) + model.add(ae.Conv2d(8, 8, (3, 3), (1, 1), padding="same")) + model.add(ae.BatchNorm()) + model.add(ae.LeakyReLU(alpha=0.1)) + model.add(ae.GlobalAvgPool()) + model.add(ae.Dense(8, self.NUM_CLASSES)) + model.configure( + loss=ae.SoftmaxCategoricalCrossEntropy(), + optimizer=ae.Adam(lr=0.01), + ) + model.to("cupy") + model.manual_seed(seed=self.SEED) + model.set_precision("float32") + model.finalize(input_shape=(8, 8, 3), fuse_kernels=fuse_kernels) + return model + + def _activations(self, layers): + return [type(layer).__name__ for layer in layers if isinstance(layer, (ae.ReLU, ae.LeakyReLU))] + + def test_fused_batchnorm_activation_matches_unfused(self): + """fuse_kernels=True skips the ReLU/LeakyReLU after each BatchNorm on CuPy, and + gives the same outputs, gradients and optimizer step as fuse_kernels=False.""" + if not config.HAS_CUPY or self.backend_name != "cupy": + self.skipTest("Kernel fusion only applies on the CuPy backend.") + + X, y = self.make_synthetic_image_data(n_samples=16, height=8, width=8) + fused = self._build_bn_model(fuse_kernels=True) + unfused = self._build_bn_model(fuse_kernels=False) + self.assertEqual(self._activations(fused._forward_layers), []) + self.assertEqual(self._activations(fused._backward_layers), []) + self.assertEqual(self._activations(unfused._forward_layers), ["ReLU", "LeakyReLU"]) + + outputs = [model.forward(X, training=True) for model in (fused, unfused)] + self.assertTrue(self.xp.allclose(*outputs, atol=1e-5, rtol=1e-4)) + + for model in (fused, unfused): + model.train(X, y, epochs=1, verbose=0) + for fused_layer, unfused_layer in zip(fused.trainable_layers, unfused.trainable_layers): + # A bias in front of BatchNorm has a true gradient of 0 (BatchNorm removes the + # mean), so its computed one is rounding noise that Adam scales up to ~lr. + idx = fused.layers.index(fused_layer) + before_norm = any(isinstance(layer, ae.BatchNorm) for layer in fused.layers[idx + 1:idx + 2]) + zero_grad = {"biases"} if before_norm else set() + with self.subTest(layer=type(fused_layer).__name__): + for name in ("dweights", "dbiases"): + self.assertTrue(self.xp.allclose( + getattr(fused_layer, name), getattr(unfused_layer, name), atol=1e-5, rtol=1e-4 + )) + for name, tensor in fused_layer.get_parameters().items(): + if name in zero_grad: + continue + self.assertTrue(self.xp.allclose( + tensor, unfused_layer.get_parameters()[name], atol=1e-5, rtol=1e-4 + )) + + def test_fused_model_runs_every_layer_after_to_numpy(self): + """Regression guard: a model fused on CuPy must run its activations again after + to('numpy'), whose BatchNorm doesn't apply them, and re-fuse on the way back.""" + if not config.HAS_CUPY or self.backend_name != "cupy": + self.skipTest("Cross-backend migration requires CuPy to be active.") + + X, y = self.make_synthetic_image_data(n_samples=16, height=8, width=8) + model = self._build_bn_model() + model.train(X, y, epochs=1, verbose=0) + gpu_logits = config.to_device(model.predict(X, return_logits=True), target="numpy") + + model.to("numpy") + X_cpu = config.to_device(X, target="numpy") + self.assertEqual(model._forward_layers, tuple(model.layers)) + cpu_logits = model.predict(X_cpu, return_logits=True) + np.testing.assert_allclose(cpu_logits, gpu_logits, rtol=1e-4, atol=1e-4) + self.assertGreaterEqual(model.layers[2].output.min(), 0.0) + + model.to("cupy") + self.assertEqual(self._activations(model._forward_layers), []) + regpu_logits = config.to_device(model.predict(X, return_logits=True), target="numpy") + np.testing.assert_allclose(regpu_logits, gpu_logits, rtol=1e-4, atol=1e-4) + + def test_loaded_model_refuses_batchnorm_activation(self): + """Fusion isn't saved: the manifest keeps every layer, and load() fuses again.""" + if not config.HAS_CUPY or self.backend_name != "cupy": + self.skipTest("Kernel fusion only applies on the CuPy backend.") + + X, y = self.make_synthetic_image_data(n_samples=16, height=8, width=8) + model = self._build_bn_model() + model.train(X, y, epochs=1, verbose=0) + expected = model.predict(X, return_logits=True) + + with tempfile.TemporaryDirectory() as temp_dir: + path = os.path.join(temp_dir, "fused.aether") + model.save(path) + loaded = ae.Model.load(path, device="cupy") + + self.assertEqual(self._activations(loaded.layers), ["ReLU", "LeakyReLU"]) + self.assertEqual(self._activations(loaded._forward_layers), []) + np.testing.assert_allclose( + loaded.predict(X, return_logits=True), expected, rtol=1e-4, atol=1e-4 + ) + register_test_suites(globals(), TestDeviceMigration) \ No newline at end of file