[xpu] support: each parameter has different lr in merged_momentum (#71212)

zhouquan32 · web-flow · commit f4e8acbdfd99 · 2025-03-05T10:26:59.000+08:00
* support: each parameter has different lr in merged_momentum

* fix code style

* support only one lr in merged_momentum

* support only one lr in merged_momentum

* limit: lr should be either 1 or len of param

* fix code style

* add some test cases for merged_moementum
diff --git a/paddle/phi/kernels/xpu/merged_momentum_kernel.cc b/paddle/phi/kernels/xpu/merged_momentum_kernel.cc
@@ -12,16 +12,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "paddle/phi/kernels/merged_momentum_kernel.h"
+
 #include <sys/syscall.h>
 #include <unistd.h>
+
 #include <iostream>
 #include <string>
 #include <vector>
 
-#include "paddle/phi/kernels/merged_momentum_kernel.h"
-
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/momentum_kernel.h"
 
 namespace phi {
 
@@ -43,9 +45,9 @@ void MergedMomentumKernel(
     std::vector<DenseTensor*> velocity_out,
     std::vector<DenseTensor*> master_param_out) {
   using XPUType = typename XPUTypeTrait<T>::Type;
-  auto lr = learning_rate[0];
   T mu = static_cast<T>(mu_in);
   int op_num = params.size();
+  int lr_len = learning_rate.size();
   PADDLE_ENFORCE_EQ(op_num,
                     params_out.size(),
                     errors::InvalidArgument(
@@ -78,11 +80,17 @@ void MergedMomentumKernel(
           "the size of Input(Grad) is %d, the size of Input(Param) is %d.",
           grad.size(),
           op_num));
+  PADDLE_ENFORCE_EQ(
+      lr_len == 1 || lr_len == op_num,
+      true,
+      errors::InvalidArgument(
+          "The len of learning_rate should be either 1 or %d.", op_num));
   std::vector<XPUType*> param_list(op_num);
   std::vector<XPUType*> velocity_list(op_num);
   std::vector<XPUType*> grad_list(op_num);
   std::vector<XPUType*> velocity_out_list(op_num);
   std::vector<XPUType*> param_out_list(op_num);
+  std::vector<const float*> lr_list(op_num);
   std::vector<int> sizes(op_num);
   std::vector<float> l2_weight_decay(op_num);
   if (op_num > 0) {
@@ -96,6 +104,9 @@ void MergedMomentumKernel(
       param_out_list[j] = reinterpret_cast<XPUType*>(params_out[j]->data<T>());
       velocity_out_list[j] =
           reinterpret_cast<XPUType*>(velocity_out[j]->data<T>());
+      if (lr_len == op_num) {
+        lr_list[j] = learning_rate[j]->data<float>();
+      }
       sizes[j] = static_cast<int>(params[j]->numel());
       if (regularization_method[j] != "l2_decay") {
         l2_weight_decay[j] = 0.0f;
@@ -116,6 +127,24 @@ void MergedMomentumKernel(
   } else {
     return;
   }
+  if (lr_len == op_num) {
+    for (int j = 0; j < op_num; j++) {
+      int r = xpu::momentum(dev_ctx.x_context(),
+                            param_list[j],
+                            velocity_list[j],
+                            grad_list[j],
+                            param_out_list[j],
+                            velocity_out_list[j],
+                            sizes[j],
+                            lr_list[j],
+                            use_nesterov,
+                            mu,
+                            l2_weight_decay[j]);
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "momentum");
+    }
+    return;
+  }
+  auto lr = learning_rate[0];
   int r = xpu::merged_momentum(dev_ctx.x_context(),
                                param_list,
                                velocity_list,
@@ -129,7 +158,6 @@ void MergedMomentumKernel(
                                use_nesterov);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "merged_momentum");
 }
-
 }  // namespace phi
 
 PD_REGISTER_KERNEL(merged_momentum,
diff --git a/test/xpu/test_merged_momentum_op_xpu.py b/test/xpu/test_merged_momentum_op_xpu.py
@@ -38,6 +38,7 @@ def setUp(self):
 
         def set_case(self):
             self.shapes = [[3, 4], [2, 7], [5, 6, 8]]
+            self.lr_shapes = [[1], [1], [1]]
             self.place = paddle.base.XPUPlace(0)
             self.seed = 1
 
@@ -47,18 +48,22 @@ def testalltype(self):
     class TestMergedMomentum1(TestMergedMomentumOp):
         def set_case(self):
             self.shapes = [[3, 4], [2, 7], [5, 6, 8]]
+            self.lr_shapes = [[1]]
 
     class TestMergedMomentum2(TestMergedMomentumOp):
         def set_case(self):
             self.shapes = [[3, 4], [2, 7]]
+            self.lr_shapes = [[1], [1]]
 
     class TestMergedMomentum3(TestMergedMomentumOp):
         def set_case(self):
             self.shapes = [[3, 4]]
+            self.lr_shapes = [[1]]
 
     class TestMergedMomentum4(TestMergedMomentumOp):
         def set_case(self):
             self.shapes = [[3, 4], [2, 7], [5, 6, 7], [9, 9], [10, 12]]
+            self.lr_shapes = [[1]]
 
 
 support_types = get_xpu_op_support_types('merged_momentum')
diff --git a/test/xpu/test_merged_momentum_op_xpu_base.py b/test/xpu/test_merged_momentum_op_xpu_base.py
@@ -29,7 +29,7 @@ def run_momentum_op(
     grads,
     velocities,
     master_params,
-    learning_rate,
+    learning_rates,
     place,
     multi_precision,
     mu=0.9,
@@ -39,6 +39,7 @@ def run_momentum_op(
 ):
     assert len(params) == len(grads)
     assert len(params) == len(velocities)
+    assert len(learning_rates) == len(params) or len(learning_rates) == 1
     if multi_precision:
         assert len(params) == len(master_params)
     op_type = 'merged_momentum' if use_merged else 'momentum'
@@ -64,11 +65,12 @@ def run_momentum_op(
                 )
                 for v in velocities
             ]
-            lr_var = helper.create_variable(
-                persistable=True,
-                shape=learning_rate.shape,
-                dtype=learning_rate.dtype,
-            )
+            lr_vars = [
+                helper.create_variable(
+                    persistable=True, shape=lr.shape, dtype=lr.dtype
+                )
+                for lr in learning_rates
+            ]
 
             feed_dict = OrderedDict()
 
@@ -98,8 +100,14 @@ def run_momentum_op(
                     ]
                 )
             )
-            feed_dict.update({lr_var.name: learning_rate})
-
+            feed_dict.update(
+                OrderedDict(
+                    [
+                        (lr_var.name, lr_val)
+                        for lr_var, lr_val in zip(lr_vars, learning_rates)
+                    ]
+                )
+            )
             if multi_precision:
                 master_param_vars = [
                     helper.create_variable(
@@ -129,6 +137,10 @@ def run_momentum_op(
                 for i, (p, g, v) in enumerate(
                     zip(param_vars, grad_vars, velocity_vars)
                 ):
+                    if len(lr_vars) == 1:
+                        lr_var = lr_vars[0]
+                    else:
+                        lr_var = lr_vars[i]
                     inputs = {
                         'Param': p,
                         'Grad': g,
@@ -158,7 +170,7 @@ def run_momentum_op(
                     'Param': param_vars,
                     'Grad': grad_vars,
                     'Velocity': velocity_vars,
-                    'LearningRate': lr_var,
+                    'LearningRate': lr_vars,
                 }
                 outputs = {'ParamOut': param_vars, 'VelocityOut': velocity_vars}
                 if multi_precision:
@@ -190,34 +202,42 @@ class TestMergedMomentumBase(unittest.TestCase):
     def setUp(self):
         paddle.enable_static()
         self.shapes = [[3, 4], [2, 7], [5, 6], [7, 8]]
+        self.lr_shapes = [[1]]
         self.seed = 10
         self.place = paddle.base.XPUPlace(0)
         self.__class__.use_xpu = True
 
     def gen_rand_data(self, shapes, dtype):
         return [np.random.random(s).astype(dtype) for s in shapes]
 
-    def prepare_data(self, shapes, multi_precision, seed, dtype, place):
+    def prepare_data(
+        self, shapes, lr_shapes, multi_precision, seed, dtype, place
+    ):
         np.random.seed(seed)
         params = self.gen_rand_data(shapes, dtype)
         grads = self.gen_rand_data(shapes, dtype)
         velocities = self.gen_rand_data(shapes, dtype)
-        learning_rate = self.gen_rand_data([[1]], np.float32)[0]
+        learning_rates = self.gen_rand_data(lr_shapes, np.float32)
         if multi_precision:
             master_params = [p.astype(dtype) for p in params]
         else:
             master_params = None
-        return params, grads, velocities, master_params, learning_rate
+        return params, grads, velocities, master_params, learning_rates
 
     def check_with_place(self, place, dtype, multi_precision=False):
         (
             params,
             grads,
             velocities,
             master_params,
-            learning_rate,
+            learning_rates,
         ) = self.prepare_data(
-            self.shapes, multi_precision, self.seed, dtype, place
+            self.shapes,
+            self.lr_shapes,
+            multi_precision,
+            self.seed,
+            dtype,
+            place,
         )
 
         def run_op(use_nesterov, use_merged):
@@ -227,7 +247,7 @@ def run_op(use_nesterov, use_merged):
                 grads,
                 velocities,
                 master_params,
-                learning_rate,
+                learning_rates,
                 place,
                 multi_precision,
                 rescale_grad=rescale_grad,
