saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

On device training offline tooling (#11520)

Browse source
This commit is contained in:
Baiju Meswani 2022-05-24 18:21:39 -07:00 committed by GitHub
parent d8a1531c37
commit 3a22a866a1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 1434 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
cmake/onnxruntime_python.cmake
View file

@ -320,6 +320,17 @@ if (onnxruntime_ENABLE_TRAINING)
file(GLOB onnxruntime_python_utils_data_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/utils/data/*"
)
if (onnxruntime_ENABLE_TRAINING_ON_DEVICE)
file(GLOB onnxruntime_python_onnxblock_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/onnxblock/*"
)
file(GLOB onnxruntime_python_onnxblock_loss_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/onnxblock/loss/*"
)
file(GLOB onnxruntime_python_onnxblock_optim_srcs CONFIGURE_DEPENDS
"${ORTTRAINING_SOURCE_DIR}/python/training/onnxblock/optim/*"
)
endif()
else()
file(GLOB onnxruntime_python_capi_training_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/python/training/*.py"
@ -644,6 +655,23 @@ if (onnxruntime_ENABLE_TRAINING)
${onnxruntime_python_utils_data_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/data/
)
if (onnxruntime_ENABLE_TRAINING_ON_DEVICE)
add_custom_command(
TARGET onnxruntime_pybind11_state POST_BUILD
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/onnxblock
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/onnxblock/loss
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/onnxblock/optim
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_onnxblock_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/onnxblock/
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_onnxblock_loss_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/onnxblock/loss/
COMMAND ${CMAKE_COMMAND} -E copy
${onnxruntime_python_onnxblock_optim_srcs}
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/onnxblock/optim/
)
endif()
endif()
if (onnxruntime_USE_DNNL)

13
orttraining/orttraining/python/training/onnxblock/__init__.py Normal file
View file

@ -0,0 +1,13 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# __init__.py
from .model import Model, TrainingModel
from .checkpoint_utils import save_checkpoint
from . import loss, optim
from .model_accessor import onnx_model
import onnx
_producer_name = "onnxblock offline tooling"
_opset_import = onnx.helper.make_opsetid("com.microsoft", 1)

146
orttraining/orttraining/python/training/onnxblock/_building_blocks.py Normal file
View file

@ -0,0 +1,146 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# _building_blocks.py
import copy
import onnx
import onnxruntime.training.onnxblock as onnxblock
import onnxruntime.training.onnxblock.model_accessor as accessor
import onnxruntime.training.onnxblock._graph_utils as graph_utils
class Sub(onnxblock.Model):
"""Adds Sub node to an onnx model."""
def __init__(self):
super(Sub, self).__init__()
def build(self, sub_input_name1, sub_input_name2):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph node for sub
sub_node_input_names = [sub_input_name1, sub_input_name2]
sub_node_output_name = graph_utils.generate_random_graph_name("sub_output")
sub_node_output_names = [sub_node_output_name]
sub_node = onnx.helper.make_node(
"Sub",
sub_node_input_names,
sub_node_output_names,
name=graph_utils.generate_random_graph_name("Sub"),
)
onnx_model.graph.node.append(sub_node)
# create the graph output for sub
graph_output = copy.deepcopy(
graph_utils.get_output_from_output_name(onnx_model, sub_input_name1)
)
graph_output.name = sub_node_output_name
del onnx_model.graph.output[:]
onnx_model.graph.output.append(graph_output)
return sub_node_output_name
class Pow(onnxblock.Model):
"""Adds Pow node to the onnx model."""
def __init__(self, exponent):
super(Pow, self).__init__()
self._exponent = exponent
def build(self, pow_input_name):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph initializer for the exponent
pow_node_exponent_name = graph_utils.generate_random_graph_name("pow_exponent")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(
pow_node_exponent_name, onnx.TensorProto.FLOAT, [1], [self._exponent]
)
)
# create the graph node for pow
pow_node_input_names = [pow_input_name, pow_node_exponent_name]
pow_node_output_name = graph_utils.generate_random_graph_name("pow_output")
pow_node_output_names = [pow_node_output_name]
pow_node = onnx.helper.make_node(
"Pow",
pow_node_input_names,
pow_node_output_names,
name=graph_utils.generate_random_graph_name("Pow"),
)
onnx_model.graph.node.append(pow_node)
# create the graph output for pow
graph_output = copy.deepcopy(
graph_utils.get_output_from_output_name(onnx_model, pow_input_name)
)
graph_output.name = pow_node_output_name
del onnx_model.graph.output[:]
onnx_model.graph.output.append(graph_output)
return pow_node_output_name
class _Reduce(onnxblock.Model):
"""Base class for the reduce blocks."""
def __init__(self):
super(_Reduce, self).__init__()
def _reduce(self, reduce_input_name, reduction_op):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph node for reduce
reduce_node_input_names = [reduce_input_name]
reduce_node_output_name = graph_utils.generate_random_graph_name(
"reduce_output"
)
reduce_node_output_names = [reduce_node_output_name]
reduce_node = onnx.helper.make_node(
reduction_op,
reduce_node_input_names,
reduce_node_output_names,
name=graph_utils.generate_random_graph_name(reduction_op),
)
onnx_model.graph.node.append(reduce_node)
# create the graph output for reduce
reduce_input = copy.deepcopy(
graph_utils.get_output_from_output_name(onnx_model, reduce_input_name)
)
output_rank = len(reduce_input.type.tensor_type.shape.dim)
graph_outputs = [
onnx.helper.make_tensor_value_info(
reduce_node_output_name, onnx.TensorProto.FLOAT, [1] * output_rank
)
]
del onnx_model.graph.output[:]
onnx_model.graph.output.extend(graph_outputs)
return reduce_node_output_name
class ReduceMean(_Reduce):
"""Adds ReduceMean node to the onnx model."""
def __init__(self):
super(ReduceMean, self).__init__()
def build(self, reduce_input_name):
return super()._reduce(reduce_input_name, "ReduceMean")
class ReduceSum(_Reduce):
"""Adds ReduceSum node to the onnx model."""
def __init__(self):
super(ReduceSum, self).__init__()
def build(self, reduce_input_name):
return super(ReduceSum, self)._reduce(reduce_input_name, "ReduceSum")

205
orttraining/orttraining/python/training/onnxblock/_graph_utils.py Normal file
View file

@ -0,0 +1,205 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# _graph_utils.py
import copy
import onnx
import random
from onnxruntime.capi._pybind_state import GradientGraphBuilder
def get_output_from_output_name(onnx_model, output_name):
"""Returns the graph output given the output name"""
# Iterate over the graph outputs looking for output_name
for output in onnx_model.graph.output:
if output.name == output_name:
return output
raise LookupError(f"The provided output name {output_name} is not a graph output.")
def get_random_number():
"""Return a random number in the range 0, 1000."""
return random.randint(0, 1000)
def generate_random_graph_name(token):
"""Return a string that can be used in the graph as a graph attribute name."""
return f"{get_random_number()}.{token}"
def build_gradient_graph(
accessor, user_args_requiring_grad, user_args_not_requiring_grad, output_names
):
"""Builds the gradient graph on top of the given input forward only graph."""
model = accessor.model
# Collect names of parameters that need gradients computed
all_args_requiring_gradient = set()
# Move all trainable and non trainable initializers to graph inputs.
# This allows training to pass in the parameters from outside the graph
# so as to share the parameters across multiple sessions.
graph_inputs = model.graph.input
initializers = []
for initializer in model.graph.initializer:
if not initializer.name[0].isdigit():
# Move only those initializers as inputs that are not local
# to the onnx model. i.e. initializers that are model parameters.
# These are tpically those initializers without any number prefixed
# to their names.
graph_inputs.append(
onnx.helper.make_tensor_value_info(
initializer.name, initializer.data_type, initializer.dims
)
)
if initializer.name not in user_args_not_requiring_grad:
all_args_requiring_gradient.add(initializer.name)
else:
# All other initializers stay where they were.
initializers.append(initializer)
# Update the initializers in the graph
del model.graph.initializer[:]
model.graph.initializer.extend(initializers)
# At this point, we have the eval model
accessor.eval_model = copy.deepcopy(model)
# Any graph input that requires gradient, should have been already added to
# args_requiring_grad. So, add these arguments to set of arguments
# whose gradient should be built.
for argument_name in user_args_requiring_grad:
all_args_requiring_gradient.add(argument_name)
# Assumption is that the first graph output is the loss output
if isinstance(output_names, str):
output_names = [output_names]
builder = GradientGraphBuilder(
model.SerializeToString(),
set(output_names),
all_args_requiring_gradient,
output_names[0],
)
builder.build()
gradient_model = onnx.load_from_string(builder.get_model())
# copy the gradient graph into the user's graph
model.graph.CopyFrom(gradient_model.graph)
del model.opset_import[:]
model.opset_import.extend(gradient_model.opset_import)
return all_args_requiring_gradient
def build_gradient_accumulation_graph(grad_model, all_args_requiring_gradient_names):
"""Builds gradient accumulation nodes on top of a training model.
Adds an InPlaceAccumulator node for every gradient so that the gradients
are accumulated in a gradient buffer (which is an input to InPlaceAccumulator).
For example, if there is a gradient in the graph called fc1.weight_grad,
an InPlaceAccumulator will be added for that gradient whose input will
be a graph input (fc1.weight_grad.accumulation.buffer) and the newly
computed gradient (fc1.weight_grad).
gradient.accumulation.buffer gradient
| | |
Ʌ v v
| |_________________________|
| |
| InPlaceAccumulator
| |
| v
|______________________|
"""
# TODO: Avoid hard coded input/output strings
gradient_output_names = {
f"{arg_name}_grad" for arg_name in all_args_requiring_gradient_names
}
graph_inputs = grad_model.graph.input
graph_nodes = grad_model.graph.node
graph_outputs = []
lazy_reset_grad_input_name = "lazy_reset_grad"
gradient_accumulation_name = "accumulation"
gradient_buffer_name_suffix = "buffer"
gradient_output_name_suffix = "out"
for idx, graph_output in enumerate(grad_model.graph.output):
if graph_output.name not in gradient_output_names:
# If graph output is not a gradient, there is no
# need to build the gradient accumulator node for it.
graph_outputs.append(graph_output)
continue
# gradient accumulation node inputs and output names
grad_name = graph_output.name
grad_accumulation_buffer_name = (
f"{grad_name}.{gradient_accumulation_name}.{gradient_buffer_name_suffix}"
)
grad_accumulation_output_name = (
f"{grad_name}.{gradient_accumulation_name}.{gradient_output_name_suffix}"
)
# Gradient accumulation node
acc_node = onnx.helper.make_node(
"InPlaceAccumulator",
[grad_accumulation_buffer_name, grad_name, lazy_reset_grad_input_name],
[grad_accumulation_output_name],
name=f"GradientAccumulator{idx}",
domain="com.microsoft",
)
graph_nodes.append(acc_node)
# grad buffer is also a graph input
grad_accumulation_buffer_input = copy.deepcopy(graph_output)
grad_accumulation_buffer_input.name = grad_accumulation_buffer_name
graph_inputs.append(grad_accumulation_buffer_input)
# accumulated gradient is also a graph output
grad_accumulation_output = copy.deepcopy(graph_output)
grad_accumulation_output.name = grad_accumulation_output_name
graph_outputs.append(grad_accumulation_output)
lazy_reset_grad_input = onnx.helper.make_tensor_value_info(
lazy_reset_grad_input_name, onnx.TensorProto.BOOL, [1]
)
graph_inputs.append(lazy_reset_grad_input)
del grad_model.graph.output[:]
grad_model.graph.output.extend(graph_outputs)
def get_model_parameters(model, args_not_requiring_gradient):
"""Returns trainable and non trainable onnx model parameters."""
trainable_params = []
non_trainable_params = []
for initializer in model.graph.initializer:
# All model parameters should have their names not begin with
# a digit. So, check to see if the initializer's first char
# is a digit. If not, it is either a trainable, or a non
# trainable parameter.
# Note that this assumption can be made because the export logic
# does not change the names of the original model parameters.
# and the original model parameters don't have their names begin
# with a digit.
# On the other hand, const initializers are generated by export
# logic and have a digit prefix.
# TODO: validate this assumption. If assumption is not valid,
# the alternative is to enforce the user to provide the parameter names.
if not initializer.name[0].isdigit():
if initializer.name in args_not_requiring_gradient:
non_trainable_params.append(initializer)
else:
trainable_params.append(initializer)
return trainable_params, non_trainable_params

21
orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py Normal file
View file

@ -0,0 +1,21 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# checkpoint_utils.py
from onnxruntime.capi._pybind_state import save_checkpoint as _internal_save_checkpoint
def save_checkpoint(parameters, path_to_checkpoint):
"""Saves the parameters to the checkpoint directory path_to_checkpoint."""
if parameters is None:
raise RuntimeError("No checkpoint parameters provided.")
# TODO: use Parameter class to pass information to backend
# Serialize the parameters and save the checkpoint
trainable_params, non_trainable_params = parameters
trainable_params = [param.SerializeToString() for param in trainable_params]
non_trainable_params = [param.SerializeToString() for param in non_trainable_params]
_internal_save_checkpoint(
trainable_params, non_trainable_params, path_to_checkpoint
)

5
orttraining/orttraining/python/training/onnxblock/loss/__init__.py Normal file
View file

@ -0,0 +1,5 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# __init__.py
from .loss import MSELoss, CrossEntropyLoss

157
orttraining/orttraining/python/training/onnxblock/loss/loss.py Normal file
View file

@ -0,0 +1,157 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# loss.py
import copy
import onnx
import onnxruntime.training.onnxblock as onnxblock
import onnxruntime.training.onnxblock.model_accessor as accessor
import onnxruntime.training.onnxblock._graph_utils as graph_utils
import onnxruntime.training.onnxblock._building_blocks as building_blocks
class MSELoss(onnxblock.Model):
"""MSELoss onnxblock for adding MSE loss to an onnx model.
Parameters:
reduction: string representing the reduction method on the loss output.
can be one of "mean" or "sum"
"""
def __init__(self, reduction="mean"):
super(MSELoss, self).__init__()
# determine the reduction type
if reduction != "mean" and reduction != "sum":
raise RuntimeError(f"Reduction {reduction} not supported.")
self._reduce = (
building_blocks.ReduceMean()
if reduction == "mean"
else building_blocks.ReduceSum()
)
self._sub = building_blocks.Sub()
self._square = building_blocks.Pow(2.0)
def build(self, loss_input_name, target_name="target"):
"""Adds an MSELoss subgraph on top of the base_model.
Creates a block that measures the mean squared error between
loss_input_name and the target_name.
Args:
loss_input_name: string input representing the loss input
target_name: string input representing the target
Returns:
Returns a string of the output name from the loss
"""
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create a new graph input. this is the target input needed to compare
# the graph output against to calculate loss.
target_input = copy.deepcopy(
graph_utils.get_output_from_output_name(onnx_model, loss_input_name)
)
target_input.name = target_name
onnx_model.graph.input.append(target_input)
# create the mse loss
# loss = reduce(square(sub(output, target)))
return self._reduce(self._square(self._sub(loss_input_name, target_name)))
class CrossEntropyLoss(onnxblock.Model):
"""CrossEntropyLoss onnxblock for adding Cross Entropy loss to an onnx model.
Parameters:
weight: boolean representing whether a manual rescaling weight given to
each class should be added to the inputs.
reduction: string representing the reduction method on the loss output.
can be one of "mean" or "sum"
ignore_index: specifies a target value that is ignored and does not
contribute to the input gradient.
"""
def __init__(self, weight=False, reduction="mean", ignore_index=None):
super(CrossEntropyLoss, self).__init__()
# determine the reduction type
if reduction != "mean" and reduction != "sum":
raise RuntimeError(f"Reduction {reduction} not supported.")
self._weight = weight
self._reduction = reduction
self._ignore_index = ignore_index
def build(self, scores_input_name, labels_name="labels", weight_name="loss_weight"):
"""Adds a CrossEntropyLoss subgraph on top of an onnx model.
Creates a block that measures the softmax cross entropy between
scores_input_name and the labels_name.
Args:
loss_input_name: string input representing the loss input
labels_name: string input representing the labels
Returns:
Returns a string of the output name from the loss
"""
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create a new graph input. this is the labels input needed to compare
# the graph output against to calculate loss.
labels_input = copy.deepcopy(
graph_utils.get_output_from_output_name(onnx_model, scores_input_name)
)
labels_input.name = labels_name
labels_input.type.tensor_type.elem_type = onnx.TensorProto.INT32
# if the predictions are (num_examples x num_classes)
# labels should be (num_examples x 1)
del labels_input.type.tensor_type.shape.dim[1]
onnx_model.graph.input.append(labels_input)
if self._weight:
weight_input = copy.deepcopy(
graph_utils.get_output_from_output_name(onnx_model, scores_input_name)
)
weight_input.name = weight_name
dim_to_keep = weight_input.type.tensor_type.shape.dim[1]
del weight_input.type.tensor_type.shape.dim[:]
weight_input.type.tensor_type.shape.dim.append(dim_to_keep)
onnx_model.graph.input.append(weight_input)
# create a new graph node for the loss
loss_node_input_names = [scores_input_name, labels_name]
if self._weight:
loss_node_input_names.append(weight_name)
loss_node_output_name = graph_utils.generate_random_graph_name("loss")
loss_node_output_names = [
loss_node_output_name,
graph_utils.generate_random_graph_name("log_prob"),
]
loss_node = onnx.helper.make_node(
"SoftmaxCrossEntropyLoss",
loss_node_input_names,
loss_node_output_names,
reduction=self._reduction,
ignore_index=self._ignore_index,
name=graph_utils.generate_random_graph_name("SoftmaxCrossEntropyLoss"),
)
onnx_model.graph.node.append(loss_node)
# create a new graph output for the loss
graph_outputs = [
onnx.helper.make_tensor_value_info(
loss_node_output_name, onnx.TensorProto.FLOAT, []
)
]
del onnx_model.graph.output[:]
onnx_model.graph.output.extend(graph_outputs)
return loss_node_output_name

123
orttraining/orttraining/python/training/onnxblock/model.py Normal file
View file

@ -0,0 +1,123 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# model.py
from abc import ABC, abstractmethod
import onnx
import onnxruntime.training.onnxblock.model_accessor as accessor
import onnxruntime.training.onnxblock._graph_utils as graph_utils
class Model(ABC):
"""Builds the forward model based on user's build method."""
def __init__(self):
...
@abstractmethod
def build(self, *args, **kwargs):
"""Customize and build the forward graph for this model.
This method is to be overriden by the user's implementation.
"""
...
def __call__(self, *args, **kwargs):
"""Calls the user's build method and runs validation on top.
The output onnx model is got by invoking the user's build method.
"""
# build the user model
output = self.build(*args, **kwargs)
# validate and check the model
onnx.checker.check_model(accessor.global_accessor.model, True)
return output
class TrainingModel(Model):
"""Builds the training model based on user's build method."""
def __init__(self):
super(TrainingModel, self).__init__()
self._arg_requiring_grad = set()
self._arg_not_requiring_grad = set()
self._parameters = None
@abstractmethod
def build(self, *args, **kwargs):
"""Customize and build the forward graph for this training model.
This method is to be overriden by the user's implementation.
"""
...
def requires_grad(self, argument_name, value=True):
"""Control whether the given graph input/parameter requires gradient."""
if value is True:
if argument_name in self._arg_not_requiring_grad:
self._arg_not_requiring_grad.remove(argument_name)
self._arg_requiring_grad.add(argument_name)
else:
if argument_name in self._arg_requiring_grad:
self._arg_requiring_grad.remove(argument_name)
self._arg_not_requiring_grad.add(argument_name)
def parameters(self):
"""Returns trainable and non trainable parameters.
Model parameters that are extracted while building the training model
are returned by this method.
Note that the parameters are not known before the training model is
built. As a result, if this method is invoked before the training model
is built, an exception will be raised.
"""
if self._parameters is None:
raise RuntimeError(
"Please build the training model first before trying to "
"retrieve the parameters."
)
return self._parameters
def __call__(self, *args, **kwargs):
"""Calls the user's build method and builds the gradient graph on top.
The onnx model contains the user's training model such that:
1. It contains the gradient graph.
2. It contains inputs in the order: user inputs, weight parameters,
gradient inputs.
3. It contains the outputs in the order: user outputs, gradient outputs.
4. Before the gradient is built, the eval model is stored in the global accessor.
Note that the model parameters are moved to be graph inputs.
"""
# build the user model
output = self.build(*args, **kwargs)
# get all the model parameters for the user_model
# and store them in self._parameters
self._parameters = graph_utils.get_model_parameters(
accessor.global_accessor.model, self._arg_not_requiring_grad
)
# TODO: Perform shape inference before building gradient graph
# build the gradient graph
all_args_requiring_gradient_names = graph_utils.build_gradient_graph(
accessor.global_accessor,
self._arg_requiring_grad,
self._arg_not_requiring_grad,
output,
)
# add gradient accumulation nodes
graph_utils.build_gradient_accumulation_graph(
accessor.global_accessor.model, all_args_requiring_gradient_names
)
# validate and check the model
onnx.checker.check_model(accessor.global_accessor.model, True)
return output

33
orttraining/orttraining/python/training/onnxblock/model_accessor.py Normal file
View file

@ -0,0 +1,33 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# global_model.py
from contextlib import contextmanager
class ModelAccessor:
"""This class stores the onnx model that is manipulated by the onnx blocks."""
def __init__(self, model):
self.model = model
self.eval_model = None
# This variable resides in the global namespace.
# Different methods can access this global model and manipulate it.
# Its construction and destruction is managed by the onnx_model contextmanager
global_accessor = None
@contextmanager
def onnx_model(model=None):
"""Context manager that is the entry point to graph manipulations on model.
Manages the construction and destruction of the global model.
"""
global global_accessor
global_accessor = ModelAccessor(model)
try:
yield global_accessor
finally:
global_accessor = None

5
orttraining/orttraining/python/training/onnxblock/optim/__init__.py Normal file
View file

@ -0,0 +1,5 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# __init__.py
from .optim import AdamW

161
orttraining/orttraining/python/training/onnxblock/optim/optim.py Normal file
View file

@ -0,0 +1,161 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# optim.py
import copy
import onnx
import onnxruntime.training.onnxblock as onnxblock
import onnxruntime.training.onnxblock.model_accessor as accessor
class AdamW(onnxblock.Model):
"""Builds AdamW optimizer onnxblock for the given training model."""
def __init__(
self, bias_correction=True, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.0
):
super(AdamW, self).__init__()
self._bias_correction = bias_correction
self._betas = betas
self._eps = eps
self._weight_decay = weight_decay
self._max_norm_clip = 1.0
def build(self, parameters):
"""Returns an AdamW optimizer model based on the input training model."""
# TODO: Avoid hard coded input/output strings
learning_rate_name = "learning_rate"
step_name = "step"
gradient_output_suffix = "_grad.accumulation.out"
first_order_moment_suffix = "exp_avg"
second_order_moment_fuffix = "exp_avg_sq"
output_name_suffix = "out"
trainable_parameters, _ = parameters
graph_nodes = []
graph_inputs = [
onnx.helper.make_tensor_value_info(
learning_rate_name, onnx.TensorProto.FLOAT, [1]
),
onnx.helper.make_tensor_value_info(step_name, onnx.TensorProto.INT64, [1]),
]
graph_outputs = []
# Iterate over all training graph outputs that are gradient outputs
for idx, param in enumerate(trainable_parameters):
param_name = param.name
grad_name = f"{param_name}{gradient_output_suffix}"
first_order_moment_name = f"{param_name}.{first_order_moment_suffix}"
second_order_moment_name = f"{param_name}.{second_order_moment_fuffix}"
# prepare node (and graph) inputs and outputs
node_input_names = [
learning_rate_name, # learning rate
step_name, # training step (used for beta correction)
param_name, # param to be updated
grad_name, # gradient of the param to be used for update
first_order_moment_name, # first order moment for this param
second_order_moment_name, # second order moment for this param
]
param_tensor_value_info = onnx.helper.make_tensor_value_info(
param_name, param.data_type, param.dims
)
grad_tensor_value_info = onnx.helper.make_tensor_value_info(
grad_name, param.data_type, param.dims
)
first_order_moment_tensor_value_info = onnx.helper.make_tensor_value_info(
first_order_moment_name, param.data_type, param.dims
)
second_order_moment_tensor_value_info = onnx.helper.make_tensor_value_info(
second_order_moment_name, param.data_type, param.dims
)
node_inputs = [
param_tensor_value_info,
grad_tensor_value_info,
first_order_moment_tensor_value_info,
second_order_moment_tensor_value_info,
]
graph_inputs.extend(node_inputs)
step_output_name = f"{param_name}.{step_name}.{output_name_suffix}"
param_output_name = f"{param_name}.{output_name_suffix}"
first_order_moment_output_name = (
f"{first_order_moment_name}.{output_name_suffix}"
)
second_order_moment_output_name = (
f"{second_order_moment_name}.{output_name_suffix}"
)
param_output_tensor_value_info = onnx.helper.make_tensor_value_info(
param_output_name, param.data_type, param.dims
)
first_order_moment_output_tensor_value_info = (
onnx.helper.make_tensor_value_info(
first_order_moment_output_name, param.data_type, param.dims
)
)
second_order_moment_output_tensor_value_info = (
onnx.helper.make_tensor_value_info(
second_order_moment_output_name, param.data_type, param.dims
)
)
node_output_names = [
step_output_name, # step out
first_order_moment_output_name, # first order moment output
second_order_moment_output_name, # second order moment output
param_output_name, # updated weights
]
node_outputs = [
onnx.helper.make_tensor_value_info(
step_output_name, onnx.TensorProto.INT64, [1]
),
first_order_moment_output_tensor_value_info,
second_order_moment_output_tensor_value_info,
param_output_tensor_value_info,
]
graph_outputs.extend(node_outputs)
# AdamOptimizer node attributes
node_attributes = {
"alpha": self._betas[0], # beta1
"beta": self._betas[1], # beta2
"lambda": self._weight_decay, # weight decay
"epsilon": self._eps, # epsilon
"do_bias_correction": 1
if self._bias_correction
else 0, # bias_correction
"weight_decay_mode": 1, # weight decay mode
"max_norm_clip": self._max_norm_clip, # used for gradient scaling
}
# make the node
optimizer_node = onnx.helper.make_node(
"AdamOptimizer",
node_input_names,
node_output_names,
name=f"AdamOptimizer{idx}",
domain="com.microsoft",
**node_attributes,
)
graph_nodes.append(optimizer_node)
# make the graph and the model
graph = onnx.helper.make_graph(
graph_nodes, "Optimizer Graph", graph_inputs, graph_outputs
)
model = onnx.helper.make_model(
graph,
producer_name=onnxblock._producer_name,
opset_imports=[onnxblock._opset_import],
)
accessor.global_accessor.model = model
return [output.name for output in graph_outputs]

526
orttraining/orttraining/test/python/orttraining_test_onnxblock.py Normal file
View file

@ -0,0 +1,526 @@
import torch
import io
import onnx
import onnxruntime.training.onnxblock as onnxblock
import onnxruntime
from onnxruntime.capi import _pybind_state as C
import pytest
import tempfile
import os
import copy
import numpy as np
# PyTorch Module definitions
class SimpleNet(torch.nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(SimpleNet, self).__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, model_input):
out = self.fc1(model_input)
out = self.relu(out)
out = self.fc2(out)
return out
# onnxblock Model definitions
class SimpleModelWithMSELoss(onnxblock.Model):
def __init__(self):
super(SimpleModelWithMSELoss, self).__init__()
self.loss = onnxblock.loss.MSELoss()
def build(self, output_name):
return self.loss(output_name)
class SimpleModelWithCrossEntropyLoss(onnxblock.Model):
def __init__(self):
super(SimpleModelWithCrossEntropyLoss, self).__init__()
self.loss = onnxblock.loss.CrossEntropyLoss()
def build(self, output_name):
return self.loss(output_name)
class SimpleTrainingModelWithMSELoss(onnxblock.TrainingModel):
def __init__(self):
super(SimpleTrainingModelWithMSELoss, self).__init__()
self.loss = onnxblock.loss.MSELoss()
def build(self, output_name):
return self.loss(output_name)
class SimpleTrainingModelWithCrossEntropyLoss(onnxblock.TrainingModel):
def __init__(self):
super(SimpleTrainingModelWithCrossEntropyLoss, self).__init__()
self.loss = onnxblock.loss.CrossEntropyLoss()
def build(self, output_name):
return self.loss(output_name)
# Test utility methods
def _get_onnx_model(torch_model, model_inputs):
model_outputs = torch_model(*model_inputs)
if isinstance(model_outputs, torch.Tensor):
model_outputs = [model_outputs]
dynamic_axes = {}
input_names = []
output_names = []
for i, model_input in enumerate(model_inputs):
input_name = f"input-{i}"
input_names.append(input_name)
dynamic_axes[input_name] = {}
for dim_idx in range(len(model_input.shape)):
dynamic_axes[input_name].update({dim_idx: f"{input_name}_dim{dim_idx}"})
for i, model_output in enumerate(model_outputs):
output_name = f"output-{i}"
output_names.append(output_name)
dynamic_axes[output_name] = {}
for dim_idx in range(len(model_output.shape)):
dynamic_axes[output_name].update({dim_idx: f"{output_name}_dim{dim_idx}"})
f = io.BytesIO()
torch.onnx.export(
torch_model,
model_inputs,
f,
input_names=input_names,
output_names=output_names,
opset_version=14,
do_constant_folding=False,
training=torch.onnx.TrainingMode.TRAINING,
dynamic_axes=dynamic_axes,
export_params=True,
keep_initializers_as_inputs=False,
)
return onnx.load_model_from_string(f.getvalue())
def _to_numpy(tensor):
return (
tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
)
def _get_models(device, N, D_in, H, D_out):
"""Returns the pt and onnx models for SimpleNet"""
pt_model = SimpleNet(D_in, H, D_out).to(device)
x = torch.randn(N, D_in, device=device)
onnx_model = _get_onnx_model(pt_model, (x,))
return pt_model, onnx_model
def _get_training_ort_output_names(pt_model, onnx_model):
"""Returns the ort output names"""
ort_output_names = [onnx_model.graph.output[0].name]
for name, _ in pt_model.named_parameters():
ort_output_names.append(f"{name}_grad.accumulation.out")
return ort_output_names
def _get_training_ort_inputs(x, target, pt_model, onnx_model, target_type=None):
"""Returns the ort inputs"""
ort_inputs = {
onnx_model.graph.input[0].name: _to_numpy(copy.deepcopy(x)),
onnx_model.graph.input[1].name: _to_numpy(copy.deepcopy(target))
if target_type is None
else _to_numpy(copy.deepcopy(target).type(target_type)),
}
if target_type is not None:
ort_inputs[onnx_model.graph.input[1].name]
for name, param in pt_model.named_parameters():
ort_inputs[name] = _to_numpy(copy.deepcopy(param))
ort_inputs[f"{name}_grad.accumulation.buffer"] = _to_numpy(
torch.zeros_like(param)
)
ort_inputs["lazy_reset_grad"] = np.full(1, True)
return ort_inputs
# All unit tests
@pytest.mark.parametrize(
"graph",
[
SimpleModelWithMSELoss,
SimpleModelWithCrossEntropyLoss,
SimpleTrainingModelWithMSELoss,
SimpleTrainingModelWithCrossEntropyLoss,
],
)
def test_loss_composition(graph):
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
_, onnx_model = _get_models(device, N, D_in, H, D_out)
# When / Then no error occurs
simple_model = graph()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
def test_mse_loss_execution():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
pt_model, onnx_model = _get_models(device, N, D_in, H, D_out)
x = torch.randn(N, D_in, device=device)
target = torch.randn(N, D_out, device=device)
# Build the onnx model with loss
simple_model = SimpleModelWithMSELoss()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
ort_output_names = [onnx_model.graph.output[0].name]
ort_inputs = {
onnx_model.graph.input[0].name: _to_numpy(copy.deepcopy(x)),
onnx_model.graph.input[1].name: _to_numpy(copy.deepcopy(target)),
}
def mse_loss(prediction, target):
loss = torch.nn.MSELoss()
return loss(prediction, target)
# When
with tempfile.NamedTemporaryFile(suffix=".onnx") as onnx_fo:
onnx.save(onnx_model, onnx_fo.name)
ort_session = onnxruntime.InferenceSession(
onnx_fo.name, providers=C.get_available_providers()
)
ort_outs = ort_session.run(ort_output_names, ort_inputs)
torch_outs = mse_loss(pt_model(x), target)
# Then
assert np.allclose(ort_outs[0], _to_numpy(torch_outs))
def test_crossentropy_loss_execution():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
pt_model, onnx_model = _get_models(device, N, D_in, H, D_out)
x = torch.randn(N, D_in, device=device)
target = torch.randint(high=D_out, size=(N,), dtype=torch.int64, device=device)
# Build the onnx model with loss
simple_model = SimpleModelWithCrossEntropyLoss()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
ort_output_names = [onnx_model.graph.output[0].name]
ort_inputs = {
onnx_model.graph.input[0].name: _to_numpy(copy.deepcopy(x)),
onnx_model.graph.input[1].name: _to_numpy(
copy.deepcopy(target).type(torch.int32)
),
}
def crossentropy_loss(prediction, target):
loss = torch.nn.CrossEntropyLoss()
return loss(prediction, target)
# When
with tempfile.NamedTemporaryFile(suffix=".onnx") as onnx_fo:
onnx.save(onnx_model, onnx_fo.name)
ort_session = onnxruntime.InferenceSession(
onnx_fo.name, providers=C.get_available_providers()
)
ort_outs = ort_session.run(ort_output_names, ort_inputs)
torch_outs = crossentropy_loss(pt_model(x), target)
# Then
assert np.allclose(ort_outs[0], _to_numpy(torch_outs))
def test_mse_loss_training_graph_execution():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
pt_model, onnx_model = _get_models(device, N, D_in, H, D_out)
x = torch.randn(N, D_in, device=device)
target = torch.randn(N, D_out, device=device)
# Build the onnx trainingmodel with loss
simple_model = SimpleTrainingModelWithMSELoss()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
ort_output_names = _get_training_ort_output_names(pt_model, onnx_model)
ort_inputs = _get_training_ort_inputs(x, target, pt_model, onnx_model)
def mse_loss(prediction, target):
loss = torch.nn.MSELoss()
return loss(prediction, target)
# When
with tempfile.NamedTemporaryFile(suffix=".onnx") as onnx_fo:
onnx.save(onnx_model, onnx_fo.name)
ort_session = onnxruntime.InferenceSession(
onnx_fo.name, providers=C.get_available_providers()
)
ort_outs = ort_session.run(ort_output_names, ort_inputs)
torch_outs = mse_loss(pt_model(x), target)
torch_outs.backward()
# Then
# assert loss is close
assert np.allclose(ort_outs[0], _to_numpy(torch_outs))
# assert all the gradients are close
for ort_grad, pt_param in zip(ort_outs[1:], pt_model.parameters()):
assert np.allclose(ort_grad, _to_numpy(pt_param.grad))
def test_crossentropy_loss_training_graph_execution():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
pt_model, onnx_model = _get_models(device, N, D_in, H, D_out)
x = torch.randn(N, D_in, device=device)
target = torch.randint(high=D_out, size=(N,), dtype=torch.int64, device=device)
# Build the onnx trainingmodel with loss
simple_model = SimpleTrainingModelWithCrossEntropyLoss()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
ort_output_names = _get_training_ort_output_names(pt_model, onnx_model)
ort_inputs = _get_training_ort_inputs(
x, target, pt_model, onnx_model, target_type=torch.int32
)
def crossentropy_loss(prediction, target):
loss = torch.nn.CrossEntropyLoss()
return loss(prediction, target)
# When
with tempfile.NamedTemporaryFile(suffix=".onnx") as onnx_fo:
onnx.save(onnx_model, onnx_fo.name)
ort_session = onnxruntime.InferenceSession(
onnx_fo.name, providers=C.get_available_providers()
)
ort_outs = ort_session.run(ort_output_names, ort_inputs)
torch_outs = crossentropy_loss(pt_model(x), target)
torch_outs.backward()
# Then
# assert loss is close
assert np.allclose(ort_outs[0], _to_numpy(torch_outs))
# assert all the gradients are close
for ort_grad, pt_param in zip(ort_outs[1:], pt_model.parameters()):
assert np.allclose(ort_grad, _to_numpy(pt_param.grad))
@pytest.mark.parametrize(
"graph", [SimpleTrainingModelWithMSELoss, SimpleTrainingModelWithCrossEntropyLoss]
)
def test_adamw_optimizer_composition(graph):
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
_, onnx_model = _get_models(device, N, D_in, H, D_out)
# When / Then no error occurs
simple_model = graph()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
optimizer = onnxblock.optim.AdamW()
with onnxblock.onnx_model() as accessor:
_ = optimizer(simple_model.parameters())
optimizer_model = accessor.model
assert optimizer_model
def test_adamw_optimizer_execution():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
pt_model, onnx_model = _get_models(device, N, D_in, H, D_out)
x = torch.randn(N, D_in, device=device)
target = torch.randn(N, D_out, device=device)
simple_model = SimpleTrainingModelWithMSELoss()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
optimizer = onnxblock.optim.AdamW()
with onnxblock.onnx_model() as accessor:
_ = optimizer(simple_model.parameters())
optimizer_model = accessor.model
learning_rate = 0.001
step = 1
ort_output_names = []
for name, _ in pt_model.named_parameters():
ort_output_names.append(f"{name}.out")
def mse_loss(prediction, target):
loss = torch.nn.MSELoss()
return loss(prediction, target)
# When
with tempfile.NamedTemporaryFile(suffix=".onnx") as onnx_fo:
onnx.save(optimizer_model, onnx_fo.name)
loss = mse_loss(pt_model(x), target)
loss.backward()
ort_inputs = {
"learning_rate": np.full(1, learning_rate, dtype=np.float32),
"step": np.full(1, step, dtype=np.int64),
}
for name, param in pt_model.named_parameters():
ort_inputs[name] = _to_numpy(copy.deepcopy(param))
ort_inputs[f"{name}_grad.accumulation.out"] = _to_numpy(
copy.deepcopy(param.grad)
)
ort_inputs[f"{name}.exp_avg"] = _to_numpy(torch.zeros_like(param))
ort_inputs[f"{name}.exp_avg_sq"] = _to_numpy(torch.zeros_like(param))
# Then no error occurs when executing the model
ort_session = onnxruntime.InferenceSession(
onnx_fo.name, providers=C.get_available_providers()
)
_ = ort_session.run(ort_output_names, ort_inputs)
def test_retrieve_parameters():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
pt_model, onnx_model = _get_models(device, N, D_in, H, D_out)
simple_model = SimpleTrainingModelWithMSELoss()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
# When
trainable_params, non_trainable_params = simple_model.parameters()
# Then
assert not non_trainable_params
for ort_param, (pt_param_name, pt_param) in zip(
trainable_params, pt_model.named_parameters()
):
assert ort_param.name == pt_param_name
assert np.allclose(
np.frombuffer(ort_param.raw_data, dtype=np.float32).reshape(pt_param.shape),
_to_numpy(pt_param),
)
def test_retrieve_parameters_before_building_gradient_graph():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
_, onnx_model = _get_models(device, N, D_in, H, D_out)
simple_model = SimpleTrainingModelWithMSELoss()
# When / Then
with pytest.raises(Exception) as ex_info:
_, _ = simple_model.parameters()
assert (
"Please build the training model first before trying to retrieve the parameters."
in str(ex_info.value)
)
def test_save_checkpoint():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
_, onnx_model = _get_models(device, N, D_in, H, D_out)
simple_model = SimpleTrainingModelWithMSELoss()
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
trainable_params, non_trainable_params = simple_model.parameters()
# When
with tempfile.TemporaryDirectory() as checkpoint_dir_name:
checkpoint_file_path = os.path.join(checkpoint_dir_name, "checkpoint")
onnxblock.save_checkpoint(
(trainable_params, non_trainable_params), checkpoint_file_path
)
# Then
assert os.path.exists(checkpoint_file_path)
def test_set_requires_grad_on_parameters():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
_, onnx_model = _get_models(device, N, D_in, H, D_out)
simple_model = SimpleTrainingModelWithMSELoss()
# When
simple_model.requires_grad("fc2.weight", False)
simple_model.requires_grad("fc1.bias", False)
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
trainable_params, non_trainable_params = simple_model.parameters()
# Then
expected_trainable_parameters = {"fc1.weight", "fc2.bias"}
expected_non_trainable_parameters = {"fc2.weight", "fc1.bias"}
for param in trainable_params:
assert param.name in expected_trainable_parameters
for param in non_trainable_params:
assert param.name in expected_non_trainable_parameters
def test_set_requires_grad_on_inputs():
# Given
device = "cuda"
N, D_in, H, D_out = 64, 784, 500, 10
_, onnx_model = _get_models(device, N, D_in, H, D_out)
# When
simple_model = SimpleTrainingModelWithMSELoss()
simple_model.requires_grad("input-0")
with onnxblock.onnx_model(onnx_model):
_ = simple_model(onnx_model.graph.output[0].name)
# Then
expected_input_gradient_buffer_name = "input-0_grad.accumulation.buffer"
expected_input_gradient_output_name = "input-0_grad.accumulation.out"
graph_input_names = {graph_input.name for graph_input in onnx_model.graph.input}
graph_output_names = {graph_output.name for graph_output in onnx_model.graph.output}
assert expected_input_gradient_buffer_name in graph_input_names
assert expected_input_gradient_output_name in graph_output_names

3
setup.py
View file

@ -329,6 +329,7 @@ requirements_file = "requirements.txt"
local_version = None
enable_training = parse_arg_remove_boolean(sys.argv, '--enable_training')
enable_training_on_device = parse_arg_remove_boolean(sys.argv, '--enable_training_on_device')
disable_auditwheel_repair = parse_arg_remove_boolean(sys.argv, '--disable_auditwheel_repair')
default_training_package_device = parse_arg_remove_boolean(sys.argv, '--default_training_package_device')
@ -374,6 +375,8 @@ if enable_training:
'onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator',
'onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.fused_ops',
'onnxruntime.training.utils.data'])
if enable_training_on_device:
packages.append('onnxruntime.training.onnxblock')
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.aten_op_executor'] = ['*.cc']
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.torch_interop_utils'] = ['*.cc']
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator'] = ['*.cc']

10
tools/ci_build/build.py
View file

@ -176,6 +176,8 @@ def parse_arguments():
"--enable_training_ops", action='store_true', help="Enable training ops in inference graph.")
parser.add_argument(
"--enable_training_torch_interop", action='store_true', help="Enable training kernels interop with torch.")
parser.add_argument(
"--enable_training_on_device", action='store_true', help="Enable on device training in ORT.")
parser.add_argument(
"--disable_nccl", action='store_true', help="Disable Nccl.")
parser.add_argument(
@ -836,6 +838,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
"-Donnxruntime_ENABLE_TRAINING=" + ("ON" if args.enable_training else "OFF"),
"-Donnxruntime_ENABLE_TRAINING_OPS=" + ("ON" if args.enable_training_ops else "OFF"),
"-Donnxruntime_ENABLE_TRAINING_TORCH_INTEROP=" + ("ON" if args.enable_training_torch_interop else "OFF"),
"-Donnxruntime_ENABLE_TRAINING_ON_DEVICE=" + ("ON" if args.enable_training_on_device else "OFF"),
# Enable advanced computations such as AVX for some traininig related ops.
"-Donnxruntime_ENABLE_CPU_FP16_OPS=" + ("ON" if args.enable_training else "OFF"),
"-Donnxruntime_USE_NCCL=" + ("OFF" if args.disable_nccl else "ON"),
@ -1801,7 +1804,7 @@ def build_python_wheel(
source_dir, build_dir, configs, use_cuda, cuda_version, use_rocm, rocm_version, use_dnnl,
use_tensorrt, use_openvino, use_nuphar, use_tvm, use_vitisai, use_acl, use_armnn, use_dml,
wheel_name_suffix, enable_training, nightly_build=False, default_training_package_device=False,
use_ninja=False, build_eager_mode=False):
use_ninja=False, build_eager_mode=False, enable_training_on_device=False):
for config in configs:
cwd = get_config_build_dir(build_dir, config)
if is_windows() and not use_ninja:
@ -1819,6 +1822,8 @@ def build_python_wheel(
args.append('--wheel_name_suffix={}'.format(wheel_name_suffix))
if enable_training:
args.append("--enable_training")
if enable_training_on_device:
args.append("--enable_training_on_device")
if build_eager_mode:
args.append("--disable_auditwheel_repair")
@ -2451,7 +2456,8 @@ def main():
nightly_build=nightly_build,
default_training_package_device=default_training_package_device,
use_ninja=(args.cmake_generator == 'Ninja'),
build_eager_mode=args.build_eager_mode
build_eager_mode=args.build_eager_mode,
enable_training_on_device=args.enable_training_on_device
)
if args.build_nuget:
build_nuget_package(