Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

【PIR Dist Op Reg No.19】 reg pull_box_sparse #62982

Merged
merged 23 commits into from
May 6, 2024
Merged
Show file tree
Hide file tree
Changes from 16 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions paddle/fluid/ir_adaptor/translator/op_compat_gen.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,9 @@ def insert_new_mutable_attributes(
op_arg_name_mappings['push_sparse_v2'].update(
{"out_grad_in": "Out@GRAD", "out_grad_out": "Out@GRAD"}
)
op_arg_name_mappings['push_box_sparse'].update(
{"out_grad_in": "Out@GRAD", "out_grad_out": "Out@GRAD"}
)

op_arg_name_mappings['push_gpups_sparse'].update(
{"out_grad": "Out@GRAD", "out_grad_grad": "Out@GRAD"}
Expand Down
4 changes: 4 additions & 0 deletions paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,10 @@
'push_dense',
'limit_by_capacity',
'global_scatter',
'pull_box_sparse',
'pull_box_sparse_',
'push_box_sparse',
'push_box_sparse_',
]


Expand Down
20 changes: 20 additions & 0 deletions paddle/fluid/pir/dialect/operator/ir/ops.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -1351,6 +1351,15 @@
func : prune_gate_by_capacity
data_type : gate_idx

- op : pull_box_sparse
args : (Tensor w, Tensor[] ids, bool is_sparse = false, bool is_distributed = false, int size = 1)
enkilee marked this conversation as resolved.
Show resolved Hide resolved
output : Tensor[](out){ids.size()}
infer_meta :
func : PullBoxSparseInferMeta
kernel :
func : pull_box_sparse
data_type : ids

- op : pull_gpups_sparse
args : (Tensor w, Tensor[] ids, int[] size={}, bool is_sparse=false, bool is_distributed=false)
output : Tensor[](out){ids.size()}
Expand All @@ -1362,6 +1371,17 @@
optional : w
backward: push_gpups_sparse

- op : push_box_sparse
enkilee marked this conversation as resolved.
Show resolved Hide resolved
args : (Tensor[] ids, Tensor[] out_grad_in, bool is_sparse = false, bool is_distributed = false, int size = 1)
output : Tensor[](out_grad_out){out_grad_in.size()}
infer_meta :
func : UnchangedMultiInferMeta
param : [out_grad_in]
kernel :
func : push_box_sparse
data_type : out_grad_in
inplace : (out_grad_in -> out_grad_out)

- op : push_dense
args : (Tensor[] ids, int table_id = -1, float scale_data_norm = -1.0f, str[] input_names = {})
output :
Expand Down
2 changes: 2 additions & 0 deletions paddle/fluid/pir/dialect/operator/utils/utils.cc
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,8 @@ const std::unordered_set<std::string> LegacyOpList = {
CReduceMinOp::name(),
CReduceProdOp::name(),
CScatterOp::name(),
PullBoxSparseOp::name(),
PushBoxSparseOp::name(),
PushSparseV2Op::name(),
PartialSendOp::name(),
PartialRecvOp::name()};
Expand Down
18 changes: 18 additions & 0 deletions paddle/phi/api/yaml/op_compat.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2671,13 +2671,31 @@
outputs :
out : Out

- op : pull_box_sparse
inputs :
{ w : W, ids: Ids}
outputs :
out : Out
attrs :
sparse : is_sparse
extra :
attrs : [bool is_sparse = false, bool is_distributed = false, int size = 1]

- op : pull_gpups_sparse
backward : push_gpups_sparse
inputs :
{w : W, ids : Ids}
outputs :
out : Out

- op : push_box_sparse
inputs :
ids: Ids
outputs :
out : Out
attrs :
sparse : is_sparse

- op : push_dense
inputs :
ids : Ids
Expand Down
26 changes: 26 additions & 0 deletions paddle/phi/infermeta/binary.cc
Original file line number Diff line number Diff line change
Expand Up @@ -2950,6 +2950,32 @@ void PruneGateByCapacityInferMeta(const MetaTensor& gate_idx,
new_gate_idx->set_dtype(gate_idx.dtype());
}

void PullBoxSparseInferMeta(const MetaTensor& w,
const std::vector<const MetaTensor*>& ids,
bool is_sparse,
bool is_distributed,
int size,
std::vector<MetaTensor*> out) {
auto hidden_size = static_cast<int64_t>(size);
const size_t n_ids = ids.size();
for (size_t i = 0; i < n_ids; ++i) {
MetaTensor* output = out[i];
auto ids_dims = ids[i]->dims();
int ids_rank = ids_dims.size();
PADDLE_ENFORCE_EQ(ids_dims[ids_rank - 1],
1UL,
phi::errors::InvalidArgument(
"Shape error in %lu id, the last dimension of the "
"'Ids' tensor must be 1.",
i));
auto out_dim =
common::vectorize(common::slice_ddim(ids_dims, 0, ids_rank - 1));
out_dim.push_back(hidden_size);
output->set_dims(common::make_ddim(out_dim));
output->share_lod(*ids[i]);
enkilee marked this conversation as resolved.
Show resolved Hide resolved
}
}

void RepeatInterleaveWithTensorIndexInferMeta(const MetaTensor& x,
const MetaTensor& repeats,
int dim,
Expand Down
7 changes: 7 additions & 0 deletions paddle/phi/infermeta/binary.h
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,13 @@ void PReluInferMeta(const MetaTensor& x,
MetaTensor* out,
MetaConfig config = MetaConfig());

void PullBoxSparseInferMeta(const MetaTensor& w,
const std::vector<const MetaTensor*>& ids,
bool is_sparse,
bool is_distributed,
int size,
std::vector<MetaTensor*> out);

void PullGpupsSparseInferMeta(const MetaTensor& w,
const std::vector<const MetaTensor*>& ids,
const std::vector<int>& size,
Expand Down
1 change: 1 addition & 0 deletions test/ir/pir/translator/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_push_dense_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_pull_box_sparse_translator)
list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_pull_gpups_sparse_translator)
list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_random_routing_translator)
list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_limit_by_capacity_translator)
Expand Down
51 changes: 51 additions & 0 deletions test/ir/pir/translator/test_pull_box_sparse_translator.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
# 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 test_op_translator

import paddle
from paddle.base.layer_helper import LayerHelper


class TestPullBoxSparseOpTranslator(
test_op_translator.TestOpWithBackwardTranslator
):
def append_op(self):
self.forward_op_type = "pull_box_sparse"
self.backward_op_type = "push_box_sparse"
ids = paddle.ones(shape=(1, 1), dtype='float32')
w = paddle.ones(shape=(1, 1), dtype='float32')
out = paddle.ones(shape=(1, 1), dtype='float32')
attrs = {
'is_sparse': False,
'is_distributed': False,
'size': 1,
}
forward_helper = LayerHelper(self.forward_op_type)
forward_helper.append_op(
type=self.forward_op_type,
inputs={"W": w, "Ids": [ids]},
outputs={"Out": [out]},
attrs=attrs,
)
return out

enkilee marked this conversation as resolved.
Show resolved Hide resolved
def test_translator(self):
self.check()


if __name__ == "__main__":
unittest.main()