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