saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

support broadcast shape for elementwise node in padding elimination (#16710)

Browse source 
With PaddingElimination optimizer, input1 of element-wise op may be
flattened like:

```
  input1 (shape:[batch_size, seq_len, ...])        input1 (shape:[valid_tokens, ...])
        \                                               \
         \               input2                          \               input2
          \                /              ----->          \               /
           \              /                                \             /
	    Element-wise Op                                Element-wise Op
```
So, the shape of input2 should be processed accordingly:
1. If input2.shape.dim_size <= input1.shape.dim_size-2, i.e. input2 has
no [batch_size, seq_len] at begining,
we needn't to process the shape of input2 because it's compatible with
the flattened shape of input1 (shape:[valid_tokens, ...]).
   
2. If the shape of input2 has the same dim_size with shape of input1 and
has [batch_size, seqlen] at begening,
to be compatible with flattened shape of input1, we need to insert
flatten pattern for input2 also,
which flatten the shape of input2 from [batch_size, seq_len, ...] to
[valida_tokens, ...].
   
  
3. (which done in this pr) In other case for shape of input2, like [1,
seq_len, ...] or [batch_size, 1, ...], we firstly need to expand it
to [batch_size, seq_len, ...] which is convenient to flatten. And then
insert flatten pattern.
This commit is contained in:
guyang3532 2023-08-10 19:07:22 +08:00 committed by GitHub
parent b4e0fc87ea
commit ef6f4a4aa1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 169 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

215
orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc
View file

@ -65,13 +65,77 @@ NodeArg* GetDimsValue(Graph& graph, NodeArg* input, NodeArg* indices_arg, Node&
return gather_out_args[0];
}
// Insert Expand to the in_index-th input of node.
// The node should have two inputs and the shape of the other input (node.InputDefs()[1-in_index]) should be
// [batch_size, seq_len, ...]. This function insert an Expand to expand shape of the in_index-th input of node with
// a shape arg of [batch_size, seq_len, 1, 1, ...] which size is equal with node.InputDefs()[1-in_index]->Shape().size.
NodeArg* InsertExpandForNodeInput(Graph& graph,
Node& node,
uint32_t in_index,
NodeArg* first_two_dims_arg,
const logging::Logger& logger) {
auto full_sized_input_shape = node.InputDefs()[1 - in_index]->Shape();
ORT_ENFORCE(full_sized_input_shape->dim_size() >= 2);
NodeArg* expand_shape_arg = nullptr;
if (full_sized_input_shape->dim_size() == 2) {
expand_shape_arg = first_two_dims_arg;
} else {
InlinedVector<int64_t> other_indices(static_cast<int64_t>(full_sized_input_shape->dim_size()) - 2, 1);
InlinedVector<NodeArg*> concat_input_args;
concat_input_args.push_back(first_two_dims_arg);
concat_input_args.push_back(
CreateInitializerFromVector(graph,
{static_cast<int64_t>(other_indices.size())},
other_indices,
graph.GenerateNodeArgName("other_shape")));
InlinedVector<NodeArg*> concat_output_args{&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("concat_shape_result"),
nullptr)};
onnxruntime::NodeAttributes attributes;
attributes["axis"] = ONNX_NAMESPACE::MakeAttribute("axis", int64_t(0));
Node& concat_node = graph.AddNode(graph.GenerateNodeName("concat_shape"), "Concat", "", concat_input_args,
concat_output_args, &attributes, kOnnxDomain);
ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(concat_node), "Failed to concat shape for " + concat_node.Name());
concat_node.SetExecutionProviderType(node.GetExecutionProviderType());
expand_shape_arg = concat_output_args[0];
}
InlinedVector<NodeArg*> expand_input_args;
expand_input_args.reserve(2);
expand_input_args.push_back(node.MutableInputDefs()[in_index]);
expand_input_args.push_back(expand_shape_arg);
InlinedVector<NodeArg*> expand_output_args;
expand_output_args.push_back(
&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("inputs_expand_result"),
node.MutableInputDefs()[1 - in_index]->TypeAsProto()));
Node* new_expand_node = InsertIntermediateNodeOnDestInput(
graph, node,
in_index,
0,
0,
graph.GenerateNodeName("ExpandPaddingShape"),
"Expand",
"Expand shape of one input arg to align the other arg.",
expand_input_args,
expand_output_args,
{},
"",
logger);
new_expand_node->SetExecutionProviderType(node.GetExecutionProviderType());
return new_expand_node->MutableOutputDefs()[0];
}
// Insert Reshape + ShrunkenGather to flatten the in_index-th input of node.
// The gather_index_arg is the indices of the elements that are not padding.
NodeArg* InsertNodesForInput(Graph& graph,
Node& node,
uint32_t in_index,
NodeArg* gather_index_arg,
const logging::Logger& logger) {
NodeArg* InsertFlattenPatternForInput(Graph& graph,
Node& node,
uint32_t in_index,
NodeArg* gather_index_arg,
const logging::Logger& logger) {
InlinedVector<NodeArg*> reshape_input_args;
reshape_input_args.reserve(2);
reshape_input_args.push_back(node.MutableInputDefs()[in_index]);
@ -216,45 +280,23 @@ void IterateSubgraphFromNode(Graph& graph,
graph_utils::IsSupportedOptypeVersionAndDomain(*cur, "Mul", {7, 13, 14})) {
ORT_ENFORCE(subgraph.find(cur->MutableInputDefs()[0]) != subgraph.end() ||
subgraph.find(cur->MutableInputDefs()[1]) != subgraph.end());
NodeArg* arg_in_subgraph = nullptr;
NodeArg* arg_not_in_subgraph = nullptr;
if (subgraph.find(cur->MutableInputDefs()[0]) != subgraph.end()) {
arg_in_subgraph = cur->MutableInputDefs()[0];
arg_not_in_subgraph = cur->MutableInputDefs()[1];
} else if (subgraph.find(cur->MutableInputDefs()[1]) != subgraph.end()) {
arg_in_subgraph = cur->MutableInputDefs()[1];
arg_not_in_subgraph = cur->MutableInputDefs()[0];
}
// arg_in_subgraph is contained in subgraph, so its shape must be [batch_size, seq_len, ...]
// Now only support cases of the two shapes are absolutely same or the other shape dim size is smaller by 2.
// For example, [batch_size, seqlen, hidden_size] and [batch_size, seqlen, hidden_size].
// [batch_size, seqlen, hidden_size] and [hidden_size].
// TODO: support other case such as:
// [batch_size, seqlen, hidden_size] and [batch_size, 1, hidden_size]
if (arg_in_subgraph->Shape() && arg_not_in_subgraph->Shape() &&
(arg_not_in_subgraph->Shape()->dim_size() <= arg_in_subgraph->Shape()->dim_size() - 2 ||
(arg_in_subgraph->Shape()->dim_size() == arg_not_in_subgraph->Shape()->dim_size() &&
arg_in_subgraph->Shape()->dim(0) == arg_not_in_subgraph->Shape()->dim(0) &&
arg_in_subgraph->Shape()->dim(1) == arg_not_in_subgraph->Shape()->dim(1)))) {
if (cur->InputDefs()[0]->Shape() && cur->InputDefs()[1]->Shape()) {
if ((subgraph.find(cur->MutableInputDefs()[0]) == subgraph.end() &&
cur->InputDefs()[0]->Shape()->dim_size() > cur->InputDefs()[1]->Shape()->dim_size()) ||
(subgraph.find(cur->MutableInputDefs()[1]) == subgraph.end() &&
cur->InputDefs()[1]->Shape()->dim_size() > cur->InputDefs()[0]->Shape()->dim_size())) {
// If the shape of one of the inputs is not in the subgraph, and it has more dimensions,
// this case is not supported now.
LOG_DEBUG_INFO(logger, "PaddingElimination::Input shapes of node:" + cur->Name() + " are not compatible." +
" arg not in subgraph has more dimensions.");
candidate_outputs.insert(cur);
continue;
}
subgraph.insert(cur->MutableOutputDefs()[0]);
PushAllOutputNode(graph, to_visit, cur, visited);
// There are two possibilities here:
// 1. The size of arg_not_in_subgraph->Shape is smaller than arg_in_subgraph->Shape by 2,
// do not need to add flatten pattern to arg_not_in_subgraph.
// 2. The size of arg_not_in_subgraph->Shape is same with arg_in_subgraph->Shape and the first
// two dims value are exactly same, then there are also two possibilities:
// <1>. The arg_not_in_subgraph is propagated from embedding_node (contained in subgraph),
// do not need to process it.
// <2>. The arg_not_in_subgraph is not propagated from embedding_node (not contained in subgraph),
// need to add flatten pattern to arg_not_in_subgraph.
// Here we just add cur node to candidate_inputs and process it (add flatten pattern to its input) after
// the graph iteration, according to whether it's contained in subgraph.
if (arg_in_subgraph->Shape()->dim_size() == arg_not_in_subgraph->Shape()->dim_size()) {
candidate_inputs.insert(cur);
}
candidate_inputs.insert(cur);
} else {
LOG_DEBUG_INFO(logger, "PaddingElimination::Input shapes of node:" + cur->Name() + "are not compatible.");
LOG_DEBUG_INFO(logger, "PaddingElimination::Input of node:" + cur->Name() + " have no shape.");
candidate_outputs.insert(cur);
continue;
}
@ -341,14 +383,15 @@ Status PaddingElimination::ApplyImpl(Graph& graph, bool& modified, int graph_lev
// Make sure each node_arg in subgraph has first two consecutive dims to be flattened.
// All node_args in subgraph is propagated from the embedding node
std::unordered_set<NodeArg*> subgraph;
// input args of nodes in candidate_inputs should be in subgraph or to be added Reshape + Gather
// record node that its input args may be input of the subgraph into candidate_inputs
// input args of nodes in candidate_inputs should be in subgraph or to be added Reshape + Gather.
// Put nodes that its input args may be input of the subgraph into candidate_inputs
std::unordered_set<Node*> candidate_inputs;
// input args of nodes in candidate_outputs, if in subgraph, should be added GatherGrad + Reshape
// record node that its input args may be output of the subgraph into candidate_outputs
std::unordered_set<Node*> candidate_outputs;
int64_t handled_input_count = 0;
int64_t handled_output_count = 0;
int64_t expanded_input_count = 0;
// Find the valid embedding node
for (auto node_index : node_topology_list) {
@ -395,25 +438,36 @@ Status PaddingElimination::ApplyImpl(Graph& graph, bool& modified, int graph_lev
return Status::OK();
}
if (!input_ids_arg->Shape()) {
LOG_DEBUG_INFO(logger, "Exit PaddingElimination optimization for not finding shape of input_ids.");
return Status::OK();
}
auto input_ids_shape = input_ids_arg->Shape();
// For now, we only support all the dims of input_ids_shape besides the first two has dim_value.
for (int k = 2; k < input_ids_shape->dim_size(); k++) {
if (!input_ids_shape->dim(k).has_dim_value()) {
LOG_DEBUG_INFO(logger, "Exit PaddingElimination optimization for shape dims of input_ids has no value.");
return Status::OK();
}
}
IterateSubgraphFromNode(graph, embedding_node, subgraph, candidate_inputs, candidate_outputs, logger);
// Add Reshape + Sub + NonZero + Squeeze to get the not padding index to be gathered
InlinedVector<NodeArg*> reshape_input_args;
reshape_input_args.reserve(2);
reshape_input_args.push_back(input_ids_arg);
std::vector<int64_t> new_input_ids_shape;
InlinedVector<int64_t> new_input_ids_shape;
new_input_ids_shape.reserve(static_cast<int64_t>(input_ids_shape->dim_size()) - 1);
new_input_ids_shape.push_back(-1); // Flatten the two leading dims
auto input_ids_shape = input_ids_arg->Shape();
for (int k = 2; k < input_ids_shape->dim_size(); k++) {
ORT_ENFORCE(input_ids_shape->dim(k).has_dim_value());
new_input_ids_shape.push_back(input_ids_shape->dim(k).dim_value());
}
ONNX_NAMESPACE::TensorProto new_shape_const_tensor;
new_shape_const_tensor.set_name(graph.GenerateNodeArgName("flattened_shape"));
new_shape_const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
new_shape_const_tensor.add_dims(new_input_ids_shape.size());
new_shape_const_tensor.set_raw_data(new_input_ids_shape.data(), new_input_ids_shape.size() * sizeof(int64_t));
NodeArg* new_input_ids_shape_arg = &graph_utils::AddInitializer(graph, new_shape_const_tensor);
reshape_input_args.push_back(new_input_ids_shape_arg);
reshape_input_args.push_back(
CreateInitializerFromVector(graph,
{static_cast<int64_t>(new_input_ids_shape.size())},
new_input_ids_shape,
graph.GenerateNodeArgName("flattened_shape")));
InlinedVector<NodeArg*> reshape_output_args;
reshape_output_args.push_back(
@ -432,30 +486,54 @@ Status PaddingElimination::ApplyImpl(Graph& graph, bool& modified, int graph_lev
NodeArg* squeeze_out_arg = InsertNodesForValidIndices(
graph, reshape_output_args[0], embedding_node->MutableInputDefs()[2], embedding_node->GetExecutionProviderType());
// Get the first two dims value of input_ids which is [batch_size, seq_len]
NodeArg* first_two_dims_arg = GetDimsValue(graph,
input_ids_arg,
CreateInitializerFromVector(graph, {2}, {0, 1}, graph.GenerateNodeArgName("first_two_indices")),
*embedding_node);
// Add flatten pattern to each input node of the subgraph
// to flattern the shape of [batch_size, seqlen, ...] to [valid_token_count, ...]
InsertNodesForInput(graph, *embedding_node, 1, squeeze_out_arg, logger);
InsertFlattenPatternForInput(graph, *embedding_node, 1, squeeze_out_arg, logger);
handled_input_count++;
modified = true;
for (auto& node : candidate_inputs) {
for (uint32_t i = 0; i < node->InputDefs().size(); ++i) {
if (subgraph.find(node->MutableInputDefs()[i]) == subgraph.end()) {
InsertNodesForInput(graph, *node, i, squeeze_out_arg, logger);
// The type of node is one of Elementwise ops.
// The input size must be 2 and there must be more than one input in the subgraph.
ORT_ENFORCE(node->InputDefs().size() == 2);
// Because candidate_inputs are nodes iterated from embedding node, each of them must have at least one arg in
// the subgraph and the i-th input of this node is not in the subgraph, so the other input must be in the subgraph
// and has shape of [batch_size, seq_len, ...]
NodeArg* arg_in_subgraph = node->MutableInputDefs()[1 - i];
NodeArg* arg_not_in_subgraph = node->MutableInputDefs()[i];
// There are three possibilities for the shape of arg_not_in_subgraph:
// 1. The size of arg_not_in_subgraph->Shape is smaller than arg_in_subgraph->Shape by 2,
// which means the shape of arg_not_in_subgraph has no [batch_size, seq_len] in beginning,
// and do not need to add flatten pattern to it.
// 2. The arg_not_in_subgraph->Shape.size == arg_in_subgraph->Shape.size or arg_in_subgraph->Shape.size - 1,
// and the first two dims do not equal [batch_size, seq_len].
// In this case we just expand the arg_not_in_subgraph->Shape to [batch_size, seq_len, ...],
// then the case becomes same with 3.
// 3. The size of arg_not_in_subgraph->Shape is equal with size of arg_in_subgraph->Shape,
// and the first two dims of arg_not_in_subgraph->Shape is [batch_size, seq_len].
// Because the shape of arg_in_subgraph will be flattened to [valid_tokens, ... ] automatically after
// the shape of input_ids is flattened, so we need to insert flatten pattern for arg_not_in_subgraph->Shape.
if (arg_not_in_subgraph->Shape()->dim_size() <= arg_in_subgraph->Shape()->dim_size() - 2) {
continue;
} else if (arg_in_subgraph->Shape()->dim_size() != arg_not_in_subgraph->Shape()->dim_size() ||
arg_in_subgraph->Shape()->dim(0) != arg_not_in_subgraph->Shape()->dim(0) ||
arg_in_subgraph->Shape()->dim(1) != arg_not_in_subgraph->Shape()->dim(1)) {
InsertExpandForNodeInput(graph, *node, i, first_two_dims_arg, logger);
expanded_input_count++;
}
InsertFlattenPatternForInput(graph, *node, i, squeeze_out_arg, logger);
handled_input_count++;
}
}
}
std::vector<int64_t> first_two_indices{0, 1};
ONNX_NAMESPACE::TensorProto first_two_indices_const_tensor;
first_two_indices_const_tensor.set_name(graph.GenerateNodeArgName("first_two_indices"));
first_two_indices_const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
first_two_indices_const_tensor.add_dims(first_two_indices.size());
first_two_indices_const_tensor.set_raw_data(first_two_indices.data(), first_two_indices.size() * sizeof(int64_t));
NodeArg* first_two_indices_arg = &graph_utils::AddInitializer(graph, first_two_indices_const_tensor);
// Get the first two dims value of input_ids which is [batch_size, seq_len]
NodeArg* first_two_dims_arg = GetDimsValue(graph, input_ids_arg, first_two_indices_arg, *embedding_node);
// Add pattern to each output node of the subgraph
// to unflatten the shape of [valid_token_count, ...] to [batch_size, seq_len, ...]
for (const auto& node : candidate_outputs) {
@ -479,8 +557,11 @@ Status PaddingElimination::ApplyImpl(Graph& graph, bool& modified, int graph_lev
edge->SetShape(flattened_shape);
}
}
LOGS(logger, INFO) << "PaddingElimination::Total handled input node count: " << handled_input_count
<< " output node count: " << handled_output_count;
if (handled_input_count > 0 || handled_output_count > 0) {
LOGS(logger, INFO) << "PaddingElimination::Total handled input node count: " << handled_input_count
<< " output node count: " << handled_output_count
<< " expanded input count: " << expanded_input_count;
}
return Status::OK();
}

40
orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
View file

@ -5741,14 +5741,17 @@ def test_runtime_inspector_label_and_embed_sparsity_detection(embed_is_sparse, l
"test_cases",
[
("Add", 0),
("Add", 1),
("Add", 2),
("Add", 3),
("Add", 4),
("Sub", 0),
("Sub", 1),
("Sub", 2),
("Sub", 3),
("Sub", 4),
("Mul", 0),
("Mul", 1),
("Mul", 2),
("Mul", 3),
("Mul", 4),
("MatMul", 0),
("MatMul", 1),
("Dropout", 0),
@ -5761,8 +5764,6 @@ def test_ops_for_padding_elimination(test_cases):
os.environ["ORTMODULE_ENABLE_EMBEDDING_SPARSE_OPTIMIZER"] = "1"
test_op = test_cases[0]
case = test_cases[1]
# test_op = "Sub"
# case = 2
class ToyModel(torch.nn.Module):
def __init__(self, vocab_size, hidden_size, pad_token_id):
@ -5776,10 +5777,13 @@ def test_ops_for_padding_elimination(test_cases):
# in case 0, the shapes of inputs of test_op are [batch_size, seqlen, hidden_size] and [hidden_size],
# the test_op should be included in padding elimination subgraph and the GatherGrad should be added to
# output of test_op.
# in case 1, the shapes of inputs of test_op are [batch_size, seqlen, hidden_size] and [batch_size, 1, hidden_size],
# this case is not support in padding elimination, so the test_op should not be included in padding
# elimination subgraph and the GatherGrad should be added before test_op.
# in case 2, the shapes of inputs of test_op are [batch_size, seqlen, hidden_size] and [batch_size, seqlen, hidden_size],
# in case 2, the shapes of inputs of test_op are [batch_size, seqlen, hidden_size] and [batch_size, 1, hidden_size],
# the test_op should be included in padding elimination subgraph and a 'Expand + Reshape + ShrunkenGather'
# pattern should be insert to the arg of [batch_size, 1, hidden_size].
# in case 3, the shapes of inputs of test_op are [batch_size, seqlen, hidden_size] and [1, hidden_size],
# the test_op should be included in padding elimination subgraph and a 'Expand + Reshape + ShrunkenGather'
# pattern should be insert to the arg of [batch_size, 1, hidden_size].
# in case 4, the shapes of inputs of test_op are [batch_size, seqlen, hidden_size] and [batch_size, seqlen, hidden_size],
# the test_op should be included in padding elimination subgraph and the GatherGrad should be added to
# output of test_op. Besides, the other input of Add should be added 'Reshape + ShrunkenGather' to
# flatten and elimination padding.
@ -5788,9 +5792,11 @@ def test_ops_for_padding_elimination(test_cases):
one_input = None
if case == 0:
one_input = torch.ones(self.hidden_size, dtype=torch.long).to(device)
elif case == 1:
one_input = torch.ones((input_shape[0], 1, self.hidden_size), dtype=torch.long).to(device)
elif case == 2:
one_input = torch.ones((input_shape[0], 1, self.hidden_size), dtype=torch.long).to(device)
elif case == 3:
one_input = torch.ones((1, self.hidden_size), dtype=torch.long).to(device)
elif case == 4:
one_input = torch.ones(input_shape, dtype=torch.long).to(device)
one_input = one_input.unsqueeze(-1).expand(-1, -1, self.hidden_size)
inputs_embeds = self.word_embeddings(input_ids)
@ -5878,7 +5884,7 @@ def test_ops_for_padding_elimination(test_cases):
assert len([node.op_type for node in training_model.graph.node if node.op_type == "NonZero"]) == 1
assert len([node.op_type for node in training_model.graph.node if node.op_type == "Squeeze"]) == 1
assert len([node.op_type for node in training_model.graph.node if node.op_type == "PadAndUnflatten"]) == 1
if case == 2:
if case >= 2:
assert len([node.op_type for node in training_model.graph.node if node.op_type == "ShrunkenGather"]) == 2
else:
assert len([node.op_type for node in training_model.graph.node if node.op_type == "ShrunkenGather"]) == 1
@ -5893,10 +5899,7 @@ def test_ops_for_padding_elimination(test_cases):
gathergrad_input_optypes = [find_input_node_type(training_model, arg) for arg in gathergrad_node.input]
if test_op == "Add" or test_op == "Mul" or test_op == "Sub":
if case == 0:
assert test_op in gathergrad_input_optypes
elif case == 1:
assert "ATen" in gathergrad_input_optypes
assert test_op in gathergrad_input_optypes
elif test_op == "MatMul":
if case == 0:
assert "ATen" in gathergrad_input_optypes
@ -5995,7 +5998,7 @@ def test_e2e_padding_elimination():
def run_optim_step(optimizer):
optimizer.step()
optimizer.zero_grad()
optimizer.zero_grad(set_to_none=False)
# Generate one batch of inputs (shape:[batch_size, max_seq_length]) and masks (shape:[batch_size, max_seq_length]).
# Each input has random length from 1 to max_seq_length*0.8 with values from 2 to vocab_size and padded with 1 at
@ -6041,8 +6044,7 @@ def test_e2e_padding_elimination():
run_optim_step(ort_optimizer)
for pt_param, ort_param in zip(pt_model.parameters(), ort_model.parameters()):
if pt_param.grad is not None:
_test_helpers.assert_values_are_close(pt_param.grad, ort_param.grad, atol=1e-4, rtol=1e-5)
_test_helpers.assert_values_are_close(pt_param.grad, ort_param.grad, atol=1e-4, rtol=1e-5)
if os.getenv("ORTMODULE_ROCM_TEST", "0") == "1":
# For ROCm EP, the difference between ORT and PyTorch is larger than CUDA EP.