【PIR Dist Op Reg No.28】 reg dgc (#62781)

xiaoyewww · web-flow · commit 6adee412cdaa · 2024-04-09T19:32:48.000+08:00
* feat(pir): reg dgc

* feat(pir): dgc

* feat(pir): dgc

* feat(pir): reg dgc

* feat(pir): reg dgc

* feat(pir): reg dgc

* feat(pir): reg dgc

* feat(pir): reg dgc

* feat(pir): reg dgc
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -146,6 +146,7 @@
     'distributed_push_sparse',
     'distributed_lookup_table',
     'dgc_momentum',
+    'dgc',
     'dpsgd',
     'embedding_grad_sparse',
     'ftrl',
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -479,6 +479,14 @@
   optional : in_accum, in_state, out_scale, out_accum, out_state
   inplace : (scale -> out_scale, in_accum -> out_accum, in_state -> out_state)
 
+- op : dgc
+  args : (Tensor u, Tensor v, Tensor grad, Tensor param, Tensor current_step, Tensor nranks, float[] sparsity, float m=0.9, bool use_nesterov=true, float rampup_begin_step=0.0, float rampup_step=0.0, float regular_coeff=0.0, int regular_type=0)
+  output : Tensor(u_out), Tensor(v_out), Tensor(encode_grad), Tensor(grad_out), Tensor(k), Tensor(gather_buff)
+  kernel :
+    func : dgc
+    param : [u, v, grad, param, current_step, nranks, m, use_nesterov, sparsity, rampup_begin_step, rampup_step, regular_coeff, regular_type]
+  optional: param
+
 - op : dgc_momentum
   args : (Tensor param, Tensor grad, Tensor velocity, Tensor learning_rate, Tensor master_param, Tensor current_step_tensor, Tensor nranks_tensor, float mu, bool use_nesterov=false, str regularization_method="", float regularization_coeff=0.0f, bool multi_precision=false, float rescale_grad=1.0f, float rampup_begin_step=-1.0f)
   output : Tensor(param_out), Tensor(velocity_out), Tensor(master_param_out), Tensor(grad_out)
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3717,6 +3717,12 @@
   outputs:
     {param_out : ParamOut, moment_out : MomentOut}
 
+- op: dgc
+  inputs:
+    {u: U, v: V, grad: Grad}
+  outputs:
+    {u_out: U_out, v_out: V_out, encode_grad: EncodeGrad, grad_out: Grad_out, gather_buff: GatherBuff}
+
 - op: distribute_fpn_proposals
   inputs :
     {fpn_rois: FpnRois, rois_num: RoisNum}
diff --git a/test/ir/pir/translator/CMakeLists.txt b/test/ir/pir/translator/CMakeLists.txt
@@ -19,8 +19,7 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
      test_distributed_lookup_table_translate)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
      test_distributed_push_sparse_translator)
-list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_distributed_fused_lamb_init)
-list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_dgc_momentum_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_dgc_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_nop_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_allgather_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_send_translator)
@@ -36,6 +35,10 @@ if(NOT WITH_DISTRIBUTE)
   list(REMOVE_ITEM TEST_INTERP_CASES ${DISTRIBUTED_OP_TRANSLATOR_TEST})
 endif()
 
+if(NOT WITH_DGC)
+  list(REMOVE_ITEM TEST_INTERP_CASES test_dgc_translator)
+endif()
+
 foreach(target ${TEST_INTERP_CASES})
   py_test_modules(${target} MODULES ${target})
 endforeach()
diff --git a/test/ir/pir/translator/test_dgc_translator.py b/test/ir/pir/translator/test_dgc_translator.py
@@ -0,0 +1,77 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import test_op_translator
+
+import paddle
+from paddle.base.layer_helper import LayerHelper
+
+
+class TestDgcOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "dgc"
+        g_array_size = 102400
+        u = paddle.ones(shape=(g_array_size,), dtype='float32')
+        v = paddle.ones(shape=(g_array_size,), dtype='float32')
+        grad = paddle.ones(shape=(g_array_size,), dtype='float32')
+        param = paddle.ones(shape=(g_array_size,), dtype='float32')
+        current_step = paddle.to_tensor([0.0], dtype='float32')
+        nranks = paddle.to_tensor([2.0], dtype='float32')
+
+        u_out = paddle.ones(shape=(g_array_size,), dtype='float32')
+        v_out = paddle.ones(shape=(g_array_size,), dtype='float32')
+        encode_grad = paddle.ones(shape=(g_array_size,), dtype='float32')
+        grad_out = paddle.ones(shape=(g_array_size,), dtype='float32')
+        k = paddle.to_tensor([0.0], dtype='float32')
+        gather_buff = paddle.ones(shape=(g_array_size,), dtype='float32')
+        attrs = {
+            'm': 0.9,
+            'use_nesterov': True,
+            'sparsity': [],
+            'padding_idx': -1,
+            'rampup_begin_step': 0.0,
+            'rampup_step': 0.0,
+            'regular_coeff': 0.0,
+            'regular_type': 0,
+        }
+        helper = LayerHelper(self.op_type)
+        helper.append_op(
+            type=self.op_type,
+            inputs={
+                "U": u,
+                "V": v,
+                "Grad": grad,
+                "Param": param,
+                "current_step": current_step,
+                "nranks": nranks,
+            },
+            outputs={
+                "U_out": u_out,
+                "V_out": v_out,
+                "EncodeGrad": encode_grad,
+                "Grad_out": grad_out,
+                "k": k,
+                "GatherBuff": gather_buff,
+            },
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()
