【PIR Dist Op Reg No.29】 reg prune_gate_by_capacity (#62494)

xiaoyewww · web-flow · commit e4de989ed706 · 2024-03-18T15:48:36.000+08:00
* feat(pir): reg prune_gate_by_capacity

* feat(pir): reg prune_gate_by_capacity

* feat(pir): reg prune_gate_by_capacity

* feat(pir): reg prune_gate_by_capacity

* feat(pir): reg prune_gate_by_capacity
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -178,6 +178,7 @@
     'c_reduce_prod',
     'c_reduce_prod_',
     'c_scatter',
+    'prune_gate_by_capacity',
     'push_sparse_v2',
     'push_sparse_v2_',
     'partial_send',
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1232,6 +1232,15 @@
   backward : prod_grad
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
+- op : prune_gate_by_capacity
+  args : (Tensor gate_idx, Tensor expert_count, int64_t n_expert, int64_t n_worker)
+  output : Tensor(new_gate_idx)
+  infer_meta :
+    func : PruneGateByCapacityInferMeta
+  kernel :
+    func : prune_gate_by_capacity
+    data_type : gate_idx
+
 - op : push_sparse_v2
   args : (Tensor[] ids, Tensor[] w, Tensor[] out_grad_in, int embeddingdim = 11, int tableid = 0, str accessorclass = "", str ctrlabelname = "", int paddingid = 0, bool scalesparsegrad = true, str[] inputnames = {}, bool is_distributed = true)
   output : Tensor[](out_grad_out){out_grad_in.size()}
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3809,6 +3809,12 @@
   inputs :
     x : X
 
+- op: prune_gate_by_capacity
+  inputs:
+    {gate_idx: GateIdx, expert_count: ExpertCount}
+  outputs:
+    new_gate_idx: NewGateIdx
+
 - op: random_routing
   inputs:
     {prob : Prob, topk_value : TopK_Value, topk_idx : TopK_Idx}
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2855,6 +2855,35 @@ void PriorBoxInferMeta(const MetaTensor& input,
   var->set_dims(common::make_ddim(dim_vec));
 }
 
+void PruneGateByCapacityInferMeta(const MetaTensor& gate_idx,
+                                  const MetaTensor& expert_count,
+                                  int64_t n_expert,
+                                  int64_t n_worker,
+                                  MetaTensor* new_gate_idx) {
+  auto expert_count_dims = expert_count.dims();
+
+  int64_t expert_count_num_ele = 1;
+  for (int i = 0; i < static_cast<int>(expert_count_dims.size()); i++) {
+    expert_count_num_ele *= expert_count_dims[i];
+  }
+
+  PADDLE_ENFORCE_EQ(
+      expert_count_num_ele,
+      n_expert * n_worker,
+      phi::errors::Unavailable(
+          "The number of elements for expert_count is ( %ld ) incorrect. "
+          "Because the number of expert_count must equal the "
+          "product of n_worker ( %ld ) and n_expert ( %ld ). "
+          "Please input appropriate expert_count again!",
+          expert_count_num_ele,
+          n_worker,
+          n_expert));
+
+  auto gate_idx_in_dims = gate_idx.dims();
+  new_gate_idx->set_dims(gate_idx_in_dims);
+  new_gate_idx->set_dtype(gate_idx.dtype());
+}
+
 void RepeatInterleaveWithTensorIndexInferMeta(const MetaTensor& x,
                                               const MetaTensor& repeats,
                                               int dim,
diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h
@@ -492,6 +492,12 @@ void PriorBoxInferMeta(const MetaTensor& input,
                        MetaTensor* out,
                        MetaTensor* var);
 
+void PruneGateByCapacityInferMeta(const MetaTensor& gate_idx,
+                                  const MetaTensor& expert_count,
+                                  int64_t n_expert,
+                                  int64_t n_worker,
+                                  MetaTensor* new_gate_idx);
+
 void SearchsortedInferMeta(const MetaTensor& sorted_sequence,
                            const MetaTensor& value,
                            bool out_int32,
diff --git a/test/ir/pir/translator/CMakeLists.txt b/test/ir/pir/translator/CMakeLists.txt
@@ -21,6 +21,8 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_distributed_fused_lamb_init)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_nop_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_send_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_recv_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
+     test_prune_gate_by_capacity_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_random_routing_translator)
 
 if(NOT WITH_DISTRIBUTE)
diff --git a/test/ir/pir/translator/test_prune_gate_by_capacity_translator.py b/test/ir/pir/translator/test_prune_gate_by_capacity_translator.py
@@ -0,0 +1,43 @@
+# Copyright (c) 2024 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 TestPruneGateByCapacityOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "prune_gate_by_capacity"
+        gate_idx = paddle.ones(shape=(200,), dtype='int64')
+        expert_count = paddle.ones(shape=(48,), dtype='int64')
+        new_gate_idx = paddle.zeros_like(expert_count)
+        attrs = {'n_expert': 24, 'n_worker': 2}
+        helper = LayerHelper(self.op_type)
+        helper.append_op(
+            type=self.op_type,
+            inputs={"GateIdx": gate_idx, "ExpertCount": expert_count},
+            outputs={"NewGateIdx": new_gate_idx},
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()
