diff --git a/cmake/onnxruntime_optimizer.cmake b/cmake/onnxruntime_optimizer.cmake index 4d4f9daf26..84e64b634a 100644 --- a/cmake/onnxruntime_optimizer.cmake +++ b/cmake/onnxruntime_optimizer.cmake @@ -78,6 +78,8 @@ if (onnxruntime_ENABLE_TRAINING_APIS) list(APPEND onnxruntime_optimizer_src_patterns "${ORTTRAINING_SOURCE_DIR}/core/optimizer/*.h" "${ORTTRAINING_SOURCE_DIR}/core/optimizer/*.cc" + "${ORTTRAINING_SOURCE_DIR}/core/optimizer/compute_optimizer/*.h" + "${ORTTRAINING_SOURCE_DIR}/core/optimizer/compute_optimizer/*.cc" ) endif() diff --git a/docs/ORTModule_Training_Guidelines.md b/docs/ORTModule_Training_Guidelines.md index defccdd90f..e638e73777 100644 --- a/docs/ORTModule_Training_Guidelines.md +++ b/docs/ORTModule_Training_Guidelines.md @@ -135,6 +135,9 @@ debugging). Disable it with following command: export ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=0 ``` + There are fine-grained flags to control different optimizations, but only applied when `ORTMODULE_ENABLE_COMPUTE_OPTIMIZER` is enabled. + - `ORTMODULE_ENABLE_LABEL_SPARSITY_OPT` is used to enable the optimization leveraging label sparsity. + #### ORTMODULE_ENABLE_INPUT_DENSITY_INSPECTOR - **Feature Area**: *ORTMODULE/Runtime inspector* diff --git a/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc b/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc index ab33c2eabe..eb0234a1c4 100644 --- a/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc +++ b/onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc @@ -47,7 +47,11 @@ Node* InsertIntermediateNodeOnDestInput(Graph& graph, for (size_t j = 0; j < new_node.MutableOutputDefs().size(); ++j) { graph.UpdateProducerNode(new_node.MutableOutputDefs()[j]->Name(), new_node.Index()); } - graph.AddConsumerNode(src_node_arg->Name(), &new_node); + + for (size_t j = 0; j < new_node.MutableInputDefs().size(); ++j) { + graph.AddConsumerNode(new_node.MutableInputDefs()[j]->Name(), &new_node); + } + const Node* src_node = graph.GetProducerNode(src_node_arg->Name()); if (src_node) { int src_out_index = optimizer_utils::IndexOfNodeOutput(*src_node, *src_node_arg); @@ -150,6 +154,41 @@ std::pair> CompareInputShapeWithOutputShape( return std::make_pair>(true, std::move(rets)); } +int GetONNXOpSetVersion(const Graph& graph) { + int onnx_opset = -1; + auto onnx_domain_it = graph.DomainToVersionMap().find(kOnnxDomain); + if (onnx_domain_it != graph.DomainToVersionMap().end()) { + onnx_opset = onnx_domain_it->second; + } else { + auto onnx_domain_alias_it = graph.DomainToVersionMap().find(kOnnxDomainAlias); + if (onnx_domain_alias_it != graph.DomainToVersionMap().end()) + onnx_opset = onnx_domain_alias_it->second; + else + ORT_THROW("ONNX domain not found in this model"); + } + return onnx_opset; +} + +NodeArg* CreateInitializerFromVector(Graph& graph, + const InlinedVector& dims, + const InlinedVector& values, + const std::string& name) { + ONNX_NAMESPACE::TensorProto const_tensor; + const_tensor.set_name(name); + const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + + int64_t total_count = 1; + for (const int64_t dim : dims) { + const_tensor.add_dims(dim); + total_count *= dim; + } + + ORT_ENFORCE(total_count == static_cast(values.size())); + + const_tensor.set_raw_data(values.data(), values.size() * sizeof(int64_t)); + return &graph_utils::AddInitializer(graph, const_tensor); +} + } // namespace onnxruntime::optimizer::compute_optimizer #endif diff --git a/onnxruntime/core/optimizer/compute_optimizer/shared_utils.h b/onnxruntime/core/optimizer/compute_optimizer/shared_utils.h index cfc38a5ae1..2f7815c72c 100644 --- a/onnxruntime/core/optimizer/compute_optimizer/shared_utils.h +++ b/onnxruntime/core/optimizer/compute_optimizer/shared_utils.h @@ -156,5 +156,23 @@ std::pair> CompareInputShapeWithOutputShape( const ONNX_NAMESPACE::TensorShapeProto* full_broadcasted_shape, const ONNX_NAMESPACE::TensorShapeProto* target_shape); +/** + * @brief Get opset version from the graph. + */ +int GetONNXOpSetVersion(const Graph& graph); + +/** + * @brief Create an initializer from given dims/value vector and name. + * + * @param dims A int vector as the shape of the created initializer. If we want to create a scalar initializer, + * we should pass an empty vector. + * @param values A int vector containing the value buffer. + */ + +NodeArg* CreateInitializerFromVector(Graph& graph, + const InlinedVector& dims, + const InlinedVector& values, + const std::string& name); + } // namespace onnxruntime::optimizer::compute_optimizer #endif diff --git a/onnxruntime/core/optimizer/compute_optimizer/upstream_gather.cc b/onnxruntime/core/optimizer/compute_optimizer/upstream_gather.cc index 85dfa5f319..25623651c0 100644 --- a/onnxruntime/core/optimizer/compute_optimizer/upstream_gather.cc +++ b/onnxruntime/core/optimizer/compute_optimizer/upstream_gather.cc @@ -307,6 +307,44 @@ std::optional IsSupportedGather(Graph& graph, Node& node, return SliceInfo(graph, &node, dim_size == 0, "axis", axis, true); } +std::optional IsSupportedShrunkenGather(Graph& graph, Node& node, + const InlinedHashSet& + compatible_execution_providers, + const logging::Logger& logger) { + if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "ShrunkenGather", {1}, kMSDomain) || + !graph_utils::IsSupportedProvider(node, compatible_execution_providers)) { + return std::nullopt; + } + + auto data_shape = node.MutableInputDefs()[0]->Shape(); + auto indices_shape = node.MutableInputDefs()[1]->Shape(); + auto gather_out_shape = node.MutableOutputDefs()[0]->Shape(); + if (data_shape == nullptr || indices_shape == nullptr || gather_out_shape == nullptr) { + LOG_DEBUG_INFO(logger, "Skip ShrunkenGather node " + node.Name() + " due to undefined shape." + + std::to_string(data_shape == nullptr) + std::to_string(indices_shape == nullptr) + + std::to_string(gather_out_shape == nullptr)); + return std::nullopt; + } + + const int data_rank = data_shape->dim_size(); + if (data_rank <= 1) { + LOG_DEBUG_INFO(logger, "Skip ShrunkenGather node " + node.Name() + " due to data rank <= 1."); + return std::nullopt; + } + + int axis = static_cast(node.GetAttributes().at("axis").i()); + axis = axis < 0 ? axis + data_rank : axis; + int dim_size = indices_shape->dim_size(); + + if (dim_size == 0) { + LOG_DEBUG_INFO(logger, "Skip ShrunkenGather node " + node.Name() + " due to unsupported dim size: " + + std::to_string(dim_size)); + return std::nullopt; + } + + return SliceInfo(graph, &node, false /*is_slice_scalar*/, "axis", axis, true); +} + } // namespace std::optional UpStreamGatherGraphTransformer::IsSupportedForUpstream( @@ -317,7 +355,9 @@ std::optional UpStreamGatherGraphTransformer::IsSupportedForUpstream( if (!gather_info.has_value()) { gather_info = IsSupportedGather(graph, node, GetCompatibleExecutionProviders(), logger); } - + if (!gather_info.has_value()) { + gather_info = IsSupportedShrunkenGather(graph, node, GetCompatibleExecutionProviders(), logger); + } return gather_info; } diff --git a/onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.cc b/onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.cc index 3fbc8e7d7c..4701bf1464 100644 --- a/onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.cc +++ b/onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.cc @@ -83,29 +83,6 @@ TensorShapeProto CreateTensorShapeInsertDimAtAxis(const TensorShapeProto* src_sh return updated_shape; } -NodeArg* CreateUnsqueezeAxesInitializer(Graph& graph, const std::vector& values) { - ONNX_NAMESPACE::TensorProto axes_const_tensor; - axes_const_tensor.set_name(graph.GenerateNodeArgName("axes")); - axes_const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - axes_const_tensor.add_dims(values.size()); - axes_const_tensor.set_raw_data(values.data(), values.size() * sizeof(int64_t)); - return &graph_utils::AddInitializer(graph, axes_const_tensor); -} - -int GetONNXOpSetVersion(const Graph& graph) { - int onnx_opset = -1; - auto onnx_domain_it = graph.DomainToVersionMap().find(kOnnxDomain); - if (onnx_domain_it != graph.DomainToVersionMap().end()) { - onnx_opset = onnx_domain_it->second; - } else { - auto onnx_domain_alias_it = graph.DomainToVersionMap().find(kOnnxDomainAlias); - if (onnx_domain_alias_it != graph.DomainToVersionMap().end()) - onnx_opset = onnx_domain_alias_it->second; - else - ORT_THROW("ONNX domain not found in this model"); - } - return onnx_opset; -} } // namespace bool UpdateSliceOutputShape(NodeArg& arg_to_update, int axis_to_update, @@ -161,7 +138,8 @@ void AdaptInputAndOutputForScalarSlice(Graph& graph, Node& current_node, int cur "Unsqueeze", "Unsqueeze node", {current_node.MutableInputDefs()[input_index], - CreateUnsqueezeAxesInitializer(graph, {pair.second.non_negative_axis})}, + CreateInitializerFromVector(graph, {1}, {pair.second.non_negative_axis}, + graph.GenerateNodeArgName("axes"))}, {&graph.GetOrCreateNodeArg( graph.GenerateNodeArgName("unsqueeze_adaptor"), current_node.MutableInputDefs()[input_index]->TypeAsProto())}, @@ -218,7 +196,7 @@ void AdaptInputAndOutputForScalarSlice(Graph& graph, Node& current_node, int cur "Squeeze", "Squeeze node", {consumer.MutableInputDefs()[index], - CreateUnsqueezeAxesInitializer(graph, {slice_axis})}, + CreateInitializerFromVector(graph, {1}, {slice_axis}, graph.GenerateNodeArgName("axes"))}, {&graph.GetOrCreateNodeArg( graph.GenerateNodeArgName("squeeze_adaptor"), consumer.MutableInputDefs()[index]->TypeAsProto())}, @@ -450,7 +428,7 @@ bool LayerNormalizationGatherActor::PreCheck(const Graph& /* graph */, auto axis = static_cast(current_node.GetAttributes().at("axis").i()); axis = axis < 0 ? axis + current_node.InputDefs()[0]->Shape()->dim_size() : axis; - // Make sure layernorm/softmax's reduction happens after the axis we want to slice. + // Make sure LayerNormalization's reduction happens after the axis we want to slice. if (axis <= info.non_negative_axis) { return false; } @@ -458,17 +436,8 @@ bool LayerNormalizationGatherActor::PreCheck(const Graph& /* graph */, const NodeArg* gather_data_input_arg = current_node.OutputDefs()[info.GetDataProducerOutputIndex()]; const auto& gather_data_input_shape = gather_data_input_arg->Shape(); - // Only handle the case where the first input is 3D. - auto ln_input_arg = current_node.InputDefs()[0]; - auto ln_input_shape = ln_input_arg->Shape(); - auto ln_input_rank = ln_input_shape->dim_size(); - if (ln_input_rank != 3) { - LOG_DEBUG_INFO(logger, "Fail LayerNormalizationGatherActor::PreCheck for node " + current_node.Name() + - ": data_input_rank is " + std::to_string(ln_input_rank)); - return false; - } - - auto [success, ret] = CompareInputShapeWithOutputShape(gather_data_input_shape, ln_input_shape); + auto [success, ret] = CompareInputShapeWithOutputShape(gather_data_input_shape, + current_node.InputDefs()[0]->Shape()); if (!success) { // This should not happen!!! LOG_DEBUG_INFO(logger, "Fail LayerNormalizationGatherActor::PreCheck for node " + current_node.Name() + @@ -481,9 +450,9 @@ bool LayerNormalizationGatherActor::PreCheck(const Graph& /* graph */, all_input_cmp_rets[0] = std::move(ret); shape_update_func = [&info](Node& node) -> void { - // Be noted: LayerNorm's 2nd and 3rd output have shape [batch, sequence, 1]. The dim is still kept even - // for reduced axes. - // So the slicing axis is same with the 1st output. + // Be noted: If LayerNorm's data input is [dim1, dim2, dim3], reduce axis is 1, + // then its 2nd and 3rd outputs have shape [dim1, dim2, 1]. The dim is still kept even + // for reduced axes, so the slicing axis is same with the 1st output. for (size_t output_idx = 0; output_idx < node.MutableOutputDefs().size(); ++output_idx) { UpdateSliceOutputShape(*node.MutableOutputDefs()[output_idx], info.non_negative_axis, info.output_dim_on_axis); @@ -501,7 +470,7 @@ bool SoftmaxGatherActor::PreCheck(const Graph& graph, const Node& current_node, auto axis = static_cast(current_node.GetAttributes().at("axis").i()); axis = axis < 0 ? axis + current_node.InputDefs()[0]->Shape()->dim_size() : axis; - // Make sure layernorm/softmax's reduction happens after the axis we want to slice. + // Make sure Softmax's reduction happens after the axis we want to slice. if (axis <= info.non_negative_axis) { return false; } @@ -613,26 +582,6 @@ bool ReshapeGatherActor::PostProcess( optimizer_utils::AppendTensorFromInitializer(graph, *current_node.InputDefs()[1], new_shape_const_values, true); const int slice_axis = info_without_node.non_negative_axis; - auto create_new_initializer_from_vector = [&graph](NodeArg* arg_to_be_replaced, - const InlinedVector& new_values) -> NodeArg* { - // Create new TensorProto. - ONNX_NAMESPACE::TensorProto constant_tensor_proto; - constant_tensor_proto.set_name(graph.GenerateNodeArgName(arg_to_be_replaced->Name())); - constant_tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - auto length = new_values.size(); - constant_tensor_proto.add_dims(length); - constant_tensor_proto.set_raw_data(new_values.data(), length * sizeof(int64_t)); - - // Add initializer into Graph. - NodeArg* new_shape_arg = &graph_utils::AddInitializer(graph, constant_tensor_proto); - // Update the output arg shape. - ONNX_NAMESPACE::TensorShapeProto new_shape; - new_shape.add_dim()->set_dim_value(length); - new_shape_arg->SetShape(new_shape); - - return new_shape_arg; - }; - // If the shape constant on slice_axis is 0, then it keeps the original dim of input. // If it is a scalar slice, then we just remove that dim. Otherwise, we don't need to update the dim value. if (new_shape_const_values[slice_axis] == 0) { @@ -645,10 +594,11 @@ bool ReshapeGatherActor::PostProcess( new_values.push_back(new_shape_const_values[i]); } } - auto new_shape_arg = create_new_initializer_from_vector(arg_to_be_replaced, new_values); + auto new_shape_arg = CreateInitializerFromVector(graph, {static_cast(new_values.size())}, new_values, + graph.GenerateNodeArgName(arg_to_be_replaced->Name())); graph_utils::ReplaceNodeInput(current_node, 1, *new_shape_arg); } else { - LOG_DEBUG_INFO(logger, "Reshape's shape has 0 specified for aixs: " + std::to_string(slice_axis) + + LOG_DEBUG_INFO(logger, "Reshape's shape has 0 specified for axis: " + std::to_string(slice_axis) + ", not need an update."); } return true; @@ -657,7 +607,9 @@ bool ReshapeGatherActor::PostProcess( // If it selected shape is a dim value, we can update the shape tensor directory. if (info_without_node.output_dim_on_axis.has_dim_value()) { new_shape_const_values[slice_axis] = info_without_node.output_dim_on_axis.dim_value(); - auto new_shape_arg = create_new_initializer_from_vector(current_node.MutableInputDefs()[1], new_shape_const_values); + auto new_shape_arg = + CreateInitializerFromVector(graph, {static_cast(new_shape_const_values.size())}, new_shape_const_values, + graph.GenerateNodeArgName(current_node.MutableInputDefs()[1]->Name())); graph_utils::ReplaceNodeInput(current_node, 1, *new_shape_arg); return true; } diff --git a/onnxruntime/core/optimizer/compute_optimizer/upstream_reshape.cc b/onnxruntime/core/optimizer/compute_optimizer/upstream_reshape.cc index 89fcfc6480..0d862dd033 100644 --- a/onnxruntime/core/optimizer/compute_optimizer/upstream_reshape.cc +++ b/onnxruntime/core/optimizer/compute_optimizer/upstream_reshape.cc @@ -140,19 +140,15 @@ ReshapeInfo UpStreamReshapeGraphTransformer::PropagateReshapeForInput( input_args.push_back(current_node.MutableInputDefs()[current_node_input_index]); // Prepare the target shape initializer. (Currently, only constant target shape is supported.) - std::vector new_shape; + InlinedVector new_shape; new_shape.push_back(-1); auto input_shape = current_node.MutableInputDefs()[current_node_input_index]->Shape(); for (int k = 2; k < input_shape->dim_size(); ++k) { ORT_ENFORCE(input_shape->dim(k).has_dim_value()); new_shape.push_back(input_shape->dim(k).dim_value()); } - ONNX_NAMESPACE::TensorProto new_shape_const_tensor; - new_shape_const_tensor.set_name(graph.GenerateNodeArgName("new_shape")); - new_shape_const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - new_shape_const_tensor.add_dims(new_shape.size()); - new_shape_const_tensor.set_raw_data(new_shape.data(), new_shape.size() * sizeof(int64_t)); - NodeArg* new_shape_arg = &graph_utils::AddInitializer(graph, new_shape_const_tensor); + NodeArg* new_shape_arg = CreateInitializerFromVector(graph, {static_cast(new_shape.size())}, new_shape, + graph.GenerateNodeArgName("new_shape")); input_args.push_back(new_shape_arg); diff --git a/onnxruntime/core/optimizer/compute_optimizer/upstream_transformer_base.cc b/onnxruntime/core/optimizer/compute_optimizer/upstream_transformer_base.cc index f0097da265..4b60d7772a 100644 --- a/onnxruntime/core/optimizer/compute_optimizer/upstream_transformer_base.cc +++ b/onnxruntime/core/optimizer/compute_optimizer/upstream_transformer_base.cc @@ -52,7 +52,7 @@ Status UpStreamGraphTransformerBase::ApplyImpl(Graph& graph, bool& modif size_t reordered_node_count = 0; // For summary size_t passthrough_count = 0; - for (auto index : order) { + for (const auto index : order) { auto* node_ptr = graph.GetNode(index); if (!node_ptr) // node was removed. @@ -91,7 +91,11 @@ Status UpStreamGraphTransformerBase::ApplyImpl(Graph& graph, bool& modif break; } - if (graph.GetConsumerNodes(input_tensor_producer_node->MutableOutputDefs()[0]->Name()).size() > 1) { + // This condition implies few things: + // 1. The node's outputs are only used once (that's the node to upstream). + // 2. If the node has more than one outputs, only one of the outputs has output edge + // and that output is used by the node to upstream. + if (input_tensor_producer_node->GetOutputEdgesCount() > 1) { LOG_DEBUG_INFO(logger, log_prefix + " stops at node " + input_tensor_producer_node->Name() + " since multiple consumers found"); continue; diff --git a/onnxruntime/test/optimizer/compute_optimizer_test.cc b/onnxruntime/test/optimizer/compute_optimizer_test.cc index b4e2db858e..9f33ee054e 100644 --- a/onnxruntime/test/optimizer/compute_optimizer_test.cc +++ b/onnxruntime/test/optimizer/compute_optimizer_test.cc @@ -21,26 +21,22 @@ #include "core/graph/graph_utils.h" #include "core/graph/graph_viewer.h" #include "core/graph/model.h" - #include "core/optimizer/common_subexpression_elimination.h" #include "core/optimizer/compute_optimizer/upstream_gather.h" #include "core/optimizer/compute_optimizer/upstream_reshape.h" #include "core/optimizer/utils.h" -#include "core/platform/env.h" -#include "core/session/inference_session.h" #include "core/util/math.h" +#include "test/common/tensor_op_test_utils.h" #include "test/compare_ortvalue.h" #include "test/framework/test_utils.h" #include "test/optimizer/graph_transform_test_builder.h" #include "test/optimizer/graph_transform_test_fixture.h" - +#include "test/optimizer/test_optimizer_utils.h" #include "test/providers/provider_test_utils.h" -#include "test/test_environment.h" #include "test/util/include/temp_dir.h" #include "test/util/include/asserts.h" #include "test/util/include/default_providers.h" -#include "test/common/tensor_op_test_utils.h" namespace onnxruntime { namespace test { @@ -131,138 +127,6 @@ TEST(ComputeOptimizerTests, GatherND_MatMul) { GatherNDComputationReductionTest("MatMul", *logger, SingleOpDefaultValidationFunc); } -/** - * @brief Class represent a input data (dimensions, data type and value). - */ -struct TestInputData { - template - TestInputData(const std::string& name, const TensorShapeVector& dims, const std::vector& values) - : name_(name), dims_(dims), values_(values) {} - - OrtValue ToOrtValue() { - OrtValue ortvalue; - std::vector dims; - dims.reserve(dims_.size()); - dims.insert(dims.end(), dims_.begin(), dims_.end()); - std::visit([&ortvalue, &dims](auto&& arg) { - using T = std::decay_t; - if constexpr (std::is_same_v> || - std::is_same_v> || - std::is_same_v>) - CreateMLValue( - TestCPUExecutionProvider()->GetAllocator(OrtMemTypeDefault), dims, arg, &ortvalue); - else - static_assert("Unspported types!"); - }, - values_); - - return ortvalue; - } - - std::string GetName() const { - return name_; - } - - private: - std::string name_; - TensorShapeVector dims_; - std::variant, std::vector, std::vector> values_; -}; - -void RandomFillFloatVector(const TensorShapeVector& shape, std::vector& data) { - static RandomValueGenerator random{1234}; - data = random.Gaussian(shape, 0.0f, 0.25f); -} - -void RandomFillHalfVector(const TensorShapeVector& shape, std::vector& data) { - std::vector data_float(TensorShape(shape).Size()); - RandomFillFloatVector(shape, data_float); - std::transform(data_float.begin(), data_float.end(), data.begin(), - [](float value) { return MLFloat16(math::floatToHalf(value)); }); -} - -void RandomMasks(int64_t batch, int64_t sequence_length, std::vector& data) { - static RandomValueGenerator random{5678}; - const std::vector num_count_to_random{batch}; - std::vector random_seq_lens = random.Uniform(num_count_to_random, 0, sequence_length); - data.resize(batch * sequence_length); // fill with zeros first. - for (int64_t i = 0; i < batch; ++i) { - for (int64_t j = 0; j < sequence_length; ++j) { - if (j > random_seq_lens[i]) { - break; - } - - data[i * sequence_length + j] = 1; - } - } -} - -struct InputContainer { - InputContainer() = default; - - template - TestInputData& AddInput(const std::string& name, const TensorShapeVector dims, const std::vector& values) { - inputs_.emplace_back(TestInputData(name, dims, values)); - return inputs_.back(); - } - - template - TestInputData& AddInput(const std::string& name, TensorShapeVector dims, - std::function< - void(const TensorShapeVector& shape, std::vector& data)> - func = nullptr) { - std::vector values(TensorShape(dims).Size()); - if (func) { - func(dims, values); - } - - inputs_.emplace_back(TestInputData(name, dims, values)); - return inputs_.back(); - } - - void ToInputMap(NameMLValMap& feeds) const { - for (auto input : inputs_) { - feeds.insert({input.GetName(), input.ToOrtValue()}); - } - } - - private: - std::vector inputs_; -}; - -static void RunModelWithData(const PathString& model_uri, const std::string session_log_id, - const std::string& provider_type, const InputContainer& input_container, - const std::vector& output_names, - std::vector& run_results) { - SessionOptions so; - // we don't want any transformation here. - so.graph_optimization_level = TransformerLevel::Default; - so.session_logid = session_log_id; - - InferenceSession session_object{so, GetEnvironment()}; - std::unique_ptr execution_provider; - if (provider_type == onnxruntime::kCpuExecutionProvider) - execution_provider = DefaultCpuExecutionProvider(); - else if (provider_type == onnxruntime::kCudaExecutionProvider) - execution_provider = DefaultCudaExecutionProvider(); - else if (provider_type == onnxruntime::kRocmExecutionProvider) - execution_provider = DefaultRocmExecutionProvider(); - EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); - - Status st; - ASSERT_TRUE((st = session_object.Load(model_uri)).IsOK()) << st.ErrorMessage(); - ASSERT_TRUE((st = session_object.Initialize()).IsOK()) << st.ErrorMessage(); - - NameMLValMap feeds; - input_container.ToInputMap(feeds); - - // Now run - RunOptions run_options; - st = session_object.Run(run_options, feeds, output_names, &run_results); - - ASSERT_TRUE(st.IsOK()) << "RunModelWithData run graph failed with error: " << st.ErrorMessage(); -} - TEST(ComputeOptimizerTests, GatherND_E2E) { const logging::Logger* logger = &logging::LoggingManager::DefaultLogger(); auto model_uri = MODEL_FOLDER "computation_reduction/gathernd/e2e.onnx"; @@ -1606,6 +1470,112 @@ TEST(ComputeOptimizerTests, GatherRobertaE2E) { } } +/* +Test graph include multiple equivalent subgraphs as below. + graph input [4, 32, 256] (float) graph input [4, 32, 256] (float) + | | + \_____________ ______________/ + \ / + Add ______ [16] + | / + ShrunkenGather, axis = 1 + | + Identity + | + graph output [4, 16, 256] (float) + +Add an Identity node because currently we don't allow ShrunkenGather generates graph output. +*/ +TEST(ComputeOptimizerTests, ShrunkenGatherElementwiseOps_PropagationOnTwoBranches) { + const logging::Logger* logger = &logging::LoggingManager::DefaultLogger(); + InlinedVector gather_indices; + auto pre_graph_checker = [&gather_indices](Graph& graph) -> Status { + auto op_count_pre = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_pre.size() == 3U); + TEST_RETURN_IF_NOT(op_count_pre["Add"] == 1); + TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.ShrunkenGather"] == 1); + TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 1); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "ShrunkenGather") { + TEST_RETURN_IF_NOT(gather_indices.empty()); + constexpr bool require_constant = true; + NodeArg* initializer_node_arg = graph.GetNodeArg(node.InputDefs()[1]->Name()); + TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, gather_indices, + require_constant)); + } + } + return Status::OK(); + }; + + auto post_graph_checker = [&gather_indices](Graph& graph) { + auto op_count_post = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_post.size() == 3U); + TEST_RETURN_IF_NOT(op_count_post["Add"] == 1); + TEST_RETURN_IF_NOT(op_count_post["com.microsoft.ShrunkenGather"] == 2); + TEST_RETURN_IF_NOT(op_count_post["Identity"] == 1); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "Add") { + const auto& input_defs = node.InputDefs(); + + { + auto producer_node = graph.GetProducerNode(input_defs[0]->Name()); + TEST_RETURN_IF_NOT(producer_node != nullptr); + TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather"); + + InlinedVector values; + constexpr bool require_constant = true; + NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name()); + TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values, + require_constant)); + for (size_t i = 0; i < values.size(); i++) { + TEST_RETURN_IF_NOT(values[i] == gather_indices[i]); + } + } + + { + auto producer_node = graph.GetProducerNode(input_defs[1]->Name()); + TEST_RETURN_IF_NOT(producer_node != nullptr); + TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather"); + + InlinedVector values; + constexpr bool require_constant = true; + NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name()); + TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values, require_constant)); + for (size_t i = 0; i < values.size(); i++) { + TEST_RETURN_IF_NOT(values[i] == gather_indices[i]); + } + } + } + } + return Status::OK(); + }; + + auto build_test_case = [](ModelTestBuilder& builder) { + auto* input1_arg = builder.MakeInput({{4, 32, 256}}); + auto* input2_arg = builder.MakeInput({{4, 32, 256}}); + auto* add_out = builder.MakeIntermediate(); + builder.AddNode("Add", {input1_arg, input2_arg}, {add_out}); + + const std::vector slice_shape{16}; + static RandomValueGenerator random{8888}; + std::vector random_slices = random.Uniform(slice_shape, 0, 32); + auto* slice_initializer = builder.MakeInitializer(slice_shape, random_slices); + auto* gather_out = builder.MakeIntermediate(); + builder.AddNode("ShrunkenGather", {add_out, slice_initializer}, {gather_out}, kMSDomain) + .AddAttribute("axis", static_cast(1)); + + auto* identity_out = builder.MakeOutput(); + builder.AddNode("Identity", {gather_out}, {identity_out}); + }; + + std::unique_ptr transformer = std::make_unique(); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger, std::move(transformer), + TransformerLevel::Level1, + 1, pre_graph_checker, post_graph_checker)); +} + /* Test graph include multiple equivalent subgraphs as below. graph input [4, 32, 256] (int64_t) graph input [4, 32, 256] (int64_t) diff --git a/onnxruntime/test/optimizer/graph_transform_test_builder.h b/onnxruntime/test/optimizer/graph_transform_test_builder.h index e6dad26a6a..3e42bc456b 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_builder.h +++ b/onnxruntime/test/optimizer/graph_transform_test_builder.h @@ -186,6 +186,11 @@ class ModelTestBuilder { return MakeInitializer({static_cast(data.size())}, data); } + NodeArg* MakeEmptyInput() { + NodeArg* empty = &graph_.GetOrCreateNodeArg("", nullptr); + return empty; + } + Node& AddNode(const std::string& op_type, const std::vector& input_args, const std::vector& output_args, diff --git a/onnxruntime/test/optimizer/test_optimizer_utils.cc b/onnxruntime/test/optimizer/test_optimizer_utils.cc new file mode 100644 index 0000000000..08cb129b1d --- /dev/null +++ b/onnxruntime/test/optimizer/test_optimizer_utils.cc @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#ifdef _MSC_VER +#pragma warning(push) +#pragma warning(disable : 4244) +#endif + +#include +#include "core/framework/ort_value.h" +#include "core/graph/model.h" + +#include "core/platform/env.h" +#include "core/session/inference_session.h" + +#include "test/test_environment.h" +#include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" +#include "test/optimizer/test_optimizer_utils.h" +#include "test/common/tensor_op_test_utils.h" + +namespace onnxruntime { +namespace test { + +void RandomFillFloatVector(const TensorShapeVector& shape, std::vector& data) { + static RandomValueGenerator random{1234}; + data = random.Gaussian(shape, 0.0f, 0.25f); +} + +void RandomFillHalfVector(const TensorShapeVector& shape, std::vector& data) { + std::vector data_float(TensorShape(shape).Size()); + RandomFillFloatVector(shape, data_float); + std::transform(data_float.begin(), data_float.end(), data.begin(), + [](float value) { return MLFloat16(math::floatToHalf(value)); }); +} + +void RandomMasks(int64_t batch, int64_t sequence_length, std::vector& data) { + static RandomValueGenerator random{5678}; + const std::vector num_count_to_random{batch}; + std::vector random_seq_lens = random.Uniform(num_count_to_random, 0, sequence_length); + data.resize(batch * sequence_length); // fill with zeros first. + for (int64_t i = 0; i < batch; ++i) { + for (int64_t j = 0; j < sequence_length; ++j) { + if (j > random_seq_lens[i]) { + break; + } + + data[i * sequence_length + j] = 1; + } + } +} + +void RunModelWithData(const PathString& model_uri, const std::string session_log_id, + const std::string& provider_type, const InputContainer& input_container, + const std::vector& output_names, + std::vector& run_results) { + SessionOptions so; + // we don't want any transformation here. + so.graph_optimization_level = TransformerLevel::Default; + so.session_logid = session_log_id; + + InferenceSession session_object{so, GetEnvironment()}; + std::unique_ptr execution_provider; + if (provider_type == onnxruntime::kCpuExecutionProvider) + execution_provider = DefaultCpuExecutionProvider(); + else if (provider_type == onnxruntime::kCudaExecutionProvider) + execution_provider = DefaultCudaExecutionProvider(); + else if (provider_type == onnxruntime::kRocmExecutionProvider) + execution_provider = DefaultRocmExecutionProvider(); + EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); + + Status st; + ASSERT_TRUE((st = session_object.Load(model_uri)).IsOK()) << st.ErrorMessage(); + ASSERT_TRUE((st = session_object.Initialize()).IsOK()) << st.ErrorMessage(); + + NameMLValMap feeds; + input_container.ToInputMap(feeds); + + // Now run + RunOptions run_options; + st = session_object.Run(run_options, feeds, output_names, &run_results); + + ASSERT_TRUE(st.IsOK()) << "RunModelWithData run graph failed with error: " << st.ErrorMessage(); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/test_optimizer_utils.h b/onnxruntime/test/optimizer/test_optimizer_utils.h new file mode 100644 index 0000000000..36df383a68 --- /dev/null +++ b/onnxruntime/test/optimizer/test_optimizer_utils.h @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/common/common.h" +#include "core/framework/tensor_shape.h" +#include "core/framework/float16.h" +#include "core/framework/framework_common.h" +#include "core/framework/ort_value.h" +#include "test/framework/test_utils.h" + +namespace onnxruntime { +namespace test { + +/** + * @brief Class represent a input data (dimensions, data type and value). + */ +struct TestInputData { + template + TestInputData(const std::string& name, const TensorShapeVector& dims, const std::vector& values) + : name_(name), dims_(dims), values_(values) {} + + OrtValue ToOrtValue() { + OrtValue ortvalue; + std::vector dims; + dims.reserve(dims_.size()); + dims.insert(dims.end(), dims_.begin(), dims_.end()); + std::visit([&ortvalue, &dims](auto&& arg) { + using T = std::decay_t; + if constexpr (std::is_same_v> || + std::is_same_v> || + std::is_same_v>) + CreateMLValue( + TestCPUExecutionProvider()->GetAllocator(OrtMemTypeDefault), dims, arg, &ortvalue); + else + static_assert("Unspported types!"); + }, + values_); + + return ortvalue; + } + + std::string GetName() const { + return name_; + } + + private: + std::string name_; + TensorShapeVector dims_; + std::variant, std::vector, std::vector> values_; +}; + +/** + * @brief A container for all input data. + * + */ +struct InputContainer { + InputContainer() = default; + + template + TestInputData& AddInput(const std::string& name, const TensorShapeVector dims, const std::vector& values) { + inputs_.emplace_back(TestInputData(name, dims, values)); + return inputs_.back(); + } + + template + TestInputData& AddInput(const std::string& name, TensorShapeVector dims, + std::function< + void(const TensorShapeVector& shape, std::vector& data)> + func = nullptr) { + std::vector values(TensorShape(dims).Size()); + if (func) { + func(dims, values); + } + + inputs_.emplace_back(TestInputData(name, dims, values)); + return inputs_.back(); + } + + void ToInputMap(NameMLValMap& feeds) const { + for (auto input : inputs_) { + feeds.insert({input.GetName(), input.ToOrtValue()}); + } + } + + private: + std::vector inputs_; +}; + +void RandomFillFloatVector(const TensorShapeVector& shape, std::vector& data); + +void RandomFillHalfVector(const TensorShapeVector& shape, std::vector& data); + +void RandomMasks(int64_t batch, int64_t sequence_length, std::vector& data); + +void RunModelWithData(const PathString& model_uri, const std::string session_log_id, + const std::string& provider_type, const InputContainer& input_container, + const std::vector& output_names, + std::vector& run_results); +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc b/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc index 77c8c2e460..e4647c58c5 100644 --- a/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc +++ b/orttraining/orttraining/core/framework/ortmodule_graph_builder.cc @@ -152,7 +152,7 @@ Status OrtModuleGraphBuilder::OptimizeForwardGraph(std::unordered_setMainGraph(); ORT_RETURN_IF_ERROR(forward_graph.Resolve()); - GraphTransformerManager graph_transformation_mgr{2}; + GraphTransformerManager graph_transformation_mgr{3}; std::unique_ptr cpu_execution_provider = std::make_unique(CPUExecutionProviderInfo()); diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc new file mode 100644 index 0000000000..4e20dc0cdc --- /dev/null +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc @@ -0,0 +1,225 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h" + +#include "core/graph/graph_utils.h" +#include "core/framework/random_seed.h" +#include "core/optimizer/initializer.h" +#include "core/optimizer/utils.h" +#include "core/optimizer/compute_optimizer/shared_utils.h" +#include "core/optimizer/compute_optimizer/upstream_gather_actors.h" + +namespace onnxruntime { + +// Put utilities in anonymous namespace. +namespace { +NodeArg* InsertNodesForValidLabelIndices(Graph& graph, Node& node, NodeArg* label_input, NodeArg* reduce_index_input) { + InlinedVector input_args{label_input, reduce_index_input}; + + InlinedVector output_args{&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("label_sub_result"), + node.MutableInputDefs()[1]->TypeAsProto())}; + + Node& sub_node = graph.AddNode(graph.GenerateNodeName("labels_sub"), "Sub", "label sub padding idx", input_args, + output_args, nullptr, kOnnxDomain); + ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(sub_node), "Failed to set op schema for " + sub_node.Name()); + sub_node.SetExecutionProviderType(node.GetExecutionProviderType()); + + auto non_zero_out_arg = &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("labels_filter_pad_result"), + node.MutableInputDefs()[1]->TypeAsProto()); + + Node& non_zero_node = graph.AddNode(graph.GenerateNodeName("labels_filter_pad"), "NonZero", + "labels filtering padding idx", + {sub_node.MutableOutputDefs()[0]}, + {non_zero_out_arg}, nullptr, kOnnxDomain); + + ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(non_zero_node), + "Failed to set op schema for " + non_zero_node.Name()); + + const std::string dim_name = MakeString("valid_label_count_", utils::GetRandomSeed()); + + // 1D input NonZero generates output of shape (1,valid_token_count). + ONNX_NAMESPACE::TensorShapeProto non_zero_output_shape; + non_zero_output_shape.add_dim()->set_dim_value(1); + non_zero_output_shape.add_dim()->set_dim_param(dim_name); + non_zero_out_arg->SetShape(non_zero_output_shape); + non_zero_node.SetExecutionProviderType(node.GetExecutionProviderType()); + + InlinedVector squeeze_input_args; + squeeze_input_args.push_back(non_zero_out_arg); + + bool opset_lower_than_13 = onnxruntime::optimizer::compute_optimizer::GetONNXOpSetVersion(graph) < 13; + onnxruntime::NodeAttributes attributes; + if (opset_lower_than_13) { + attributes["axes"] = ONNX_NAMESPACE::MakeAttribute("axes", std::vector{0}); + } else { + squeeze_input_args.push_back(onnxruntime::optimizer::compute_optimizer::CreateInitializerFromVector( + graph, {1}, {0}, graph.GenerateNodeArgName("axes"))); + } + + auto squeeze_out_arg = &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("squeeze_adaptor"), + non_zero_out_arg->TypeAsProto()); + Node& squeeze_node = graph.AddNode(graph.GenerateNodeName("squeeze_adaptor"), "Squeeze", "nonzero_squeezer", + squeeze_input_args, {squeeze_out_arg}, &attributes, kOnnxDomain); + ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(squeeze_node), + "Failed to set op schema for " + squeeze_node.Name()); + + // After Squeeze, the shape becomes (valid_token_count). + ONNX_NAMESPACE::TensorShapeProto squeeze_output_shape; + squeeze_output_shape.add_dim()->set_dim_param(dim_name); + squeeze_out_arg->SetShape(squeeze_output_shape); + squeeze_node.SetExecutionProviderType(node.GetExecutionProviderType()); + + return squeeze_out_arg; +} +} // namespace + +Status InsertGatherBeforeSceLoss::ApplyImpl(Graph& graph, bool& modified, int /*graph_level*/, + const logging::Logger& logger) const { + LOG_DEBUG_INFO(logger, "Enter InsertGatherBeforeSceLoss"); + + GraphViewer graph_viewer(graph); + size_t handled_sce_node_count = 0; // For summary + const auto& order = graph_viewer.GetNodesInTopologicalOrder(); + for (const auto index : order) { + auto* node_ptr = graph.GetNode(index); + if (!node_ptr) + // node was removed. + continue; + + auto& node = *node_ptr; + + bool is_onnx_sce = graph_utils::IsSupportedOptypeVersionAndDomain(node, "SoftmaxCrossEntropyLoss", {12, 13}); + bool is_internal_sce = graph_utils::IsSupportedOptypeVersionAndDomain(node, "SoftmaxCrossEntropyLossInternal", {1}, + kMSDomain); + + if ((!is_onnx_sce && !is_internal_sce) || + !graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders())) { + continue; + } + + // Check whether this SCE node is handled or not. + const Node* labels_producer = graph.GetProducerNode(node.MutableInputDefs()[1]->Name()); + // Skip if already inserted a ShrunkenGather node. + if (labels_producer && graph_utils::IsSupportedOptypeVersionAndDomain( + *labels_producer, "ShrunkenGather", {1}, kMSDomain)) { + LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + + ") due to labels input is already consumed by a ShrunkenGather node."); + continue; + } + + // Check shape requirements. + auto logits_shape = node.MutableInputDefs()[0]->Shape(); + auto labels_shape = node.MutableInputDefs()[1]->Shape(); + if (logits_shape == nullptr || labels_shape == nullptr) { + LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + + ") due to undefined input shapes."); + continue; + } + + if (logits_shape->dim_size() != 2 || labels_shape->dim_size() != 1) { + LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + + ") due to unsupported input shape ranks."); + continue; + } + + // Check attribute and input requirements. + std::string reduction = node.GetAttributes().at("reduction").s(); + if (reduction.compare("mean") == 0 || reduction.compare("sum") == 0) { + // loss output is a scalar, don't need reset shape. + } else { + LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + + ") due to loss [reduction=" + reduction + "]."); + continue; + } + + NodeArg* ignore_index_node_arg = nullptr; + if (is_internal_sce) { + if (node.InputDefs().size() < 4 || !graph_utils::IsConstantInitializer( + graph, node.InputDefs()[3]->Name(), /* check_outer_scope */ false)) { + LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + + ") due to target padding idx is non-constant initializer. Input count: " + std::to_string(node.InputDefs().size())); + continue; + } + ignore_index_node_arg = node.MutableInputDefs()[3]; + } else { + const auto ignore_index_attr = node.GetAttributes().find("ignore_index"); + if (ignore_index_attr != node.GetAttributes().end()) { + int64_t ignore_index_value = (*ignore_index_attr).second.i(); + ignore_index_node_arg = onnxruntime::optimizer::compute_optimizer::CreateInitializerFromVector( + graph, {}, {ignore_index_value}, graph.GenerateNodeArgName("ignore_index")); + } else { + LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + + ") due to missing ignore_index attribute."); + continue; + } + } + + std::vector sce_out1_consumers = graph.GetConsumerNodes(node.OutputDefs()[1]->Name()); + if (sce_out1_consumers.size() != 0 || graph.IsOutput(node.OutputDefs()[1])) { + LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + + ") due to log_prob output is graph output or consumed by other nodes."); + continue; + } + + // SoftmaxCrossEntropyLossInternal op definition guarantees that the the first dimension of both inputs must match, + // we don't do the check explicitly here. + + LOG_DEBUG_INFO(logger, "Inserting Sub+NonZero nodes for filtering valid tokens"); + + // It is possible a label input is used by multiple SoftmaxCrossEntropyLossInternal nodes, here we will create a + // subgraph retrieving valid tokens for each SoftmaxCrossEntropyLossInternal node. + // The duplication will be removed by CSE graph transformers. + NodeArg* valid_labels_input_arg = + InsertNodesForValidLabelIndices(graph, node, node.MutableInputDefs()[1], ignore_index_node_arg); + + // Insert the ShrunkenGather node on the two inputs. + for (int i = 0; i < 2; ++i) { + InlinedVector input_args; + input_args.reserve(2); + input_args.push_back(node.MutableInputDefs()[i]); + input_args.push_back(valid_labels_input_arg); + + InlinedVector output_args; + output_args.push_back( + &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("label_filter_result"), + node.MutableInputDefs()[i]->TypeAsProto())); + + /* new node input index to connect to node's input node*/ + int new_gather_input_index_to_connect = 0; + /* new node output index to connect to node*/ + int new_gather_output_index_to_connect = 0; + Node* new_gather_node = onnxruntime::optimizer::compute_optimizer::InsertIntermediateNodeOnDestInput( + graph, node, + i, + new_gather_input_index_to_connect, + new_gather_output_index_to_connect, + graph.GenerateNodeName("LabelsFilter"), + "ShrunkenGather", + "ShrunkenGather node to filter invalid tokens.", + input_args, + output_args, + {}, + kMSDomain, + logger); + + new_gather_node->SetExecutionProviderType(node.GetExecutionProviderType()); + auto gather_out_arg = new_gather_node->MutableOutputDefs()[0]; + + onnxruntime::optimizer::compute_optimizer::UpdateSliceOutputShape( + *gather_out_arg, 0, valid_labels_input_arg->Shape()->dim(0)); + } + + modified = true; + handled_sce_node_count += 1; + } + + LOG_DEBUG_INFO(logger, "Exit InsertGatherBeforeSceLoss, handled " + std::to_string(handled_sce_node_count) + + " SCE nodes"); + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h new file mode 100644 index 0000000000..0abb3c74a0 --- /dev/null +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +/** + * @brief Graph transformer that inserts ShrunkenGather before SCE nodes (e.g. + * SoftmaxCrossEntropyLossInternal/SoftmaxCrossEntropyLoss nodes). + * + * For label sparsity, we can remove them from the compute of loss function, by inserting ShrunkenGather + * operators for its two inputs. + * + * logits (float) [token_count, classes] labels (int64), [token_count] + * \ / + * SCE Node(ignore_index=-100) + * / \ + * loss (shape is scalar or [token_count]) log_prob [token_count, classes] + * + * Be noted in Transformer-based models: + * > `token_count` usually equals with `batch size` x `sequence length`. + * > `classes` usually equals with `vocabulary`. + * + * Only insert ShrunkGather if all following conditions are met for SCE nodes`: + * 1. Its reduction attribute value is 'sum' or 'mean', to make sure loss is a scalar. + * Otherwise, the loss is in shape [token_count], changing on `token_count` will affect subsequent computations. + * 2. Its 2nd output (log_prob) MUST NOT be graph output and not consumed by other other nodes. + * 3. Its ignore_index exists and is a constant scalar value. + * 4. Its 2nd input label's input node is not `ShrunkGather` node (to avoid this transformer duplicated applied). + * + * + * After the transformation: + * labels [token_count] + * \_______ + * \ \ + * \ Sub(ignore_index) + * \ \ + * | | + * | | + * | NonZero + * | | + * | Squeeze + * | | + * | indices of valid token [valid_token_count] + * \ | + * logits [token_count, classes] _________________\ _ _____/ + * \ / \ / + * \ / \ / + * \ / \ / + * ShrunkenGather ShrunkenGather + * [valid_token_count, classes] [valid_token_count] + * \ / + * SCE Node (ignore_index=-100) + * / \ + * / \ + * loss (shape is scalar or [valid_token_count]) log_prob [valid_token_count, classes] + * + * In this specific scenario, it is easy to infer that valid_token_count <= token_count. + * After insertion, loss computation flop is reduced. Additionally, upstream Gather graph optimization + * will try to reduce flop for other ops further. + */ +class InsertGatherBeforeSceLoss : public GraphTransformer { + public: + InsertGatherBeforeSceLoss(const InlinedHashSet& compatible_execution_providers = {}) noexcept + : GraphTransformer("InsertGatherBeforeSceLoss", compatible_execution_providers) { + } + + /** + * @brief + */ + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_config.h b/orttraining/orttraining/core/optimizer/graph_transformer_config.h index 755d878029..a73830ef1d 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_config.h +++ b/orttraining/orttraining/core/optimizer/graph_transformer_config.h @@ -24,6 +24,9 @@ struct TrainingGraphTransformerConfiguration : public GraphTransformerConfigurat // Enable compute optimizer. bool enable_compute_optimizer{false}; + + // Enable label sparsity compute optimization. + bool enable_label_sparsity_optimization{false}; }; } // namespace training diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc index 110cab2a43..e8744b93dd 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc +++ b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc @@ -11,8 +11,10 @@ #include "core/optimizer/cast_elimination.h" #include "core/optimizer/common_subexpression_elimination.h" #include "core/optimizer/compute_optimizer/upstream_gather.h" +#include "core/optimizer/compute_optimizer/upstream_reshape.h" #include "core/optimizer/concat_slice_elimination.h" #include "core/optimizer/constant_folding.h" +#include "core/optimizer/constant_sharing.h" #include "core/optimizer/conv_activation_fusion.h" #include "core/optimizer/conv_add_fusion.h" #include "core/optimizer/conv_bn_fusion.h" @@ -51,6 +53,7 @@ #include "orttraining/core/framework/distributed_run_context.h" #include "orttraining/core/optimizer/batchnorm_replacement.h" #include "orttraining/core/optimizer/bitmask_dropout_replacement.h" +#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h" #include "orttraining/core/optimizer/concat_replacement.h" #include "orttraining/core/optimizer/graph_transformer_registry.h" #include "orttraining/core/optimizer/insert_output_rewriter.h" @@ -96,6 +99,10 @@ std::vector> GeneratePreTrainingTransformers( ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique())); + // Put ConstantSharing before CommonSubexpressionElimination by intention as it can create more opportunities for + // CSE. For example, if A and B nodes both do Add operation with a same value but different initializers, by + // default, CSE will not merge them, because the different initializers are represented by different NodeArg. + transformers.emplace_back(std::make_unique(compatible_eps)); // Remove duplicate nodes. Must be applied before any recompute transformations. if (config.gelu_recompute || config.attn_dropout_recompute || config.transformer_layer_recompute) { transformers.emplace_back(std::make_unique(compatible_eps)); @@ -136,9 +143,6 @@ std::vector> GeneratePreTrainingTransformers( transformers.emplace_back(std::make_unique(compatible_eps)); transformers.emplace_back(std::make_unique(compatible_eps)); - if (config.enable_compute_optimizer) { - transformers.emplace_back(std::make_unique(compatible_eps)); - } if (config.gelu_recompute) { transformers.emplace_back(std::make_unique()); } @@ -150,12 +154,22 @@ std::vector> GeneratePreTrainingTransformers( config.number_recompute_layers, compatible_eps)); } if (config.propagate_cast_ops_config.level >= 0) { - const InlinedHashSet cuda_execution_provider = {onnxruntime::kCudaExecutionProvider, onnxruntime::kRocmExecutionProvider}; + const InlinedHashSet cuda_execution_provider = {onnxruntime::kCudaExecutionProvider, + onnxruntime::kRocmExecutionProvider}; transformers.emplace_back(std::make_unique(config.propagate_cast_ops_config.strategy, static_cast(config.propagate_cast_ops_config.level), config.propagate_cast_ops_config.allow, cuda_execution_provider)); } + + if (config.enable_compute_optimizer) { + transformers.emplace_back(std::make_unique(compatible_eps)); + if (config.enable_label_sparsity_optimization) { + transformers.emplace_back(std::make_unique(compatible_eps)); + transformers.emplace_back(std::make_unique(compatible_eps)); + } + } + } break; case TransformerLevel::Level2: { diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index 415dd7dd97..b9d5d78417 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -737,6 +737,7 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn .def_readwrite("transformer_layer_recompute", &TrainingGraphTransformerConfiguration::transformer_layer_recompute) .def_readwrite("number_recompute_layers", &TrainingGraphTransformerConfiguration::number_recompute_layers) .def_readwrite("enable_compute_optimizer", &TrainingGraphTransformerConfiguration::enable_compute_optimizer) + .def_readwrite("enable_label_sparsity_optimization", &TrainingGraphTransformerConfiguration::enable_label_sparsity_optimization) .def_readwrite("propagate_cast_ops_config", &TrainingGraphTransformerConfiguration::GraphTransformerConfiguration::propagate_cast_ops_config); py::class_ module_graph_builder_config( diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 1453234615..cbb23be541 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -185,6 +185,10 @@ class GraphExecutionManager(GraphExecutionInterface): self._enable_compute_optimizer = ( ortmodule._defined_from_envvar("ORTMODULE_ENABLE_COMPUTE_OPTIMIZER", 1, warn=True) == 1 ) + self._enable_label_sparsity_optimization = ( + self._enable_compute_optimizer + and ortmodule._defined_from_envvar("ORTMODULE_ENABLE_LABEL_SPARSITY_OPT", 0, warn=True) == 1 + ) # Flag to re-export the model due to attribute change on the original module. # Re-export will be avoided if _skip_check is enabled. @@ -459,6 +463,7 @@ class GraphExecutionManager(GraphExecutionInterface): graph_transformer_config.propagate_cast_ops_config.allow = self._propagate_cast_ops_allow graph_transformer_config.propagate_cast_ops_config.strategy = self._propagate_cast_ops_strategy graph_transformer_config.enable_compute_optimizer = self._enable_compute_optimizer + graph_transformer_config.enable_label_sparsity_optimization = self._enable_label_sparsity_optimization return graph_transformer_config def _initialize_graph_builder(self): diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/torch_interop_utils.cc b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/torch_interop_utils.cc index a8445bf64f..32d458bb39 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/torch_interop_utils.cc +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/torch_interop_utils.cc @@ -8,15 +8,15 @@ // In Torch forward run (e.g. THPFunction_apply), ctx of type THPFunction* (which is also a PyObject*) // is created (https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/python_function.cpp#L673). // The ctx is used to run user-defined forward function and backward function as the first -// parameter. The same time, a cdata of type std::shared_ptr is created +// parameter. The same time, a cdata of type std::shared_ptr is created // (https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/python_function.cpp#L677), // cdata is owned by: -// a). forward run output tensors as grad_fn_ property. (The full hierarchy is: Tensor owns +// a). forward run output tensors as grad_fn_ property. (The full hierarchy is: Tensor owns // shared_pointer; TensorImpl owns std::unique_ptr; AutogradMeta // manages grad_/grad_fn_/grad_accumulator_. Among them, grad_fn_ is std::shared_ptr, // e.g, the so called gradient function.) // https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/variable.h#L194 -// b). the consumer operator of forward run outputs, will let its own PyNode/Node (gradident function) +// b). the consumer operator of forward run outputs, will let its own PyNode/Node (gradient function) // owns the grad_fn_ (of type std::shared_ptr) of all inputs that require grad. // https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/function.h#L263 // BUT, if we run torch computation within PythonOp, b) is lost. SO, for some cases, where forward outputs @@ -25,7 +25,7 @@ // Then when PythonOpGrad runs, segment fault. // // So we add b)'s reference in this Pool when forward run returns; dereference from this Pool when backward -// completes, then ~PyNode() is called, which subsquently calls ~THPFunction() destorying ctx. +// completes, then ~PyNode() is called, which subsequently calls ~THPFunction() destroying ctx. class PyNodeSharedPointerPool { public: static PyNodeSharedPointerPool& GetInstance() { @@ -33,7 +33,7 @@ class PyNodeSharedPointerPool { return pool; }; - void RegisterGradFunc(const size_t& ctx_address, torch::autograd::AutogradMeta* autograd_meta){ + void RegisterGradFunc(const size_t& ctx_address, torch::autograd::AutogradMeta* autograd_meta) { auto it = grad_fns_.find(ctx_address); TORCH_CHECK(it == grad_fns_.end(), "should not register grad_fn twice for ctx ", ctx_address); @@ -41,21 +41,20 @@ class PyNodeSharedPointerPool { grad_fns_.emplace(ctx_address, std::move(autograd_meta->grad_fn_)); }; - void UnRegisterGradFunc(const size_t& ctx_address){ + void UnRegisterGradFunc(const size_t& ctx_address) { auto it = grad_fns_.find(ctx_address); TORCH_CHECK(it != grad_fns_.end(), "fail to find grad_fn for ctx ", ctx_address); grad_fns_.erase(ctx_address); }; - void ClearAll(){ + void ClearAll() { grad_fns_.clear(); } private: PyNodeSharedPointerPool(){}; - ~PyNodeSharedPointerPool(){ - }; + ~PyNodeSharedPointerPool(){}; PyNodeSharedPointerPool(const PyNodeSharedPointerPool&) = delete; PyNodeSharedPointerPool& operator=(const PyNodeSharedPointerPool&) = delete; @@ -65,20 +64,19 @@ class PyNodeSharedPointerPool { std::unordered_map> grad_fns_; }; - void clear_grad_fns_for_next_edges(at::Tensor target, std::vector saved_tensors) { - // For leaf tensor, there will be a AccumulateGrad (gradident function) created, which owns a - // reference to the tensor. + // For leaf tensor, there will be a AccumulateGrad (gradient function) created, which owns a + // reference to the tensor. // For any user saved tensors (with save_for_backward), if the tensor is leaf, we put the map // {AccumulateGrad*, Tensor*} into grad_fn_to_tensor_map. - std::unordered_map grad_fn_to_tensor_map; - for (auto& t: saved_tensors) { + std::unordered_map grad_fn_to_tensor_map; + for (auto& t : saved_tensors) { auto grad_fn = t.grad_fn(); if (!grad_fn) { grad_fn = torch::autograd::impl::try_get_grad_accumulator(t); if (grad_fn) { - TORCH_CHECK(grad_fn_to_tensor_map.find(grad_fn.get()) == grad_fn_to_tensor_map.end(), - "found AccumulateGrad* is used by more than one tensors."); + TORCH_CHECK(grad_fn_to_tensor_map.find(grad_fn.get()) == grad_fn_to_tensor_map.end(), + "found AccumulateGrad* is used by more than one tensors."); grad_fn_to_tensor_map.insert({grad_fn.get(), &t}); } } @@ -103,14 +101,12 @@ void clear_grad_fns_for_next_edges(at::Tensor target, std::vector sa } } -void register_grad_fn(size_t ctx_address, at::Tensor target) -{ +void register_grad_fn(size_t ctx_address, at::Tensor target) { torch::autograd::AutogradMeta* autograd_meta = torch::autograd::impl::get_autograd_meta(target); PyNodeSharedPointerPool::GetInstance().RegisterGradFunc(ctx_address, autograd_meta); } -void unregister_grad_fn(size_t ctx_address) -{ +void unregister_grad_fn(size_t ctx_address) { PyNodeSharedPointerPool::GetInstance().UnRegisterGradFunc(ctx_address); } @@ -118,7 +114,7 @@ void unregister_grad_fn(size_t ctx_address) // When training program exits, PyNodeSharedPointerPool destructor is called, if grad_fns_ is not empty, // PyNode::release_variables() will be called. // (https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/python_function.cpp#L168) -// The other hand, there is known issue when acquiring GIL in pybind11 destructors, there will be probabbly deadlock issue. +// The other hand, there is known issue when acquiring GIL in pybind11 destructors, there will be probably deadlock issue. // (https://github.com/pybind/pybind11/issues/1446) // The resolution here, we remove all maintained states before program exits. @@ -126,13 +122,13 @@ void unregister_grad_fn(size_t ctx_address) // grad functions keeps accumulating without releasing, there might be memory (bound to those gradient function) leaks. // Ideally this usually won't happen in real training case, so it should be fine. -// We CANNOT explictly clear grad functions before each forward pass to mitigate the known issue above. +// We CANNOT explicitly clear grad functions before each forward pass to mitigate the known issue above. // For example: // loss1 = forward_run(inputs1) // loss2 = forward_run(inputs2) // loss = loss1 + loss2 // loss.backward() -// If we clear grad functions in the beggining of the second `forward_run`, when `loss.backward()` runs, +// If we clear grad functions in the beginning of the second `forward_run`, when `loss.backward()` runs, // the backward path of `loss1` will fail to run PythonOpGrad ops (if there is any). void clear_all_grad_fns() { PyNodeSharedPointerPool::GetInstance().ClearAll(); @@ -140,7 +136,8 @@ void clear_all_grad_fns() { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("register_grad_fn", ®ister_grad_fn, "increase grad_fn shared pointer reference."); - m.def("unregister_grad_fn", &unregister_grad_fn, "release grad_fn shared pointer referece."); - m.def("clear_all_grad_fns", &clear_all_grad_fns, "clear all grad_fn shared pointer refereces."); - m.def("clear_grad_fns_for_next_edges", &clear_grad_fns_for_next_edges, "remove reference on next edges' gradident funtions."); + m.def("unregister_grad_fn", &unregister_grad_fn, "release grad_fn shared pointer reference."); + m.def("clear_all_grad_fns", &clear_all_grad_fns, "clear all grad_fn shared pointer references."); + m.def("clear_grad_fns_for_next_edges", &clear_grad_fns_for_next_edges, + "remove reference on next edges' gradient functions."); } diff --git a/orttraining/orttraining/test/optimizer/compute_optimizer_test.cc b/orttraining/orttraining/test/optimizer/compute_optimizer_test.cc new file mode 100644 index 0000000000..82bc750e42 --- /dev/null +++ b/orttraining/orttraining/test/optimizer/compute_optimizer_test.cc @@ -0,0 +1,428 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include "core/graph/onnx_protobuf.h" + +#include "gtest/gtest.h" +#include "gmock/gmock.h" + +#include "asserts.h" +#include "core/common/span_utils.h" +#include "core/framework/data_types.h" +#include "core/framework/ort_value.h" +#include "core/graph/graph_utils.h" +#include "core/graph/graph_viewer.h" +#include "core/graph/model.h" +#include "core/optimizer/utils.h" +#include "core/util/math.h" +#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h" + +#include "test/common/tensor_op_test_utils.h" +#include "test/compare_ortvalue.h" +#include "test/framework/test_utils.h" +#include "test/optimizer/graph_transform_test_builder.h" +#include "test/optimizer/graph_transform_test_fixture.h" +#include "test/optimizer/test_optimizer_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/temp_dir.h" +#include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { + +#define MODEL_FOLDER ORT_TSTR("testdata/transform/") + +/* +Test graph include multiple equivalent subgraphs as below. + graph input [32, 256] (float) graph input [32] (int64_t) + | | + \_____________ _______/ graph input -1, scalar (int64_t) + \ / _______/ + \ / / + SCE Node, reduction = 'mean', output_type=1 + | + | + graph output, loss, scalar (float) +*/ +TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_Allowed) { + const logging::Logger* logger = &logging::LoggingManager::DefaultLogger(); + + for (const bool is_sce_internal : {true, false}) { + auto pre_graph_checker = [is_sce_internal](Graph& graph) -> Status { + auto op_count_pre = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_pre.size() == 1U); + + if (is_sce_internal) + TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); + else + TEST_RETURN_IF_NOT(op_count_pre["SoftmaxCrossEntropyLoss"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [is_sce_internal](Graph& graph) { + auto op_count_post = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_post.size() == 5U); + TEST_RETURN_IF_NOT(op_count_post["Sub"] == 1); + TEST_RETURN_IF_NOT(op_count_post["NonZero"] == 1); + TEST_RETURN_IF_NOT(op_count_post["Squeeze"] == 1); + TEST_RETURN_IF_NOT(op_count_post["com.microsoft.ShrunkenGather"] == 2); + + if (is_sce_internal) + TEST_RETURN_IF_NOT(op_count_post["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); + else + TEST_RETURN_IF_NOT(op_count_post["SoftmaxCrossEntropyLoss"] == 1); + + const NodeArg* squeeze_output_arg = nullptr; + for (Node& node : graph.Nodes()) { + if (node.OpType() == "Squeeze") { + squeeze_output_arg = node.OutputDefs()[0]; + break; + } + } + + TEST_RETURN_IF_NOT(squeeze_output_arg != nullptr); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "SoftmaxCrossEntropyLossInternal" || node.OpType() == "SoftmaxCrossEntropyLoss") { + const auto& input_defs = node.InputDefs(); + + { + auto producer_node = graph.GetProducerNode(input_defs[0]->Name()); + TEST_RETURN_IF_NOT(producer_node != nullptr); + TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather"); + TEST_RETURN_IF_NOT(producer_node->InputDefs()[1] == squeeze_output_arg); + } + + { + auto producer_node = graph.GetProducerNode(input_defs[1]->Name()); + TEST_RETURN_IF_NOT(producer_node != nullptr); + TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather"); + TEST_RETURN_IF_NOT(producer_node->InputDefs()[1] == squeeze_output_arg); + } + } + } + return Status::OK(); + }; + + auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) { + auto* input1_arg = builder.MakeInput({{32, 256}}); + auto* input2_arg = builder.MakeInput({{32}}); + auto* sce_out1 = builder.MakeOutput(); + NodeArg* empty = builder.MakeEmptyInput(); + auto* sce_out2 = builder.MakeIntermediate(); + + if (is_sce_internal) { + auto* ignore_index_arg = builder.MakeScalarInitializer(-100); + Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal", + {input1_arg, input2_arg, empty, ignore_index_arg}, + {sce_out1, sce_out2}, kMSDomain); + sce.AddAttribute("reduction", "mean"); + sce.AddAttribute("output_type", static_cast(1)); + } else { + Node& sce = builder.AddNode("SoftmaxCrossEntropyLoss", + {input1_arg, input2_arg, empty}, + {sce_out1, sce_out2}); + sce.AddAttribute("reduction", "mean"); + sce.AddAttribute("ignore_index", static_cast(-100)); + } + }; + + std::vector opsets{12, 13, 14, 15}; + for (auto opset : opsets) { + std::unique_ptr transformer = std::make_unique(); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer), + TransformerLevel::Level1, + 1, pre_graph_checker, post_graph_checker)); + } + } +} + +/* +Test graph include multiple equivalent subgraphs as below. + graph input [32, 256] (float) graph input [32] (int64_t) + | | + \_____________ _______/ graph input -1, scalar (int64_t) + \ / _______/ + \ / / + SCE Node, reduction = 'none', output_type=1 + | + | + graph output, loss, [32] (float) +*/ +TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_ReduceNone) { + const logging::Logger* logger = &logging::LoggingManager::DefaultLogger(); + + for (const bool is_sce_internal : {true, false}) { + auto pre_graph_checker = [is_sce_internal](Graph& graph) -> Status { + auto op_count_pre = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_pre.size() == 1U); + if (is_sce_internal) + TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); + else + TEST_RETURN_IF_NOT(op_count_pre["SoftmaxCrossEntropyLoss"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [is_sce_internal](Graph& graph) { + auto op_count_post = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_post.size() == 1U); + if (is_sce_internal) + TEST_RETURN_IF_NOT(op_count_post["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); + else + TEST_RETURN_IF_NOT(op_count_post["SoftmaxCrossEntropyLoss"] == 1); + return Status::OK(); + }; + + auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) { + auto* input1_arg = builder.MakeInput({{32, 256}}); + auto* input2_arg = builder.MakeInput({{32}}); + auto* sce_out1 = builder.MakeOutput(); + + NodeArg* empty = builder.MakeEmptyInput(); + auto* sce_out2 = builder.MakeIntermediate(); + + if (is_sce_internal) { + auto* ignore_index_arg = builder.MakeScalarInitializer(-100); + Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal", + {input1_arg, input2_arg, empty, ignore_index_arg}, + {sce_out1, sce_out2}, kMSDomain); + sce.AddAttribute("reduction", "none"); + sce.AddAttribute("output_type", static_cast(1)); + } else { + Node& sce = builder.AddNode("SoftmaxCrossEntropyLoss", + {input1_arg, input2_arg, empty}, + {sce_out1, sce_out2}); + sce.AddAttribute("reduction", "none"); + sce.AddAttribute("ignore_index", static_cast(-100)); + } + }; + + std::vector opsets{12, 13, 14, 15}; + for (auto opset : opsets) { + std::unique_ptr transformer = std::make_unique(); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer), + TransformerLevel::Level1, + 1, pre_graph_checker, post_graph_checker)); + } + } +} + +/* +Test graph include multiple equivalent subgraphs as below. + graph input [32, 256] (float) graph input [32] (int64_t) + | | + \_____________ _______/ graph input -1, scalar (int64_t) + \ / _______/ + \ / / + SCE Node, reduction = 'none', output_type=1 + | + | + graph output, loss, scalar (float) +*/ +TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_NoIgnoreIndex) { + const logging::Logger* logger = &logging::LoggingManager::DefaultLogger(); + + for (const bool is_sce_internal : {true, false}) { + auto pre_graph_checker = [is_sce_internal](Graph& graph) -> Status { + auto op_count_pre = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_pre.size() == 1U); + if (is_sce_internal) + TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); + else + TEST_RETURN_IF_NOT(op_count_pre["SoftmaxCrossEntropyLoss"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [is_sce_internal](Graph& graph) { + auto op_count_post = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_count_post.size() == 1U); + if (is_sce_internal) + TEST_RETURN_IF_NOT(op_count_post["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); + else + TEST_RETURN_IF_NOT(op_count_post["SoftmaxCrossEntropyLoss"] == 1); + return Status::OK(); + }; + + auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) { + auto* input1_arg = builder.MakeInput({{32, 256}}); + auto* input2_arg = builder.MakeInput({{32}}); + auto* sce_out1 = builder.MakeOutput(); + auto* sce_out2 = builder.MakeIntermediate(); + + if (is_sce_internal) { + Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal", + {input1_arg, input2_arg}, + {sce_out1, sce_out2}, kMSDomain); + sce.AddAttribute("reduction", "sum"); + sce.AddAttribute("output_type", static_cast(1)); + } else { + Node& sce = builder.AddNode("SoftmaxCrossEntropyLoss", + {input1_arg, input2_arg}, + {sce_out1, sce_out2}); + sce.AddAttribute("reduction", "sum"); + } + }; + + std::vector opsets{12, 13, 14, 15}; + for (auto opset : opsets) { + std::unique_ptr transformer = std::make_unique(); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer), + TransformerLevel::Level1, + 1, pre_graph_checker, post_graph_checker)); + } + } +} + +TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_MlmBertE2E) { + const logging::Logger* logger = &logging::LoggingManager::DefaultLogger(); + // Be noted all dropout have a ratio be 0.0, to make it easier to compare when running with the session. + // This did not affect the transformer tests, because we did not remove the Dropout of ratio 0. in the middle. + auto model_uri = MODEL_FOLDER "computation_reduction/reshape/mlm_bert_e2e.onnx"; + std::shared_ptr model; + ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger)); + Graph& graph = model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{3}; + ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique(), + TransformerLevel::Level1)); + ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger)); + + { + const NodeArg* squeeze_output_arg = nullptr; + for (Node& node : graph.Nodes()) { + if (node.OpType() == "NonZero") { + const std::vector& consumers = graph.GetConsumerNodes(node.OutputDefs()[0]->Name()); + ASSERT_TRUE(consumers.size() == 1); + ASSERT_TRUE(consumers[0]->OpType() == "Squeeze"); + squeeze_output_arg = consumers[0]->OutputDefs()[0]; + break; + } + } + + ASSERT_TRUE(squeeze_output_arg != nullptr); + + for (Node& node : graph.Nodes()) { + if (node.OpType() == "SoftmaxCrossEntropyLossInternal") { + const auto& input_defs = node.InputDefs(); + + { + auto producer_node = graph.GetProducerNode(input_defs[0]->Name()); + ASSERT_TRUE(producer_node != nullptr); + ASSERT_TRUE(producer_node->OpType() == "ShrunkenGather"); + ASSERT_TRUE(producer_node->InputDefs()[1] == squeeze_output_arg); + } + + { + auto producer_node = graph.GetProducerNode(input_defs[1]->Name()); + ASSERT_TRUE(producer_node != nullptr); + ASSERT_TRUE(producer_node->OpType() == "ShrunkenGather"); + ASSERT_TRUE(producer_node->InputDefs()[1] == squeeze_output_arg); + } + } + } + } + + onnxruntime::test::TemporaryDirectory tmp_dir{ORT_TSTR("compute_optimizer_test_tmp_dir")}; + PathString new_model_uri{ConcatPathComponent( + tmp_dir.Path(), + ORT_TSTR("insert_gather_before_sceloss_bert_e2e_optimized.onnx"))}; + ASSERT_STATUS_OK(Model::Save(*model, new_model_uri)); + + int64_t batch_size = 8; + int64_t sequence_length = 16; + int64_t hidden_size = 1024; + + InputContainer input_container; + + input_container.AddInput("input", {batch_size, sequence_length, hidden_size}, RandomFillFloatVector); + + const TensorShapeVector dims_mask = {batch_size, sequence_length}; + std::vector attention_mask(TensorShape(dims_mask).Size(), 0); + RandomMasks(batch_size, sequence_length, attention_mask); + input_container.AddInput("attention_mask", dims_mask, attention_mask); + + input_container.AddInput("matmul1.weight", {hidden_size, 1024}, RandomFillHalfVector); + input_container.AddInput("add1.bias", {1024}, RandomFillHalfVector); + + input_container.AddInput("matmul2.weight", {hidden_size, 1024}, RandomFillHalfVector); + input_container.AddInput("add2.bias", {1024}, RandomFillHalfVector); + + input_container.AddInput("matmul3.weight", {hidden_size, 1024}, RandomFillHalfVector); + input_container.AddInput("add3.bias", {1024}, RandomFillHalfVector); + + input_container.AddInput("matmul4.weight", {hidden_size, 1024}, RandomFillHalfVector); + input_container.AddInput("add4.bias", {1024}, RandomFillHalfVector); + + input_container.AddInput("layer_norm1.weight", {hidden_size}, RandomFillFloatVector); + input_container.AddInput("layer_norm1.bias", {hidden_size}, RandomFillFloatVector); + + input_container.AddInput("matmul7.weight", {hidden_size, hidden_size * 4}, RandomFillHalfVector); + input_container.AddInput("add7.bias", {hidden_size * 4}, RandomFillHalfVector); + + input_container.AddInput("matmul8.weight", {hidden_size * 4, hidden_size}, RandomFillHalfVector); + input_container.AddInput("add8.bias", {hidden_size}, RandomFillHalfVector); + + input_container.AddInput("layer_norm2.weight", {hidden_size}, RandomFillFloatVector); + input_container.AddInput("layer_norm2.bias", {hidden_size}, RandomFillFloatVector); + + input_container.AddInput("matmul9.weight", {hidden_size, hidden_size}, RandomFillHalfVector); + input_container.AddInput("add9.bias", {hidden_size}, RandomFillHalfVector); + + input_container.AddInput("layer_norm3.weight", {hidden_size}, RandomFillFloatVector); + input_container.AddInput("layer_norm3.bias", {hidden_size}, RandomFillFloatVector); + + input_container.AddInput("matmul10.weight", {hidden_size, 30522}, RandomFillHalfVector); + input_container.AddInput("add10.bias", {30522}, RandomFillHalfVector); + + const TensorShapeVector dims_labels = {batch_size * sequence_length}; + static RandomValueGenerator random{8910}; + std::vector labels = random.Uniform(dims_labels, 0, 30522); + const std::vector num_count_to_random{batch_size}; + std::vector random_seq_lens = random.Uniform(num_count_to_random, 0, sequence_length); + for (int64_t i = 0; i < batch_size; ++i) { + for (int64_t j = 0; j < sequence_length; ++j) { + if (j > random_seq_lens[i]) { + labels[i * sequence_length + j] = -100; + } + } + } + + input_container.AddInput("labels", dims_labels, labels); + + const std::string all_provider_types[] = { + onnxruntime::kCpuExecutionProvider, +#ifdef USE_CUDA + onnxruntime::kCudaExecutionProvider, +#elif USE_ROCM + onnxruntime::kRocmExecutionProvider, +#endif + }; + + const std::vector output_names = {"output-1"}; + + for (auto& provider_type : all_provider_types) { + std::vector expected_ort_values; + RunModelWithData(model_uri, std::string("RawGraphRun"), provider_type, + input_container, output_names, expected_ort_values); + + std::vector actual_ort_values; + RunModelWithData(ToPathString(new_model_uri), std::string("OptimizedGraphRun"), + provider_type, input_container, output_names, actual_ort_values); + + ASSERT_TRUE(expected_ort_values.size() == actual_ort_values.size()); + + constexpr double per_sample_tolerance = 1e-4; + constexpr double relative_per_sample_tolerance = 1e-4; + for (size_t i = 0; i < expected_ort_values.size(); i++) { + auto ret = CompareOrtValue(actual_ort_values[i], expected_ort_values[i], + per_sample_tolerance, relative_per_sample_tolerance, false); + EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second; + } + } +} + +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py index c6bf327be0..33821843e4 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py @@ -126,7 +126,7 @@ def test_gelu_custom_func_rets_not_as_module_output(): def forward(self, model_input): out = self.relu(model_input, self.bias) # add * 9 by intention to make custom function's output - # NOT as module outputs (which are consumed by subsquent computations). + # NOT as module outputs (which are consumed by subsequent computations). # This aims to trigger a GC for "out", saying, out is released, # the underlying std::shared still have other references. # Otherwise, a segment fault will be triggered.