mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
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.
This commit is contained in:
guyang3532
GitHub
parent
e124cf8e76
commit
cfe830b248
16 changed files with 306 additions and 610 deletions
|
|
@ -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*
|
||||
|
|
|
|||
|
|
@ -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<Node*>& q, Node* node, std::unordered_set<Node*>& 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<std::string>(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<std::string>(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());
|
||||
|
|
|
|||
|
|
@ -127,10 +127,15 @@ namespace onnxruntime {
|
|||
*/
|
||||
class PaddingElimination : public GraphTransformer {
|
||||
public:
|
||||
explicit PaddingElimination(const InlinedHashSet<std::string_view>& compatible_execution_providers = {}) noexcept
|
||||
: GraphTransformer("PaddingElimination", compatible_execution_providers) {}
|
||||
explicit PaddingElimination(const InlinedHashSet<std::string_view>& 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
|
||||
|
|
|
|||
|
|
@ -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<std::string>(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.
|
||||
|
|
|
|||
|
|
@ -68,9 +68,9 @@ namespace onnxruntime {
|
|||
class InsertGatherBeforeSceLoss : public GraphTransformer {
|
||||
public:
|
||||
InsertGatherBeforeSceLoss(const InlinedHashSet<std::string_view>& compatible_execution_providers = {},
|
||||
const std::vector<std::string>& 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<std::string> sparse_label_input_names_;
|
||||
bool print_density_ = false;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -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<std::string> 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;
|
||||
|
|
|
|||
|
|
@ -195,11 +195,12 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
|
|||
transformers.emplace_back(std::make_unique<UpStreamGatherGraphTransformer>(compatible_eps));
|
||||
transformers.emplace_back(std::make_unique<UpStreamReshapeGraphTransformer>(compatible_eps));
|
||||
transformers.emplace_back(std::make_unique<InsertGatherBeforeSceLoss>(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<PaddingElimination>(compatible_eps));
|
||||
transformers.emplace_back(std::make_unique<PaddingElimination>(compatible_eps,
|
||||
config.print_input_density));
|
||||
transformers.emplace_back(std::make_unique<Conv1dReplacement>(compatible_eps));
|
||||
#endif
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 <ONNX graph input, padding_idx> 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 <ONNX graph input, <ignore_index, label_preprocess_function>> 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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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<std::string_view> 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<float>({{32, 256}});
|
||||
auto* input2_arg = builder.MakeInput<int64_t>({{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<int64_t>{7});
|
||||
python_op.AddAttribute("input_tensor_ranks", std::vector<int64_t>{1});
|
||||
python_op.AddAttribute("output_tensor_types", std::vector<int64_t>{7});
|
||||
python_op.AddAttribute("output_tensor_ranks", std::vector<int64_t>{1});
|
||||
if (is_sce_internal) {
|
||||
auto* ignore_index_arg = builder.MakeScalarInitializer<int64_t>(-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<int64_t>(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<int64_t>(-100));
|
||||
|
|
@ -138,7 +154,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_Allowed) {
|
|||
std::vector<int> opsets{12, 13, 14, 15, 17};
|
||||
for (auto opset : opsets) {
|
||||
std::unique_ptr<GraphTransformer> transformer =
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(compatible_eps, std::vector<std::string>{"label"});
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(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<float>({{32, 256}});
|
||||
auto* input2_arg = builder.MakeInput<int64_t>({{32}}, "label111");
|
||||
auto* input2_arg = builder.MakeInput<int64_t>({{32}}, "label");
|
||||
auto* sce_out1 = builder.MakeOutput();
|
||||
|
||||
NodeArg* empty = builder.MakeEmptyInput();
|
||||
|
|
@ -209,7 +225,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_LabelNameNotMat
|
|||
std::vector<int> opsets{12, 13, 14, 15, 17};
|
||||
for (auto opset : opsets) {
|
||||
std::unique_ptr<GraphTransformer> transformer =
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(compatible_eps, std::vector<std::string>{"label"});
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(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<float>({{32, 256}});
|
||||
auto* input2_arg = builder.MakeInput<int64_t>({{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<int64_t>{7});
|
||||
python_op.AddAttribute("input_tensor_ranks", std::vector<int64_t>{1});
|
||||
python_op.AddAttribute("output_tensor_types", std::vector<int64_t>{7});
|
||||
python_op.AddAttribute("output_tensor_ranks", std::vector<int64_t>{1});
|
||||
if (is_sce_internal) {
|
||||
auto* ignore_index_arg = builder.MakeScalarInitializer<int64_t>(-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<int64_t>(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<int64_t>(-100));
|
||||
|
|
@ -280,7 +307,7 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_ReduceNone) {
|
|||
std::vector<int> opsets{12, 13, 14, 15, 17};
|
||||
for (auto opset : opsets) {
|
||||
std::unique_ptr<GraphTransformer> transformer =
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(compatible_eps, std::vector<std::string>{"label"});
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(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<float>({{32, 256}});
|
||||
auto* input2_arg = builder.MakeInput<int64_t>({{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<int64_t>{7});
|
||||
python_op.AddAttribute("input_tensor_ranks", std::vector<int64_t>{1});
|
||||
python_op.AddAttribute("output_tensor_types", std::vector<int64_t>{7});
|
||||
python_op.AddAttribute("output_tensor_ranks", std::vector<int64_t>{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<int64_t>(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<int> opsets{12, 13, 14, 15, 17};
|
||||
for (auto opset : opsets) {
|
||||
std::unique_ptr<GraphTransformer> transformer =
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(compatible_eps, std::vector<std::string>{"label"});
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(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> model;
|
||||
ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger));
|
||||
Graph& graph = model->MainGraph();
|
||||
std::map<std::string, int> 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<NodeArg*> python_op_node_input_args;
|
||||
python_op_node_input_args.reserve(1);
|
||||
python_op_node_input_args.push_back(node.MutableInputDefs()[1]);
|
||||
InlinedVector<NodeArg*> 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<int64_t>{7});
|
||||
attributes["input_tensor_ranks"] = onnxruntime::utils::MakeAttribute("input_tensor_ranks", std::vector<int64_t>{1});
|
||||
attributes["output_tensor_types"] = onnxruntime::utils::MakeAttribute("output_tensor_types", std::vector<int64_t>{7});
|
||||
attributes["output_tensor_ranks"] = onnxruntime::utils::MakeAttribute("output_tensor_ranks", std::vector<int64_t>{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<InsertGatherBeforeSceLoss>(compatible_eps, std::vector<std::string>{"labels"}),
|
||||
std::make_unique<InsertGatherBeforeSceLoss>(compatible_eps),
|
||||
TransformerLevel::Level1));
|
||||
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger));
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue