fix test_lars_momentum (#60287)

xingmingyyj · web-flow · commit 20273340ff90 · 2023-12-26T15:32:51.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -126,6 +126,7 @@
     'fused_scale_bias_add_relu',
     'fused_dconv_drelu_dbn',
     'fused_dot_product_attention',
+    'lars_momentum',
     'recv_v2',
     'rnn_',
     'row_conv',
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1443,6 +1443,18 @@
   optional: dropout1_seed, dropout2_seed, linear1_bias, linear2_bias, ln1_scale, ln1_bias, ln2_scale, ln2_bias, ln2_mean, ln2_variance, ln1_mean, ln1_variance, ln1_out
   backward: fused_feedforward_grad
 
+- op: lars_momentum
+  args: (Tensor param,  Tensor velocity, Tensor grad, Tensor learning_rate, Tensor master_param, float mu, float lars_coeff=0.001f, float[] lars_weight_decay={0.0005}, float epsilon=0, bool multi_precision=false, float rescale_grad=1.0f)
+  output: Tensor(param_out), Tensor(velocity_out), Tensor(master_param_out)
+  infer_meta:
+    func: SparseMomentumInferMeta
+    param: [param, learning_rate, velocity]
+  kernel:
+    func: lars_momentum
+    param: [param, velocity, learning_rate, grad, master_param, lars_weight_decay, mu, lars_coeff, epsilon, multi_precision, rescale_grad]
+    data_type: param
+  optional: master_param, master_param_out
+
 - op: number_count
   args: (Tensor numbers, int upper_range)
   output: Tensor(out)
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3460,6 +3460,12 @@
   outputs :
     out : Out
 
+- op: lars_momentum
+  inputs:
+    {param : Param, grad : Grad, velocity : Velocity, learning_rate : LearningRate, master_param : MasterParam}
+  outputs :
+    {param_out: ParamOut, velocity_out: VelocityOut, master_param_out: MasterParamOut}
+
 - op: lod_array_length
   inputs :
     {x: X}
diff --git a/test/white_list/pir_op_test_white_list b/test/white_list/pir_op_test_white_list
@@ -203,6 +203,7 @@ test_mean_op
 test_memcpy_op
 test_meshgrid_op
 test_mode_op
+test_momentum_op
 test_mul_int8_mkldnn_op
 test_mul_op
 test_multi_dot_op
