diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake
index 4cd0d676c7..839857abfb 100644
--- a/cmake/onnxruntime_python.cmake
+++ b/cmake/onnxruntime_python.cmake
@@ -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)
diff --git a/orttraining/orttraining/python/training/onnxblock/__init__.py b/orttraining/orttraining/python/training/onnxblock/__init__.py
new file mode 100644
index 0000000000..e97e540103
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/__init__.py
@@ -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)
diff --git a/orttraining/orttraining/python/training/onnxblock/_building_blocks.py b/orttraining/orttraining/python/training/onnxblock/_building_blocks.py
new file mode 100644
index 0000000000..10a8f04c57
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/_building_blocks.py
@@ -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")
diff --git a/orttraining/orttraining/python/training/onnxblock/_graph_utils.py b/orttraining/orttraining/python/training/onnxblock/_graph_utils.py
new file mode 100644
index 0000000000..b7bcc99dfa
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/_graph_utils.py
@@ -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
diff --git a/orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py b/orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py
new file mode 100644
index 0000000000..411a8e266c
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/checkpoint_utils.py
@@ -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
+    )
diff --git a/orttraining/orttraining/python/training/onnxblock/loss/__init__.py b/orttraining/orttraining/python/training/onnxblock/loss/__init__.py
new file mode 100644
index 0000000000..1566b461ac
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/loss/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# __init__.py
+
+from .loss import MSELoss, CrossEntropyLoss
diff --git a/orttraining/orttraining/python/training/onnxblock/loss/loss.py b/orttraining/orttraining/python/training/onnxblock/loss/loss.py
new file mode 100644
index 0000000000..a027889bef
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/loss/loss.py
@@ -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
diff --git a/orttraining/orttraining/python/training/onnxblock/model.py b/orttraining/orttraining/python/training/onnxblock/model.py
new file mode 100644
index 0000000000..081aab54e7
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/model.py
@@ -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
diff --git a/orttraining/orttraining/python/training/onnxblock/model_accessor.py b/orttraining/orttraining/python/training/onnxblock/model_accessor.py
new file mode 100644
index 0000000000..324eae441d
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/model_accessor.py
@@ -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
diff --git a/orttraining/orttraining/python/training/onnxblock/optim/__init__.py b/orttraining/orttraining/python/training/onnxblock/optim/__init__.py
new file mode 100644
index 0000000000..58a2a41f5d
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/optim/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# __init__.py
+
+from .optim import AdamW
diff --git a/orttraining/orttraining/python/training/onnxblock/optim/optim.py b/orttraining/orttraining/python/training/onnxblock/optim/optim.py
new file mode 100644
index 0000000000..480c1ace67
--- /dev/null
+++ b/orttraining/orttraining/python/training/onnxblock/optim/optim.py
@@ -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]
diff --git a/orttraining/orttraining/test/python/orttraining_test_onnxblock.py b/orttraining/orttraining/test/python/orttraining_test_onnxblock.py
new file mode 100644
index 0000000000..8c43cdf8a0
--- /dev/null
+++ b/orttraining/orttraining/test/python/orttraining_test_onnxblock.py
@@ -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
diff --git a/setup.py b/setup.py
index a9749a9317..d1bc9cc2ef 100644
--- a/setup.py
+++ b/setup.py
@@ -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']
diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py
index 432f2d0917..83cbe548c2 100644
--- a/tools/ci_build/build.py
+++ b/tools/ci_build/build.py
@@ -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(