Fix(BN mudnn api use bug) (#291)

Co-authored-by: Albert <albert@mthreads.com>

Author: real-zhangzhe

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<<<grid_size, block_size, 0, stream>>>(
+      variance, inv_std, epsilon, count);
+}
+
 }  // extern "C"

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<float>().data(), saved_inv_std.flat<float>().data(),
+        epsilon_, static_cast<int>(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."));

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.