fuj/rms-norm (DeepLink-org#861)

* impl rms_norm_backward function, fix rms_norm forward function, fix rms_norm case config

* fix rms norm backward

* rms_norm config for ascend

Author: jingguo-st

diopi_test/python/configs/diopi_configs.py

@@ -8189,10 +8189,12 @@
             args=[
                 {
                     "ins": ['input'],
+                    "requires_grad": [True],
                     "shape": ((5, 5), (35, 125, 32), (16, 64, 64), (1, 32, 32, 8)),
                 },
                 {
                     "ins": ['weight'],
+                    "requires_grad": [True],
                     "shape": ((5, ), (32, ), (64, ), (8, )),
                 },
                 {
@@ -8201,6 +8203,9 @@
                 },
             ],
         ),
+        # saved_args=dict(grad_outputs=0, inv_rms=1),
+        saved_args=dict(inv_rms=1),
+        requires_backward=[0],
     ),
 
     # 'multihead_attention_forward': dict(

diopi_test/python/conformance/diopi_functions.py

@@ -5151,7 +5151,9 @@ def rms_norm(input, normalized_shape, weight, bias, eps):
     func = check_function(call)
     size = list(input.size().data)
     out = Tensor(size, input.get_dtype())
-    inv_rms = Tensor(size, input.get_dtype())
+    inv_rms_size = size.copy()
+    inv_rms_size[-1] = 1
+    inv_rms = Tensor(inv_rms_size, input.get_dtype())
     normalized_shape = Sizes(list(normalized_shape))
     ret = func(
         input.context(),
@@ -5164,7 +5166,21 @@ def rms_norm(input, normalized_shape, weight, bias, eps):
         eps,
     )
     check_returncode(ret)
-    return out
+    return (out, inv_rms)
+
+
+def rms_norm_backward(grad_outputs, input, weight, bias, inv_rms, normalized_shape, eps):
+    call = "diopiRMSNormBackward"
+    func = check_function(call)
+    grad_input = Tensor(list(input.size().data), input.get_dtype())
+    grad_weight = Tensor(list(weight.size().data), weight.get_dtype())
+    grad_bias = Tensor(list(bias.size().data), bias.get_dtype())
+    normalized_shape = Sizes(list(normalized_shape))
+
+    ret = func(input.context(), grad_input, grad_weight, grad_bias, grad_outputs[0], input, weight, bias, inv_rms,
+               normalized_shape, eps)
+    check_returncode(ret)
+    return {'input': grad_input, 'weight': grad_weight}
 
 
 def multihead_attention_forward(

diopi_test/python/conformance/gen_output.py

@@ -226,10 +226,12 @@ def rotary_emb(input, cos, sin, conj):
         return out
 
     def rms_norm(input, normalized_shape, weight, bias, eps):
-        variance = input.to(torch.float32).pow(2).mean(-1, keepdim=True)
-        input = input * torch.rsqrt(variance + eps)
-        out = weight * input
-        return out
+        var = input.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        inv_rms = torch.rsqrt(var + eps)
+        inp = input * inv_rms
+        out = weight * inp
+
+        return (out, inv_rms)
 
     def multihead_attention_forward(q, k, v, dropout_p, is_causal, return_debug_mask, scale):
         # 为了保证精度，因此在test的时候不使用dropout
@@ -298,8 +300,11 @@ class GenOutputData(object):
     db_case_items = {}
 
     @staticmethod
-    def run(diopi_item_config_path='diopi_case_items.cfg', input_path='data/inputs/',
-            output_path='data/outputs/', fname='all_ops', model_name='diopi'):
+    def run(diopi_item_config_path='diopi_case_items.cfg',
+            input_path='data/inputs/',
+            output_path='data/outputs/',
+            fname='all_ops',
+            model_name='diopi'):
         if not os.path.exists(input_path):
             logger.error("Input data is not generated!")
             sys.exit(0)
@@ -332,9 +337,11 @@ def run(diopi_item_config_path='diopi_case_items.cfg', input_path='data/inputs/'
                 output, saved_grads = gen_tensor_obj.gen_data(input_)
                 item['result'] = 'passed'
             except Exception as err_msg:
-                raise GenDataFailedException(f'Generate output data for diopi_functions.{func_name} [{case_name}] failed, cause by \n{err_msg}')
+                raise GenDataFailedException(
+                    f'Generate output data for diopi_functions.{func_name} [{case_name}] failed, cause by \n{err_msg}')
             GenOutputData.db_case_items[case_name] = item
             if output is not None:
+                # import pdb; pdb.set_trace()
                 with open(os.path.join(output_path, case_name), "wb") as f:
                     pickle.dump(GenOutputData.to_numpy(output), f, protocol=4)
                     logger_str = "output"

impl/ascend/device_configs.py

@@ -780,6 +780,14 @@
         ),
     ),
 
+    'rms_norm': dict(
+        name=['rms_norm'],
+        atol=1e-3,
+        rtol=1e-3,
+        atol_half=1e-2,
+        rtol_half=1e-2,
+    ),
+
     'smooth_l1_loss': dict(
         name=['smooth_l1_loss'],
         tensor_para=dict(

impl/ascend/functions_ext/rms_norm.cpp

@@ -22,7 +22,14 @@ diopiError_t diopiRMSNormBackward(diopiContextHandle_t ctx, diopiTensorHandle_t
                                   diopiConstTensorHandle_t bias, diopiConstTensorHandle_t invRms, diopiSize_t normalizedShape, double eps) {
     AscendTensor inputTensor(input);
     ASCEND_CHECK_ABORT(1 == normalizedShape.len && normalizedShape.data[0] == inputTensor.shape()[inputTensor.dim() - 1], "normalized shape error!");
-    AclOpRunner<4, 2>("RmsNorm", ctx).addInput(gradOutput).addInput(input).addInput(invRms).addInput(weight).addOutput(gradInput).addOutput(gradWeight).run();
+    AclOpRunner<4, 2>("RmsNormGrad", ctx)
+        .addInput(gradOutput)
+        .addInput(input)
+        .addInput(invRms)
+        .addInput(weight)
+        .addOutput(gradInput)
+        .addOutput(gradWeight)
+        .run();
     return diopiSuccess;
 }