From c0b6c6c94bcd7cb1ed6bb0fc8edfa8b606229411 Mon Sep 17 00:00:00 2001 From: Adam Louly Date: Fri, 18 Aug 2023 10:50:39 -0700 Subject: [PATCH] Add SGDOptimizer in the on-device training offline tooling (onnxblock) (#17085) ### Description Adding SGDOptimizer to on device training onnxblock --- .../test/testdata/training_api/adamw.onnx | Bin 580 -> 544 bytes .../core/graph/training_op_defs.cc | 2 +- .../orttraining/python/training/artifacts.py | 4 +- .../training/onnxblock/optim/__init__.py | 4 +- .../python/training/onnxblock/optim/optim.py | 309 ++++++++++++------ .../test/python/orttraining_test_onnxblock.py | 26 ++ .../orttraining_test_python_bindings.py | 26 +- .../training_ops/cuda/optimizer/common.cc | 4 +- .../orttraining/training_api/optimizer.cc | 2 +- .../training_ops/cpu/optimizer/adamw/adamw.cc | 6 +- .../cuda/optimizer/adamw/adamw.cc | 6 +- 11 files changed, 266 insertions(+), 123 deletions(-) diff --git a/onnxruntime/test/testdata/training_api/adamw.onnx b/onnxruntime/test/testdata/training_api/adamw.onnx index eed8ef5e4883b70a98fedeabcb187530ccd97c66..9195778176c71d40cd7486403fe2d015cd342aeb 100644 GIT binary patch delta 126 zcmX@YvVcW|gHwn*zo?|7C^0iHGcVohD>K*piAtf2vXhvk&E77YEQ>7A__ZMkhfo a-sJpTz1+;?qWt3gv=RYECl)RS0bu|u%^|h` delta 139 zcmZ3$a)d>RgF{F%KQFH$DJMTUTOmI!EhjTCRiPw5A4sQLb+d3So+uK^C_nigqm+`R znUR5klr%`Wm6c;kVs5y9K}lwAW>spDm6frjg{A4lMU%xu!7>VPDFxsBl+>Kb){G07 Yb-BO>af1yK;^N}qV4uv$B*(}H05+8=ga7~l diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 0a37906589..60867accb8 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -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) diff --git a/orttraining/orttraining/python/training/artifacts.py b/orttraining/orttraining/python/training/artifacts.py index b98164ec38..3d6a8e8248 100644 --- a/orttraining/orttraining/python/training/artifacts.py +++ b/orttraining/orttraining/python/training/artifacts.py @@ -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) diff --git a/orttraining/orttraining/python/training/onnxblock/optim/__init__.py b/orttraining/orttraining/python/training/onnxblock/optim/__init__.py index 5d1af763f0..4384ecf054 100644 --- a/orttraining/orttraining/python/training/onnxblock/optim/__init__.py +++ b/orttraining/orttraining/python/training/onnxblock/optim/__init__.py @@ -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"] diff --git a/orttraining/orttraining/python/training/onnxblock/optim/optim.py b/orttraining/orttraining/python/training/onnxblock/optim/optim.py index 94d4c2791d..d14b2efefe 100644 --- a/orttraining/orttraining/python/training/onnxblock/optim/optim.py +++ b/orttraining/orttraining/python/training/onnxblock/optim/optim.py @@ -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 diff --git a/orttraining/orttraining/test/python/orttraining_test_onnxblock.py b/orttraining/orttraining/test/python/orttraining_test_onnxblock.py index 72eafc7be3..c6e8b98d35 100644 --- a/orttraining/orttraining/test/python/orttraining_test_onnxblock.py +++ b/orttraining/orttraining/test/python/orttraining_test_onnxblock.py @@ -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" diff --git a/orttraining/orttraining/test/python/orttraining_test_python_bindings.py b/orttraining/orttraining/test/python/orttraining_test_python_bindings.py index 78cb9c6016..56338ddbaf 100644 --- a/orttraining/orttraining/test/python/orttraining_test_python_bindings.py +++ b/orttraining/orttraining/test/python/orttraining_test_python_bindings.py @@ -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) diff --git a/orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc b/orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc index 7bbf413e4b..a914cd1409 100644 --- a/orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc +++ b/orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc @@ -186,7 +186,7 @@ void AdamWTestLoop( // Add test outputs as baseline. if (update_signal == nullptr || *update_signal) { - test.AddOutput("updated_flag", {}, {1}); + test.AddOutput("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("updated_flag", {}, {0}); + test.AddOutput("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); diff --git a/orttraining/orttraining/training_api/optimizer.cc b/orttraining/orttraining/training_api/optimizer.cc index 26565fdd98..a6b82f1d50 100644 --- a/orttraining/orttraining/training_api/optimizer.cc +++ b/orttraining/orttraining/training_api/optimizer.cc @@ -272,7 +272,7 @@ Status Optimizer::Step() { ORT_THROW_IF_ERROR(status); // Extract step output and update - if (utils::GetScalarFromOrtValue(outputs[0]) == 1LL) { + if (utils::GetScalarFromOrtValue(outputs[0]) == true) { optimizer_state_->step++; } diff --git a/orttraining/orttraining/training_ops/cpu/optimizer/adamw/adamw.cc b/orttraining/orttraining/training_ops/cpu/optimizer/adamw/adamw.cc index 8fe0cb8925..753a167457 100644 --- a/orttraining/orttraining/training_ops/cpu/optimizer/adamw/adamw.cc +++ b/orttraining/orttraining/training_ops/cpu/optimizer/adamw/adamw.cc @@ -137,7 +137,7 @@ Status AdamWOptimizer::Compute(OpKernelContext* ctx) const { AdamWOptimizerBase::Prepare p; ORT_RETURN_IF_ERROR(PrepareForCompute(ctx, p)); - int64_t* updated_flag_ptr = p.updated_flag->template MutableData(); + bool* updated_flag_ptr = p.updated_flag->template MutableData(); const Tensor* update_signal = ctx->Input(6); if (update_signal == nullptr || *update_signal->template Data()) { @@ -182,9 +182,9 @@ Status AdamWOptimizer::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) { diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc index 855754f961..9ff54bfe7e 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc @@ -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(); + bool* updated_flag_ptr = p.updated_flag->template MutableData(); // Currently placed on CPU, need revisit when we had mixed precision training requirement. const Tensor* update_signal = ctx->Input(6); @@ -51,9 +51,9 @@ Status AdamWOptimizer::ComputeInternal(OpKernelContext* ctx) const { launch_multi_tensor_functor( 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) {