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

Add SGDOptimizer in the on-device training offline tooling (onnxblock) (#17085)

Browse source 
### Description
Adding SGDOptimizer to on device training onnxblock
This commit is contained in:
Adam Louly 2023-08-18 10:50:39 -07:00 committed by GitHub
parent ee09a5ff35
commit c0b6c6c94b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 266 additions and 123 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

BIN
onnxruntime/test/testdata/training_api/adamw.onnx vendored
View file

Binary file not shown.

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

@ -1537,7 +1537,7 @@ void RegisterTrainingOpSchemas() {
"This signal indicates if weight updates are skipped, applicable to gradient infinity check"
" in mixed precision training. ",
"T_BOOL", OpSchema::Optional)
.Output(0, "updated_flag", "Whether gradient is applied or not.", "T2")
.Output(0, "updated_flag", "Whether gradient is applied or not.", "T_BOOL")
.Output(1, "updated_weights", "Sequence of weights after optimize.", "S_WEIGHT", OpSchema::Optional)
.Output(2, "updated_momentums_1", "Sequence of momentum_1 after optimize.", "S_MOMENT", OpSchema::Optional)
.Output(3, "updated_momentums_2", "Sequence of momentum_2 after optimize.", "S_MOMENT", OpSchema::Optional)

4
orttraining/orttraining/python/training/artifacts.py
View file

@ -33,6 +33,7 @@ class OptimType(Enum):
"""
AdamW = 1
SGD = 2
def generate_artifacts(
@ -192,7 +193,8 @@ def generate_artifacts(
logging.info("Optimizer enum provided: %s", optimizer.name)
optim_model = None
optim_blocks = {OptimType.AdamW: onnxblock.optim.AdamW}
optim_blocks = {OptimType.AdamW: onnxblock.optim.AdamW, OptimType.SGD: onnxblock.optim.SGD}
optim_block = optim_blocks[optimizer]()
with onnxblock.empty_base():
_ = optim_block(model_params)

4
orttraining/orttraining/python/training/onnxblock/optim/__init__.py
View file

@ -1,6 +1,6 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
from onnxruntime.training.onnxblock.optim.optim import AdamW, ClipGradNorm
from onnxruntime.training.onnxblock.optim.optim import SGD, AdamW, ClipGradNorm
__all__ = ["AdamW", "ClipGradNorm"]
__all__ = ["AdamW", "ClipGradNorm", "SGD"]

309
orttraining/orttraining/python/training/onnxblock/optim/optim.py
View file

@ -1,7 +1,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
from typing import Optional, Tuple
from typing import Dict, List, Optional, Tuple
import onnx
@ -10,71 +10,6 @@ import onnxruntime.training.onnxblock.blocks as blocks
import onnxruntime.training.onnxblock.onnxblock as onnxblock_module
class AdamWOptimizer(blocks.Block):
"""Adds an AdamWOptimizer node to the onnx model."""
def __init__(
self,
bias_correction: Optional[bool] = True,
betas: Tuple[float, float] = (0.9, 0.999),
eps: Optional[float] = 1e-6,
weight_decay: Optional[float] = 0.0,
):
super().__init__()
self._bias_correction = bias_correction
self._betas = betas
self._eps = eps
self._weight_decay = weight_decay
def build( # pylint: disable=too-many-arguments
self,
learning_rate_name: str,
step_name: str,
parameter_sequence_name: str,
gradient_sequence_name: str,
first_order_moment_sequence_name: str,
second_order_moment_sequence_name: str,
):
"""Adds the AdamWOptimizer node to the model."""
# get the model to manipulate
onnx_model = self.base
# define the node attributes
node_attributes = {
"alpha": self._betas[0], # beta1
"beta": self._betas[1], # beta2
"epsilon": self._eps, # epsilon
"weight_decay": self._weight_decay, # weight decay
"correct_bias": 1 if self._bias_correction else 0, # bias_correction
"adam_mode": 1, # adam mode (1 for hf/transformers/AdamW)
}
# add the adamw node to the onnx model
adamw_input_names = [
learning_rate_name, # learning rate
step_name, # training step
parameter_sequence_name, # param to be updated
gradient_sequence_name, # gradient of the param to be used for update
first_order_moment_sequence_name, # first order moment for this param
second_order_moment_sequence_name, # second order moment for this param
]
adamw_output_name = _graph_utils.generate_graph_name("adamw.updated_flag")
adamw_output_names = [adamw_output_name]
adamw_node = onnx.helper.make_node(
"AdamWOptimizer",
adamw_input_names,
adamw_output_names,
name=_graph_utils.generate_graph_name("AdamWOptimizer"),
domain="com.microsoft",
**node_attributes,
)
onnx_model.graph.node.append(adamw_node)
return adamw_output_name
class ClipGradNorm(blocks.Block):
"""Builds a gradient clipping by norm sub graph for the onnx model.
@ -125,59 +60,162 @@ class ClipGradNorm(blocks.Block):
return cgn_node_output_name
class AdamW(onnxblock_module.ForwardBlock):
"""Builds AdamW optimizer onnxblock for the given training parameters.
class _OptimizerBase(blocks.Block):
def __init__(self):
super().__init__()
Creates a block that updates the model parameters based on the calculated
gradient following the AdamW algorithm.
def _build_optimizer_node(
self,
input_names: List[str],
output_name: str,
node_name: str,
node_attributes: Dict,
) -> str:
"""
Build and append an optimizer node to the ONNX graph.
Args:
bias_correction: bool indicating whether to perform bias correction.
betas: AdamW decay rate hyperparameters.
eps: term added to the denominator for computing the moments.
weight_decay: AdamW weight decay
clip_grad (optional): an instance of the ClipGradNorm. If not provided,
gradient clipping will not be done.
Args:
input_names (list): List of input tensor names for the optimizer node.
output_name (str): Output tensor name of the optimizer node.
node_name (str): Name of the optimizer node.
node_attributes (dict): Additional attributes for the optimizer node.
Returns:
Returns a string of the output names from this optimizer node.
"""
Returns:
str: The output tensor name of the optimizer node.
"""
onnx_model = self.base
# add the optimizer node to the onnx model
optimizer_node = onnx.helper.make_node(
node_name,
input_names,
[output_name],
name=_graph_utils.generate_graph_name(node_name),
domain="com.microsoft",
**node_attributes,
)
onnx_model.graph.node.append(optimizer_node)
return output_name
class SGDOptimizer(_OptimizerBase):
def __init__(self):
super().__init__()
def build(
self,
learning_rate_name: str,
gradients_name: str,
params_name: str,
) -> str:
"""
Build an SGD optimizer node.
Args:
learning_rate_name (str): Name of the learning rate input tensor.
gradients_name (str): Name of the gradients input tensor.
params_name (str): Name of the weights input tensor.
Returns:
str: The output tensor name of the SGD optimizer node.
"""
input_names = [learning_rate_name, gradients_name, params_name]
return self._build_optimizer_node(
input_names,
_graph_utils.generate_graph_name("update_completed"),
"SGDOptimizerV2",
{},
)
class AdamWOptimizer(_OptimizerBase):
def __init__(
self,
bias_correction: Optional[bool] = True,
betas: Tuple[float, float] = (0.9, 0.999),
eps: Optional[float] = 1e-6,
weight_decay: Optional[float] = 0.0,
clip_grad=None,
): # pylint: disable=too-many-arguments
):
super().__init__()
self._adamw = AdamWOptimizer(
bias_correction=bias_correction,
betas=betas,
eps=eps,
weight_decay=weight_decay,
self._bias_correction = bias_correction
self._betas = betas
self._eps = eps
self._weight_decay = weight_decay
def build(
self,
learning_rate_name: str,
step_name: str,
parameter_sequence_name: str,
gradient_sequence_name: str,
first_order_moment_sequence_name: str,
second_order_moment_sequence_name: str,
) -> str:
"""
Build an AdamW optimizer node.
Args:
learning_rate_name (str): Name of the learning rate input tensor.
step_name (str): Name of the step input tensor.
parameter_sequence_name (str): Name of the parameter sequence input tensor.
gradient_sequence_name (str): Name of the gradient sequence input tensor.
first_order_moment_sequence_name (str): Name of the first order moment sequence input tensor.
second_order_moment_sequence_name (str): Name of the second order moment sequence input tensor.
Returns:
str: The output tensor name of the AdamW optimizer node.
"""
input_names = [
learning_rate_name,
step_name,
parameter_sequence_name,
gradient_sequence_name,
first_order_moment_sequence_name,
second_order_moment_sequence_name,
]
# define the node attributes
node_attributes = {
"alpha": self._betas[0], # beta1
"beta": self._betas[1], # beta2
"epsilon": self._eps, # epsilon
"weight_decay": self._weight_decay, # weight decay
"correct_bias": 1 if self._bias_correction else 0, # bias_correction
"adam_mode": 1, # adam mode (1 for hf/transformers/AdamW)
}
return self._build_optimizer_node(
input_names,
_graph_utils.generate_graph_name("adamw.updated_flag"),
"AdamWOptimizer",
node_attributes,
)
class _Optimizer(onnxblock_module.ForwardBlock):
"""Base class for building optimizer onnxblocks."""
def __init__(self, clip_grad=None):
super().__init__()
self._clip_grad = clip_grad
def build(self, parameters):
"""Returns an AdamW optimizer model based on the input parameters."""
# get the model to manipulate and update its namespace
onnx_model = self.base
# TODO: Avoid hard coded input/output strings
learning_rate_name = "learning_rate"
step_name = "step"
params_name = "params"
first_order_moments_name = "first_order_moments"
second_order_moments_name = "second_order_moments"
gradients_name = "gradients"
step_name = "step"
first_order_moments_name = "first_order_moments"
trainable_parameters, _ = parameters
# create the graph inputs for the lr, step, params, grads, moments
onnx_model.graph.input.extend(
[
onnx.helper.make_tensor_value_info(learning_rate_name, onnx.TensorProto.FLOAT, [1]),
@ -185,17 +223,61 @@ class AdamW(onnxblock_module.ForwardBlock):
]
)
# Prepare the tensor sequence inputs for params and moments
for input_name in [params_name, gradients_name, first_order_moments_name, second_order_moments_name]:
for input_name in [params_name, gradients_name, first_order_moments_name]:
onnx_model.graph.input.append(
onnx.helper.make_tensor_sequence_value_info(input_name, trainable_parameters[0].data_type, None)
)
# Clip the gradients if needed
if self._clip_grad is not None:
gradients_name = self._clip_grad(gradients_name)
# Run multi tensor AdamWOptimizer
updated_flag_name = self._optimizer_specific_logic(
learning_rate_name, params_name, gradients_name, trainable_parameters
)
return updated_flag_name
def _optimizer_specific_logic(
self,
learning_rate_name: str,
params_name: str,
gradients_name: str,
trainable_parameters: Tuple[List[onnx.TensorProto], List[onnx.TensorProto]],
) -> str:
raise NotImplementedError("Subclasses must implement _optimizer_specific_logic method.")
class AdamW(_Optimizer):
"""Builds AdamW optimizer onnxblock for the given training parameters."""
def __init__(self, bias_correction=True, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.0, clip_grad=None):
super().__init__(clip_grad)
self._adamw = AdamWOptimizer(
bias_correction=bias_correction,
betas=betas,
eps=eps,
weight_decay=weight_decay,
)
def _optimizer_specific_logic(
self,
learning_rate_name: str,
params_name: str,
gradients_name: str,
trainable_parameters: Tuple[List[onnx.TensorProto], List[onnx.TensorProto]],
) -> str:
onnx_model = self.base
step_name = "step"
first_order_moments_name = "first_order_moments"
second_order_moments_name = "second_order_moments"
# Prepare the tensor sequence inputs for moments
onnx_model.graph.input.append(
onnx.helper.make_tensor_sequence_value_info(
second_order_moments_name, trainable_parameters[0].data_type, None
)
)
updated_flag_name = self._adamw(
learning_rate_name,
step_name,
@ -205,9 +287,34 @@ class AdamW(onnxblock_module.ForwardBlock):
second_order_moments_name,
)
# Create the graph outputs
# Create graph outputs for AdamW
onnx_model.graph.output.append(
onnx.helper.make_tensor_value_info(updated_flag_name, onnx.TensorProto.INT64, [1])
onnx.helper.make_tensor_value_info(updated_flag_name, onnx.TensorProto.BOOL, [1])
)
return updated_flag_name
class SGD(_Optimizer):
"""Builds SGD optimizer onnxblock for the given training parameters."""
def __init__(self, clip_grad=None):
super().__init__(clip_grad)
self._sgd = SGDOptimizer()
def _optimizer_specific_logic(
self,
learning_rate_name: str,
params_name: str,
gradients_name: str,
trainable_parameters: Tuple[List[onnx.TensorProto], List[onnx.TensorProto]],
) -> str:
onnx_model = self.base
updated_flag_name = self._sgd(learning_rate_name, params_name, gradients_name)
# Create graph outputs for SGD
onnx_model.graph.output.append(
onnx.helper.make_tensor_value_info(updated_flag_name, onnx.TensorProto.BOOL, [1])
)
return updated_flag_name

26
orttraining/orttraining/test/python/orttraining_test_onnxblock.py
View file

@ -554,6 +554,32 @@ def test_adamw_optimizer_execution():
_ = ort_session.run(ort_output_names, ort_inputs)
@pytest.mark.parametrize(
"block",
[SimpleTrainingBlockWithMSELoss, SimpleTrainingBlockWithCrossEntropyLoss, SimpleTrainingBlockWithBCEWithLogitsLoss],
)
@pytest.mark.parametrize("grad_clipping", [None, onnxblock.optim.ClipGradNorm(2.5)])
def test_sgd_optimizer_composition(block, grad_clipping):
# Given
device = "cpu"
batch_size, input_size, hidden_size, output_size = 64, 784, 500, 10
pt_model, base_model = _get_models(device, batch_size, input_size, hidden_size, output_size)
# When / Then no error occurs
simple_block = block()
for name, _ in pt_model.named_parameters():
simple_block.requires_grad(name)
with onnxblock.base(base_model):
_ = simple_block(base_model.graph.output[0].name)
optimizer = onnxblock.optim.SGD(clip_grad=grad_clipping)
with onnxblock.empty_base() as accessor:
_ = optimizer(simple_block.parameters())
optimizer_model = accessor.model
assert optimizer_model
def test_retrieve_parameters():
# Given
device = "cuda"

26
orttraining/orttraining/test/python/orttraining_test_python_bindings.py
View file

@ -32,6 +32,7 @@ def _create_training_artifacts(
artifact_directory: str | os.PathLike,
requires_grad: list[str] | None = None,
frozen_params: list[str] | None = None,
optimizer_type=artifacts.OptimType.AdamW,
):
device = "cpu"
batch_size, input_size, hidden_size, output_size = 64, 784, 500, 10
@ -45,7 +46,7 @@ def _create_training_artifacts(
artifacts.generate_artifacts(
onnx_model,
optimizer=artifacts.OptimType.AdamW,
optimizer=optimizer_type,
loss=artifacts.LossType.CrossEntropyLoss,
requires_grad=requires_grad,
frozen_params=frozen_params,
@ -113,7 +114,8 @@ def test_eval_step():
assert fetches
def test_optimizer_step():
@pytest.mark.parametrize("optimizer_type", [artifacts.OptimType.SGD, artifacts.OptimType.AdamW])
def test_optimizer_step(optimizer_type):
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
@ -125,7 +127,7 @@ def test_optimizer_step():
_,
optimizer_model_file_path,
_,
) = _create_training_artifacts(temp_dir)
) = _create_training_artifacts(temp_dir, optimizer_type=optimizer_type)
# Create Checkpoint State.
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module and Optimizer.
@ -142,7 +144,8 @@ def test_optimizer_step():
assert not np.array_equal(old_flatten_params.numpy(), new_params.numpy())
def test_get_and_set_lr():
@pytest.mark.parametrize("optimizer_type", [artifacts.OptimType.SGD, artifacts.OptimType.AdamW])
def test_get_and_set_lr(optimizer_type):
with tempfile.TemporaryDirectory() as temp_dir:
(
checkpoint_file_path,
@ -150,7 +153,7 @@ def test_get_and_set_lr():
_,
optimizer_model_file_path,
_,
) = _create_training_artifacts(temp_dir)
) = _create_training_artifacts(temp_dir, optimizer_type=optimizer_type)
# Create Checkpoint State.
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module and Optimizer.
@ -168,7 +171,8 @@ def test_get_and_set_lr():
assert lr != new_lr
def test_scheduler_step():
@pytest.mark.parametrize("optimizer_type", [artifacts.OptimType.SGD, artifacts.OptimType.AdamW])
def test_scheduler_step(optimizer_type):
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
@ -180,7 +184,7 @@ def test_scheduler_step():
_,
optimizer_model_file_path,
_,
) = _create_training_artifacts(temp_dir)
) = _create_training_artifacts(temp_dir, optimizer_type=optimizer_type)
# Create Checkpoint State.
state = CheckpointState.load_checkpoint(checkpoint_file_path)
# Create a Module and Optimizer.
@ -240,8 +244,9 @@ def test_training_module_checkpoint():
assert np.array_equal(old_flatten_params.numpy(), new_params.numpy())
@pytest.mark.parametrize("optimizer_type", [artifacts.OptimType.SGD, artifacts.OptimType.AdamW])
@pytest.mark.parametrize("trainable_only", [True, False])
def test_copy_buffer_to_parameters(trainable_only):
def test_copy_buffer_to_parameters(trainable_only, optimizer_type):
# Generating random data for testing.
inputs = torch.randn(64, 784).numpy()
labels = torch.randint(high=10, size=(64,), dtype=torch.int64).numpy()
@ -254,7 +259,10 @@ def test_copy_buffer_to_parameters(trainable_only):
optimizer_model_file_path,
_,
) = _create_training_artifacts(
temp_dir, requires_grad=["fc2.weight", "fc2.bias"], frozen_params=["fc1.weight", "fc1.bias"]
temp_dir,
requires_grad=["fc2.weight", "fc2.bias"],
frozen_params=["fc1.weight", "fc1.bias"],
optimizer_type=optimizer_type,
)
state = CheckpointState.load_checkpoint(checkpoint_file_path)

