diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc index 37dcc47512..758ec7b8eb 100644 --- a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc @@ -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 other_indices(static_cast(full_sized_input_shape->dim_size()) - 2, 1); + InlinedVector concat_input_args; + concat_input_args.push_back(first_two_dims_arg); + concat_input_args.push_back( + CreateInitializerFromVector(graph, + {static_cast(other_indices.size())}, + other_indices, + graph.GenerateNodeArgName("other_shape"))); + + InlinedVector 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 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 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 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 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 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 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 reshape_input_args; + reshape_input_args.reserve(2); reshape_input_args.push_back(input_ids_arg); - std::vector new_input_ids_shape; + InlinedVector new_input_ids_shape; + new_input_ids_shape.reserve(static_cast(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(new_input_ids_shape.size())}, + new_input_ids_shape, + graph.GenerateNodeArgName("flattened_shape"))); InlinedVector 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 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(); } diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index e0dbc44741..8d2bd19bff 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -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.