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