4
orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc
View file

@ -186,7 +186,7 @@ void AdamWTestLoop(
// Add test outputs as baseline.
if (update_signal == nullptr || *update_signal) {
test.AddOutput<int64_t>("updated_flag", {}, {1});
test.AddOutput<bool>("updated_flag", {}, {1});
test.AddSeqOutput("updated_weights", data.UpdatedWeightSeq(), weight_tolerance.first, weight_tolerance.second);
test.AddSeqOutput("updated_momentums_1", data.UpdatedMomentum_1_Seq(), momentum_1_tolerance.first,
momentum_1_tolerance.second);
@ -195,7 +195,7 @@ void AdamWTestLoop(
} else {
// No update happens.
test.AddOutput<int64_t>("updated_flag", {}, {0});
test.AddOutput<bool>("updated_flag", {}, {0});
test.AddSeqOutput("updated_weights", data.WeightSeq(), weight_tolerance.first, weight_tolerance.second);
test.AddSeqOutput("updated_momentums_1", data.Momentum_1_Seq(), momentum_1_tolerance.first,
momentum_1_tolerance.second);

2
orttraining/orttraining/training_api/optimizer.cc
View file

@ -272,7 +272,7 @@ Status Optimizer::Step() {
ORT_THROW_IF_ERROR(status);
// Extract step output and update
if (utils::GetScalarFromOrtValue<int64_t>(outputs[0]) == 1LL) {
if (utils::GetScalarFromOrtValue<bool>(outputs[0]) == true) {
optimizer_state_->step++;
}

6
orttraining/orttraining/training_ops/cpu/optimizer/adamw/adamw.cc
View file

@ -137,7 +137,7 @@ Status AdamWOptimizer<T>::Compute(OpKernelContext* ctx) const {
AdamWOptimizerBase::Prepare p;
ORT_RETURN_IF_ERROR(PrepareForCompute(ctx, p));
int64_t* updated_flag_ptr = p.updated_flag->template MutableData<int64_t>();
bool* updated_flag_ptr = p.updated_flag->template MutableData<bool>();
const Tensor* update_signal = ctx->Input<Tensor>(6);
if (update_signal == nullptr || *update_signal->template Data<bool>()) {
@ -182,9 +182,9 @@ Status AdamWOptimizer<T>::Compute(OpKernelContext* ctx) const {
}
}
*updated_flag_ptr = 1;
*updated_flag_ptr = true;
} else {
*updated_flag_ptr = 0;
*updated_flag_ptr = false;
}
if (p.updated_weights != nullptr) {

6
orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc
View file

@ -35,7 +35,7 @@ Status AdamWOptimizer::ComputeInternal(OpKernelContext* ctx) const {
AdamWOptimizerBase::Prepare p;
ORT_RETURN_IF_ERROR(PrepareForCompute(ctx, p));
int64_t* updated_flag_ptr = p.updated_flag->template MutableData<int64_t>();
bool* updated_flag_ptr = p.updated_flag->template MutableData<bool>();
// Currently placed on CPU, need revisit when we had mixed precision training requirement.
const Tensor* update_signal = ctx->Input<Tensor>(6);
@ -51,9 +51,9 @@ Status AdamWOptimizer::ComputeInternal(OpKernelContext* ctx) const {
launch_multi_tensor_functor<MTA_ADAMW_GROUP_SIZE, TFunctor>(
Stream(ctx), MTA_ADAMW_CHUNK_SIZE, p.grouped_tensor_sizes, p.grouped_tensor_pointers, functor,
alpha_, beta_, epsilon_, *lr_ptr, weight_decay_, adam_mode_, correct_bias_, *step_ptr);
*updated_flag_ptr = 1;
*updated_flag_ptr = true;
} else {
*updated_flag_ptr = 0;
*updated_flag_ptr = false;
}
if (p.updated_weights != nullptr) {