From 763989aeecdf6b0a7f4667e61b52842b856d6e28 Mon Sep 17 00:00:00 2001 From: Albert Date: Fri, 17 Jul 2026 09:11:05 +0800 Subject: [PATCH] Fix(BN mudnn api use bug) --- .../kernels/nn/musa_fused_batchnorm_kernel.mu | 19 +++++++++ .../kernels/nn/musa_fused_batchnorm_op.cc | 18 ++++++++- test/ops/fused_batchnorm_op_test.py | 40 +++++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/musa_ext/kernels/nn/musa_fused_batchnorm_kernel.mu b/musa_ext/kernels/nn/musa_fused_batchnorm_kernel.mu index a0e6a262..10ecf9e5 100644 --- a/musa_ext/kernels/nn/musa_fused_batchnorm_kernel.mu +++ b/musa_ext/kernels/nn/musa_fused_batchnorm_kernel.mu @@ -20,6 +20,14 @@ __global__ void BesselCorrectionKernel(float* data, float factor, int count) { } } +__global__ void VarianceToInvStdKernel(const float* variance, float* inv_std, + float epsilon, int count) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < count) { + inv_std[idx] = rsqrtf(variance[idx] + epsilon); + } +} + extern "C" { // Launch Bessel correction kernel on the given stream. @@ -36,4 +44,15 @@ void LaunchBesselCorrection(float* data, float factor, int count, data, factor, count); } +// TensorFlow exposes population variance as reserve_space_2, while muDNN's +// training backward kernel consumes reciprocal standard deviation. +void LaunchVarianceToInvStd(const float* variance, float* inv_std, + float epsilon, int count, musaStream_t stream) { + if (count <= 0) return; + int block_size = 256; + int grid_size = (count + block_size - 1) / block_size; + VarianceToInvStdKernel<<>>( + variance, inv_std, epsilon, count); +} + } // extern "C" diff --git a/musa_ext/kernels/nn/musa_fused_batchnorm_op.cc b/musa_ext/kernels/nn/musa_fused_batchnorm_op.cc index e29fbe38..bf92a99d 100644 --- a/musa_ext/kernels/nn/musa_fused_batchnorm_op.cc +++ b/musa_ext/kernels/nn/musa_fused_batchnorm_op.cc @@ -16,6 +16,8 @@ extern "C" { void LaunchBesselCorrection(float* data, float factor, int count, musaStream_t stream); +void LaunchVarianceToInvStd(const float* variance, float* inv_std, + float epsilon, int count, musaStream_t stream); } namespace tensorflow { @@ -260,7 +262,19 @@ class MusaFusedBatchNormGradOp : public MusaOpKernel { mTensor mt_scale = CreateMTensor(scale, mFormat::NCHW); mTensor mt_saved_mean = CreateMTensor(saved_mean, mFormat::NCHW); - mTensor mt_saved_var = CreateMTensor(saved_var, mFormat::NCHW); + // TensorFlow's reserve_space_2 contains population variance. muDNN's + // training backward kernels interpret their `v` input as reciprocal + // standard deviation and use it directly (including inv_std^3 terms). + // Passing variance happens to look plausible when variance is near one, + // but makes dx scale catastrophically for high-variance inputs. + Tensor saved_inv_std; + OP_REQUIRES_OK(ctx, + ctx->allocate_temp(DT_FLOAT, saved_var.shape(), + &saved_inv_std)); + LaunchVarianceToInvStd( + saved_var.flat().data(), saved_inv_std.flat().data(), + epsilon_, static_cast(saved_var.NumElements()), stream); + mTensor mt_saved_inv_std = CreateMTensor(saved_inv_std, mFormat::NCHW); mTensor mt_d_scale = CreateMTensor(*d_scale, mFormat::NCHW); mTensor mt_d_offset = CreateMTensor(*d_offset, mFormat::NCHW); @@ -274,7 +288,7 @@ class MusaFusedBatchNormGradOp : public MusaOpKernel { mStatus status = bn_op.RunBwd( handle, mt_dx, mt_d_mean, mt_d_var, mt_d_scale, mt_d_offset, mt_x, - mt_dy, mt_saved_mean, mt_saved_var, mt_scale, maintainer); + mt_dy, mt_saved_mean, mt_saved_inv_std, mt_scale, maintainer); OP_REQUIRES(ctx, status == mStatus::SUCCESS, errors::Internal("MUSA BN Backward failed.")); diff --git a/test/ops/fused_batchnorm_op_test.py b/test/ops/fused_batchnorm_op_test.py index ef26c756..e6f45015 100755 --- a/test/ops/fused_batchnorm_op_test.py +++ b/test/ops/fused_batchnorm_op_test.py @@ -215,6 +215,46 @@ def grad_dx_nchw(x, scale, offset, mean, var): atol=1e-3, ) + def testLayerNormGradNCHWLargeVariance(self): + """LayerNorm's NCHW fused-BN backward uses variance correctly. + + Keras implements rank-2 LayerNormalization with FusedBatchNormV3 by + reshaping [batch, features] to NCHW [1, batch, features, 1]. TensorFlow + passes population variance to FusedBatchNormGradV3, whereas muDNN's + training backward API consumes reciprocal standard deviation. Passing + variance directly makes dx grow by many orders of magnitude when the + input variance is far from one, as seen in Wukong's FMB blocks. + """ + rng = np.random.default_rng(42) + x_np = (10.0 * rng.normal(size=(128, 1536))).astype(np.float32) + dy_np = rng.normal(scale=1e-3, size=x_np.shape).astype(np.float32) + + def run(device, dtype): + layer_dtype = ( + "mixed_bfloat16" if dtype == tf.bfloat16 else "float32" + ) + with tf.device(device): + x = tf.cast(tf.constant(x_np), dtype) + dy = tf.cast(tf.constant(dy_np), dtype) + layer = tf.keras.layers.LayerNormalization( + axis=-1, epsilon=1e-3, dtype=layer_dtype + ) + with tf.GradientTape() as tape: + tape.watch(x) + y = layer(x) + dx = tape.gradient(y, x, output_gradients=tf.cast(dy, y.dtype)) + return tf.cast(dx, tf.float32) + + for dtype in (tf.float32, tf.bfloat16): + dx_cpu = run("/CPU:0", dtype) + dx_musa = run("/device:MUSA:0", dtype) + self.assertAllClose( + dx_musa.numpy(), + dx_cpu.numpy(), + rtol=0.15, + atol=5e-5, + ) + def testFusedBatchNormGradDXNotAffectedByPriorMemory(self): """BN backward dx must not inherit values from a previously freed buffer.