From 0ecfd69c6e6be153270205c062a34831a6c9183f Mon Sep 17 00:00:00 2001 From: wangyiheng02 <928430299@qq.com> Date: Wed, 31 Jul 2024 15:04:09 +0800 Subject: [PATCH 01/29] no test --- .../interface/infer_symbolic_shape/same_operands_result.cc | 1 + .../interface/infer_symbolic_shape/same_operands_result.h | 1 + paddle/phi/ops/yaml/inconsistent/static_ops.yaml | 1 + 3 files changed, 3 insertions(+) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index 4d601d15d8548..09257bab2f0ed 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -121,6 +121,7 @@ OP_SAME_OPERANDS_AND_RESULT(ScatterNdAdd) OP_SAME_OPERANDS_AND_RESULT(Scatter) OP_SAME_OPERANDS_AND_RESULT(Scatter_) OP_SAME_OPERANDS_AND_RESULT(Select) +OP_SAME_OPERANDS_AND_RESULT(ShadowFeedTensors) OP_SAME_OPERANDS_AND_RESULT(Sign) OP_SAME_OPERANDS_AND_RESULT(Sin) OP_SAME_OPERANDS_AND_RESULT(Sin_) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h index ae4b285148fe3..26f9a799c2141 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h @@ -114,6 +114,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(ScatterNdAdd) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Scatter) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Scatter_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Select) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShadowFeedTensors) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sign) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin_) diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml index 148f572379367..7bf48c70b54e6 100644 --- a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml +++ b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml @@ -829,6 +829,7 @@ kernel: func: shadow_feed_tensors param: [x] + interfaces : paddle::dialect::InferSymbolicShapeInterface - op : share_data_ args : (Tensor x) From 8a582f33b912e50a8d7d75e48ad425069c097a2e Mon Sep 17 00:00:00 2001 From: wangyiheng02 <928430299@qq.com> Date: Mon, 5 Aug 2024 11:02:14 +0800 Subject: [PATCH 02/29] unchanged vector --- .../infer_symbolic_shape/same_operands_result.cc | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index 09257bab2f0ed..aa0388ff348bc 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -189,6 +189,21 @@ bool ArgsortOpInferSymbolicShape( return true; } +bool UnchangedVectorOpInferSymbolicShape( + pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { + const symbol::ShapeOrDataDimExprs &operand_shape_or_data = + infer_context->GetShapeOrDataForValue(op->operand_source(0)); + size_t size = operand_shape_or_data.size(); + std::vector shape(size); + for (size_t i = 0; i < size; ++i) { + shape[i] = operand_shape_or_data.shape()[i]; + } + infer_context->SetShapeOrDataForValue( + op->result(0), + symbol::TensorShapeOrDataDimExprs(shape, operand_shape_or_data.data())); + return true; +} + } // namespace paddle::dialect namespace cinn::dialect {} // namespace cinn::dialect From d7cb8730fc97d39c841e430dc4590d9ba2b5d877 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Tue, 6 Aug 2024 15:48:20 +0800 Subject: [PATCH 03/29] moe --- .../multiary_infer_sym.cc | 77 ++++++++++++++++--- .../infer_symbolic_shape/multiary_infer_sym.h | 4 +- paddle/phi/ops/yaml/ops.yaml | 2 + 3 files changed, 70 insertions(+), 13 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 84c3bc82ee3aa..40c9233afa4d7 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1199,18 +1199,73 @@ bool RoiAlignOpInferSymbolicShape( // return MergedMomentumOpInferSymbolicShape(op, infer_context); // } -// bool MoeOpInferSymbolicShape(pir::Operation *op, -// pir::InferSymbolicShapeContext *infer_context) { -// // pass -// return true; -// } +bool MoeOpInferSymbolicShape(pir::Operation *op, + pir::InferSymbolicShapeContext *infer_context) { + const auto &x_shape_or_data = + infer_context->GetShapeOrDataForValue(op->operand_source(0)); + infer_context->SetShapeOrDataForValue(op->result(0), x_shape_or_data); + return true; +} -// bool MulticlassNMS3OpInferSymbolicShape(pir::Operation *op, -// pir::InferSymbolicShapeContext -// *infer_context) { -// // pass -// return true; -// } +bool MulticlassNMS3OpInferSymbolicShape( + pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { + const auto &bboxes_shape_or_data = + infer_context->GetShapeOrDataForValue(op->operand_source(0)); + const auto &scores_shape_or_data = + infer_context->GetShapeOrDataForValue(op->operand_source(1)); + const auto &rois_num_shape_or_data = + infer_context->GetShapeOrDataForValue(op->operand_source(2)); + + const auto &bboxes_shape = bboxes_shape_or_data.shape(); + const auto &scores_shape = scores_shape_or_data.shape(); + + PADDLE_ENFORCE_EQ( + scores_shape.size() == 2 || scores_shape.size() == 3, + true, + errors::InvalidArgument("The rank of Input(Scores) must be 2 or 3.")); + + PADDLE_ENFORCE_EQ( + bboxes_shape.size(), + 3, + errors::InvalidArgument("The rank of Input(BBoxes) must be 3.")); + + if (scores_shape.size() == 3) { + PADDLE_ENFORCE_EQ( + bboxes_shape[2] == 4 || bboxes_shape[2] == 8 || bboxes_shape[2] == 16 || + bboxes_shape[2] == 24 || bboxes_shape[2] == 32, + true, + errors::InvalidArgument("The last dimension of Input(BBoxes) must be " + "4, 8, 16, 24, or 32.")); + infer_context->AddEqualCstr(bboxes_shape[1], scores_shape[2]); + } else { + PADDLE_ENFORCE_EQ(bboxes_shape[2], + 4, + errors::InvalidArgument( + "The last dimension of Input(BBoxes) must be 4.")); + infer_context->AddEqualCstr(bboxes_shape[1], scores_shape[1]); + } + + std::vector out_shape = {symbol::DimExpr(-1), + bboxes_shape[2] + 2}; + infer_context->SetShapeOrDataForValue( + op->result(0), + symbol::ShapeOrDataDimExprs{ + symbol::TensorShapeOrDataDimExprs(out_shape)}); + + std::vector index_shape = {symbol::DimExpr(-1), 1}; + infer_context->SetShapeOrDataForValue( + op->result(1), + symbol::ShapeOrDataDimExprs{ + symbol::TensorShapeOrDataDimExprs(index_shape)}); + + std::vector nms_rois_num_shape = {symbol::DimExpr(-1)}; + infer_context->SetShapeOrDataForValue( + op->result(2), + symbol::ShapeOrDataDimExprs{ + symbol::TensorShapeOrDataDimExprs(nms_rois_num_shape)}); + + return true; +} bool MeshgridOpInferSymbolicShape( pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h index a06c91cfde013..e02b74b82a3c2 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h @@ -58,8 +58,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam_) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum_) -// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) -// OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNMS3) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNMS3) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MemoryEfficientAttention) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Meshgrid) OP_DECLARE_INFER_SYMBOLIC_SHAPE(NearestInterp) diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml index 319d8bea6d9cd..8314e38a2f4df 100755 --- a/paddle/phi/ops/yaml/ops.yaml +++ b/paddle/phi/ops/yaml/ops.yaml @@ -3250,6 +3250,7 @@ func : multiclass_nms3 data_type : scores optional : rois_num, nms_rois_num + interfaces : paddle::dialect::InferSymbolicShapeInterface - op : multinomial args : (Tensor x, Scalar(int) num_samples = 1, bool replacement = false) @@ -5038,6 +5039,7 @@ func: MoeInferMeta kernel: func: moe + interfaces: paddle::dialect::InferSymbolicShapeInterface - op: number_count args: (Tensor numbers, int upper_range) From 2a4181391886df6b5d443416f07d3e4a24a4845e Mon Sep 17 00:00:00 2001 From: JeffWangYH <78253082+Jeff114514@users.noreply.github.com> Date: Tue, 6 Aug 2024 15:52:06 +0800 Subject: [PATCH 04/29] Update same_operands_result.cc --- .../interface/infer_symbolic_shape/same_operands_result.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index b3602f248a846..39d7ff91dfbfa 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -134,7 +134,6 @@ OP_SAME_OPERANDS_AND_RESULT(ScatterNdAdd) OP_SAME_OPERANDS_AND_RESULT(Scatter) OP_SAME_OPERANDS_AND_RESULT(Scatter_) OP_SAME_OPERANDS_AND_RESULT(Select) -OP_SAME_OPERANDS_AND_RESULT(ShadowFeedTensors) OP_SAME_OPERANDS_AND_RESULT(Sign) OP_SAME_OPERANDS_AND_RESULT(Sin) OP_SAME_OPERANDS_AND_RESULT(Sin_) From 3913b9a0bac3981f79c18dd36d97994d7eb4e38f Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Tue, 6 Aug 2024 17:59:15 +0800 Subject: [PATCH 05/29] unchange --- .../infer_symbolic_shape/same_operands_result.cc | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index 5d5323c128e7b..5ad13f0319858 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -134,6 +134,7 @@ OP_SAME_OPERANDS_AND_RESULT(ScatterNdAdd) OP_SAME_OPERANDS_AND_RESULT(Scatter) OP_SAME_OPERANDS_AND_RESULT(Scatter_) OP_SAME_OPERANDS_AND_RESULT(Select) +OP_SAME_OPERANDS_AND_RESULT(ShadowFeedTensors) OP_SAME_OPERANDS_AND_RESULT(Sign) OP_SAME_OPERANDS_AND_RESULT(Sin) OP_SAME_OPERANDS_AND_RESULT(Sin_) @@ -224,21 +225,6 @@ bool ArgsortOpInferSymbolicShape( return true; } -bool UnchangedVectorOpInferSymbolicShape( - pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { - const symbol::ShapeOrDataDimExprs &operand_shape_or_data = - infer_context->GetShapeOrDataForValue(op->operand_source(0)); - size_t size = operand_shape_or_data.size(); - std::vector shape(size); - for (size_t i = 0; i < size; ++i) { - shape[i] = operand_shape_or_data.shape()[i]; - } - infer_context->SetShapeOrDataForValue( - op->result(0), - symbol::TensorShapeOrDataDimExprs(shape, operand_shape_or_data.data())); - return true; -} - } // namespace paddle::dialect namespace cinn::dialect {} // namespace cinn::dialect From 7d00b39d4fe29af1389f3920c9e98cb1f9b093b2 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 7 Aug 2024 11:54:50 +0800 Subject: [PATCH 06/29] name --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 40c9233afa4d7..d21659a376121 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1207,7 +1207,7 @@ bool MoeOpInferSymbolicShape(pir::Operation *op, return true; } -bool MulticlassNMS3OpInferSymbolicShape( +bool MulticlassNms3OpInferSymbolicShape( pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { const auto &bboxes_shape_or_data = infer_context->GetShapeOrDataForValue(op->operand_source(0)); From ce0552d2a9e9387358dae45891ffcf446818b3e9 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 7 Aug 2024 11:58:22 +0800 Subject: [PATCH 07/29] name --- .../interface/infer_symbolic_shape/multiary_infer_sym.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h index e02b74b82a3c2..03dacb83787dc 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h @@ -59,7 +59,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) -OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNMS3) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNms3) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MemoryEfficientAttention) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Meshgrid) OP_DECLARE_INFER_SYMBOLIC_SHAPE(NearestInterp) From b180bb96edba591994b317cb817845b852603e92 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Thu, 8 Aug 2024 12:48:25 +0800 Subject: [PATCH 08/29] test --- .../same_operands_result.cc | 1 + .../same_operands_result.h | 1 + test/legacy_test/test_moe_op.py | 193 ++++++++++++++++++ 3 files changed, 195 insertions(+) create mode 100644 test/legacy_test/test_moe_op.py diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index 5ad13f0319858..8a80cbfb2ce23 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -160,6 +160,7 @@ OP_SAME_OPERANDS_AND_RESULT(Sigmoid) OP_SAME_OPERANDS_AND_RESULT(Sigmoid_) OP_SAME_OPERANDS_AND_RESULT(LeakyRelu) OP_SAME_OPERANDS_AND_RESULT(LeakyRelu_) +OP_SAME_OPERANDS_AND_RESULT(Moe) OP_SAME_OPERANDS_AND_RESULT(ThresholdedRelu) OP_SAME_OPERANDS_AND_RESULT(ThresholdedRelu_) OP_SAME_OPERANDS_AND_RESULT(SquareSr) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h index 6ae226dda7ccb..94d0d7b295804 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h @@ -99,6 +99,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logit) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logit_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsigmoid) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsigmoid_) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Poisson) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pow) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pow_) diff --git a/test/legacy_test/test_moe_op.py b/test/legacy_test/test_moe_op.py new file mode 100644 index 0000000000000..c5aedc67003b7 --- /dev/null +++ b/test/legacy_test/test_moe_op.py @@ -0,0 +1,193 @@ +# 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 +# +# http://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 numpy as np +from op_test import OpTest +from test_sparse_attention_op import get_cuda_version + +import paddle +import paddle.nn.functional as F +from paddle import nn +from paddle.incubate.nn.functional import moe, swiglu +from paddle.nn.layer.common import Linear + +paddle.seed(42) + + +class Expert(nn.Layer): + def __init__(self, d_model, d_feedforward): + super().__init__() + self.fc1 = nn.Linear( + d_model, d_feedforward * 2 + ) # Swiglu expects twice the hidden_dim + self.swiglu = swiglu + self.fc2 = nn.Linear(d_feedforward, d_model) + + def forward(self, x): + x = self.fc1(x) + x = self.swiglu(x) + x = self.fc2(x) + return x + + +@unittest.skipIf( + not paddle.is_compiled_with_cuda() + or get_cuda_version() < 11030 + or paddle.device.cuda.get_device_capability()[0] < 8, + "MoE requires CUDA >= 11.2 and CUDA_ARCH >= 8", +) +class TestMoEOp(OpTest): + def setUp(self): + self.config() + self.rtol = 1e-2 + self.atol = 1e-2 + + paddle.set_default_dtype(self.x_type) + self.__class__.op_type = "moe" + self.__class__.no_need_check_grad = True + + self.experts = nn.LayerList( + [ + Expert(self.d_model, self.d_feedforward) + for _ in range(self.num_expert) + ] + ) + + self.bmm_w0 = paddle.to_tensor( + np.array([expert.fc1.weight.numpy() for expert in self.experts]), + dtype=self.x_type, + ) + self.bmm_b0 = paddle.to_tensor( + np.array( + [expert.fc1.bias.numpy() for expert in self.experts] + ).reshape(self.num_expert, 1, -1), + dtype=self.x_type, + ) + + self.bmm_w1 = paddle.to_tensor( + np.array([expert.fc2.weight.numpy() for expert in self.experts]), + dtype=self.x_type, + ) + self.bmm_b1 = paddle.to_tensor( + np.array( + [expert.fc2.bias.numpy() for expert in self.experts] + ).reshape(self.num_expert, 1, -1), + dtype=self.x_type, + ) + + self.tensor_x = paddle.to_tensor( + np.random.randn(self.batch_size, self.seq_len, self.d_model) * 0.1, + dtype=self.x_type, + ) + + self.bmm_w0.stop_gradient = True + self.bmm_b0.stop_gradient = True + self.bmm_w1.stop_gradient = True + self.bmm_b1.stop_gradient = True + self.tensor_x.stop_gradient = True + + self.gate = Linear(self.d_model, self.num_expert) + self.gate_weight = self.gate.weight.cast(paddle.float32) + self.gate_weight.stop_gradient = True + + paddle.set_default_dtype(self.x_type) + self.activation = swiglu + + def config(self): + self.x_type = np.float16 + self.batch_size = 10 + self.seq_len = 128 + self.num_expert = 32 + self.d_model = 768 + self.d_feedforward = 3072 + self.top_k = 2 + self.quant_method = "None" + + def GetMoEOut(self, tensor_x): + paddle.disable_static(place=paddle.CUDAPlace(0)) + moe_out = moe( + tensor_x, + self.gate_weight, + self.bmm_w0, + self.bmm_b0, + self.bmm_w1, + self.bmm_b1, + None if self.quant_method == "None" else self.scale0, + None if self.quant_method == "None" else self.scale1, + self.quant_method, + 2, + ) + return moe_out + + def GetBaselineOut(self, hidden_states): + paddle.disable_static(place=paddle.CUDAPlace(0)) + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = paddle.reshape(hidden_states, [-1, hidden_dim]) + router_logits = self.gate(hidden_states).cast(paddle.float32) + + routing_weights = F.softmax(router_logits, axis=-1, dtype='float32') + routing_weights, selected_experts = paddle.topk( + routing_weights, self.top_k, axis=-1 + ) + routing_weights /= paddle.sum(routing_weights, axis=-1, keepdim=True) + routing_weights = routing_weights.cast(np.float32) + + final_hidden_states = paddle.zeros_like(hidden_states) + + expert_mask = paddle.transpose( + F.one_hot(selected_experts, num_classes=self.num_expert), [2, 1, 0] + ) + + for expert_idx in range(self.num_expert): + expert_layer = self.experts[expert_idx] + idx, top_x = paddle.where(expert_mask[expert_idx]) + + current_state = paddle.index_select( + hidden_states, top_x, axis=0 + ).reshape([-1, hidden_dim]) + current_hidden_states = ( + expert_layer(current_state) * routing_weights[top_x, idx] + ) + paddle.index_add_( + x=final_hidden_states, + index=top_x.squeeze(), + axis=0, + value=current_hidden_states.to(hidden_states.dtype), + ) + + final_hidden_states = paddle.reshape( + final_hidden_states, [batch_size, sequence_length, hidden_dim] + ) + return final_hidden_states + + def test_moe_op(self): + ref_out = self.GetBaselineOut(self.tensor_x).cast(np.float32) + moe_out = self.GetMoEOut(self.tensor_x).cast(np.float32) + np.testing.assert_allclose( + ref_out, moe_out, rtol=self.rtol, atol=self.atol + ) + + +class TestMoEOpBf16(TestMoEOp): + def config(self): + super().config() + self.x_type = paddle.bfloat16 + self.rtol = 1e-2 + self.atol = 1e-2 + + +if __name__ == "__main__": + unittest.main() From fa9ad91a6b7b3bc59a863ee4761e7a0ed6fde0bf Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Sun, 11 Aug 2024 12:48:17 +0800 Subject: [PATCH 09/29] err --- .../multiary_infer_sym.cc | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 69aaafdec6faf..bf8d2ada2bf62 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1294,28 +1294,29 @@ bool MulticlassNms3OpInferSymbolicShape( const auto &bboxes_shape = bboxes_shape_or_data.shape(); const auto &scores_shape = scores_shape_or_data.shape(); - PADDLE_ENFORCE_EQ( - scores_shape.size() == 2 || scores_shape.size() == 3, - true, - errors::InvalidArgument("The rank of Input(Scores) must be 2 or 3.")); + PADDLE_ENFORCE_EQ(scores_shape.size() == 2 || scores_shape.size() == 3, + true, + common::errors::InvalidArgument( + "The rank of Input(Scores) must be 2 or 3.")); PADDLE_ENFORCE_EQ( bboxes_shape.size(), 3, - errors::InvalidArgument("The rank of Input(BBoxes) must be 3.")); + common::errors::InvalidArgument("The rank of Input(BBoxes) must be 3.")); if (scores_shape.size() == 3) { - PADDLE_ENFORCE_EQ( - bboxes_shape[2] == 4 || bboxes_shape[2] == 8 || bboxes_shape[2] == 16 || - bboxes_shape[2] == 24 || bboxes_shape[2] == 32, - true, - errors::InvalidArgument("The last dimension of Input(BBoxes) must be " - "4, 8, 16, 24, or 32.")); + PADDLE_ENFORCE_EQ(bboxes_shape[2] == 4 || bboxes_shape[2] == 8 || + bboxes_shape[2] == 16 || bboxes_shape[2] == 24 || + bboxes_shape[2] == 32, + true, + common::errors::InvalidArgument( + "The last dimension of Input(BBoxes) must be " + "4, 8, 16, 24, or 32.")); infer_context->AddEqualCstr(bboxes_shape[1], scores_shape[2]); } else { PADDLE_ENFORCE_EQ(bboxes_shape[2], 4, - errors::InvalidArgument( + common::errors::InvalidArgument( "The last dimension of Input(BBoxes) must be 4.")); infer_context->AddEqualCstr(bboxes_shape[1], scores_shape[1]); } From b7d57d076153d92730959744046a454391186d50 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Sun, 11 Aug 2024 12:58:06 +0800 Subject: [PATCH 10/29] same same same next sym next sym next sym densetensor densetensortype densetensortype densetensortype rois_num dyn_cast ? dyn_cast --- .../multiary_infer_sym.cc | 128 +++++++++++++----- .../infer_symbolic_shape/multiary_infer_sym.h | 1 - .../same_operands_result.cc | 13 +- 3 files changed, 104 insertions(+), 38 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index bf8d2ada2bf62..c4971ea9c42ed 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1274,14 +1274,6 @@ bool RoiAlignOpInferSymbolicShape( // return MergedMomentumOpInferSymbolicShape(op, infer_context); // } -bool MoeOpInferSymbolicShape(pir::Operation *op, - pir::InferSymbolicShapeContext *infer_context) { - const auto &x_shape_or_data = - infer_context->GetShapeOrDataForValue(op->operand_source(0)); - infer_context->SetShapeOrDataForValue(op->result(0), x_shape_or_data); - return true; -} - bool MulticlassNms3OpInferSymbolicShape( pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { const auto &bboxes_shape_or_data = @@ -1291,50 +1283,113 @@ bool MulticlassNms3OpInferSymbolicShape( const auto &rois_num_shape_or_data = infer_context->GetShapeOrDataForValue(op->operand_source(2)); - const auto &bboxes_shape = bboxes_shape_or_data.shape(); - const auto &scores_shape = scores_shape_or_data.shape(); - - PADDLE_ENFORCE_EQ(scores_shape.size() == 2 || scores_shape.size() == 3, - true, - common::errors::InvalidArgument( - "The rank of Input(Scores) must be 2 or 3.")); + const std::vector &box_dims = bboxes_shape_or_data.shape(); + const std::vector &score_dims = scores_shape_or_data.shape(); + const size_t score_size = score_dims.size(); PADDLE_ENFORCE_EQ( - bboxes_shape.size(), + score_size == 2 || score_size == 3, + true, + common::errors::InvalidArgument( + "The rank of Input(Scores) must be 2 or 3. But received rank = %d", + score_size)); + PADDLE_ENFORCE_EQ( + box_dims.size(), 3, - common::errors::InvalidArgument("The rank of Input(BBoxes) must be 3.")); + common::errors::InvalidArgument( + "The rank of Input(BBoxes) must be 3. But received rank = %d", + box_dims.size())); - if (scores_shape.size() == 3) { - PADDLE_ENFORCE_EQ(bboxes_shape[2] == 4 || bboxes_shape[2] == 8 || - bboxes_shape[2] == 16 || bboxes_shape[2] == 24 || - bboxes_shape[2] == 32, - true, - common::errors::InvalidArgument( - "The last dimension of Input(BBoxes) must be " - "4, 8, 16, 24, or 32.")); - infer_context->AddEqualCstr(bboxes_shape[1], scores_shape[2]); + if (score_size == 3) { + PADDLE_ENFORCE_EQ( + box_dims[2] == 4 || box_dims[2] == 8 || box_dims[2] == 16 || + box_dims[2] == 24 || box_dims[2] == 32, + true, + common::errors::InvalidArgument("The last dimension of Input(BBoxes) " + "must be 4 or 8 or 16 or 24 or 32")); + PADDLE_ENFORCE_EQ( + box_dims[1], + score_dims[2], + common::errors::InvalidArgument( + "The 2nd dimension of Input(BBoxes) must be equal to " + "last dimension of Input(Scores), which represents the " + "predicted bboxes." + "But received box_dims[1](%s) != socre_dims[2](%s)", + box_dims[1], + score_dims[2])); } else { - PADDLE_ENFORCE_EQ(bboxes_shape[2], + PADDLE_ENFORCE_EQ(box_dims[2], 4, common::errors::InvalidArgument( - "The last dimension of Input(BBoxes) must be 4.")); - infer_context->AddEqualCstr(bboxes_shape[1], scores_shape[1]); + "The last dimension of Input(BBoxes) must be 4. But " + "received dimension = %d", + box_dims[2])); + PADDLE_ENFORCE_EQ( + box_dims[1], + score_dims[1], + common::errors::InvalidArgument( + "The 2nd dimension of Input" + "(BBoxes) must be equal to the 2nd dimension of Input(Scores). " + "But received box dimension = %d, score dimension = %d", + box_dims[1], + score_dims[1])); } - std::vector out_shape = {symbol::DimExpr(-1), - bboxes_shape[2] + 2}; + const auto &next_symbol_out = infer_context->GetNextSymName(); + const auto &next_symbol_index = infer_context->GetNextSymName(); + infer_context->AddEqualCstr(next_symbol_out, next_symbol_index); + + int rois_num_numel = [&] { + auto rois_num_shape = rois_num_shape_or_data.shape(); + int numel = 0; + for (size_t i = 0; i < rois_num_shape.size(); ++i) { + numel += rois_num_shape[i].dyn_cast(); + } + return numel; + }(); + + int n = [&] { + auto bboxes = bboxes_shape_or_data.dyn_cast(); + bool has_roisnum = + !rois_num_shape_or_data.isa(); + int n = 0; + if (has_roisnum) { + auto rois_num_tensor = + rois_num_shape_or_data.dyn_cast(); + if (score_size == 3) { + n = score_dims[0].dyn_cast(); + } else { + n = rois_num_numel; + } + } else { + if (score_size == 3) { + n = score_dims[0].dyn_cast(); + } else { + n = bboxes.lod().back().size() - 1; + } + } + return n; + }(); + + std::vector out_shape; + out_shape.emplace_back(next_symbol_out); + out_shape.emplace_back(box_dims[2] + 2); + + std::vector index_shape; + index_shape.emplace_back(next_symbol_index); + index_shape.emplace_back(1); + + std::vector nms_rois_num_shape; + nms_rois_num_shape.emplace_back(n); + infer_context->SetShapeOrDataForValue( op->result(0), symbol::ShapeOrDataDimExprs{ symbol::TensorShapeOrDataDimExprs(out_shape)}); - - std::vector index_shape = {symbol::DimExpr(-1), 1}; infer_context->SetShapeOrDataForValue( op->result(1), symbol::ShapeOrDataDimExprs{ symbol::TensorShapeOrDataDimExprs(index_shape)}); - - std::vector nms_rois_num_shape = {symbol::DimExpr(-1)}; infer_context->SetShapeOrDataForValue( op->result(2), symbol::ShapeOrDataDimExprs{ @@ -1401,7 +1456,8 @@ bool MovingAverageAbsMaxScale_OpInferSymbolicShape( } // bool NceOpInferSymbolicShape(pir::Operation *op, -// pir::InferSymbolicShapeContext *infer_context) { +// pir::InferSymbolicShapeContext *infer_context) +// { // // pass // return true; // } diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h index 0f4c067a7a39b..f276e05325aab 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h @@ -64,7 +64,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam_) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum_) -OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNms3) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MemoryEfficientAttention) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Meshgrid) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index 26d86a91ede41..6d11fa7575170 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -142,7 +142,6 @@ OP_SAME_OPERANDS_AND_RESULT(Scatter) OP_SAME_OPERANDS_AND_RESULT(Scatter_) OP_SAME_OPERANDS_AND_RESULT(Select) OP_SAME_OPERANDS_AND_RESULT(ShadowFeed) -OP_SAME_OPERANDS_AND_RESULT(ShadowFeedTensors) OP_SAME_OPERANDS_AND_RESULT(ShareData_) OP_SAME_OPERANDS_AND_RESULT(Sign) OP_SAME_OPERANDS_AND_RESULT(Sin) @@ -241,6 +240,18 @@ bool ArgsortOpInferSymbolicShape( return true; } +bool ShadowFeedTensorsOpInferSymbolicShape( + pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { + pir::Value operand_source = op->operand_source(0); + const symbol::TensorListShapeOrDataDimExprs &shape_or_data_list = + infer_context->GetShapeOrDataForValue(operand_source) + .dyn_cast(); + for (size_t i = 0; i < shape_or_data_list.size(); ++i) { + infer_context->SetShapeOrDataForValue(op->result(i), shape_or_data_list[i]); + } + return true; +} + } // namespace paddle::dialect namespace cinn::dialect {} // namespace cinn::dialect From 60880cbf7b5fa9e7443c778ab93a150442136c22 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 14:50:21 +0800 Subject: [PATCH 11/29] dyn_cast --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index c4971ea9c42ed..626af02b20f5b 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1349,13 +1349,11 @@ bool MulticlassNms3OpInferSymbolicShape( }(); int n = [&] { - auto bboxes = bboxes_shape_or_data.dyn_cast(); + auto bboxes = bboxes_shape_or_data.dyn_cast(); bool has_roisnum = !rois_num_shape_or_data.isa(); int n = 0; if (has_roisnum) { - auto rois_num_tensor = - rois_num_shape_or_data.dyn_cast(); if (score_size == 3) { n = score_dims[0].dyn_cast(); } else { From ff9ab662b038304524442fecef75420356319212 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 15:14:23 +0800 Subject: [PATCH 12/29] new sym --- .../multiary_infer_sym.cc | 33 ++----------------- 1 file changed, 2 insertions(+), 31 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 626af02b20f5b..fa01cce83338c 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1337,38 +1337,9 @@ bool MulticlassNms3OpInferSymbolicShape( const auto &next_symbol_out = infer_context->GetNextSymName(); const auto &next_symbol_index = infer_context->GetNextSymName(); + const auto &next_symbol_nms_rois_num = infer_context->GetNextSymName(); infer_context->AddEqualCstr(next_symbol_out, next_symbol_index); - int rois_num_numel = [&] { - auto rois_num_shape = rois_num_shape_or_data.shape(); - int numel = 0; - for (size_t i = 0; i < rois_num_shape.size(); ++i) { - numel += rois_num_shape[i].dyn_cast(); - } - return numel; - }(); - - int n = [&] { - auto bboxes = bboxes_shape_or_data.dyn_cast(); - bool has_roisnum = - !rois_num_shape_or_data.isa(); - int n = 0; - if (has_roisnum) { - if (score_size == 3) { - n = score_dims[0].dyn_cast(); - } else { - n = rois_num_numel; - } - } else { - if (score_size == 3) { - n = score_dims[0].dyn_cast(); - } else { - n = bboxes.lod().back().size() - 1; - } - } - return n; - }(); - std::vector out_shape; out_shape.emplace_back(next_symbol_out); out_shape.emplace_back(box_dims[2] + 2); @@ -1378,7 +1349,7 @@ bool MulticlassNms3OpInferSymbolicShape( index_shape.emplace_back(1); std::vector nms_rois_num_shape; - nms_rois_num_shape.emplace_back(n); + nms_rois_num_shape.emplace_back(next_symbol_nms_rois_num); infer_context->SetShapeOrDataForValue( op->result(0), From 43925394e9108be11c8e078df521deab78993124 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 15:45:24 +0800 Subject: [PATCH 13/29] no moe optest --- test/legacy_test/test_moe_op.py | 193 -------------------------------- 1 file changed, 193 deletions(-) delete mode 100644 test/legacy_test/test_moe_op.py diff --git a/test/legacy_test/test_moe_op.py b/test/legacy_test/test_moe_op.py deleted file mode 100644 index c5aedc67003b7..0000000000000 --- a/test/legacy_test/test_moe_op.py +++ /dev/null @@ -1,193 +0,0 @@ -# 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 -# -# http://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 numpy as np -from op_test import OpTest -from test_sparse_attention_op import get_cuda_version - -import paddle -import paddle.nn.functional as F -from paddle import nn -from paddle.incubate.nn.functional import moe, swiglu -from paddle.nn.layer.common import Linear - -paddle.seed(42) - - -class Expert(nn.Layer): - def __init__(self, d_model, d_feedforward): - super().__init__() - self.fc1 = nn.Linear( - d_model, d_feedforward * 2 - ) # Swiglu expects twice the hidden_dim - self.swiglu = swiglu - self.fc2 = nn.Linear(d_feedforward, d_model) - - def forward(self, x): - x = self.fc1(x) - x = self.swiglu(x) - x = self.fc2(x) - return x - - -@unittest.skipIf( - not paddle.is_compiled_with_cuda() - or get_cuda_version() < 11030 - or paddle.device.cuda.get_device_capability()[0] < 8, - "MoE requires CUDA >= 11.2 and CUDA_ARCH >= 8", -) -class TestMoEOp(OpTest): - def setUp(self): - self.config() - self.rtol = 1e-2 - self.atol = 1e-2 - - paddle.set_default_dtype(self.x_type) - self.__class__.op_type = "moe" - self.__class__.no_need_check_grad = True - - self.experts = nn.LayerList( - [ - Expert(self.d_model, self.d_feedforward) - for _ in range(self.num_expert) - ] - ) - - self.bmm_w0 = paddle.to_tensor( - np.array([expert.fc1.weight.numpy() for expert in self.experts]), - dtype=self.x_type, - ) - self.bmm_b0 = paddle.to_tensor( - np.array( - [expert.fc1.bias.numpy() for expert in self.experts] - ).reshape(self.num_expert, 1, -1), - dtype=self.x_type, - ) - - self.bmm_w1 = paddle.to_tensor( - np.array([expert.fc2.weight.numpy() for expert in self.experts]), - dtype=self.x_type, - ) - self.bmm_b1 = paddle.to_tensor( - np.array( - [expert.fc2.bias.numpy() for expert in self.experts] - ).reshape(self.num_expert, 1, -1), - dtype=self.x_type, - ) - - self.tensor_x = paddle.to_tensor( - np.random.randn(self.batch_size, self.seq_len, self.d_model) * 0.1, - dtype=self.x_type, - ) - - self.bmm_w0.stop_gradient = True - self.bmm_b0.stop_gradient = True - self.bmm_w1.stop_gradient = True - self.bmm_b1.stop_gradient = True - self.tensor_x.stop_gradient = True - - self.gate = Linear(self.d_model, self.num_expert) - self.gate_weight = self.gate.weight.cast(paddle.float32) - self.gate_weight.stop_gradient = True - - paddle.set_default_dtype(self.x_type) - self.activation = swiglu - - def config(self): - self.x_type = np.float16 - self.batch_size = 10 - self.seq_len = 128 - self.num_expert = 32 - self.d_model = 768 - self.d_feedforward = 3072 - self.top_k = 2 - self.quant_method = "None" - - def GetMoEOut(self, tensor_x): - paddle.disable_static(place=paddle.CUDAPlace(0)) - moe_out = moe( - tensor_x, - self.gate_weight, - self.bmm_w0, - self.bmm_b0, - self.bmm_w1, - self.bmm_b1, - None if self.quant_method == "None" else self.scale0, - None if self.quant_method == "None" else self.scale1, - self.quant_method, - 2, - ) - return moe_out - - def GetBaselineOut(self, hidden_states): - paddle.disable_static(place=paddle.CUDAPlace(0)) - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = paddle.reshape(hidden_states, [-1, hidden_dim]) - router_logits = self.gate(hidden_states).cast(paddle.float32) - - routing_weights = F.softmax(router_logits, axis=-1, dtype='float32') - routing_weights, selected_experts = paddle.topk( - routing_weights, self.top_k, axis=-1 - ) - routing_weights /= paddle.sum(routing_weights, axis=-1, keepdim=True) - routing_weights = routing_weights.cast(np.float32) - - final_hidden_states = paddle.zeros_like(hidden_states) - - expert_mask = paddle.transpose( - F.one_hot(selected_experts, num_classes=self.num_expert), [2, 1, 0] - ) - - for expert_idx in range(self.num_expert): - expert_layer = self.experts[expert_idx] - idx, top_x = paddle.where(expert_mask[expert_idx]) - - current_state = paddle.index_select( - hidden_states, top_x, axis=0 - ).reshape([-1, hidden_dim]) - current_hidden_states = ( - expert_layer(current_state) * routing_weights[top_x, idx] - ) - paddle.index_add_( - x=final_hidden_states, - index=top_x.squeeze(), - axis=0, - value=current_hidden_states.to(hidden_states.dtype), - ) - - final_hidden_states = paddle.reshape( - final_hidden_states, [batch_size, sequence_length, hidden_dim] - ) - return final_hidden_states - - def test_moe_op(self): - ref_out = self.GetBaselineOut(self.tensor_x).cast(np.float32) - moe_out = self.GetMoEOut(self.tensor_x).cast(np.float32) - np.testing.assert_allclose( - ref_out, moe_out, rtol=self.rtol, atol=self.atol - ) - - -class TestMoEOpBf16(TestMoEOp): - def config(self): - super().config() - self.x_type = paddle.bfloat16 - self.rtol = 1e-2 - self.atol = 1e-2 - - -if __name__ == "__main__": - unittest.main() From fe37df4e2ff29bf3e67566eb9c100eeaae98fd57 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 16:06:51 +0800 Subject: [PATCH 14/29] new test --- test/legacy_test/test_moe_op.py | 273 ++++++++++++++++++++++++++++++++ 1 file changed, 273 insertions(+) create mode 100644 test/legacy_test/test_moe_op.py diff --git a/test/legacy_test/test_moe_op.py b/test/legacy_test/test_moe_op.py new file mode 100644 index 0000000000000..614cc1f2b6962 --- /dev/null +++ b/test/legacy_test/test_moe_op.py @@ -0,0 +1,273 @@ +# 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 +# +# http://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 numpy as np +from op_test import OpTest +from test_sparse_attention_op import get_cuda_version + +import paddle +import paddle.nn.functional as F +from paddle import _C_ops, nn +from paddle.base.layer_helper import LayerHelper +from paddle.framework import in_dynamic_mode +from paddle.incubate.nn.functional import swiglu +from paddle.nn.layer.common import Linear + + +def moe( + x, + gate_weight, + ffn1_weight, + ffn1_bias, + ffn2_weight, + ffn2_bias, + ffn1_scale=None, + ffn2_scale=None, + quant_method="None", + moe_topk=2, +): + if in_dynamic_mode(): + final_out = _C_ops.moe( + x, + gate_weight, + ffn1_weight, + ffn1_scale, + ffn1_bias, + ffn2_weight, + ffn2_scale, + ffn2_bias, + quant_method, + moe_topk, + ) + return final_out + else: + helper = LayerHelper('moe', **locals()) + final_out = helper.create_variable_for_type_inference(dtype=x.dtype) + helper.append_op( + type='moe', + inputs={ + 'x': x, + 'gate_weight': gate_weight, + 'ffn1_weight': ffn1_weight, + 'ffn1_scale': ffn1_scale, + 'ffn1_bias': ffn1_bias, + 'ffn2_weight': ffn2_weight, + 'ffn2_scale': ffn2_scale, + 'ffn2_bias': ffn2_bias, + }, + outputs={'out': final_out}, + attrs={ + 'quant_method': quant_method, + 'moe_topk': moe_topk, + }, + ) + return final_out + + +# 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 +# +# http://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. + +paddle.seed(42) + + +class Expert(nn.Layer): + def __init__(self, d_model, d_feedforward): + super().__init__() + self.fc1 = nn.Linear( + d_model, d_feedforward * 2 + ) # Swiglu expects twice the hidden_dim + self.swiglu = swiglu + self.fc2 = nn.Linear(d_feedforward, d_model) + + def forward(self, x, idx): + x = self.fc1(x) + x = swiglu(x) + x = self.fc2(x) + return x + + +@unittest.skipIf( + not paddle.is_compiled_with_cuda() + or get_cuda_version() < 11030 + or paddle.device.cuda.get_device_capability()[0] < 8, + "MoE requires CUDA >= 11.2 and CUDA_ARCH >= 8", +) +class TestMoEOp(OpTest): + def setUp(self): + self.config() + self.rtol = 1e-2 + self.atol = 1e-2 + + paddle.set_default_dtype(self.x_type) + self.__class__.op_type = "moe" + # Since it's only used in inference. + self.__class__.no_need_check_grad = True + + self.experts = nn.LayerList( + [ + Expert(self.d_model, self.d_feedforward) + for _ in range(self.num_expert) + ] + ) + + self.bmm_w0 = paddle.to_tensor( + np.array([expert.fc1.weight.numpy() for expert in self.experts]), + dtype=self.x_type, + ) + self.bmm_b0 = paddle.to_tensor( + np.array( + [expert.fc1.bias.numpy() for expert in self.experts] + ).reshape(self.num_expert, 1, -1), + dtype=self.x_type, + ) + + # d_model*2 for swiglu + self.bmm_w1 = paddle.to_tensor( + np.array([expert.fc2.weight.numpy() for expert in self.experts]), + dtype=self.x_type, + ) + self.bmm_b1 = paddle.to_tensor( + np.array( + [expert.fc2.bias.numpy() for expert in self.experts] + ).reshape(self.num_expert, 1, -1), + dtype=self.x_type, + ) + + self.tensor_x = paddle.to_tensor( + np.random.randn(self.batch_size, self.seq_len, self.d_model) * 0.1, + dtype=self.x_type, + ) + + self.bmm_w0.stop_gradient = True + self.bmm_b0.stop_gradient = True + self.bmm_w1.stop_gradient = True + self.bmm_b1.stop_gradient = True + self.tensor_x.stop_gradient = True + + self.gate = Linear(self.d_model, self.num_expert) + + self.gate_weight = self.gate.weight.cast(paddle.float32) + self.gate_weight.stop_gradient = True + + paddle.set_default_dtype(self.x_type) + self.activation = swiglu + + def config(self): + self.x_type = np.float16 + self.batch_size = 10 + self.seq_len = 128 + self.num_expert = 32 + self.d_model = 768 + self.d_feedforward = 3072 + self.top_k = 2 + self.quant_method = "None" + + def GetMoeOut(self, tensor_x): + paddle.disable_static(place=paddle.CUDAPlace(0)) + moe_out = moe( + tensor_x, + self.gate_weight, + self.bmm_w0, + self.bmm_b0, + self.bmm_w1, + self.bmm_b1, + None if self.quant_method == "None" else self.scale0, + None if self.quant_method == "None" else self.scale1, + self.quant_method, + 2, + ) + + return moe_out + + def GetBaselineOut(self, hidden_states): + paddle.disable_static(place=paddle.CUDAPlace(0)) + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = paddle.reshape(hidden_states, [-1, hidden_dim]) + router_logits = self.gate(hidden_states).cast(paddle.float32) + + routing_weights = F.softmax(router_logits, axis=-1, dtype='float32') + routing_weights, selected_experts = paddle.topk( + routing_weights, self.top_k, axis=-1 + ) + routing_weights /= paddle.sum(routing_weights, axis=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.cast(np.float32) + + final_hidden_states = paddle.zeros_like(hidden_states) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = paddle.transpose( + F.one_hot(selected_experts, num_classes=self.num_expert), [2, 1, 0] + ) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_expert): + expert_layer = self.experts[expert_idx] + idx, top_x = paddle.where(expert_mask[expert_idx]) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = paddle.index_select( + hidden_states, top_x, axis=0 + ).reshape([-1, hidden_dim]) + current_hidden_states = ( + expert_layer(current_state, expert_idx) + * routing_weights[top_x, idx] + ) + # Use scatter to accumulate the results + paddle.index_add_( + x=final_hidden_states, + index=top_x.squeeze(), + axis=0, + value=current_hidden_states.to(hidden_states.dtype), + ) + + final_hidden_states = paddle.reshape( + final_hidden_states, [batch_size, sequence_length, hidden_dim] + ) + return final_hidden_states + + def test_moe_op_new(self): + ref_out = self.GetBaselineOut(self.tensor_x).cast(np.float32) + moe_out = self.GetMoeOut(self.tensor_x).cast(np.float32) + np.testing.assert_allclose( + ref_out, moe_out, rtol=self.rtol, atol=self.atol + ) + + +class TestMoEOpBf16(TestMoEOp): + def config(self): + super().config() + self.x_type = paddle.bfloat16 + self.rtol = 1e-2 + self.atol = 1e-2 + + +if __name__ == "__main__": + unittest.main() From 6fef51b7229c5532c1c714eb7f664c0faf3051b9 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 17:27:58 +0800 Subject: [PATCH 15/29] mkddim --- .../infer_symbolic_shape/multiary_infer_sym.cc | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index fa01cce83338c..1798a5d961601 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1351,18 +1351,22 @@ bool MulticlassNms3OpInferSymbolicShape( std::vector nms_rois_num_shape; nms_rois_num_shape.emplace_back(next_symbol_nms_rois_num); + VLOG(3) << "out_shape: " << common::make_ddim(out_shape); + VLOG(3) << "index_shape: " << common::make_ddim(index_shape); + VLOG(3) << "nms_rois_num_shape: " << common::make_ddim(nms_rois_num_shape); + infer_context->SetShapeOrDataForValue( op->result(0), symbol::ShapeOrDataDimExprs{ - symbol::TensorShapeOrDataDimExprs(out_shape)}); + symbol::TensorShapeOrDataDimExprs(common::make_ddim(out_shape))}); infer_context->SetShapeOrDataForValue( op->result(1), symbol::ShapeOrDataDimExprs{ - symbol::TensorShapeOrDataDimExprs(index_shape)}); + symbol::TensorShapeOrDataDimExprs(common::make_ddim(index_shape))}); infer_context->SetShapeOrDataForValue( op->result(2), - symbol::ShapeOrDataDimExprs{ - symbol::TensorShapeOrDataDimExprs(nms_rois_num_shape)}); + symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs( + common::make_ddim(nms_rois_num_shape))}); return true; } From 8190ce6a3769acba4f866372d3fe344fb65a2db9 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 17:33:37 +0800 Subject: [PATCH 16/29] no mkddim --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 1798a5d961601..040507fc5484b 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1351,9 +1351,9 @@ bool MulticlassNms3OpInferSymbolicShape( std::vector nms_rois_num_shape; nms_rois_num_shape.emplace_back(next_symbol_nms_rois_num); - VLOG(3) << "out_shape: " << common::make_ddim(out_shape); - VLOG(3) << "index_shape: " << common::make_ddim(index_shape); - VLOG(3) << "nms_rois_num_shape: " << common::make_ddim(nms_rois_num_shape); + VLOG(3) << "out_shape: " << out_shape; + VLOG(3) << "index_shape: " << index_shape; + VLOG(3) << "nms_rois_num_shape: " << nms_rois_num_shape; infer_context->SetShapeOrDataForValue( op->result(0), From e35dfb6d4e141a2eec36ba2d8e7e789cfb864953 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 17:39:08 +0800 Subject: [PATCH 17/29] no mkddim --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 040507fc5484b..db331210726aa 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1358,15 +1358,15 @@ bool MulticlassNms3OpInferSymbolicShape( infer_context->SetShapeOrDataForValue( op->result(0), symbol::ShapeOrDataDimExprs{ - symbol::TensorShapeOrDataDimExprs(common::make_ddim(out_shape))}); + symbol::TensorShapeOrDataDimExprs(out_shape)}); infer_context->SetShapeOrDataForValue( op->result(1), symbol::ShapeOrDataDimExprs{ - symbol::TensorShapeOrDataDimExprs(common::make_ddim(index_shape))}); + symbol::TensorShapeOrDataDimExprs(index_shape)}); infer_context->SetShapeOrDataForValue( op->result(2), - symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs( - common::make_ddim(nms_rois_num_shape))}); + symbol::ShapeOrDataDimExprs{ + symbol::TensorShapeOrDataDimExprs(nms_rois_num_shape)}); return true; } From d9c191b80adbbe6d36b2b4120886754c17f94a25 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Mon, 12 Aug 2024 23:10:51 +0800 Subject: [PATCH 18/29] pass --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 4 ---- 1 file changed, 4 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index db331210726aa..fa01cce83338c 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1351,10 +1351,6 @@ bool MulticlassNms3OpInferSymbolicShape( std::vector nms_rois_num_shape; nms_rois_num_shape.emplace_back(next_symbol_nms_rois_num); - VLOG(3) << "out_shape: " << out_shape; - VLOG(3) << "index_shape: " << index_shape; - VLOG(3) << "nms_rois_num_shape: " << nms_rois_num_shape; - infer_context->SetShapeOrDataForValue( op->result(0), symbol::ShapeOrDataDimExprs{ From e62de69bf85f46dc56fb315989712d36f3d4de37 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Tue, 13 Aug 2024 11:04:55 +0800 Subject: [PATCH 19/29] remove unused --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index fa01cce83338c..acfe49b6192fc 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1280,8 +1280,6 @@ bool MulticlassNms3OpInferSymbolicShape( infer_context->GetShapeOrDataForValue(op->operand_source(0)); const auto &scores_shape_or_data = infer_context->GetShapeOrDataForValue(op->operand_source(1)); - const auto &rois_num_shape_or_data = - infer_context->GetShapeOrDataForValue(op->operand_source(2)); const std::vector &box_dims = bboxes_shape_or_data.shape(); const std::vector &score_dims = scores_shape_or_data.shape(); From 5686d2bf43f1a9561f98a9811f41ce33bff35b55 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 11:03:24 +0800 Subject: [PATCH 20/29] fix --- .../multiary_infer_sym.cc | 39 ++++--------------- test/legacy_test/test_moe_op.py | 34 +++++++++------- 2 files changed, 27 insertions(+), 46 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index acfe49b6192fc..d1c7bc6a9753e 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1305,45 +1305,21 @@ bool MulticlassNms3OpInferSymbolicShape( true, common::errors::InvalidArgument("The last dimension of Input(BBoxes) " "must be 4 or 8 or 16 or 24 or 32")); - PADDLE_ENFORCE_EQ( - box_dims[1], - score_dims[2], - common::errors::InvalidArgument( - "The 2nd dimension of Input(BBoxes) must be equal to " - "last dimension of Input(Scores), which represents the " - "predicted bboxes." - "But received box_dims[1](%s) != socre_dims[2](%s)", - box_dims[1], - score_dims[2])); + infer_context->AddEqualCstr(box_dims[1], score_dims[2]); } else { - PADDLE_ENFORCE_EQ(box_dims[2], - 4, - common::errors::InvalidArgument( - "The last dimension of Input(BBoxes) must be 4. But " - "received dimension = %d", - box_dims[2])); - PADDLE_ENFORCE_EQ( - box_dims[1], - score_dims[1], - common::errors::InvalidArgument( - "The 2nd dimension of Input" - "(BBoxes) must be equal to the 2nd dimension of Input(Scores). " - "But received box dimension = %d, score dimension = %d", - box_dims[1], - score_dims[1])); + infer_context->AddEqualCstr(box_dims[2], symbol::DimExpr(4)); + infer_context->AddEqualCstr(box_dims[1], score_dims[1]); } - const auto &next_symbol_out = infer_context->GetNextSymName(); - const auto &next_symbol_index = infer_context->GetNextSymName(); + const auto &next_symbol_out_and_index = infer_context->GetNextSymName(); const auto &next_symbol_nms_rois_num = infer_context->GetNextSymName(); - infer_context->AddEqualCstr(next_symbol_out, next_symbol_index); std::vector out_shape; - out_shape.emplace_back(next_symbol_out); + out_shape.emplace_back(next_symbol_out_and_index); out_shape.emplace_back(box_dims[2] + 2); std::vector index_shape; - index_shape.emplace_back(next_symbol_index); + index_shape.emplace_back(next_symbol_out_and_index); index_shape.emplace_back(1); std::vector nms_rois_num_shape; @@ -1423,8 +1399,7 @@ bool MovingAverageAbsMaxScale_OpInferSymbolicShape( } // bool NceOpInferSymbolicShape(pir::Operation *op, -// pir::InferSymbolicShapeContext *infer_context) -// { +// pir::InferSymbolicShapeContext *infer_context){ // // pass // return true; // } diff --git a/test/legacy_test/test_moe_op.py b/test/legacy_test/test_moe_op.py index 614cc1f2b6962..1e02900cc6490 100644 --- a/test/legacy_test/test_moe_op.py +++ b/test/legacy_test/test_moe_op.py @@ -77,20 +77,6 @@ def moe( return final_out -# 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 -# -# http://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. - paddle.seed(42) @@ -176,6 +162,23 @@ def setUp(self): paddle.set_default_dtype(self.x_type) self.activation = swiglu + self.inputs = { + self.tensor_x, + self.gate_weight, + self.bmm_w0, + self.bmm_b0, + self.bmm_w1, + self.bmm_b1, + None if self.quant_method == "None" else self.scale0, + None if self.quant_method == "None" else self.scale1, + self.quant_method, + 2, + } + self.python_api = moe + self.outputs = { + "Out": self.GetMoeOut(self.tensor_x), + } + def config(self): self.x_type = np.float16 self.batch_size = 10 @@ -260,6 +263,9 @@ def test_moe_op_new(self): ref_out, moe_out, rtol=self.rtol, atol=self.atol ) + def test_check_output(self): + self.check_output(check_pir=True) + class TestMoEOpBf16(TestMoEOp): def config(self): From 39d024bc6834565f58e67fe29dc62a3cba345426 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 12:34:06 +0800 Subject: [PATCH 21/29] fix --- test/legacy_test/test_moe_op.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/legacy_test/test_moe_op.py b/test/legacy_test/test_moe_op.py index 1e02900cc6490..de4e044e53612 100644 --- a/test/legacy_test/test_moe_op.py +++ b/test/legacy_test/test_moe_op.py @@ -262,6 +262,7 @@ def test_moe_op_new(self): np.testing.assert_allclose( ref_out, moe_out, rtol=self.rtol, atol=self.atol ) + self.test_check_output() def test_check_output(self): self.check_output(check_pir=True) From bf708f023cd68d44a2efcb3126481ec1670bc042 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 19:35:13 +0800 Subject: [PATCH 22/29] fix --- .../infer_symbolic_shape/multiary_infer_sym.cc | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 86bee28d6a1f4..9da8b32a042fd 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1460,6 +1460,16 @@ bool RoiAlignOpInferSymbolicShape( // return MergedMomentumOpInferSymbolicShape(op, infer_context); // } +bool MoeOpInferSymbolicShape(pir::Operation *op, + pir::InferSymbolicShapeContext *infer_context) { + const auto &x_shape = + infer_context->GetShapeOrDataForValue(op->operand_source(0)); + infer_context->SetShapeOrDataForValue( + op->result(0), + symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)}); + return true; +} + bool MulticlassNms3OpInferSymbolicShape( pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { const auto &bboxes_shape_or_data = From 7283fff321882bb0e079a1374c2fcef8978ee49e Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 19:35:59 +0800 Subject: [PATCH 23/29] fix --- .../operator/interface/infer_symbolic_shape/multiary_infer_sym.h | 1 + .../interface/infer_symbolic_shape/same_operands_result.cc | 1 - .../interface/infer_symbolic_shape/same_operands_result.h | 1 - 3 files changed, 1 insertion(+), 2 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h index 6a816bec07de0..fc65d076bfba1 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h @@ -70,6 +70,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNms3) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MemoryEfficientAttention) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Meshgrid) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MovingAverageAbsMaxScale) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MovingAverageAbsMaxScale_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(NearestInterp) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index 9e092ba61c180..40b661c7f6ec3 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -171,7 +171,6 @@ OP_SAME_OPERANDS_AND_RESULT(Sigmoid) OP_SAME_OPERANDS_AND_RESULT(Sigmoid_) OP_SAME_OPERANDS_AND_RESULT(LeakyRelu) OP_SAME_OPERANDS_AND_RESULT(LeakyRelu_) -OP_SAME_OPERANDS_AND_RESULT(Moe) OP_SAME_OPERANDS_AND_RESULT(ThresholdedRelu) OP_SAME_OPERANDS_AND_RESULT(ThresholdedRelu_) OP_SAME_OPERANDS_AND_RESULT(SquareSr) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h index b6419fea24317..9e903c730f14f 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h @@ -109,7 +109,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logit_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsigmoid) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsigmoid_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mish) -OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Poisson) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pow) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pow_) From 27af0b030ba5ba4852d2d53798e6d930c6474114 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 20:47:03 +0800 Subject: [PATCH 24/29] fix --- .../multiary_infer_sym.cc | 11 +- .../same_operands_result.cc | 12 - .../same_operands_result.h | 1 - test/legacy_test/test_moe_op.py | 280 ------------------ 4 files changed, 3 insertions(+), 301 deletions(-) delete mode 100644 test/legacy_test/test_moe_op.py diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 9da8b32a042fd..8787a62d98a6c 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1487,12 +1487,8 @@ bool MulticlassNms3OpInferSymbolicShape( common::errors::InvalidArgument( "The rank of Input(Scores) must be 2 or 3. But received rank = %d", score_size)); - PADDLE_ENFORCE_EQ( - box_dims.size(), - 3, - common::errors::InvalidArgument( - "The rank of Input(BBoxes) must be 3. But received rank = %d", - box_dims.size())); + + infer_context->AddEqualCstr(box_dims.size(), symbol::DimExpr(3)); if (score_size == 3) { PADDLE_ENFORCE_EQ( @@ -1508,7 +1504,6 @@ bool MulticlassNms3OpInferSymbolicShape( } const auto &next_symbol_out_and_index = infer_context->GetNextSymName(); - const auto &next_symbol_nms_rois_num = infer_context->GetNextSymName(); std::vector out_shape; out_shape.emplace_back(next_symbol_out_and_index); @@ -1519,7 +1514,7 @@ bool MulticlassNms3OpInferSymbolicShape( index_shape.emplace_back(1); std::vector nms_rois_num_shape; - nms_rois_num_shape.emplace_back(next_symbol_nms_rois_num); + nms_rois_num_shape.emplace_back(infer_context->GetNextSymName()); infer_context->SetShapeOrDataForValue( op->result(0), diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc index 40b661c7f6ec3..fbb455340db62 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc @@ -242,18 +242,6 @@ bool ArgsortOpInferSymbolicShape( return true; } -bool ShadowFeedTensorsOpInferSymbolicShape( - pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { - pir::Value operand_source = op->operand_source(0); - const symbol::TensorListShapeOrDataDimExprs &shape_or_data_list = - infer_context->GetShapeOrDataForValue(operand_source) - .dyn_cast(); - for (size_t i = 0; i < shape_or_data_list.size(); ++i) { - infer_context->SetShapeOrDataForValue(op->result(i), shape_or_data_list[i]); - } - return true; -} - } // namespace paddle::dialect namespace cinn::dialect {} // namespace cinn::dialect diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h index 9e903c730f14f..5f5127a98f608 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h @@ -138,7 +138,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Scatter) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Scatter_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Select) OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShadowFeed) -OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShadowFeedTensors) OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShareData_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sign) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Silu) diff --git a/test/legacy_test/test_moe_op.py b/test/legacy_test/test_moe_op.py deleted file mode 100644 index de4e044e53612..0000000000000 --- a/test/legacy_test/test_moe_op.py +++ /dev/null @@ -1,280 +0,0 @@ -# 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 -# -# http://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 numpy as np -from op_test import OpTest -from test_sparse_attention_op import get_cuda_version - -import paddle -import paddle.nn.functional as F -from paddle import _C_ops, nn -from paddle.base.layer_helper import LayerHelper -from paddle.framework import in_dynamic_mode -from paddle.incubate.nn.functional import swiglu -from paddle.nn.layer.common import Linear - - -def moe( - x, - gate_weight, - ffn1_weight, - ffn1_bias, - ffn2_weight, - ffn2_bias, - ffn1_scale=None, - ffn2_scale=None, - quant_method="None", - moe_topk=2, -): - if in_dynamic_mode(): - final_out = _C_ops.moe( - x, - gate_weight, - ffn1_weight, - ffn1_scale, - ffn1_bias, - ffn2_weight, - ffn2_scale, - ffn2_bias, - quant_method, - moe_topk, - ) - return final_out - else: - helper = LayerHelper('moe', **locals()) - final_out = helper.create_variable_for_type_inference(dtype=x.dtype) - helper.append_op( - type='moe', - inputs={ - 'x': x, - 'gate_weight': gate_weight, - 'ffn1_weight': ffn1_weight, - 'ffn1_scale': ffn1_scale, - 'ffn1_bias': ffn1_bias, - 'ffn2_weight': ffn2_weight, - 'ffn2_scale': ffn2_scale, - 'ffn2_bias': ffn2_bias, - }, - outputs={'out': final_out}, - attrs={ - 'quant_method': quant_method, - 'moe_topk': moe_topk, - }, - ) - return final_out - - -paddle.seed(42) - - -class Expert(nn.Layer): - def __init__(self, d_model, d_feedforward): - super().__init__() - self.fc1 = nn.Linear( - d_model, d_feedforward * 2 - ) # Swiglu expects twice the hidden_dim - self.swiglu = swiglu - self.fc2 = nn.Linear(d_feedforward, d_model) - - def forward(self, x, idx): - x = self.fc1(x) - x = swiglu(x) - x = self.fc2(x) - return x - - -@unittest.skipIf( - not paddle.is_compiled_with_cuda() - or get_cuda_version() < 11030 - or paddle.device.cuda.get_device_capability()[0] < 8, - "MoE requires CUDA >= 11.2 and CUDA_ARCH >= 8", -) -class TestMoEOp(OpTest): - def setUp(self): - self.config() - self.rtol = 1e-2 - self.atol = 1e-2 - - paddle.set_default_dtype(self.x_type) - self.__class__.op_type = "moe" - # Since it's only used in inference. - self.__class__.no_need_check_grad = True - - self.experts = nn.LayerList( - [ - Expert(self.d_model, self.d_feedforward) - for _ in range(self.num_expert) - ] - ) - - self.bmm_w0 = paddle.to_tensor( - np.array([expert.fc1.weight.numpy() for expert in self.experts]), - dtype=self.x_type, - ) - self.bmm_b0 = paddle.to_tensor( - np.array( - [expert.fc1.bias.numpy() for expert in self.experts] - ).reshape(self.num_expert, 1, -1), - dtype=self.x_type, - ) - - # d_model*2 for swiglu - self.bmm_w1 = paddle.to_tensor( - np.array([expert.fc2.weight.numpy() for expert in self.experts]), - dtype=self.x_type, - ) - self.bmm_b1 = paddle.to_tensor( - np.array( - [expert.fc2.bias.numpy() for expert in self.experts] - ).reshape(self.num_expert, 1, -1), - dtype=self.x_type, - ) - - self.tensor_x = paddle.to_tensor( - np.random.randn(self.batch_size, self.seq_len, self.d_model) * 0.1, - dtype=self.x_type, - ) - - self.bmm_w0.stop_gradient = True - self.bmm_b0.stop_gradient = True - self.bmm_w1.stop_gradient = True - self.bmm_b1.stop_gradient = True - self.tensor_x.stop_gradient = True - - self.gate = Linear(self.d_model, self.num_expert) - - self.gate_weight = self.gate.weight.cast(paddle.float32) - self.gate_weight.stop_gradient = True - - paddle.set_default_dtype(self.x_type) - self.activation = swiglu - - self.inputs = { - self.tensor_x, - self.gate_weight, - self.bmm_w0, - self.bmm_b0, - self.bmm_w1, - self.bmm_b1, - None if self.quant_method == "None" else self.scale0, - None if self.quant_method == "None" else self.scale1, - self.quant_method, - 2, - } - self.python_api = moe - self.outputs = { - "Out": self.GetMoeOut(self.tensor_x), - } - - def config(self): - self.x_type = np.float16 - self.batch_size = 10 - self.seq_len = 128 - self.num_expert = 32 - self.d_model = 768 - self.d_feedforward = 3072 - self.top_k = 2 - self.quant_method = "None" - - def GetMoeOut(self, tensor_x): - paddle.disable_static(place=paddle.CUDAPlace(0)) - moe_out = moe( - tensor_x, - self.gate_weight, - self.bmm_w0, - self.bmm_b0, - self.bmm_w1, - self.bmm_b1, - None if self.quant_method == "None" else self.scale0, - None if self.quant_method == "None" else self.scale1, - self.quant_method, - 2, - ) - - return moe_out - - def GetBaselineOut(self, hidden_states): - paddle.disable_static(place=paddle.CUDAPlace(0)) - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = paddle.reshape(hidden_states, [-1, hidden_dim]) - router_logits = self.gate(hidden_states).cast(paddle.float32) - - routing_weights = F.softmax(router_logits, axis=-1, dtype='float32') - routing_weights, selected_experts = paddle.topk( - routing_weights, self.top_k, axis=-1 - ) - routing_weights /= paddle.sum(routing_weights, axis=-1, keepdim=True) - # we cast back to the input dtype - routing_weights = routing_weights.cast(np.float32) - - final_hidden_states = paddle.zeros_like(hidden_states) - - # One hot encode the selected experts to create an expert mask - # this will be used to easily index which expert is going to be sollicitated - expert_mask = paddle.transpose( - F.one_hot(selected_experts, num_classes=self.num_expert), [2, 1, 0] - ) - - # Loop over all available experts in the model and perform the computation on each expert - for expert_idx in range(self.num_expert): - expert_layer = self.experts[expert_idx] - idx, top_x = paddle.where(expert_mask[expert_idx]) - - # Index the correct hidden states and compute the expert hidden state for - # the current expert. We need to make sure to multiply the output hidden - # states by `routing_weights` on the corresponding tokens (top-1 and top-2) - current_state = paddle.index_select( - hidden_states, top_x, axis=0 - ).reshape([-1, hidden_dim]) - current_hidden_states = ( - expert_layer(current_state, expert_idx) - * routing_weights[top_x, idx] - ) - # Use scatter to accumulate the results - paddle.index_add_( - x=final_hidden_states, - index=top_x.squeeze(), - axis=0, - value=current_hidden_states.to(hidden_states.dtype), - ) - - final_hidden_states = paddle.reshape( - final_hidden_states, [batch_size, sequence_length, hidden_dim] - ) - return final_hidden_states - - def test_moe_op_new(self): - ref_out = self.GetBaselineOut(self.tensor_x).cast(np.float32) - moe_out = self.GetMoeOut(self.tensor_x).cast(np.float32) - np.testing.assert_allclose( - ref_out, moe_out, rtol=self.rtol, atol=self.atol - ) - self.test_check_output() - - def test_check_output(self): - self.check_output(check_pir=True) - - -class TestMoEOpBf16(TestMoEOp): - def config(self): - super().config() - self.x_type = paddle.bfloat16 - self.rtol = 1e-2 - self.atol = 1e-2 - - -if __name__ == "__main__": - unittest.main() From a2ae50534e461bfddbe77e0000a4d7df032cd1d3 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 20:50:31 +0800 Subject: [PATCH 25/29] fixx --- paddle/phi/ops/yaml/inconsistent/static_ops.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml index efd2ac108a5b0..6453d67cb869b 100644 --- a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml +++ b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml @@ -810,7 +810,6 @@ kernel: func: shadow_feed_tensors param: [x] - interfaces : paddle::dialect::InferSymbolicShapeInterface - op : share_data_ args : (Tensor x) From 776a427d8751dfdfa51f728a3c41cbab8713fbf9 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 20:58:46 +0800 Subject: [PATCH 26/29] fixx --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index 8787a62d98a6c..da5d52b77e4a2 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1463,7 +1463,7 @@ bool RoiAlignOpInferSymbolicShape( bool MoeOpInferSymbolicShape(pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { const auto &x_shape = - infer_context->GetShapeOrDataForValue(op->operand_source(0)); + infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape(); infer_context->SetShapeOrDataForValue( op->result(0), symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)}); @@ -1488,7 +1488,7 @@ bool MulticlassNms3OpInferSymbolicShape( "The rank of Input(Scores) must be 2 or 3. But received rank = %d", score_size)); - infer_context->AddEqualCstr(box_dims.size(), symbol::DimExpr(3)); + infer_context->AddEqualCstr(box_dims.size(), 3); if (score_size == 3) { PADDLE_ENFORCE_EQ( From 6b9205c9ca586f3482c0451365232f3bed996744 Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Wed, 14 Aug 2024 21:10:40 +0800 Subject: [PATCH 27/29] fixx --- .../interface/infer_symbolic_shape/multiary_infer_sym.cc | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index da5d52b77e4a2..b947238dce27e 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1487,8 +1487,12 @@ bool MulticlassNms3OpInferSymbolicShape( common::errors::InvalidArgument( "The rank of Input(Scores) must be 2 or 3. But received rank = %d", score_size)); - - infer_context->AddEqualCstr(box_dims.size(), 3); + PADDLE_ENFORCE_EQ( + box_dims.size(), + 3, + common::errors::InvalidArgument( + "The rank of Input(BBoxes) must be 3. But received rank = %d", + box_dims.size())); if (score_size == 3) { PADDLE_ENFORCE_EQ( From 8adfc221896cc98fce5e05f865ea90cfe9ac082b Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Thu, 15 Aug 2024 13:09:41 +0800 Subject: [PATCH 28/29] todo --- .../multiary_infer_sym.cc | 26 +++++-------------- test/legacy_test/test_multinomial_op.py | 2 ++ 2 files changed, 8 insertions(+), 20 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc index b947238dce27e..0fae5cb439bc7 100644 --- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc +++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc @@ -1480,27 +1480,13 @@ bool MulticlassNms3OpInferSymbolicShape( const std::vector &box_dims = bboxes_shape_or_data.shape(); const std::vector &score_dims = scores_shape_or_data.shape(); const size_t score_size = score_dims.size(); - - PADDLE_ENFORCE_EQ( - score_size == 2 || score_size == 3, - true, - common::errors::InvalidArgument( - "The rank of Input(Scores) must be 2 or 3. But received rank = %d", - score_size)); - PADDLE_ENFORCE_EQ( - box_dims.size(), - 3, - common::errors::InvalidArgument( - "The rank of Input(BBoxes) must be 3. But received rank = %d", - box_dims.size())); - + // TODO(Jeff114514): + // add constraint for score_size == 2 || score_size == 3 + // add constraint for box_dims.size() == 3 if (score_size == 3) { - PADDLE_ENFORCE_EQ( - box_dims[2] == 4 || box_dims[2] == 8 || box_dims[2] == 16 || - box_dims[2] == 24 || box_dims[2] == 32, - true, - common::errors::InvalidArgument("The last dimension of Input(BBoxes) " - "must be 4 or 8 or 16 or 24 or 32")); + // TODO(Jeff114514): + // add constraint for box_dims[2] == 4 || box_dims[2] == 8 || box_dims[2] == + // 16 || box_dims[2] == 24 || box_dims[2] == 32 infer_context->AddEqualCstr(box_dims[1], score_dims[2]); } else { infer_context->AddEqualCstr(box_dims[2], symbol::DimExpr(4)); diff --git a/test/legacy_test/test_multinomial_op.py b/test/legacy_test/test_multinomial_op.py index ff4d136ce57f1..471a57b474599 100644 --- a/test/legacy_test/test_multinomial_op.py +++ b/test/legacy_test/test_multinomial_op.py @@ -64,6 +64,7 @@ def init_data(self): def test_check_output(self): self.check_output_customized(self.verify_output, check_pir=True) + self.check_output(check_pir=True) def sample_output(self, out): return sample_output_one_dimension(out, 4) @@ -72,6 +73,7 @@ def verify_output(self, outs): # normalize the input to get the probability prob = self.input_np / self.input_np.sum(axis=-1, keepdims=True) sample_prob = self.sample_output(np.array(outs[0])) + print(self.outputs) np.testing.assert_allclose( sample_prob, prob, From 276e829bb3ee27030d06dc9787e7fe1a6de862ad Mon Sep 17 00:00:00 2001 From: Jeff114514 <928430299@qq.com> Date: Thu, 15 Aug 2024 13:11:59 +0800 Subject: [PATCH 29/29] todo --- test/legacy_test/test_multinomial_op.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/legacy_test/test_multinomial_op.py b/test/legacy_test/test_multinomial_op.py index 471a57b474599..ff4d136ce57f1 100644 --- a/test/legacy_test/test_multinomial_op.py +++ b/test/legacy_test/test_multinomial_op.py @@ -64,7 +64,6 @@ def init_data(self): def test_check_output(self): self.check_output_customized(self.verify_output, check_pir=True) - self.check_output(check_pir=True) def sample_output(self, out): return sample_output_one_dimension(out, 4) @@ -73,7 +72,6 @@ def verify_output(self, outs): # normalize the input to get the probability prob = self.input_np / self.input_np.sum(axis=-1, keepdims=True) sample_prob = self.sample_output(np.array(outs[0])) - print(self.outputs) np.testing.assert_allclose( sample_prob, prob,