saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-28 20:11:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Statistics tool for ORTModule convergence parity (#15020)

Browse source 
### Statistics tool for ORTModule convergence parity

As ORTModule get more and more validated, it is pretty fast to
intergrade PyTorch based model with ORT.

The same time, we need make sure once there is convergence issue, we
don't spend months of time to investigate. As part of this efforts, this
PR is introducing a tool to dump activation statistics without much
involvement from users. The dumping results contains only some statistic
numbers plus sampled data, which is not big, compared with dumping all
the tensors, it is much faster and space efficient.

For us to use it, two single lines are needed before wrapping ORTModule.
For baseline run, need also apply the same trick.

```
+	from onnxruntime.training.utils.hooks import SubscriberManager, StatisticsSubscriber
+	SubscriberManager.subscribe(model, [StatisticsSubscriber("pt_out", override_output_dir=True)])
```

Once you run the steps, following command can be used to merge result
into per-step-summary respectively for ORT and baseline runs.
 
```bash
python -m onnxruntime.training.utils.hooks.merge_activation_summary --pt_dir pt_out --ort_dir ort_out --output_dir /tmp/output
```

Docs is added here as part of this PR [convergence investigation
notes](https://github.com/microsoft/onnxruntime/blob/pengwa/conv_tool/docs/ORTModule_Convergence_Notes.md)

Based on the generated merged files, we can compare them with tools. 


![image](https://user-images.githubusercontent.com/10530022/224653929-4e4480bd-bb02-4bbe-bd44-2672bdf91a87.png)

### Design and Implementation

This PR introduced a common mechanism registering custom logic for
nn.Module's post forward hooks. And statistics for activation
(StatisticsSubscriber) is one of the implementations. If there is other
needs, we can define another XXSubscriber to do the customized things.
This commit is contained in:
pengwa 2023-03-23 20:34:24 +08:00 committed by GitHub
parent 039ca10822
commit 1d32285536
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 973 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cmake/onnxruntime_python.cmake
View file

@ -416,6 +416,9 @@ if (onnxruntime_ENABLE_TRAINING)
file(GLOB onnxruntime_python_utils_data_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/utils/data/*"
)
file(GLOB onnxruntime_python_utils_hooks_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/utils/hooks/*"
)
if (onnxruntime_ENABLE_TRAINING_APIS)
file(GLOB onnxruntime_python_onnxblock_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/onnxblock/*"
@ -720,6 +723,7 @@ if (onnxruntime_ENABLE_TRAINING)
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/fused_ops
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/data/
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/hooks/
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_capi_training_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/capi/training/
@ -771,6 +775,9 @@ if (onnxruntime_ENABLE_TRAINING)
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_utils_data_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/data/
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_utils_hooks_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/hooks/
)
if (onnxruntime_ENABLE_TRAINING_APIS)
add_custom_command(

49
docs/ORTModule_Convergence_Notes.md Normal file
View file

@ -0,0 +1,49 @@
# ORTModule Training Convergence Investigation
## 1. Discovering
Convergence issues can be identified by:
- Large discrepancy on core training metrics including training loss, evaluation loss, model specific AUC metrics.
- Runtime failures (for example loss scaler reach the minimum triggering an exception).
Before looking into further, we should clarify few things (if possible):
- If we change seed for baseline run, whether the metric diff is big?
(Make sure the discrepancy is not introduced by random)
- What's the very first steps we see obvious diverges?
- Still repro once remove randomness?
- Set same seeds
- Set dropout ratio to 0
- Set compute to be deterministic and torch-comparable (TODO(pengwa): need a flag for this).
## 2. Collect Activation Statistics
Add codes:
```diff
+ from onnxruntime.training.utils.hooks import SubscriberManager, StatisticsSubscriber
+ SubscriberManager.subscribe(model, [StatisticsSubscriber("pt_out", override_output_dir=True)])
```
Run training script to the steps that triggered the divergence. A folder named `pt_out` is created in current working directory. For each step, there is a folder containing summaries for every activation tensor.
Add few lines of code:
```diff
from onnxruntime.training.ortmodule import ORTModule
from onnxruntime.training.utils.hooks import SubscriberManager, StatisticsSubscriber
model = ORTModule(model)
+ SubscriberManager.subscribe(model, [StatisticsSubscriber("ort_out", override_output_dir=True)])
```
> `StatisticsSubscriber` can be initialized before OR after wrapping ORTModule.
Run training script to the steps that triggered the divergence. Similarly, a folder named `ort_out` is created in current working directory.
Run command to generate per step summary
```bash
python -m onnxruntime.training.utils.hooks.merge_activation_summary --pt_dir pt_out --ort_dir ort_out --output_dir /tmp/output
```
Manual diff the generate per-step summary to find the where is the first big diff happens.

41
onnxruntime/python/tools/symbolic_shape_infer.py
View file

@ -331,7 +331,7 @@ class SymbolicShapeInference:
return unique_dims[int_dim]
else:
if self.verbose_ > 0:
logger.debug("dim {} has been mergd with dim {}".format(unique_dims[1:], unique_dims[0]))
logger.debug("dim {} has been merged with dim {}".format(unique_dims[1:], unique_dims[0]))
return dims[0]
else:
return None
@ -2199,7 +2199,7 @@ class SymbolicShapeInference:
vi.CopyFrom(helper.make_tensor_value_info(node.output[1], onnx.TensorProto.INT32, mask_index_shape))
if len(node.output) > 2:
# Optional output of add before layer nomalization is done
# Optional output of add before layer normalization is done
# shape is same as the output
vi = self.known_vi_[node.output[2]]
vi.CopyFrom(helper.make_tensor_value_info(node.output[2], word_embedding_dtype, output_shape))
@ -2234,20 +2234,35 @@ class SymbolicShapeInference:
output_tensor_ranks = get_attribute(node, "output_tensor_ranks")
assert output_tensor_ranks
# set the context output seperately.
# set the context output separately.
# The first output is autograd's context.
vi = self.known_vi_[node.output[0]]
vi.CopyFrom(helper.make_tensor_value_info(node.output[0], onnx.TensorProto.INT64, []))
# Outputs after autograd's context are tensors.
# We assume their ranks are fixed for different model inputs.
for i in range(len(node.output) - 1):
# Process the i-th tensor outputs.
vi = self.known_vi_[node.output[i + 1]]
sympy_shape = self._new_symbolic_shape(output_tensor_ranks[i], node)
shape = get_shape_from_sympy_shape(sympy_shape)
value_info = helper.make_tensor_value_info(node.output[i + 1], output_tensor_types[i], shape)
vi.CopyFrom(value_info)
if get_attribute(node, "name").decode() in ["_InspectActivation", "_IncrementStep"]:
# PythonOp with name being "_InspectActivation" or "_IncrementStep" will behave exactly same as a normal
# PythonOp when execution. The only difference is that
# 1). those ops having same number of tensor inputs and tensor outputs;
# 2). and the i-th output tensor's shape is same as i-th input tensor's shape.
# Be noted, the count of custom autograd function might be bigger than output count, because there might
# be other non-tensor constant inputs (string, object, int, tuple, etc). But we did not make those constant
# inputs as ONNX op's input, instead they are stored in the attributes.
assert len(node.output) == len(node.input) + 1 # The output contains one extra context info.
for input_index in range(len(node.output) - 1):
# Process the i-th tensor outputs.
vi = self.known_vi_[node.output[input_index + 1]]
shape = self._get_shape(node, input_index)
output_dtype = self.known_vi_[node.input[input_index]].type.tensor_type.elem_type
vi.CopyFrom(helper.make_tensor_value_info(node.output[input_index + 1], output_dtype, shape))
else:
# Outputs after autograd's context are tensors.
# We assume their ranks are fixed for different model inputs.
for i in range(len(node.output) - 1):
# Process the i-th tensor outputs.
vi = self.known_vi_[node.output[i + 1]]
sympy_shape = self._new_symbolic_shape(output_tensor_ranks[i], node)
shape = get_shape_from_sympy_shape(sympy_shape)
value_info = helper.make_tensor_value_info(node.output[i + 1], output_tensor_types[i], shape)
vi.CopyFrom(value_info)
def _propagate_shape_and_type(self, node, input_index=0, output_index=0):
shape = self._get_shape(node, input_index)

4
orttraining/orttraining/core/graph/gradient_builder.cc
View file

@ -1825,6 +1825,10 @@ IMPLEMENT_GRADIENT_BUILDER(GetPythonOpGradient) {
"PythonOpGrad requiring gradient output count mismatch.");
attrs.push_back(MakeAttribute("output_tensor_requires_grads", bw_tensor_output_requires_grads));
if (src_attrs.find("comment") != src_attrs.end() && utils::HasString(src_attrs.at("comment"))) {
attrs.push_back(MakeAttribute("comment", src_attrs.at("comment").s()));
}
std::vector<ArgDef> output_args;
for (int i = 0; i < GetSrcNodeInputSize(); ++i) {
if (IsGradientRequiredForSrcNodeInput(i)) {

114
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -3871,6 +3871,10 @@ Return true if all elements are true and false otherwise.
"Indicate if the model is exported in training_mode, by default, False.",
AttributeProto::INT,
static_cast<int64_t>(0))
.Attr(
"comment",
"comment only for debugging purposes.",
AttributeProto::STRING, false)
.TypeConstraint(
"T",
OpSchema::all_tensor_types(),
@ -3891,16 +3895,16 @@ Return true if all elements are true and false otherwise.
"PythonOp's input list and \"input_tensor_types\" attribute should have the same length.");
for (auto i = 0; i < input_tensor_types_count; ++i) {
const auto inferred_input_type = ctx.getInputType(i);
ORT_ENFORCE(inferred_input_type, "PythonOp's ", i, "th input type is missing.");
ORT_ENFORCE(inferred_input_type, "PythonOp's ", i, "-th input type is missing.");
ORT_ENFORCE(inferred_input_type->value_case() == TypeProto::kTensorType,
"PythonOp's ", i, "th input type must be a tensor.");
"PythonOp's ", i, "-th input type must be a tensor.");
ORT_ENFORCE(inferred_input_type->tensor_type().elem_type() == input_tensor_types_proto->ints().at(i),
"PythonOp's ", i, "th input type must be ",
"PythonOp's ", i, "-th input type must be ",
TensorProto_DataType_Name(input_tensor_types_proto->ints().at(i)), " but got ",
TensorProto_DataType_Name(inferred_input_type->tensor_type().elem_type()));
}
// The first output is a pointer which points to
// The first output is a pointer that points to
// a Python object created by torch.autograd.Function.apply.
// For details, see how we interpret it (the 1st input of PythonOpGrad)
// in PythonOpGrad's implementation.
@ -3913,22 +3917,39 @@ Return true if all elements are true and false otherwise.
// This is a required field.
ORT_ENFORCE(output_tensor_types_proto, "PythonOp's must have \"output_tensor_types\" attribute.");
size_t rank_count = 0;
// Set inferred output types.
for (auto i = 1; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i - 1)));
// Create symbolic shape.
const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks");
ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
for (int64_t j = 0; j < output_tensor_ranks->ints().at(i - 1); ++j) {
std::stringstream ss;
ss << "PythonOp_unknown_rank_" << rank_count++;
rank_only_shape.add_dim()->set_dim_param(ss.str());
std::string func_name = getAttribute(ctx, "name", "");
if (func_name == "_InspectActivation" || func_name == "_IncrementStep") {
// PythonOp with the name attribute being "_InspectActivation" or "_IncrementStep" will behave exactly the
// same as a normal PythonOp when execution. The only difference is that:
// 1). those ops having the same number of tensor inputs and tensor outputs;
// 2). and the i-th output tensor's shape is the same as i-th input tensor's shape.
// Be noted, the count of custom autograd function might be bigger than the output count, because there might
// be other non-tensor constant inputs (string, object, int, tuple, etc). But we did not make those constant
// inputs as ONNX op's input, instead they are stored in the attributes.
ORT_ENFORCE(ctx.getNumOutputs() == ctx.getNumInputs() + 1); // The output contains one extra context info.
// Set inferred output types.
for (size_t i = 1; i < static_cast<size_t>(ctx.getNumOutputs()); ++i) {
size_t input_idx = i - static_cast<size_t>(1);
propagateElemTypeFromInputToOutput(ctx, input_idx, i);
propagateShapeFromInputToOutput(ctx, input_idx, i);
}
} else {
size_t rank_count = 0;
// Create a symbolic shape.
const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks")->ints();
// Set inferred output types.
for (auto i = 1; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i - 1)));
ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
for (int64_t j = 0; j < output_tensor_ranks.at(i - 1); ++j) {
std::stringstream ss;
ss << "PythonOp_unknown_rank_" << rank_count++;
rank_only_shape.add_dim()->set_dim_param(ss.str());
}
// Assign symbolic shape.
updateOutputShape(ctx, i, rank_only_shape);
// Assign symbolic shape.
updateOutputShape(ctx, i, rank_only_shape);
}
}
});
@ -3964,13 +3985,13 @@ Return true if all elements are true and false otherwise.
AttributeProto::STRING)
.Attr(
"inplace",
"Indicate if the output should reuse input memory. Todo(pengwa): do we really need it?",
"Indicate if the output should reuse input memory. Todo(pengwa): do we need it?",
AttributeProto::INT,
static_cast<int64_t>(0))
.Attr(
"input_tensor_types",
"Input types of autograd.Function.backward (including only tensor inputs)."
"This attribute is mostly used for input checks for better robustnes.",
"This attribute is mostly used for input checks for better robustness.",
AttributeProto::INTS,
false)
.Attr(
@ -3998,6 +4019,11 @@ Return true if all elements are true and false otherwise.
"A string inidicating autograd.Function.backward outputs's type."
"value 'c' - non-tensor output; value 'd' - tensor output.",
AttributeProto::STRING)
.Attr(
"comment",
"comment only for debugging purposes.",
AttributeProto::STRING,
false)
.TypeConstraint(
"T",
OpSchema::all_tensor_types(),
@ -4022,11 +4048,11 @@ Return true if all elements are true and false otherwise.
// For details, see how we interpret it in PythonOpGrad implementation.
for (auto i = 1; i < input_count; ++i) {
const auto inferred_input_type = ctx.getInputType(i);
ORT_ENFORCE(inferred_input_type, "PythonOpGrad's ", i, "th input type is missing.");
ORT_ENFORCE(inferred_input_type, "PythonOpGrad's ", i, "-th input type is missing.");
ORT_ENFORCE(inferred_input_type->value_case() == TypeProto::kTensorType,
"PythonOpGrad's ", i, "th input type must be a tensor.");
"PythonOpGrad's ", i, "-th input type must be a tensor.");
ORT_ENFORCE(inferred_input_type->tensor_type().elem_type() == input_tensor_types_proto->ints().at(i - 1),
"PythonOpGrad's ", i, "th input type must be ", input_tensor_types_proto->ints().at(i - 1));
"PythonOpGrad's ", i, "-th input type must be ", input_tensor_types_proto->ints().at(i - 1));
}
// Load expected output types.
@ -4035,20 +4061,36 @@ Return true if all elements are true and false otherwise.
"PythonOpGrad's output list and \"output_tensor_types\" attribute should have the same length.");
// This is a required field.
ORT_ENFORCE(output_tensor_types_proto, "PythonOpGrad's must have \"output_tensor_types\" attribute.");
// Set inferred output types.
size_t rank_count = 0;
for (auto i = 0; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i)));
const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks");
ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
for (int64_t j = 0; j < output_tensor_ranks->ints().at(i); ++j) {
std::stringstream ss;
ss << "PythonOpGrad_unknown_rank_" << rank_count++;
rank_only_shape.add_dim()->set_dim_param(ss.str());
}
// Assign symbolic shape.
updateOutputShape(ctx, i, rank_only_shape);
std::string func_name = getAttribute(ctx, "name", "");
if (func_name == "_InspectActivation" || func_name == "_IncrementStep") {
// PythonOpGrad with name attribute being "_InspectActivation" or "_IncrementStep" will behave exactly
// the same as a normal PythonOpGrad when execution. The only difference is that:
// 1). those ops having the same number of tensor inputs and tensor outputs;
// 2). and the i-th output tensor's shape is same as i-th input tensor's shape.
ORT_ENFORCE(ctx.getNumOutputs() == ctx.getNumInputs() - 1); // inputs contains one extra context input
for (size_t i = 0; i < static_cast<size_t>(ctx.getNumOutputs()); ++i) {
size_t input_idx = i + static_cast<size_t>(1);
propagateElemTypeFromInputToOutput(ctx, input_idx, i);
propagateShapeFromInputToOutput(ctx, input_idx, i);
}
} else {
// Set inferred output types.
size_t rank_count = 0;
const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks")->ints();
for (auto i = 0; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i)));
ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
for (int64_t j = 0; j < output_tensor_ranks.at(i); ++j) {
std::stringstream ss;
ss << "PythonOpGrad_unknown_rank_" << rank_count++;
rank_only_shape.add_dim()->set_dim_param(ss.str());
}
// Assign symbolic shape.
updateOutputShape(ctx, i, rank_only_shape);
}
}
});

7
orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py
View file

@ -118,6 +118,7 @@ def _export_pt_1_10(g, n, *args, **kwargs):
input_pointer_scalar_positions = []
tensor_args = []
debug_comment = ""
# Encode inputs to autograd.Function.
for i, arg, call_type in zip(range(len(args)), args, cconv):
if call_type == "d":
@ -155,6 +156,11 @@ def _export_pt_1_10(g, n, *args, **kwargs):
ORTModuleONNXModelException, Exception(f"Unknown argument type found: {type(arg)}.")
)
else:
if name == "_InspectActivation" and isinstance(arg, str):
# _InspectActivation is a special case where the first argument is a string
# that is used to determine the activation name to be inspected.
debug_comment += arg
# All other inputs are accessed via "pointers".
input_pointer_scalar_positions.append(i)
input_pointer_scalars.append(id(arg))
@ -187,6 +193,7 @@ def _export_pt_1_10(g, n, *args, **kwargs):
"output_tensor_types_i": output_tensor_types,
"output_tensor_ranks_i": output_tensor_ranks,
"training_mode_i": 1 if training_mode else 0,
"comment_s": debug_comment,
}
if len(input_int_scalars) > 0:

9
orttraining/orttraining/python/training/utils/hooks/__init__.py Normal file
View file

@ -0,0 +1,9 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from ._subscriber_manager import SubscriberManager
from ._statistics_subscriber import StatisticsSubscriber

113
orttraining/orttraining/python/training/utils/hooks/_statistics_subscriber.py Normal file
View file

@ -0,0 +1,113 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from pathlib import Path
from typing import Union
import os
import warnings
import shutil
import torch
from ._subscriber_base import SubscriberBase
class StatisticsSubscriber(SubscriberBase):
"""
This subscriber is used to dump the activation statistics into files.
Each activation will be summarized into 1 or 2 files, depending on whether it is used in the backward pass.
> In the forward pass, summarize the tensor's statistics and write to a file.
> In the backward pass, summarize the tensor's gradient statistics and write it into another file.
So for each run step, there will be many files.
Currently, the statistics mainly include:
> Number of inf/nan values.
> Common statistics: tensor shape, data type, total element size, min/max/mean/std of the tensor elements.
> Number of zero/negative/positive elements.
> Sampled elements (taking the first 128 elements for example).
> To be extended...
`merge_activation_summary.py` can be used to merge the files into one file per training step.
"""
def __init__(
self,
output_dir: str,
start_step: Union[None, int] = None,
end_step: Union[None, int] = None,
override_output_dir: bool = False,
):
"""
Steps in [start_step, end_step) will run subscriber actions.
Args:
output_dir: the directory in all activation statistics files will be stored.
start_step: the first step that runs subscriber actions.
end_step: the end step (exclusively) that runs subscriber actions.
override_output_dir: whether `output_dir` can be overridden if it already exists.
"""
super().__init__(start_step=start_step, end_step=end_step)
self._output_dir = output_dir
if os.path.exists(self._output_dir):
if override_output_dir:
warnings.warn(f"Output directory {self._output_dir} already exists, overriding it.")
shutil.rmtree(self._output_dir)
else:
raise FileExistsError(
f"Output directory {self._output_dir} already exists. "
"Set override_output_dir=True for StatisticsSubscriber if this is the intention."
)
def module_post_forward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
output_file_path = os.path.join(f"{self._output_dir}", f"step_{step}")
return self._summarize_activations(activation, depth, name, output_file_path, True)
def module_pre_backward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
output_file_path = os.path.join(f"{self._output_dir}", f"step_{step}")
return self._summarize_activations(activation, depth, name, output_file_path, False)
def _summarize_activations(self, tensor: torch.Tensor, depth: int, name: str, step_folder: str, is_forward: bool):
display_name = name + " forward run" if is_forward is True else name + " backward run"
output_file_name = name + "_forward" if is_forward is True else name + "_backward"
if tensor is None or not isinstance(tensor, torch.Tensor):
print(f"{display_name} not a torch tensor, value: {tensor}")
return
step_path = Path(step_folder)
if not step_path.exists():
step_path.mkdir(parents=True, exist_ok=False)
order_file_path = step_path / "order.txt"
tensor_file_path = step_path / output_file_name
# This is to try the best effort to align the count of numbers per line for easier comparison in diff views,
# though it does not always guarantee to do this way.
torch.set_printoptions(precision=6, linewidth=128)
flatten_array = tensor.flatten()
zero_tensor = torch.tensor(0, dtype=flatten_array.dtype, device=flatten_array.device)
num_nan = torch.isnan(flatten_array).sum()
num_inf = torch.isinf(flatten_array).sum()
with order_file_path.open(mode="a", encoding="utf-8") as f:
f.write(f"{output_file_name}\n")
with tensor_file_path.open(mode="w", encoding="utf-8") as f:
f.write(
f"{'>'*depth + display_name} shape: {tensor.shape} dtype: {tensor.dtype} size: {flatten_array.size()} \n"
f"min: {flatten_array.min()} max: {flatten_array.max()}, mean: {flatten_array.mean()}, "
f"std: {flatten_array.std()} \n"
f"nan: {num_nan}, inf: {num_inf}\n"
)
f.write(f"samples(top 128): {flatten_array[:128]}\n")
f.write(
f"neg: {torch.less(flatten_array, zero_tensor).to(torch.int64).sum()}, "
f"pos: {torch.greater(flatten_array, zero_tensor).to(torch.int64).sum()}, "
f"zero: {torch.eq(flatten_array, zero_tensor).to(torch.int64).sum()},\n"
)
f.write(f"{'='*16}\n")

73
orttraining/orttraining/python/training/utils/hooks/_subscriber_base.py Normal file
View file

@ -0,0 +1,73 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from typing import Union
import sys
import torch
class SubscriberBase:
"""
Base class for all module hook subscribers.
Currently, the hook here only means post-forward hook and pre-backward hook.
A module hook subscriber is a class that implements the `module_post_forward_impl` and `module_pre_backward_impl`
function.
> The post_forward hook is called after the activation is generated in the forward path.
> The pre_backward hook is called before the activation gradient is computed.
The post_forward path:
Module_A generates activation tensor_a --> Post forward hook (calling subscribers' forward one by one) -->
Module_B generates activation tensor_b --> ...
The pre_backward path:
Module_B backward run, tensor_b's gradient is computed as tensor_b_grad -->
Pre-backward hook (calling subscribers' backward one by one) -->
Module_A backward run, tensor_a's gradient is computed as tensor_a_grad
Be noted: the "Pre"/"Post" is described from the perspective of Module_A.
"""
def __init__(self, start_step: Union[None, int], end_step: Union[None, int]):
"""
Steps in [start_step, end_step) will run the subscriber's actions, and other steps will skip.
If start_step is None, 0 is given; if end_step is None, sys.maxsize is given.
"""
self._start_step: int = start_step if start_step is not None else 0
self._end_step: int = end_step if end_step is not None else sys.maxsize
def module_post_forward(self, activation: torch.Tensor, depth: int, name: str, step: int):
"""
This function will be run after the torch Module forward is completed.
Args:
activation: Tensor to be inspected.
depth: The indent level of the torch Module generating `activation`.
name: The unique name for the `activation`.
step: Current execution step.
"""
if self._start_step <= step < self._end_step:
self.module_post_forward_impl(activation, depth, name, step)
def module_pre_backward(self, activation: torch.Tensor, depth: int, name: str, step: int):
"""
This function will be run before the torch Module backward run.
Args:
activation: Tensor to be inspected.
depth: The indent level of the torch Module generating `activation`.
name: The unique name for the `activation`.
step: Current execution step.
"""
if self._start_step <= step < self._end_step:
self.module_pre_backward_impl(activation, depth, name, step)
def module_post_forward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
raise NotImplementedError()
def module_pre_backward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
raise NotImplementedError()

300
orttraining/orttraining/python/training/utils/hooks/_subscriber_manager.py Normal file
View file

@ -0,0 +1,300 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from collections import abc
from typing import Callable, List, Union
import torch
from onnxruntime.training.ortmodule import ORTModule
from ._subscriber_base import SubscriberBase
class _RuntimeStates:
"""
A data struct holding states for runtime context. Tho kinds of states are included:
> Global states that are one-time collected during model hook registration. A global execution step is
also initialized to reflect how many steps have been executed, it will get updated after each step
completes its forward path.
> Intra-execution step states, initialized and cleaned up intended only for current execution step.
Usually, it carries intermediate information during the model execution.
"""
class _GlobalStates:
def __init__(self):
# Used to track current execution step, e.g. how many forward/backward path is called.
self.execution_step = 0
# Used to store the depth of each module, which indicate the indentation level of the module.
self.module_index_to_depth = {}
# Used to store the unique id of each sequential activation.
self.module_to_module_index = {}
self.subscribers = set()
class _ExecutionStepStates:
def __init__(self):
# Used to store the activation tensor names, if already handled, then skipped.
# Need to clear after each step.
self.observed_activation_names = {}
def __init__(self):
self.global_states = _RuntimeStates._GlobalStates()
self.reset_step_states()
def reset_step_states(self):
self.execution_step_states = _RuntimeStates._ExecutionStepStates()
class _InspectActivation(torch.autograd.Function):
"""
This class is used to run the subscriber's forward and backward functions.
"""
@staticmethod
def forward(ctx, activation_name: str, module_idx: int, run_ctx: _RuntimeStates, input_tensor):
"""
Make sure there is a same number of `tensor` type inputs and outputs.
This is enforced by ORT's PythonOp's schema check.
"""
depth = run_ctx.global_states.module_index_to_depth[module_idx]
input_tensor_copied = None
if input_tensor is None or not isinstance(input_tensor, torch.Tensor):
input_tensor_copied = input_tensor
else:
input_tensor_copied = input_tensor.detach().clone()
ctx.current_step = run_ctx.global_states.execution_step
ctx.name = activation_name
ctx.id = module_idx
ctx.depth = depth
ctx.subscribers = run_ctx.global_states.subscribers
# Run subscribers sequentially.
for subscriber in run_ctx.global_states.subscribers:
subscriber.module_post_forward(input_tensor_copied, depth, activation_name, ctx.current_step)
return input_tensor.detach() if input_tensor is not None else None
@staticmethod
def backward(ctx, grad_output):
val = None
if grad_output is None or not isinstance(grad_output, torch.Tensor):
val = grad_output
else:
val = grad_output.detach().clone()
for subscriber in ctx.subscribers:
subscriber.module_pre_backward(val, ctx.depth, ctx.name, ctx.current_step)
return None, None, None, grad_output.detach() if grad_output is not None else None
class _IncrementStep(torch.autograd.Function):
"""
This class is used to manage the global execution step, e.g.
global step increment by one, once a full forward path is completed and the state clear.
This autograd Function is registered as a post-forward hook to the root module. So once the root
module's forward path is completed, this backward function will be called immediately, triggering
global step increment and state clear.
"""
@staticmethod
def forward(ctx, run_ctx, input_tensor):
"""
Make sure there is a same number of `tensor` inputs and outputs.
This is enforced by ORT's PythonOp's schema check.
"""
ctx.current_step = run_ctx.global_states.execution_step
ctx.run_ctx = run_ctx
# We cannot do the step incremental here. Imagine the outside-most module has multiple outputs,
# we need to increase the step only at the very last output handling.
# We avoid the complexity to probe the last output handling, and instead, we assume once
# the very first backward of the outside-most module is called, then the forward pass MUST be completed.
# Be noted: it is not safe to register _IncrementStep only for one of the outputs of the outside-most module,
# because we are not sure which output branch is executed earlier, for example.
# OuterMostModuleOutputs
# / \
# OuterMostModuleOutputs_0_0th_output OuterMostModuleOutputs_0_1th_output
# | |
# PythonOp(_InspectActivation) PythonOp(_InspectActivation)
# | |
# PythonOp(_IncrementStep) graph output
# |
# graph output
# The PythonOp(_InspectActivation) (who relies on global step) after 1th output is possible
# to run before or after PythonOp(_IncrementStep), so increasing the step is not safe.
return input_tensor.detach() if isinstance(input_tensor, torch.Tensor) else input_tensor
@staticmethod
def backward(ctx, grad_output):
# In case there are multiple backward calls for multiple outputs of the outside-most module.
if ctx.current_step == ctx.run_ctx.global_states.execution_step:
if ctx.current_step >= 0:
print(f"{'='*6} Completed forward pass for STEP {ctx.current_step} {'='*6}")
ctx.run_ctx.global_states.execution_step += 1
ctx.run_ctx.reset_step_states()
return None, grad_output.detach() if isinstance(grad_output, torch.Tensor) else grad_output
class SubscriberManager:
"""
This class is used to manage all the subscribers and register the post-forward hook to the root module.
`subscribe()` is used to register a list of subscribers.
Currently, the hook handled here is post forward hook for nn.Module. The hook is registered for all nn.Modules
recursively. Each hook inserts a PythonOp for every tensor output generated by the corresponding module.
Each subscriber implementation is called in the PythonOp's forward function, and backward function.
There is one special handling for global step increment and state clear. A post-forward hook is registered
for the outside-most module, which is the root module. In that hook, _IncrementStep is called, which will
increase the step by 1 once the very first time its backward is called (check _IncrementStep for details).
"""
def __init__(self):
self._run_ctx: _RuntimeStates = _RuntimeStates()
def subscribe(self, module: Union[torch.nn.Module, ORTModule], subscribers: List[SubscriberBase]):
"""
The API is called externally to register hooks that are implicitly defined by subscribers.
Each time all global states will be cleaned up once called.
"""
if not isinstance(module, torch.nn.Module):
raise ValueError("module must be a torch.nn.Module instance")
self._reset_all_states()
if isinstance(module, ORTModule):
module = module.module
for subscriber in subscribers:
if not isinstance(subscriber, SubscriberBase):
raise ValueError("subscriber must be a SubscriberBase instance")
self._run_ctx.global_states.subscribers.add(subscriber)
self._initialize(module)
def _reset_all_states(self):
self._run_ctx = _RuntimeStates()
def _initialize(self, module: torch.nn.Module):
"""
Register hooks for the specified module.
"""
if len(self._run_ctx.global_states.subscribers) == 0:
raise RuntimeError("No subscribers are registered.")
next_module_index = [0]
# Register post forward hook for every module, inside the hook, we loop every tensor output of the module,
# and wrap it with an autograd Function called _InspectActivation (which takes in a tensor and returns the same
# tensor). In this way, we keep ORT and PyTorch run have the same boundary to check activation equality.
self._register_hooks_recursively(module, 1, next_module_index)
# Register post forward hook for the outside-most module, then we increase the dump step.
# Be noted, if backward is not triggered, the global dump step remains the original number,
# which means the subsequent run will override the previous dump files. This indeed happens to imagine ORTModule
# firstly export graph (run the forward only), after the gradient graph is built, another forward+backward is
# triggered, override the previous dump files.
def _post_forward_outmost_module_hook(module, _, module_outputs):
def _apply_to_tensors_func(_, outputs):
return _IncrementStep.apply(self._run_ctx, outputs)
return self._apply_function_to_tensors(module, module_outputs, _apply_to_tensors_func)
module.register_forward_hook(_post_forward_outmost_module_hook)
def _register_hooks_recursively(self, module: torch.nn.Module, depth: int, next_module_index: List[int]):
"""
Called to register hooks for every `torch.nn.Module`. Due to `Module` can contain child `Module`s,
this function is called recursively by passing in `next_module_index` - a list of int to maintain a
global incremental unique module id.
Args:
module: torch.nn.Module to register hook.
depth: the indent of the module compared with the outside-most Module.
next_module_index: list of int, carrying a global unique module index that can be used next.
"""
module_index = next_module_index[0]
self._run_ctx.global_states.module_index_to_depth[module_index] = depth
self._run_ctx.global_states.module_to_module_index[module] = module_index
for child in module.children():
if isinstance(child, torch.nn.Module) and child not in self._run_ctx.global_states.module_to_module_index:
next_module_index[0] += 1
self._register_hooks_recursively(child, depth + 1, next_module_index)
def _post_forward_module_hook(module, _, module_outputs):
if module in self._run_ctx.global_states.module_to_module_index and isinstance(module, torch.nn.Module):
module_index = self._run_ctx.global_states.module_to_module_index[module]
def _apply_to_tensors_func(index, activation_tensor):
name = f"{module.__class__.__name__}_{module_index}_{index}th_output"
if name not in self._run_ctx.execution_step_states.observed_activation_names:
self._run_ctx.execution_step_states.observed_activation_names[name] = True
return _InspectActivation.apply(name, module_index, self._run_ctx, activation_tensor)
return activation_tensor
return self._apply_function_to_tensors(module, module_outputs, _apply_to_tensors_func)
return module_outputs
module.register_forward_hook(_post_forward_module_hook)
def _is_builtin_type(self, obj):
# https://stackoverflow.com/a/17795199
return obj.__class__.__module__ in ["__builtin__", "builtins"]
def _apply_function_to_tensors(self, module: torch.nn.Module, data, func: Callable):
"""
Apply func to all tensors in the given object.
Args:
module: the module that generates the tensors.
data: the object that contains activation tensors.
func: the function to apply to the tensors.
"""
tensor_output_idx: List[int] = [0]
def _apply_to_tensors_by_flatten(
module: torch.nn.Module,
index_for_tensor_output: List[int],
outputs,
func: Callable,
):
if isinstance(outputs, abc.Sequence):
touched_outputs = []
for output in outputs:
touched_output = _apply_to_tensors_by_flatten(module, index_for_tensor_output, output, func)
touched_outputs.append(touched_output)
return outputs.__class__(touched_outputs)
if isinstance(outputs, abc.Mapping):
# apply inplace to avoid recreating dict inherited objects
for key in outputs:
outputs[key] = _apply_to_tensors_by_flatten(
module,
index_for_tensor_output,
outputs[key],
func,
)
return outputs
if isinstance(outputs, torch.Tensor):
cur_id = index_for_tensor_output[0]
index_for_tensor_output[0] += 1
return func(cur_id, outputs)
if not self._is_builtin_type(outputs):
raise RuntimeError(f"Unknown type {type(outputs)}")
return outputs
return _apply_to_tensors_by_flatten(module, tensor_output_idx, data, func)

156
orttraining/orttraining/python/training/utils/hooks/merge_activation_summary.py Normal file
View file

@ -0,0 +1,156 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
"""
This merges convergence debugging per activation summary files into per-step summary files.
Both PyTorch and ORT run result directories are needed, check `StatisticsSubscriber` usage from
[ORTModule_Convergence_Notes](docs/ORTModule_Convergence_Notes.md) for how to export the results.
Since ORT and PyTorch run the compute in different orders (e.g. different typological orders sort
implementations), when we generate a per-step summary, we want the summary to be comparable between
ORT and PyTorch run. So during the merge, the same typological order is used.
Example:
python merge_activation_summary.py --pt_dir pt_out --ort_dir ort_out --output_dir /tmp/output
"""
from pathlib import Path
import argparse
import logging
import os
import shutil
logger = logging.getLogger(__name__)
def generate_summaries_per_step(args):
pt_dir = args.pt_dir
ort_dir = args.ort_dir
output_path = Path(args.output_dir)
output_path.mkdir(parents=True, exist_ok=False)
# We should use the order.txt generated by PyTorch run, which means, we follow the PyTorch typological order to compare
# activation results. Here we assume to get the order.txt from pt_dir/step_0/order.txt
topo_order_file_path = Path(f"{pt_dir}/step_0/order.txt")
src_ort_path = Path(ort_dir)
src_pt_path = Path(pt_dir)
merge_ort_path = output_path / "merge_ort"
merge_pt_path = output_path / "merge_pt"
def generate_summary_per_step(topo_order_file_path: Path, dump_src_path: Path, merge_dest_path: Path):
logger.warning(
"Start generating summary per step for [%s] following typological order in [%s]",
dump_src_path.as_posix(),
topo_order_file_path.as_posix(),
)
with topo_order_file_path.open(mode="r", encoding="utf-8") as order_file:
tensor_name_in_order = order_file.readlines()
if merge_dest_path.exists():
shutil.rmtree(merge_dest_path.as_posix())
merge_dest_path.mkdir(parents=True, exist_ok=False)
for dump_step_path in dump_src_path.iterdir():
if dump_step_path.is_dir():
step_name = dump_step_path.name
merge_filename_for_sub_dir = merge_dest_path / f"{step_name}_.txt"
# Open merge_filename_for_sub_dir in write mode
with merge_filename_for_sub_dir.open(mode="w", encoding="utf-8") as outfile:
for filename in tensor_name_in_order:
filename = filename.rstrip("\n")
full_filename = dump_step_path / filename
if not full_filename.exists():
# Be noted that some tensor handled in PyTorch might be missing in ORT graph
# (if the activation is not used by others, which is pruned during export)
logger.warning("tensor %s not exist", full_filename)
continue
with full_filename.open(mode="r", encoding="utf-8") as infile:
outfile.write(infile.read())
outfile.write("\n")
logger.warning(
"Finish generating summary per step for [%s] following typological order in [%s], merged files are in [%s]",
dump_src_path.as_posix(),
topo_order_file_path.as_posix(),
merge_dest_path.as_posix(),
)
generate_summary_per_step(topo_order_file_path, src_pt_path, merge_pt_path)
generate_summary_per_step(topo_order_file_path, src_ort_path, merge_ort_path)
def parse_arguments():
"""Parse arguments
Merge per-step summary from dumped files.
Data are collected with `StatisticsSubscriber` for ORT and PyTorch runs.
Required parameters include
> dump root folders for ORT and PyTorch.
> output folder.
Returns:
Namespace: arguments
"""
parser = argparse.ArgumentParser()
parser.add_argument(
"--pt_dir",
required=True,
type=str,
help="Root of input directory of PyTorch run result of activation dump.",
)
parser.add_argument(
"--ort_dir",
required=True,
type=str,
help="Root of input directory of ORTModule run result of activation dump.",
)
parser.add_argument(
"--output_dir",
required=True,
type=str,
help="Root of output directory for generated PyTorch/ORTModule per-step summaries.",
)
parser.add_argument(
"--overwrite",
required=False,
action="store_true",
help="Overwrite exists output files.",
)
args = parser.parse_args()
return args
def main():
args = parse_arguments()
if os.path.exists(args.output_dir):
if args.overwrite:
logger.warning("Output directory %s already exists, overwriting it.", args.output_dir)
shutil.rmtree(args.output_dir)
else:
raise FileExistsError(
f"Output directory {args.output_dir} already exists. " "Enable --overwrite to allow overwriting it."
)
logger.info("Arguments: %s", str(args))
generate_summaries_per_step(args)
if __name__ == "__main__":
main()

10
orttraining/orttraining/test/python/orttraining_ortmodule_tests.py
View file

@ -142,6 +142,14 @@ def run_data_sampler_tests(cwd, log):
run_subprocess(command, cwd=cwd, log=log).check_returncode()
def run_hooks_tests(cwd, log):
log.debug("Running: Data hooks tests")
command = [sys.executable, "-m", "pytest", "-sv", "orttraining_test_hooks.py"]
run_subprocess(command, cwd=cwd, log=log).check_returncode()
def run_pytorch_export_contrib_ops_tests(cwd, log):
log.debug("Running: PyTorch Export Contrib Ops Tests")
@ -192,6 +200,8 @@ def main():
run_data_sampler_tests(cwd, log)
run_hooks_tests(cwd, log)
run_experimental_gradient_graph_tests(cwd, log)
# TODO(bmeswani): Enable this test once it can run with latest pytorch

138
orttraining/orttraining/test/python/orttraining_test_hooks.py Normal file
View file

@ -0,0 +1,138 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
import os
import tempfile
import pytest
import torch
from onnxruntime.training.ortmodule import ORTModule
from onnxruntime.training.utils.hooks import SubscriberManager, StatisticsSubscriber
class NeuralNetSingleOutput(torch.nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super().__init__()
self.fc1 = torch.nn.Linear(input_size, hidden_size)
self.relu = torch.nn.ReLU()
self.fc2 = torch.nn.Linear(hidden_size, num_classes)
def forward(self, input1, input2):
model_input = input1 + input2
out = self.fc1(model_input)
out = self.relu(out)
out = self.fc2(out)
return out
class NeuralNetMultipleOutputs(torch.nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super().__init__()
self.fc1 = torch.nn.Linear(input_size, hidden_size)
self.relu = torch.nn.ReLU()
self.fc2 = torch.nn.Linear(hidden_size, num_classes)
def forward(self, input1, input2):
model_input = input1 + input2
out = self.fc1(model_input)
out = self.relu(out)
out = self.fc2(out)
return out, model_input
@pytest.mark.parametrize("device", ["cpu", "cuda"])
@pytest.mark.parametrize("backend", ["torch", "ortmodule"])
def test_statistic_subscriber_single_output(device, backend):
input_size = 8
hidden_size = 16
num_classes = 32
model = NeuralNetSingleOutput(input_size, hidden_size, num_classes)
model.to(device)
model.train()
with tempfile.TemporaryDirectory() as temporary_dir:
output_dir_path = os.path.join(temporary_dir, f"{backend}_out")
sub_manager = SubscriberManager()
sub_manager.subscribe(model, [StatisticsSubscriber(output_dir_path, override_output_dir=True)])
if backend == "ortmodule":
model = ORTModule(model)
batch_size = 4
input1_tensor = torch.randn(batch_size, input_size, device=device)
input2_tensor = torch.randn(batch_size, input_size, device=device)
for _ in range(5):
y = model(input1_tensor, input2_tensor)
y.sum().backward()
assert os.path.exists(output_dir_path)
expected_files = [
"order.txt",
"Linear_1_0th_output_forward",
"Linear_1_0th_output_backward",
"NeuralNetSingleOutput_0_0th_output_forward",
"NeuralNetSingleOutput_0_0th_output_backward",
"ReLU_2_0th_output_forward",
"ReLU_2_0th_output_backward",
"Linear_3_0th_output_forward",
"Linear_3_0th_output_backward",
]
for i in range(5):
step_dir = os.path.join(output_dir_path, f"step_{i}")
for file in expected_files:
assert os.path.exists(os.path.join(step_dir, file))
@pytest.mark.parametrize("device", ["cpu", "cuda"])
@pytest.mark.parametrize("backend", ["torch", "ortmodule"])
def test_statistic_subscriber_multiple_outputs(device, backend):
input_size = 8
hidden_size = 16
num_classes = 32
model = NeuralNetMultipleOutputs(input_size, hidden_size, num_classes)
model.to(device)
model.train()
with tempfile.TemporaryDirectory() as temporary_dir:
output_dir_path = os.path.join(temporary_dir, f"{backend}_out")
sub_manager = SubscriberManager()
sub_manager.subscribe(model, [StatisticsSubscriber(output_dir_path, override_output_dir=True)])
if backend == "ortmodule":
model = ORTModule(model)
batch_size = 4
input1_tensor = torch.randn(batch_size, input_size, device=device).requires_grad_(True)
input2_tensor = torch.randn(batch_size, input_size, device=device)
for _ in range(5):
y_output1, y_output2 = model(input1_tensor, input2_tensor)
y = y_output1.sum() + y_output2.sum()
y.backward()
assert os.path.exists(output_dir_path)
expected_files = [
"order.txt",
"NeuralNetMultipleOutputs_0_0th_output_forward",
"NeuralNetMultipleOutputs_0_0th_output_backward",
"NeuralNetMultipleOutputs_0_1th_output_forward",
"NeuralNetMultipleOutputs_0_1th_output_backward",
"Linear_1_0th_output_forward",
"Linear_1_0th_output_backward",
"ReLU_2_0th_output_forward",
"ReLU_2_0th_output_backward",
"Linear_3_0th_output_forward",
"Linear_3_0th_output_backward",
]
for i in range(5):
step_dir = os.path.join(output_dir_path, f"step_{i}")
assert len(os.listdir(step_dir)) == len(expected_files)
for file in expected_files:
assert os.path.exists(os.path.join(step_dir, file))

1
setup.py
View file

@ -545,6 +545,7 @@ if enable_training or enable_training_apis:
"onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator",
"onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.fused_ops",
"onnxruntime.training.utils.data",
"onnxruntime.training.utils.hooks",
]
)