From cfe830b248c8913ec989dfb1eb7b8310f0630efe Mon Sep 17 00:00:00 2001 From: guyang3532 <62738430+guyang3532@users.noreply.github.com> Date: Fri, 10 May 2024 21:55:43 +0800 Subject: [PATCH] Generalize label input sparsity check and refactor (#20636) ### Description The InsertGatherBeforeSceLoss optimization is enabled when the density of label padding less than 90%. We need to check the density of the label padding to decide whether enable the optimization. Before this pr, we just check the inputs of graph and correlate one with the SCE node by iterate graph from the SCE node back to one graph input. This is hard to be general because there may be complicated pattern between graph input and SCE node. This pr check padding density by the direct input of SCE module rather than the input of graph at the first graph execution when exporting onnx graph. And if the density < 90%, insert a flag PythonOp after the SCE node as: ``` SoftmaxCrossEntropy | PythonOp (func_name: FlagAndPrintDensity) (insert if density < 90%) | Following graph ``` When the InsertGatherBeforeSceLoss is invoked, it check if there is the flag PythonOp(func_name: FlagAndPrintDensity) after the SCE node and if it is, remove it and do the padding elimination optimization. If the env of ORTMODULE_PRINT_INPUT_DENSITY is 1, we will print input density each step by the PythonOp (func_name: FlagAndPrintDensity). In this case the PythonOp will not be removed. --- docs/ORTModule_Training_Guidelines.md | 24 +- .../compute_optimizer/padding_elimination.cc | 34 +- .../compute_optimizer/padding_elimination.h | 9 +- .../sceloss_compute_optimization.cc | 43 +- .../sceloss_compute_optimization.h | 6 +- .../core/optimizer/graph_transformer_config.h | 3 +- .../core/optimizer/graph_transformer_utils.cc | 5 +- .../python/orttraining_pybind_state.cc | 2 +- .../ortmodule/_graph_execution_manager.py | 169 +++---- .../training/ortmodule/_inference_manager.py | 7 +- .../python/training/ortmodule/_io.py | 11 +- .../training/ortmodule/_runtime_inspector.py | 426 +----------------- .../training/ortmodule/_training_manager.py | 7 +- .../python/training/ortmodule/options.py | 21 +- .../test/optimizer/compute_optimizer_test.cc | 138 ++++-- .../python/orttraining_test_ortmodule_api.py | 11 +- 16 files changed, 306 insertions(+), 610 deletions(-) diff --git a/docs/ORTModule_Training_Guidelines.md b/docs/ORTModule_Training_Guidelines.md index 1609aa7ae0..8d5472ba30 100644 --- a/docs/ORTModule_Training_Guidelines.md +++ b/docs/ORTModule_Training_Guidelines.md @@ -208,19 +208,6 @@ debugging). export ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=0 # Disable ``` -#### ORTMODULE_ENABLE_SPARSE_OPTIMIZER - -- **Feature Area**: *ORTMODULE/Optimizations* -- **Description**: By default, this is enabled. This env var can be used for enabling or disabling the input data sparsity -based performance optimizations, including embedding sparsity and label sparsity. -This optimization is applicable when using optimum, which has an implementation of the ModuleWithLoss class that wraps the HuggingFace Training that allows loss computation inside ONNX Runtime (ORT). -If you're not using optimum but want to implement a similar wrapper in your codebase to compute the loss inside ONNX Runtime (ORT), you can refer to this [Link](ORTModule_ModuleWithLoss_Wrapper.md) for detailed steps and guidelines on how to achieve this. - - ```bash - export ORTMODULE_ENABLE_SPARSE_OPTIMIZER=1 # Enable - export ORTMODULE_ENABLE_SPARSE_OPTIMIZER=0 # Disable - ``` - #### ORTMODULE_PRINT_INPUT_DENSITY - **Feature Area**: *ORTMODULE/RuntimeInspector* @@ -254,6 +241,17 @@ data sparsity based performance optimizations. export ORTMODULE_ENABLE_EMBEDDING_SPARSE_OPTIMIZER=0 # Disable ``` +#### ORTMODULE_ENABLE_LABEL_SPARSE_OPTIMIZER + +- **Feature Area**: *ORTMODULE/Optimizations* +- **Description**: By default, this is enabled. This env var can be used for enabling or disabling the label input +data sparsity based performance optimizations. + + ```bash + export ORTMODULE_ENABLE_LABEL_SPARSE_OPTIMIZER=1 # Enable + export ORTMODULE_ENABLE_LABEL_SPARSE_OPTIMIZER=0 # Disable + ``` + #### ORTMODULE_CACHE_DIR - **Feature Area**: *ORTMODULE/RuntimeOptions* diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc index 07ec7e17b2..9209b5fdea 100644 --- a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc @@ -21,8 +21,8 @@ constexpr const char* kInspectActivationFuncName = "onnxruntime.training.utils.hooks._statistics_subscriber._InspectActivation"; constexpr const char* kIncrementStepFuncName = "onnxruntime.training.utils.hooks._subscriber_manager._IncrementStep"; -constexpr const char* kFlagPaddingEliminationFuncName = - "onnxruntime.training.ortmodule._runtime_inspector.FlagPaddingElimination"; +constexpr const char* kFlagAndPrintDensityFuncName = + "onnxruntime.training.ortmodule._runtime_inspector.FlagAndPrintDensity"; void PushAllOutputNode(Graph& graph, std::queue& q, Node* node, std::unordered_set& visited) { for (auto iter = node->OutputNodesBegin(); iter != node->OutputNodesEnd(); ++iter) { @@ -396,26 +396,28 @@ Status PaddingElimination::ApplyImpl(Graph& graph, bool& modified, int graph_lev if (outputNodeCount != 1) { continue; } - auto embedding_output_node = graph.GetNode(node.OutputNodesBegin()->Index()); - if (embedding_output_node == nullptr || - !graph_utils::IsSupportedOptypeVersionAndDomain(*embedding_output_node, "PythonOp", {1}, kMSDomain) || - static_cast(embedding_output_node->GetAttributes().at("func_name").s()) != - kFlagPaddingEliminationFuncName) { + Node* embedding_input_node = graph.GetMutableProducerNode(node.MutableInputDefs()[1]->Name()); + if (embedding_input_node == nullptr || + !graph_utils::IsSupportedOptypeVersionAndDomain(*embedding_input_node, "PythonOp", {1}, kMSDomain) || + static_cast(embedding_input_node->GetAttributes().at("func_name").s()) != + kFlagAndPrintDensityFuncName) { LOG_DEBUG_INFO(logger, "not find PythonOp of flagPaddingElimination after embedding node"); continue; } - if (graph_utils::CanRemoveNode(graph, *embedding_output_node, logger)) { - if (graph_utils::RemoveNode(graph, *embedding_output_node)) { - modified = true; + if (!print_density_) { + if (graph_utils::CanRemoveNode(graph, *embedding_input_node, logger)) { + if (graph_utils::RemoveNode(graph, *embedding_input_node)) { + modified = true; + } else { + LOG_DEBUG_INFO(logger, "Failed to remove node " + embedding_input_node->Name() + + "(" + embedding_input_node->OpType() + ")"); + continue; + } } else { - LOG_DEBUG_INFO(logger, "Failed to remove node " + embedding_output_node->Name() + - "(" + embedding_output_node->OpType() + ")"); + LOG_DEBUG_INFO(logger, "Can not remove node " + embedding_input_node->Name() + + "(" + embedding_input_node->OpType() + ")"); continue; } - } else { - LOG_DEBUG_INFO(logger, "Can not remove node " + embedding_output_node->Name() + - "(" + embedding_output_node->OpType() + ")"); - continue; } const ONNX_NAMESPACE::TensorProto* padding_initializer = graph_utils::GetConstantInitializer(graph, node.InputDefs()[2]->Name()); diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.h b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.h index 0dd62be142..cc3c90dac2 100644 --- a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.h +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.h @@ -127,10 +127,15 @@ namespace onnxruntime { */ class PaddingElimination : public GraphTransformer { public: - explicit PaddingElimination(const InlinedHashSet& compatible_execution_providers = {}) noexcept - : GraphTransformer("PaddingElimination", compatible_execution_providers) {} + explicit PaddingElimination(const InlinedHashSet& compatible_execution_providers = {}, + const bool print_input_density = false) noexcept + : GraphTransformer("PaddingElimination", compatible_execution_providers), + print_density_(print_input_density) {} Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; + + private: + bool print_density_ = false; }; } // namespace onnxruntime diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc index 71aa1366bc..3c809ce801 100644 --- a/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc @@ -16,15 +16,16 @@ namespace onnxruntime { +namespace { + +constexpr const char* kFlagAndPrintDensityFuncName = + "onnxruntime.training.ortmodule._runtime_inspector.FlagAndPrintDensity"; +} // namespace + Status InsertGatherBeforeSceLoss::ApplyImpl(Graph& graph, bool& modified, int /*graph_level*/, const logging::Logger& logger) const { LOG_DEBUG_INFO(logger, "Enter InsertGatherBeforeSceLoss"); - if (sparse_label_input_names_.size() == 0) { - LOG_DEBUG_INFO(logger, "Exit InsertGatherBeforeSceLoss, no sparse label input names."); - return Status::OK(); - } - GraphViewer graph_viewer(graph); [[maybe_unused]] size_t handled_sce_node_count = 0; // For summary const auto& order = graph_viewer.GetNodesInTopologicalOrder(); @@ -48,7 +49,7 @@ Status InsertGatherBeforeSceLoss::ApplyImpl(Graph& graph, bool& modified, int /* const NodeArg* label_input_arg = node.InputDefs()[1]; // Check whether this SCE node is handled or not. - const Node* labels_producer = graph.GetProducerNode(label_input_arg->Name()); + Node* labels_producer = graph.GetMutableProducerNode(label_input_arg->Name()); // Skip if already inserted a ShrunkenGather node. if (labels_producer && graph_utils::IsSupportedOptypeVersionAndDomain( *labels_producer, "ShrunkenGather", {1}, kMSDomain)) { @@ -57,18 +58,28 @@ Status InsertGatherBeforeSceLoss::ApplyImpl(Graph& graph, bool& modified, int /* continue; } - // Label input can be a graph input or from a Reshape node taking a graph input as its data input. - if (labels_producer && graph_utils::IsSupportedOptypeVersionAndDomain( - *labels_producer, "Reshape", {1, 5, 13, 14}, kOnnxDomain)) { - label_input_arg = labels_producer->InputDefs()[0]; - } - // Then check if the label input is graph input and in the sparse label input list. - if (!graph.IsInputsIncludingInitializers(label_input_arg) || - std::find(sparse_label_input_names_.begin(), sparse_label_input_names_.end(), - label_input_arg->Name()) == sparse_label_input_names_.end()) { + if (labels_producer == nullptr || + !graph_utils::IsSupportedOptypeVersionAndDomain(*labels_producer, "PythonOp", {1}, kMSDomain) || + static_cast(labels_producer->GetAttributes().at("func_name").s()) != + kFlagAndPrintDensityFuncName) { LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() + - ") due to labels input is not a graph input or not in the sparse label input list."); + ") due to labels input is not produced by a PythonOp node with flag " + + kFlagAndPrintDensityFuncName + "."); continue; + } else if (!print_density_) { + if (graph_utils::CanRemoveNode(graph, *labels_producer, logger)) { + if (graph_utils::RemoveNode(graph, *labels_producer)) { + modified = true; + } else { + LOG_DEBUG_INFO(logger, "Failed to remove node " + labels_producer->Name() + + "(" + labels_producer->OpType() + ")"); + continue; + } + } else { + LOG_DEBUG_INFO(logger, "Can not remove node " + labels_producer->Name() + + "(" + labels_producer->OpType() + ")"); + continue; + } } // Check shape requirements. diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h index 94787aca7c..cf34706115 100644 --- a/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h @@ -68,9 +68,9 @@ namespace onnxruntime { class InsertGatherBeforeSceLoss : public GraphTransformer { public: InsertGatherBeforeSceLoss(const InlinedHashSet& compatible_execution_providers = {}, - const std::vector& sparse_label_input_names = {}) noexcept + const bool print_input_density = false) noexcept : GraphTransformer("InsertGatherBeforeSceLoss", compatible_execution_providers), - sparse_label_input_names_{sparse_label_input_names} { + print_density_(print_input_density) { } /** @@ -79,7 +79,7 @@ class InsertGatherBeforeSceLoss : public GraphTransformer { Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; private: - std::vector sparse_label_input_names_; + bool print_density_ = false; }; } // namespace onnxruntime diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_config.h b/orttraining/orttraining/core/optimizer/graph_transformer_config.h index f6c1450397..f72dbfa3fd 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_config.h +++ b/orttraining/orttraining/core/optimizer/graph_transformer_config.h @@ -25,8 +25,7 @@ struct TrainingGraphTransformerConfiguration : public GraphTransformerConfigurat // Enable compute optimizer. bool enable_compute_optimizer{false}; - // Enable label sparsity compute optimization for the input names in the below list. - std::vector sparse_label_input_names; + bool print_input_density{false}; // Path for serialization of the transformed optimized model. If empty, serialization is disabled. std::string optimized_pre_grad_filepath; diff --git a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc index 5a3db9454c..436a24c34d 100644 --- a/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc +++ b/orttraining/orttraining/core/optimizer/graph_transformer_utils.cc @@ -195,11 +195,12 @@ std::vector> GeneratePreTrainingTransformers( transformers.emplace_back(std::make_unique(compatible_eps)); transformers.emplace_back(std::make_unique(compatible_eps)); transformers.emplace_back(std::make_unique(compatible_eps, - config.sparse_label_input_names)); + config.print_input_density)); #if defined(USE_CUDA) || defined(USE_ROCM) // Put this under CUDA/ROCM guard as it depends on PadAndUnflatten CUDA/ROCM kernel. // Once we have a CPU kernel for PadAndUnflatten, we can remove the guard. - transformers.emplace_back(std::make_unique(compatible_eps)); + transformers.emplace_back(std::make_unique(compatible_eps, + config.print_input_density)); transformers.emplace_back(std::make_unique(compatible_eps)); #endif } diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index c83f540f3f..3f91dc1065 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -487,7 +487,7 @@ void addObjectMethodsForTraining(py::module& m) { .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("sparse_label_input_names", &TrainingGraphTransformerConfiguration::sparse_label_input_names) + .def_readwrite("print_input_density", &TrainingGraphTransformerConfiguration::print_input_density) .def_readwrite("optimized_pre_grad_filepath", &TrainingGraphTransformerConfiguration::optimized_pre_grad_filepath) .def_readwrite("propagate_cast_ops_config", &TrainingGraphTransformerConfiguration::GraphTransformerConfiguration::propagate_cast_ops_config); diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 5316c0cd5f..ad726a82ef 100755 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -9,6 +9,7 @@ import io import logging import os from abc import ABC, abstractmethod # noqa: F401 +from functools import partial from hashlib import md5 as hash_fn from typing import Dict, List, Optional, Tuple @@ -34,7 +35,7 @@ from ._gradient_accumulation_manager import GradientAccumulationManager from ._graph_execution_interface import GraphExecutionInterface from ._io import _FlattenedModule, _InputInfo from ._logger import LogColor -from ._runtime_inspector import FlagPaddingElimination, RuntimeInspector +from ._runtime_inspector import FlagAndPrintDensity, RuntimeInspector from ._utils import check_function_has_param, get_rank from .options import DebugOptions, LogLevel, _MemoryOptimizationLevel, _RuntimeOptions from .torch_cpp_extensions.cpu.aten_op_executor import load_aten_op_executor_cpp_extension @@ -310,8 +311,8 @@ class GraphExecutionManager(GraphExecutionInterface): return False self._set_device_from_module(inputs, kwargs) - # TODO: move it into runtime_inspector embedding_hook_handles = self._add_check_embedding_sparsity_hook() + label_hook_handles = self._add_check_label_sparsity_hook() from onnxruntime.training.utils.hooks._subscriber_manager import no_increase_global_step @@ -320,6 +321,8 @@ class GraphExecutionManager(GraphExecutionInterface): for hook in embedding_hook_handles: hook.remove() + for hook in label_hook_handles: + hook.remove() if self._debug_options.save_onnx_models.save: self._onnx_models.save_exported_model( @@ -547,6 +550,7 @@ class GraphExecutionManager(GraphExecutionInterface): graph_transformer_config.propagate_cast_ops_config.allow = self._runtime_options.propagate_cast_ops_allow graph_transformer_config.propagate_cast_ops_config.strategy = self._runtime_options.propagate_cast_ops_strategy graph_transformer_config.enable_compute_optimizer = self._runtime_options.enable_compute_optimizer + graph_transformer_config.print_input_density = self._runtime_options.print_input_density if self._debug_options.save_onnx_models.save: graph_transformer_config.optimized_pre_grad_filepath = os.path.join( @@ -686,21 +690,17 @@ class GraphExecutionManager(GraphExecutionInterface): """ Add hook to check embedding sparsity and enable padding elimination if applicable. 1. Iterate through all modules to find Embedding modules with padding_idx >= 0. - 2. Register forward hook to the Embedding module and the hook will check sparsity of the embedding input. - 3. If the sparsity is below a threshold, enable padding elimination by adding FlagPaddingElimination after the - output. GraphTransformer of PaddingElimination will check the FlagPaddingElimination and do the actual + 2. Register forward pre hook to the Embedding module and the hook will check sparsity of the embedding input. + 3. If the sparsity is below a threshold, enable padding elimination by adding FlagAndPrintDensity after the + output. GraphTransformer of PaddingElimination will check the FlagAndPrintDensity and do the actual padding elimination graph modification. 4. Return the hook handles for later removal. """ - if ( - not self._runtime_options.enable_sparse_optimizer - or not self._runtime_options.enable_embedding_sparse_optimizer - or self._device.type != "cuda" - ): + if not self._runtime_options.enable_embedding_sparse_optimizer or self._device.type != "cuda": return [] - def _embedding_hook(module, args, output): + def _embedding_hook(name, module, args): ebd_input = args[0] if ebd_input is None or not isinstance(ebd_input, torch.Tensor): self._logger.warning("Embedding input is not a tensor.") @@ -709,19 +709,11 @@ class GraphExecutionManager(GraphExecutionInterface): valid_token = torch.count_nonzero(ebd_input - module.padding_idx) total_token = ebd_input.numel() embed_density = float(valid_token) / float(total_token) * 100 - if module not in self._runtime_inspector._embedding_module_to_padding_density_map: - self._logger.warning("Found Embedding module not in the map. %s", module) - return None if embed_density < 90: self._logger.info("Embedding sparsity-based optimization is ON for density: %.0f%%", embed_density) - if self._runtime_inspector._embedding_module_to_padding_density_map[module][1] != -1: - self._logger.warning( - "Found duplicate Embedding module. %s", - self._runtime_inspector._embedding_module_to_padding_density_map[module][0], - ) - self._runtime_inspector._embedding_module_to_padding_density_map[module][1] = embed_density - return FlagPaddingElimination.apply(output) + self._runtime_inspector._embedding_module_to_padding_density_map[name] = embed_density + return FlagAndPrintDensity.apply(args[0], module.padding_idx, "embedding") else: self._logger.info("Embedding sparsity-based optimization is OFF for density: %.0f%%", embed_density) return None @@ -730,15 +722,49 @@ class GraphExecutionManager(GraphExecutionInterface): for name, sub_module in self._flattened_module.named_modules(): if isinstance(sub_module, torch.nn.modules.sparse.Embedding): if sub_module.padding_idx is not None and sub_module.padding_idx >= 0: - self._runtime_inspector._embedding_module_to_padding_density_map[sub_module] = [name, -1] - embedding_hook_handles.append(sub_module.register_forward_hook(_embedding_hook)) + embedding_hook_handles.append(sub_module.register_forward_pre_hook(partial(_embedding_hook, name))) return embedding_hook_handles + def _add_check_label_sparsity_hook(self): + """ + Add hook to check label sparsity and enable sceloss compute optimization if applicable. + 1. Register forward pre hook to the sceloss module in the model and the hook will check sparsity of the label input. + 2. If the sparsity is below a threshold, enable sceloss compute optimization by adding FlagAndPrintDensity after the + output. GraphTransformer of InsertGatherBeforeSceLoss will check the FlagAndPrintDensity and do the actual + sceloss compute optimization graph modification. + + """ + if not self._runtime_options.enable_label_sparse_optimizer: + return None + + def _label_hook(name, module, args): + label_input = args[1] + if label_input is None or not isinstance(label_input, torch.Tensor): + self._logger.warning("Label input is not a tensor.") + return None + + valid_token = torch.count_nonzero(label_input - module.ignore_index) + total_token = label_input.numel() + label_density = float(valid_token) / float(total_token) * 100 + + if label_density < 90: + self._logger.info("Label sparsity-based optimization is ON for density: %.0f%%", label_density) + self._runtime_inspector._sceloss_module_to_ignore_density_map[name] = label_density + return (args[0], FlagAndPrintDensity.apply(args[1], module.ignore_index, "label")) + else: + self._logger.info("Label sparsity-based optimization is OFF for density: %.0f%%", label_density) + return None + + label_check_hook_handles = [] + for name, sub_module in self._flattened_module.named_modules(): + if isinstance(sub_module, torch.nn.modules.loss.CrossEntropyLoss): + label_check_hook_handles.append(sub_module.register_forward_pre_hook(partial(_label_hook, name))) + + return label_check_hook_handles + @_logger.TrackTime(_logger.ORTModuleInitPhase.DETECTION) - def _enable_conditional_optimizations( - self, graph_transformer_config: C.TrainingGraphTransformerConfiguration, inputs: Tuple, kwargs: Dict - ): + def _detect_from_inputs(self, inputs: Tuple, kwargs: Dict): """ Based on runtime inspection, enable conditional optimizations if applicable. @@ -749,63 +775,44 @@ class GraphExecutionManager(GraphExecutionInterface): enable sparsity-based optimization. """ - # Enable data sparsity inspection if sparse optimizer is ON or user wants to print input density. - if self._runtime_options.enable_sparse_optimizer or self._runtime_options.print_input_density: - self._runtime_inspector.enable_input_inspector( - self._onnx_models.processed_exported_model, self._graph_builder.get_graph_info().user_input_names + detected_device = _utils.get_device_from_module(self._original_module) or _utils.get_device_from_inputs( + inputs, kwargs + ) + + if self._runtime_options.enable_zero_stage3_support or self._mem_efficient_grad_management_is_enabled: + self._append_pull_weight_trigger_as_input(kwargs, detected_device) + + param_to_append_as_onnx_graph_inputs = [] + if self._mem_efficient_grad_management_is_enabled: + from ._mem_efficient_grad_mgmt import get_params_not_connected_to_pull_param_trigger + + param_to_append_as_onnx_graph_inputs = get_params_not_connected_to_pull_param_trigger( + self._flattened_module.named_parameters(), self._onnx_models.exported_model + ) + else: + param_to_append_as_onnx_graph_inputs = self._graph_initializers + + _io._combine_input_buffers_initializers( + param_to_append_as_onnx_graph_inputs, + self._graph_builder.get_graph_info().user_input_names, + self._input_info, + self._flattened_module.named_buffers(), + inputs, + kwargs, + detected_device, + self._runtime_inspector, + self._zero_stage3_param_map, + ) + + if self._runtime_inspector._sceloss_module_to_ignore_density_map: + self._runtime_options.label_sparsity_ratio = ",".join( + [f"{k}:{v:.0f}%" for k, v in self._runtime_inspector._sceloss_module_to_ignore_density_map.items()] ) - if self._runtime_options.enable_sparse_optimizer: - detected_device = _utils.get_device_from_module(self._original_module) or _utils.get_device_from_inputs( - inputs, kwargs - ) - - if self._runtime_options.enable_zero_stage3_support or self._mem_efficient_grad_management_is_enabled: - self._append_pull_weight_trigger_as_input(kwargs, detected_device) - - param_to_append_as_onnx_graph_inputs = [] - if self._mem_efficient_grad_management_is_enabled: - from ._mem_efficient_grad_mgmt import get_params_not_connected_to_pull_param_trigger - - param_to_append_as_onnx_graph_inputs = get_params_not_connected_to_pull_param_trigger( - self._flattened_module.named_parameters(), self._onnx_models.exported_model - ) - else: - param_to_append_as_onnx_graph_inputs = self._graph_initializers - - _, _, label_sparsity_results = _io._combine_input_buffers_initializers( - param_to_append_as_onnx_graph_inputs, - self._graph_builder.get_graph_info().user_input_names, - self._input_info, - self._flattened_module.named_buffers(), - inputs, - kwargs, - detected_device, - self._runtime_inspector, - self._zero_stage3_param_map, - ) - - # Enable sparsity-based optimization when applicable. - if len(label_sparsity_results) > 0: - graph_transformer_config.sparse_label_input_names = list(label_sparsity_results.keys()) - self._logger.info("Label sparsity-based optimization is ON for %s", label_sparsity_results) - self._runtime_options.label_sparsity_ratio = ",".join( - [f"{k}:{v:.0f}%" for k, v in label_sparsity_results.items()] - ) - - if self._runtime_inspector._embedding_module_to_padding_density_map: - self._runtime_options.embed_sparsity_ratio = ",".join( - [ - f"{v[0]}:{v[1]:.0f}%" - for v in self._runtime_inspector._embedding_module_to_padding_density_map.values() - if v[1] != -1 - ] - ) - - # If users don't want to print input density, disable the input density observer to avoid overhead - # when looping through inputs during training. - if not self._runtime_options.print_input_density: - self._runtime_inspector.disable_input_inspector() + if self._runtime_inspector._embedding_module_to_padding_density_map: + self._runtime_options.embed_sparsity_ratio = ",".join( + [f"{k}:{v:.0f}%" for k, v in self._runtime_inspector._embedding_module_to_padding_density_map.items()] + ) def _append_pull_weight_trigger_as_input(self, kwargs: Dict, device: torch.device): if self._runtime_options.enable_zero_stage3_support: diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index 13145c7c79..642dc9b0f4 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -121,10 +121,9 @@ class InferenceManager(GraphExecutionManager): # Build the inference graph if build_graph: - graph_transformer_config = self._get_graph_transformer_config() - # Set the config according to input inspection. - self._enable_conditional_optimizations(graph_transformer_config, inputs, kwargs) + self._detect_from_inputs(inputs, kwargs) + graph_transformer_config = self._get_graph_transformer_config() # Build the graph self._build_graph(graph_transformer_config) @@ -161,7 +160,7 @@ class InferenceManager(GraphExecutionManager): if self._runtime_options.enable_zero_stage3_support: self._append_pull_weight_trigger_as_input(kwargs, self._device) - prepared_input_list, _, _ = _io._combine_input_buffers_initializers( + prepared_input_list = _io._combine_input_buffers_initializers( self._graph_initializers, self._graph_info.user_input_names, self._input_info, diff --git a/orttraining/orttraining/python/training/ortmodule/_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py index 7534cc46a2..1ba62194bf 100644 --- a/orttraining/orttraining/python/training/ortmodule/_io.py +++ b/orttraining/orttraining/python/training/ortmodule/_io.py @@ -208,8 +208,6 @@ def _combine_input_buffers_initializers( _expand_inputs(kwargs, flattened_kwargs_inputs) buffer_names_dict = None result = [] - embed_sparsity_results = OrderedDict() - label_sparsity_results = OrderedDict() onnx_input_to_value_map = OrderedDict() for input_idx, name in enumerate(onnx_input_names): @@ -245,14 +243,7 @@ def _combine_input_buffers_initializers( if PrimitiveType.is_primitive_type(inp): inp = PrimitiveType.get_tensor(inp, device) - found, embedding_density, label_density = rt_inspector.inspect_input(name, inp) - if found: - if embedding_density < 100: - embed_sparsity_results[name] = embedding_density - if label_density < 100: - label_sparsity_results[name] = label_density result.append(inp) - onnx_input_to_value_map[name] = inp else: raise wrap_exception( @@ -271,7 +262,7 @@ def _combine_input_buffers_initializers( rt_inspector.memory_ob.collect_symbolic_dim_values(input_info.dynamic_axes, onnx_input_to_value_map) rt_inspector.memory_ob.symbolic_dim_collecting_completed = True - return result, embed_sparsity_results, label_sparsity_results + return result def deepcopy_model_input( diff --git a/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py b/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py index 74f27d6fe4..fb1a26661b 100644 --- a/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py +++ b/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py @@ -9,8 +9,6 @@ from typing import Dict, List, Optional, Tuple, Union import onnx import torch -from onnx import ModelProto, helper -from onnx import onnx_pb as onnx_proto from sympy import Symbol, simplify from sympy.parsing.sympy_parser import parse_expr @@ -56,415 +54,9 @@ class RuntimeInspector: """ self._logger = logger - self.input_density_ob: Union[InputDensityObserver, None] = None self.memory_ob = MemoryObserver(module, self._logger, training) self._embedding_module_to_padding_density_map = {} - - def enable_input_inspector(self, model: ModelProto, user_input_names: List[str]) -> None: - """Initialize input inspector from the given ONNX model and user input names. - - Args: - model: ONNX model. - user_input_names: User input names in the ONNX model. - - """ - if self.input_density_ob is None: - self.input_density_ob = InputDensityObserver(self._logger) - else: - raise RuntimeError("Input density observer is already enabled.") - - return self.input_density_ob.initialize(model, user_input_names) - - def inspect_input(self, input_name, input_data) -> Tuple[bool, float, float]: - """Inspect input data and print statistics. - - Args: - input_name: User input name. - input_data: User input tensor. - - Returns: - found: Whether the input name is found in `_embedding_graph_input_to_padding_idx_map` and - `_loss_label_graph_input_to_ignore_idx_map`. - embed_input_density: Density for the inspected embedding input if found to be True; otherwise, return 100. - label_input_density: Density for the inspected label input if found to be True; otherwise, return 100. - """ - if self.input_density_ob is not None: - return self.input_density_ob.inspect_from_input_data(input_name, input_data) - - return (False, 100, 100) - - def disable_input_inspector(self) -> None: - """Disable input density inspector.""" - self.input_density_ob = None - - -class InputDensityObserver: - """Training input data observer for ORTModule. - - Data observer is used to collect data/compute sparsity information for embedding and label inputs. It needs to be - firstly initialized with the ONNX model and user input names. Then, it can be used to inspect the input data - through `inspect_from_input_data()` method given user input name and input tensor. Inspection results will be - printed per `log_steps`. - - """ - - def __init__(self, logger: Logger, log_steps=1): - self._logger = logger - self._embedding_graph_input_to_padding_idx_map = {} - self._loss_label_graph_input_to_ignore_idx_map = {} - self._stats = [] - self._is_initialized = False - - self._last_step = 0 - self._current_step = 0 - self._log_steps = log_steps - - self._tensor_to_node_map = {} - - def initialize(self, model: ModelProto, user_input_names: List[str]) -> None: - """Initialize data observer from the given ONNX model and user input names. - - For embedding input (e.g. ATen embedding), try to parse the padding_idx from the ONNX model, if padding_idx is - valid, register it in _embedding_graph_input_to_padding_idx_map. - For label input (e.g. SoftmaxCrossEntropyLossInternal), try to parse the ignore_index from the ONNX model, if - ignore_index is valid, register it in _loss_label_graph_input_to_ignore_idx_map. - - Args: - model: ONNX model. - user_input_names: User input names in the ONNX model. - - """ - if self._is_initialized: - return - - try: - self._tensor_to_node_map.clear() - for node in model.graph.node: - for output_name in node.output: - if output_name != "": - self._tensor_to_node_map[output_name] = node - - self._initialize_embedding_padding_inspector(model, user_input_names) - self._initialize_loss_label_padding_inspector(model, user_input_names) - - self._is_initialized = True - except Exception as e: - self._is_initialized = False - self._logger.warning(f"Failed to initialize InputDensityObserver due to {e}") - - def _initialize_embedding_padding_inspector(self, model, user_input_names): - """Register embedding input padding inspector. - - Iterate all ATen embedding nodes, and check if the following conditions are met: - > 1. the data input is ONNX graph input; - > 2. and the padding_idx is a non-negative scalar constant. - - If yes, append into _embedding_graph_input_to_padding_idx_map, which is later - used for collecting data/compute sparsity information for embedding layer. - """ - - self._embedding_graph_input_to_padding_idx_map.clear() - - for node in model.graph.node: - if not (node.domain == "org.pytorch.aten" and node.op_type == "ATen" and len(node.input) >= 3): - continue - - found = [attr for attr in node.attribute if attr.name == "operator"] - if not found or helper.get_attribute_value(found[0]).decode() != "embedding": - continue - - tensor = None - padding_const_node = self._try_get_node_from_its_output(node.input[2]) - if padding_const_node is None: - padding_initializer_name = node.input[2] - tensor = self._try_get_initializer(model, padding_initializer_name) - - elif padding_const_node.op_type == "Constant": - found = [attr for attr in padding_const_node.attribute if attr.name == "value"] - tensor = found[0].t - else: - continue - - if tensor is None or tensor.data_type not in [onnx_proto.TensorProto.INT32, onnx_proto.TensorProto.INT64]: - continue - - value = onnx.numpy_helper.to_array(tensor) - if value.ndim != 0: - self._logger.warning(f"Embedding padding_idx must be a scalar, but got a tensor of shape {value.shape}") - continue - - padding_idx = value.item() - # Negative padding_idx in ATen embedding means there is no padding. - if padding_idx < 0: - continue - - # Given the input arg of embedding node, find the corresponding user input that feeds into the data. - # Will iterate the args recursively if some subgraph pattern is found between the input and the embedding, - # such as Input -> Cast -> Cast -> Embedding. - # TODO: This is a workaround for the case that the input of embedding is a list of Cast nodes which is found - # in Llama-2. We need to find a general way to handle all types of subgraph parttern between input and embedding. - def _get_embedding_graph_input(node_arg): - if node_arg in user_input_names: - return node_arg - input_node = self._try_get_node_from_its_output(node_arg) - if input_node.op_type == "Cast": - return _get_embedding_graph_input(input_node.input[0]) - else: - self._logger.warning(f"Cannot find embedding input {node_arg}") - return None - - embedding_graph_input = _get_embedding_graph_input(node.input[1]) - if embedding_graph_input is None: - continue - - if embedding_graph_input not in self._embedding_graph_input_to_padding_idx_map: - self._embedding_graph_input_to_padding_idx_map[embedding_graph_input] = set() - - self._embedding_graph_input_to_padding_idx_map[embedding_graph_input].add(padding_idx) - - def _initialize_loss_label_padding_inspector(self, model, user_input_names): - """Register loss label input padding inspector. - - Iterate all SoftmaxCrossEntropyLossInternal nodes, and check if the following conditions are met: - > 1. ignore_index (the 4th input) is a non-negative scalar constant; - > 2. label input (the 2nd input) is either a). ONNX graph input or b). a Reshape node with a Slice node as its - input. In the case of b), the Slice node must be a pattern defined in Bloom model. - - If yes, append > into - _loss_label_graph_input_to_ignore_idx_map, which is later used for collecting data/compute sparsity information - for labels. - """ - - def _default_label_preprocess(labels): - return labels - - self._loss_label_graph_input_to_ignore_idx_map.clear() - for node in model.graph.node: - if not ( - node.domain == "com.microsoft" - and node.op_type == "SoftmaxCrossEntropyLossInternal" - and len(node.input) == 4 - ): - continue - - tensor = None - padding_const_node = self._try_get_node_from_its_output(node.input[3]) - if padding_const_node is None: - padding_initializer_name = node.input[3] - tensor = self._try_get_initializer(model, padding_initializer_name) - - elif padding_const_node.op_type == "Constant": - found = [attr for attr in padding_const_node.attribute if attr.name == "value"] - tensor = found[0].t - else: - continue - - if tensor is None or tensor.data_type not in [onnx_proto.TensorProto.INT32, onnx_proto.TensorProto.INT64]: - continue - - value = onnx.numpy_helper.to_array(tensor) - if value.ndim != 0: - self._logger.warning( - f"SoftmaxCrossEntropyLossInternal ignore_index must be a scalar, but got a tensor of shape {value.shape}" - ) - continue - - ignore_index = value.item() - - # Check label inputs - label_graph_input = None - - label_preprocess_func = _default_label_preprocess - reshape_node = self._try_get_node_from_its_output(node.input[1]) - # The label input comes from graph input or a Reshape node consuming a graph input, which is aligned with - # orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc. - if reshape_node is None: - if node.input[1] not in user_input_names: - continue - label_graph_input = node.input[1] - else: - if reshape_node.op_type != "Reshape": - continue - - reshape_input = reshape_node.input[0] - if reshape_input in user_input_names: - label_graph_input = reshape_input - else: # Pattern defined in Bloom model. - slice_node = self._try_get_node_from_its_output(reshape_input) - if slice_node is None: - continue - - if slice_node.op_type != "Slice": - continue - - slice_input = slice_node.input[0] - starts = self._try_get_initializer_value(model, slice_node.input[1]) - ends = self._try_get_initializer_value(model, slice_node.input[2]) - axes = self._try_get_initializer_value(model, slice_node.input[3]) - steps = self._try_get_initializer_value(model, slice_node.input[4]) - if ( - slice_input in user_input_names - and starts is not None - and ends is not None - and axes is not None - and steps is not None - and len(axes) == 1 - and axes[0] == 1 - ): - label_graph_input = slice_input - - def _slice_label_preprocess(labels, s=starts[0], e=ends[0], st=steps[0]): - return labels[:, s:e:st] - - label_preprocess_func = _slice_label_preprocess - - if label_graph_input is None: - continue - - if label_graph_input not in self._loss_label_graph_input_to_ignore_idx_map: - self._loss_label_graph_input_to_ignore_idx_map[label_graph_input] = [] - - self._loss_label_graph_input_to_ignore_idx_map[label_graph_input].append( - [ignore_index, label_preprocess_func] - ) - - def inspect_from_input_data(self, name: str, inp) -> Tuple[bool, float, float]: - """Inspect input data and print statistics. - - Args: - name: User input name. - inp: User input tensor. - Returns: - found: Whether the input name is found in `_embedding_graph_input_to_padding_idx_map` and - `_loss_label_graph_input_to_ignore_idx_map`. - embed_input_density: Density for the inspected embedding input if found to be True; otherwise, return 100. - label_input_density: Density for the inspected label input if found to be True; otherwise, return 100. - """ - if not self._is_initialized: - return (False, 100, 100) - - try: - data = inp.clone() - found, embed_input_density, label_input_density = self._inspect_embed_label_input(name, data) - if found: - self._current_step += 1 - - if self._current_step - self._last_step >= self._log_steps: - self._last_step = self._current_step - self._print_embed_label_stats() - - return (found, embed_input_density, label_input_density) - except Exception as e: - self._logger.warning(f"Failed to inspect input {name} due to {e}", UserWarning) - return (False, 100, 100) - - def _inspect_embed_label_input(self, name, data): - found = False - min_embed_density = 100 - min_label_density = 100 - if ( - len(self._embedding_graph_input_to_padding_idx_map) > 0 - and name in self._embedding_graph_input_to_padding_idx_map - and isinstance(data, torch.Tensor) - ): - for padding_idx in self._embedding_graph_input_to_padding_idx_map[name]: - valid_token = torch.count_nonzero(data - padding_idx) - valid_token_per_batch = "N/A" - if data.dim() > 1: - valid_token_per_batch = str(torch.count_nonzero(data - padding_idx, dim=1).tolist()) - total_token = data.numel() - embed_density = float(valid_token) / float(total_token) * 100 - if embed_density < 90: - min_embed_density = min(min_embed_density, embed_density) - self._stats.append( - [ - self._current_step, - "EMBED", - name, - padding_idx, - embed_density, - valid_token, - total_token, - valid_token_per_batch, - ] - ) - found = True - - if ( - len(self._loss_label_graph_input_to_ignore_idx_map) > 0 - and name in self._loss_label_graph_input_to_ignore_idx_map - and isinstance(data, torch.Tensor) - ): - for ignore_index, preprocess_func in self._loss_label_graph_input_to_ignore_idx_map[name]: - data_preprocessed = preprocess_func(data) - valid_token = torch.count_nonzero(data_preprocessed - ignore_index) - total_token = data_preprocessed.numel() - label_density = float(valid_token) / float(total_token) * 100 - if label_density < 90: - min_label_density = min(min_label_density, label_density) - self._stats.append( - [ - self._current_step, - "LABEL", - name, - ignore_index, - label_density, - valid_token, - total_token, - "N/A", - ] - ) - found = True - - return found, min_embed_density, min_label_density - - def _print_embed_label_stats(self): - if len(self._stats) > 0: - stat = f">>>Valid token/label density (e.g. valid/total) in passing {self._log_steps} steps:\n" - stat += "\t| {:<10} | {:<10} | {:<15} | {:<10} | {:<10} | {:<15} | {:<15} | {:<15} |\n".format( - "STEP", - "INPUT TYPE", - "INPUT NAME", - "PAD IDX", - "DENSITY", - "VALID TOKENS", - "TOTAL TOKENS", - "VALID TOKENS/BATCH", - ) - for ( - step, - input_type, - input_name, - padding_idx, - density, - valid_token, - total_token, - valid_token_per_batch, - ) in self._stats: - stat += f"\t| {step:<10} | {input_type:<10} | {input_name:<15} | {padding_idx:<10} | {density:<9.2f}% | {valid_token:<15} | {total_token:<15} | {valid_token_per_batch:<15} |\n" - stat += "<<<\n" - self._logger.info(stat) - self._stats.clear() - - def _try_get_node_from_its_output(self, name): - if name == "" or name not in self._tensor_to_node_map: - return None - - return self._tensor_to_node_map[name] - - def _try_get_initializer(self, model, name): - for tensor in model.graph.initializer: - if tensor.name == name: - return tensor - - return None - - def _try_get_initializer_value(self, model, name): - tensor = self._try_get_initializer(model, name) - if tensor is None: - return None - value = onnx.numpy_helper.to_array(tensor) - return value + self._sceloss_module_to_ignore_density_map = {} class MemoryOptimizationSummary: @@ -750,15 +342,19 @@ class MemoryObserver: return [], None -class FlagPaddingElimination(torch.autograd.Function): +class FlagAndPrintDensity(torch.autograd.Function): """ - FlagPaddingElimination is a PyTorch autograd function that does nothing in forward pass and backward pass. - It is used as a flag to tell the GraphTransformer of PaddingElimination to modify the graph to eliminate - the embedding padding. + FlagAndPrintDensity is a PyTorch autograd function that print input density for embedding or label. + It is also used as a flag to tell the GraphTransformer of PaddingElimination and InsertGatherBeforeSceLoss + to modify the graph to eliminate the padding. """ @staticmethod - def forward(ctx, input): + def forward(ctx, input, padding_value, type_name): + valid_token = torch.count_nonzero(input - padding_value) + total_token = input.numel() + density = float(valid_token) / float(total_token) * 100 + print(type_name + " tensor density: ", density) return input @staticmethod @@ -776,5 +372,5 @@ class FlagPaddingElimination(torch.autograd.Function): @staticmethod def alias_input(node_proto_str: str): fw_alias_map = [0] - bw_alias_map = [0] + bw_alias_map = [0, -1, -1] return fw_alias_map, bw_alias_map diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index a7426bce38..476c169fd1 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -271,10 +271,9 @@ class TrainingManager(GraphExecutionManager): # Build the gradient graph if build_gradient_graph: - graph_transformer_config = self._get_graph_transformer_config() - # Set the config according to input inspection. - self._enable_conditional_optimizations(graph_transformer_config, inputs, kwargs) + self._detect_from_inputs(inputs, kwargs) + graph_transformer_config = self._get_graph_transformer_config() # Build the gradient graph self._build_graph(graph_transformer_config) @@ -324,7 +323,7 @@ class TrainingManager(GraphExecutionManager): else: param_to_append_as_onnx_graph_inputs = self._graph_initializers - prepared_input_list, _, _ = _io._combine_input_buffers_initializers( + prepared_input_list = _io._combine_input_buffers_initializers( param_to_append_as_onnx_graph_inputs, self._graph_info.user_input_names, self._input_info, diff --git a/orttraining/orttraining/python/training/ortmodule/options.py b/orttraining/orttraining/python/training/ortmodule/options.py index 1bde07dc29..eabeacee25 100644 --- a/orttraining/orttraining/python/training/ortmodule/options.py +++ b/orttraining/orttraining/python/training/ortmodule/options.py @@ -274,10 +274,10 @@ class _RuntimeOptions: # Configuration for compute optimization. self.enable_compute_optimizer = True - self.enable_sparse_optimizer = True + self.enable_embedding_sparse_optimizer = True + self.enable_label_sparse_optimizer = True self.label_sparsity_ratio = "" self.embed_sparsity_ratio = "" - self.enable_embedding_sparse_optimizer = True # Configuration for memory optimization. self.memory_optimization_level = ( @@ -335,15 +335,18 @@ class _RuntimeOptions: self.enable_compute_optimizer = int(os.getenv("ORTMODULE_ENABLE_COMPUTE_OPTIMIZER")) == 1 compute_optimizer_reset = True - if "ORTMODULE_ENABLE_SPARSE_OPTIMIZER" in os.environ or compute_optimizer_reset: - if "ORTMODULE_ENABLE_SPARSE_OPTIMIZER" in os.environ: - self.enable_sparse_optimizer = int(os.getenv("ORTMODULE_ENABLE_SPARSE_OPTIMIZER")) == 1 - self.enable_sparse_optimizer = self.enable_compute_optimizer and self.enable_sparse_optimizer + if "ORTMODULE_ENABLE_LABEL_SPARSE_OPTIMIZER" in os.environ or compute_optimizer_reset: + if "ORTMODULE_ENABLE_LABEL_SPARSE_OPTIMIZER" in os.environ: + self.enable_label_sparse_optimizer = int(os.getenv("ORTMODULE_ENABLE_LABEL_SPARSE_OPTIMIZER")) == 1 + self.enable_label_sparse_optimizer = self.enable_compute_optimizer and self.enable_label_sparse_optimizer - # TODO(pengwa): remove once validation on more models are done. - if "ORTMODULE_ENABLE_EMBEDDING_SPARSE_OPTIMIZER" in os.environ: + if "ORTMODULE_ENABLE_EMBEDDING_SPARSE_OPTIMIZER" in os.environ or compute_optimizer_reset: + if "ORTMODULE_ENABLE_EMBEDDING_SPARSE_OPTIMIZER" in os.environ: + self.enable_embedding_sparse_optimizer = ( + int(os.getenv("ORTMODULE_ENABLE_EMBEDDING_SPARSE_OPTIMIZER")) == 1 + ) self.enable_embedding_sparse_optimizer = ( - self.enable_sparse_optimizer and int(os.getenv("ORTMODULE_ENABLE_EMBEDDING_SPARSE_OPTIMIZER")) == 1 + self.enable_compute_optimizer and self.enable_embedding_sparse_optimizer ) # Configuration for memory optimization. diff --git a/orttraining/orttraining/test/optimizer/compute_optimizer_test.cc b/orttraining/orttraining/test/optimizer/compute_optimizer_test.cc index 509937bdd0..c860788d3d 100644 --- a/orttraining/orttraining/test/optimizer/compute_optimizer_test.cc +++ b/orttraining/orttraining/test/optimizer/compute_optimizer_test.cc @@ -16,6 +16,8 @@ #include "core/graph/graph_utils.h" #include "core/graph/graph_viewer.h" #include "core/graph/model.h" +#include "core/graph/node_attr_utils.h" +#include "core/optimizer/compute_optimizer/shared_utils.h" #include "core/optimizer/utils.h" #include "core/util/math.h" #include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h" @@ -31,6 +33,8 @@ #include "test/util/include/asserts.h" #include "test/util/include/default_providers.h" +using namespace onnxruntime::optimizer::compute_optimizer; + namespace onnxruntime { namespace test { @@ -44,7 +48,9 @@ const InlinedHashSet compatible_eps = {}; Test graph includes multiple equivalent subgraphs as below. graph input [32, 256] (float) graph input [32] (int64_t) | | - \_____________ _______/ graph input -1, scalar (int64_t) + \___________ ______/ graph input -1, scalar (int64_t) + \ / / + \ PythonOp (flag) / \ / _______/ \ / / SCE Node, reduction = 'mean', output_type=1 @@ -58,7 +64,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_Allowed) { 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); + TEST_RETURN_IF_NOT(op_count_pre.size() == 2U); if (is_sce_internal) TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); @@ -115,20 +121,30 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_Allowed) { auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) { auto* input1_arg = builder.MakeInput({{32, 256}}); auto* input2_arg = builder.MakeInput({{32}}, "label"); + auto* python_op_out1 = builder.MakeIntermediate(); + auto* python_op_out2 = builder.MakeIntermediate(); auto* sce_out1 = builder.MakeOutput(); NodeArg* empty = builder.MakeEmptyInput(); auto* sce_out2 = builder.MakeIntermediate(); + Node& python_op = builder.AddNode("PythonOp", {input2_arg}, {python_op_out1, python_op_out2}, kMSDomain); + python_op.AddAttribute("func_name", "onnxruntime.training.ortmodule._runtime_inspector.FlagAndPrintDensity"); + python_op.AddAttribute("input_convention", "dcc"); + python_op.AddAttribute("input_tensor_types", std::vector{7}); + python_op.AddAttribute("input_tensor_ranks", std::vector{1}); + python_op.AddAttribute("output_tensor_types", std::vector{7}); + python_op.AddAttribute("output_tensor_ranks", std::vector{1}); if (is_sce_internal) { auto* ignore_index_arg = builder.MakeScalarInitializer(-100); + Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal", - {input1_arg, input2_arg, empty, ignore_index_arg}, + {input1_arg, python_op_out2, 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}, + {input1_arg, python_op_out2, empty}, {sce_out1, sce_out2}); sce.AddAttribute("reduction", "mean"); sce.AddAttribute("ignore_index", static_cast(-100)); @@ -138,7 +154,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_Allowed) { std::vector opsets{12, 13, 14, 15, 17}; for (auto opset : opsets) { std::unique_ptr transformer = - std::make_unique(compatible_eps, std::vector{"label"}); + std::make_unique(compatible_eps); ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer), TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); @@ -158,7 +174,7 @@ Test graph includes multiple equivalent subgraphs as below. | graph output, loss, scalar (float) */ -TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_LabelNameNotMatch) { +TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_WithoutFlagPythonOp) { const logging::Logger* logger = &logging::LoggingManager::DefaultLogger(); for (const bool is_sce_internal : {true, false}) { @@ -184,7 +200,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_LabelNameNotMat auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) { auto* input1_arg = builder.MakeInput({{32, 256}}); - auto* input2_arg = builder.MakeInput({{32}}, "label111"); + auto* input2_arg = builder.MakeInput({{32}}, "label"); auto* sce_out1 = builder.MakeOutput(); NodeArg* empty = builder.MakeEmptyInput(); @@ -209,7 +225,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_LabelNameNotMat std::vector opsets{12, 13, 14, 15, 17}; for (auto opset : opsets) { std::unique_ptr transformer = - std::make_unique(compatible_eps, std::vector{"label"}); + std::make_unique(compatible_eps); ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer), TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); @@ -219,11 +235,13 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_LabelNameNotMat /* Test graph includes multiple equivalent subgraphs as below. - graph input [32, 256] (float) graph input [32] (int64_t) - | | - \_____________ _______/ graph input -1, scalar (int64_t) - \ / _______/ - \ / / + graph input [32, 256] (float) graph input [32] (int64_t) + | | + \_____________ _______/ graph input -1, scalar (int64_t) + \ / _______/ + \ / / + \ PythonOp (flag) / + \ / / SCE Node, reduction = 'none', output_type=1 | | @@ -235,7 +253,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_ReduceNone) { 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); + TEST_RETURN_IF_NOT(op_count_pre.size() == 2U); if (is_sce_internal) TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); else @@ -256,21 +274,30 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_ReduceNone) { auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) { auto* input1_arg = builder.MakeInput({{32, 256}}); auto* input2_arg = builder.MakeInput({{32}}, "label"); + auto* python_op_out1 = builder.MakeIntermediate(); + auto* python_op_out2 = builder.MakeIntermediate(); auto* sce_out1 = builder.MakeOutput(); NodeArg* empty = builder.MakeEmptyInput(); auto* sce_out2 = builder.MakeIntermediate(); + Node& python_op = builder.AddNode("PythonOp", {input2_arg}, {python_op_out1, python_op_out2}, kMSDomain); + python_op.AddAttribute("func_name", "onnxruntime.training.ortmodule._runtime_inspector.FlagAndPrintDensity"); + python_op.AddAttribute("input_convention", "dcc"); + python_op.AddAttribute("input_tensor_types", std::vector{7}); + python_op.AddAttribute("input_tensor_ranks", std::vector{1}); + python_op.AddAttribute("output_tensor_types", std::vector{7}); + python_op.AddAttribute("output_tensor_ranks", std::vector{1}); if (is_sce_internal) { auto* ignore_index_arg = builder.MakeScalarInitializer(-100); Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal", - {input1_arg, input2_arg, empty, ignore_index_arg}, + {input1_arg, python_op_out2, 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}, + {input1_arg, python_op_out2, empty}, {sce_out1, sce_out2}); sce.AddAttribute("reduction", "none"); sce.AddAttribute("ignore_index", static_cast(-100)); @@ -280,7 +307,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_ReduceNone) { std::vector opsets{12, 13, 14, 15, 17}; for (auto opset : opsets) { std::unique_ptr transformer = - std::make_unique(compatible_eps, std::vector{"label"}); + std::make_unique(compatible_eps); ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer), TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); @@ -290,11 +317,13 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_ReduceNone) { /* Test graph include multiple equivalent subgraphs as below. - graph input [32, 256] (float) graph input [32] (int64_t) - | | - \_____________ _______/ graph input -1, scalar (int64_t) - \ / _______/ - \ / / + graph input [32, 256] (float) graph input [32] (int64_t) + | | + \_____________ _______/ graph input -1, scalar (int64_t) + \ / _______/ + \ / / + \ PythonOp (flag) / + \ / / SCE Node, reduction = 'none', output_type=1 | | @@ -306,7 +335,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_NoIgnoreIndex) 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); + TEST_RETURN_IF_NOT(op_count_pre.size() == 2U); if (is_sce_internal) TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1); else @@ -327,18 +356,27 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_NoIgnoreIndex) auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) { auto* input1_arg = builder.MakeInput({{32, 256}}); auto* input2_arg = builder.MakeInput({{32}}, "label"); + auto* python_op_out1 = builder.MakeIntermediate(); + auto* python_op_out2 = builder.MakeIntermediate(); auto* sce_out1 = builder.MakeOutput(); auto* sce_out2 = builder.MakeIntermediate(); + Node& python_op = builder.AddNode("PythonOp", {input2_arg}, {python_op_out1, python_op_out2}, kMSDomain); + python_op.AddAttribute("func_name", "onnxruntime.training.ortmodule._runtime_inspector.FlagAndPrintDensity"); + python_op.AddAttribute("input_convention", "dcc"); + python_op.AddAttribute("input_tensor_types", std::vector{7}); + python_op.AddAttribute("input_tensor_ranks", std::vector{1}); + python_op.AddAttribute("output_tensor_types", std::vector{7}); + python_op.AddAttribute("output_tensor_ranks", std::vector{1}); if (is_sce_internal) { Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal", - {input1_arg, input2_arg}, + {input1_arg, python_op_out2}, {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}, + {input1_arg, python_op_out2}, {sce_out1, sce_out2}); sce.AddAttribute("reduction", "sum"); } @@ -347,7 +385,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_NoIgnoreIndex) std::vector opsets{12, 13, 14, 15, 17}; for (auto opset : opsets) { std::unique_ptr transformer = - std::make_unique(compatible_eps, std::vector{"label"}); + std::make_unique(compatible_eps); ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer), TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); @@ -363,11 +401,55 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_MlmBertE2E) { 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); + + // Insert a PythonOp Flag to enabel the optimization + GraphViewer graph_viewer(graph); + const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); + for (auto node_index : node_topology_list) { + auto& node = *graph.GetNode(node_index); + bool is_internal_sce = graph_utils::IsSupportedOptypeVersionAndDomain(node, "SoftmaxCrossEntropyLossInternal", {1}, + kMSDomain); + if (!is_internal_sce) { + continue; + } + InlinedVector python_op_node_input_args; + python_op_node_input_args.reserve(1); + python_op_node_input_args.push_back(node.MutableInputDefs()[1]); + InlinedVector python_op_node_output_args; + python_op_node_output_args.push_back( + &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("python_op_ctx"), + nullptr)); + python_op_node_output_args.push_back( + &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("python_op_out"), + nullptr)); + onnxruntime::NodeAttributes attributes; + attributes["func_name"] = onnxruntime::utils::MakeAttribute( + "func_name", + "onnxruntime.training.ortmodule._runtime_inspector.FlagAndPrintDensity"); + attributes["input_convention"] = onnxruntime::utils::MakeAttribute("input_convention", "dcc"); + attributes["input_tensor_types"] = onnxruntime::utils::MakeAttribute("input_tensor_types", std::vector{7}); + attributes["input_tensor_ranks"] = onnxruntime::utils::MakeAttribute("input_tensor_ranks", std::vector{1}); + attributes["output_tensor_types"] = onnxruntime::utils::MakeAttribute("output_tensor_types", std::vector{7}); + attributes["output_tensor_ranks"] = onnxruntime::utils::MakeAttribute("output_tensor_ranks", std::vector{1}); + Node* python_op_node = InsertIntermediateNodeOnDestInput( + graph, node, + 1, + 0 /* new_node_input_index*/, + 1 /* new_node_output_index*/, + graph.GenerateNodeName("PaddingFlag"), + "PythonOp", + "", + python_op_node_input_args, + python_op_node_output_args, + attributes, + kMSDomain, + *logger); + python_op_node->SetExecutionProviderType(node.GetExecutionProviderType()); + } onnxruntime::GraphTransformerManager graph_transformation_mgr{3}; ASSERT_STATUS_OK(graph_transformation_mgr.Register( - std::make_unique(compatible_eps, std::vector{"labels"}), + std::make_unique(compatible_eps), TransformerLevel::Level1)); ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger)); diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 24c637bd77..9e226958dd 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -5756,9 +5756,6 @@ def test_runtime_inspector_label_and_embed_sparsity_detection(embed_is_sparse, l loss.backward() return loss - # batch_size = 3 - # sequence = 4 - if embed_is_sparse: input = torch.tensor([[0, 2, 3, 4], [2, 3, 1, 1], [1, 1, 1, 1]], device=device) else: @@ -5778,10 +5775,14 @@ def test_runtime_inspector_label_and_embed_sparsity_detection(embed_is_sparse, l found_embed_is_sparse = False found_embed_is_dense = False found_label_is_sparse = False + found_label_is_dense = False for record in caplog.records: if "Label sparsity-based optimization is ON for" in record.getMessage(): found_label_is_sparse = True + if "Label sparsity-based optimization is OFF for" in record.getMessage(): + found_label_is_dense = True + if "Embedding sparsity-based optimization is OFF for" in record.getMessage(): found_embed_is_dense = True @@ -5789,7 +5790,9 @@ def test_runtime_inspector_label_and_embed_sparsity_detection(embed_is_sparse, l found_embed_is_sparse = True if label_is_sparse: - assert found_label_is_sparse + assert found_label_is_sparse and not found_label_is_dense + else: + assert not found_label_is_sparse and found_label_is_dense if embed_is_sparse: assert found_embed_is_sparse and not found_embed_is_dense