diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake
index 2788d326b9..2a57689591 100644
--- a/cmake/onnxruntime_python.cmake
+++ b/cmake/onnxruntime_python.cmake
@@ -416,6 +416,9 @@ if (onnxruntime_ENABLE_TRAINING)
   file(GLOB onnxruntime_python_utils_data_srcs CONFIGURE_DEPENDS
   "${ORTTRAINING_SOURCE_DIR}/python/training/utils/data/*"
   )
+  file(GLOB onnxruntime_python_utils_hooks_srcs CONFIGURE_DEPENDS
+  "${ORTTRAINING_SOURCE_DIR}/python/training/utils/hooks/*"
+  )
   if (onnxruntime_ENABLE_TRAINING_APIS)
     file(GLOB onnxruntime_python_onnxblock_srcs CONFIGURE_DEPENDS
     "${ORTTRAINING_SOURCE_DIR}/python/training/onnxblock/*"
@@ -720,6 +723,7 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/fused_ops
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/data/
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/hooks/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_capi_training_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/capi/training/
@@ -771,6 +775,9 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_utils_data_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/data/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_utils_hooks_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/utils/hooks/
   )
   if (onnxruntime_ENABLE_TRAINING_APIS)
     add_custom_command(
diff --git a/docs/ORTModule_Convergence_Notes.md b/docs/ORTModule_Convergence_Notes.md
new file mode 100644
index 0000000000..ff3a915efe
--- /dev/null
+++ b/docs/ORTModule_Convergence_Notes.md
@@ -0,0 +1,49 @@
+# ORTModule Training Convergence Investigation
+
+## 1. Discovering
+
+Convergence issues can be identified by:
+- Large discrepancy on core training metrics including training loss, evaluation loss, model specific AUC metrics.
+- Runtime failures (for example loss scaler reach the minimum triggering an exception).
+
+Before looking into further, we should clarify few things (if possible):
+- If we change seed for baseline run, whether the metric diff is big?
+  (Make sure the discrepancy is not introduced by random)
+- What's the very first steps we see obvious diverges?
+- Still repro once remove randomness?
+	- Set same seeds
+	- Set dropout ratio to 0
+	- Set compute to be deterministic and torch-comparable (TODO(pengwa): need a flag for this).
+
+
+## 2. Collect Activation Statistics
+
+Add codes:
+
+```diff
++	from onnxruntime.training.utils.hooks import SubscriberManager, StatisticsSubscriber
++	SubscriberManager.subscribe(model, [StatisticsSubscriber("pt_out", override_output_dir=True)])
+
+```
+Run training script to the steps that triggered the divergence. A folder named `pt_out` is created in current working directory. For each step, there is a folder containing summaries for every activation tensor.
+
+
+Add few lines of code:
+```diff
+	from onnxruntime.training.ortmodule import ORTModule
+	from onnxruntime.training.utils.hooks import SubscriberManager, StatisticsSubscriber
+	model = ORTModule(model)
++	SubscriberManager.subscribe(model, [StatisticsSubscriber("ort_out", override_output_dir=True)])
+```
+
+> `StatisticsSubscriber` can be initialized before OR after wrapping ORTModule.
+
+Run training script to the steps that triggered the divergence. Similarly, a folder named `ort_out` is created in current working directory.
+
+Run command to generate per step summary
+
+```bash
+python -m onnxruntime.training.utils.hooks.merge_activation_summary --pt_dir pt_out --ort_dir ort_out --output_dir /tmp/output
+```
+
+Manual diff the generate per-step summary to find the where is the first big diff happens.
diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py
index 581a61ef3f..ce221cb2eb 100755
--- a/onnxruntime/python/tools/symbolic_shape_infer.py
+++ b/onnxruntime/python/tools/symbolic_shape_infer.py
@@ -331,7 +331,7 @@ class SymbolicShapeInference:
                     return unique_dims[int_dim]
                 else:
                     if self.verbose_ > 0:
-                        logger.debug("dim {} has been mergd with dim {}".format(unique_dims[1:], unique_dims[0]))
+                        logger.debug("dim {} has been merged with dim {}".format(unique_dims[1:], unique_dims[0]))
                     return dims[0]
             else:
                 return None
@@ -2199,7 +2199,7 @@ class SymbolicShapeInference:
         vi.CopyFrom(helper.make_tensor_value_info(node.output[1], onnx.TensorProto.INT32, mask_index_shape))
 
         if len(node.output) > 2:
-            # Optional output of add before layer nomalization is done
+            # Optional output of add before layer normalization is done
             # shape is same as the output
             vi = self.known_vi_[node.output[2]]
             vi.CopyFrom(helper.make_tensor_value_info(node.output[2], word_embedding_dtype, output_shape))
@@ -2234,20 +2234,35 @@ class SymbolicShapeInference:
         output_tensor_ranks = get_attribute(node, "output_tensor_ranks")
         assert output_tensor_ranks
 
-        # set the context output seperately.
+        # set the context output separately.
         # The first output is autograd's context.
         vi = self.known_vi_[node.output[0]]
         vi.CopyFrom(helper.make_tensor_value_info(node.output[0], onnx.TensorProto.INT64, []))
-
-        # Outputs after autograd's context are tensors.
-        # We assume their ranks are fixed for different model inputs.
-        for i in range(len(node.output) - 1):
-            # Process the i-th tensor outputs.
-            vi = self.known_vi_[node.output[i + 1]]
-            sympy_shape = self._new_symbolic_shape(output_tensor_ranks[i], node)
-            shape = get_shape_from_sympy_shape(sympy_shape)
-            value_info = helper.make_tensor_value_info(node.output[i + 1], output_tensor_types[i], shape)
-            vi.CopyFrom(value_info)
+        if get_attribute(node, "name").decode() in ["_InspectActivation", "_IncrementStep"]:
+            # PythonOp with name being "_InspectActivation" or "_IncrementStep" will behave exactly same as a normal
+            # PythonOp when execution. The only difference is that
+            # 1). those ops having same number of tensor inputs and tensor outputs;
+            # 2). and the i-th output tensor's shape is same as i-th input tensor's shape.
+            # Be noted, the count of custom autograd function might be bigger than output count, because there might
+            # be other non-tensor constant inputs (string, object, int, tuple, etc). But we did not make those constant
+            # inputs as ONNX op's input, instead they are stored in the attributes.
+            assert len(node.output) == len(node.input) + 1  # The output contains one extra context info.
+            for input_index in range(len(node.output) - 1):
+                # Process the i-th tensor outputs.
+                vi = self.known_vi_[node.output[input_index + 1]]
+                shape = self._get_shape(node, input_index)
+                output_dtype = self.known_vi_[node.input[input_index]].type.tensor_type.elem_type
+                vi.CopyFrom(helper.make_tensor_value_info(node.output[input_index + 1], output_dtype, shape))
+        else:
+            # Outputs after autograd's context are tensors.
+            # We assume their ranks are fixed for different model inputs.
+            for i in range(len(node.output) - 1):
+                # Process the i-th tensor outputs.
+                vi = self.known_vi_[node.output[i + 1]]
+                sympy_shape = self._new_symbolic_shape(output_tensor_ranks[i], node)
+                shape = get_shape_from_sympy_shape(sympy_shape)
+                value_info = helper.make_tensor_value_info(node.output[i + 1], output_tensor_types[i], shape)
+                vi.CopyFrom(value_info)
 
     def _propagate_shape_and_type(self, node, input_index=0, output_index=0):
         shape = self._get_shape(node, input_index)
diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc
index 0ba18fbc6d..c46f03f37e 100755
--- a/orttraining/orttraining/core/graph/gradient_builder.cc
+++ b/orttraining/orttraining/core/graph/gradient_builder.cc
@@ -1825,6 +1825,10 @@ IMPLEMENT_GRADIENT_BUILDER(GetPythonOpGradient) {
               "PythonOpGrad requiring gradient output count mismatch.");
   attrs.push_back(MakeAttribute("output_tensor_requires_grads", bw_tensor_output_requires_grads));
 
+  if (src_attrs.find("comment") != src_attrs.end() && utils::HasString(src_attrs.at("comment"))) {
+    attrs.push_back(MakeAttribute("comment", src_attrs.at("comment").s()));
+  }
+
   std::vector<ArgDef> output_args;
   for (int i = 0; i < GetSrcNodeInputSize(); ++i) {
     if (IsGradientRequiredForSrcNodeInput(i)) {
diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc
index 4527e6a875..b2f87213bd 100644
--- a/orttraining/orttraining/core/graph/training_op_defs.cc
+++ b/orttraining/orttraining/core/graph/training_op_defs.cc
@@ -3871,6 +3871,10 @@ Return true if all elements are true and false otherwise.
           "Indicate if the model is exported in training_mode, by default, False.",
           AttributeProto::INT,
           static_cast<int64_t>(0))
+      .Attr(
+          "comment",
+          "comment only for debugging purposes.",
+          AttributeProto::STRING, false)
       .TypeConstraint(
           "T",
           OpSchema::all_tensor_types(),
@@ -3891,16 +3895,16 @@ Return true if all elements are true and false otherwise.
                     "PythonOp's input list and \"input_tensor_types\" attribute should have the same length.");
         for (auto i = 0; i < input_tensor_types_count; ++i) {
           const auto inferred_input_type = ctx.getInputType(i);
-          ORT_ENFORCE(inferred_input_type, "PythonOp's ", i, "th input type is missing.");
+          ORT_ENFORCE(inferred_input_type, "PythonOp's ", i, "-th input type is missing.");
           ORT_ENFORCE(inferred_input_type->value_case() == TypeProto::kTensorType,
-                      "PythonOp's ", i, "th input type must be a tensor.");
+                      "PythonOp's ", i, "-th input type must be a tensor.");
           ORT_ENFORCE(inferred_input_type->tensor_type().elem_type() == input_tensor_types_proto->ints().at(i),
-                      "PythonOp's ", i, "th input type must be ",
+                      "PythonOp's ", i, "-th input type must be ",
                       TensorProto_DataType_Name(input_tensor_types_proto->ints().at(i)), " but got ",
                       TensorProto_DataType_Name(inferred_input_type->tensor_type().elem_type()));
         }
 
-        // The first output is a pointer which points to
+        // The first output is a pointer that points to
         // a Python object created by torch.autograd.Function.apply.
         // For details, see how we interpret it (the 1st input of PythonOpGrad)
         // in PythonOpGrad's implementation.
@@ -3913,22 +3917,39 @@ Return true if all elements are true and false otherwise.
         // This is a required field.
         ORT_ENFORCE(output_tensor_types_proto, "PythonOp's must have \"output_tensor_types\" attribute.");
 
-        size_t rank_count = 0;
-        // Set inferred output types.
-        for (auto i = 1; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
-          updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i - 1)));
-
-          // Create symbolic shape.
-          const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks");
-          ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
-          for (int64_t j = 0; j < output_tensor_ranks->ints().at(i - 1); ++j) {
-            std::stringstream ss;
-            ss << "PythonOp_unknown_rank_" << rank_count++;
-            rank_only_shape.add_dim()->set_dim_param(ss.str());
+        std::string func_name = getAttribute(ctx, "name", "");
+        if (func_name == "_InspectActivation" || func_name == "_IncrementStep") {
+          // PythonOp with the name attribute being "_InspectActivation" or "_IncrementStep" will behave exactly the
+          // same as a normal PythonOp when execution. The only difference is that:
+          // 1). those ops having the same number of tensor inputs and tensor outputs;
+          // 2). and the i-th output tensor's shape is the same as i-th input tensor's shape.
+          // Be noted, the count of custom autograd function might be bigger than the output count, because there might
+          // be other non-tensor constant inputs (string, object, int, tuple, etc). But we did not make those constant
+          // inputs as ONNX op's input, instead they are stored in the attributes.
+          ORT_ENFORCE(ctx.getNumOutputs() == ctx.getNumInputs() + 1);  // The output contains one extra context info.
+          // Set inferred output types.
+          for (size_t i = 1; i < static_cast<size_t>(ctx.getNumOutputs()); ++i) {
+            size_t input_idx = i - static_cast<size_t>(1);
+            propagateElemTypeFromInputToOutput(ctx, input_idx, i);
+            propagateShapeFromInputToOutput(ctx, input_idx, i);
           }
+        } else {
+          size_t rank_count = 0;
+          // Create a symbolic shape.
+          const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks")->ints();
+          // Set inferred output types.
+          for (auto i = 1; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
+            updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i - 1)));
+            ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
+            for (int64_t j = 0; j < output_tensor_ranks.at(i - 1); ++j) {
+              std::stringstream ss;
+              ss << "PythonOp_unknown_rank_" << rank_count++;
+              rank_only_shape.add_dim()->set_dim_param(ss.str());
+            }
 
-          // Assign symbolic shape.
-          updateOutputShape(ctx, i, rank_only_shape);
+            // Assign symbolic shape.
+            updateOutputShape(ctx, i, rank_only_shape);
+          }
         }
       });
 
@@ -3964,13 +3985,13 @@ Return true if all elements are true and false otherwise.
           AttributeProto::STRING)
       .Attr(
           "inplace",
-          "Indicate if the output should reuse input memory. Todo(pengwa): do we really need it?",
+          "Indicate if the output should reuse input memory. Todo(pengwa): do we need it?",
           AttributeProto::INT,
           static_cast<int64_t>(0))
       .Attr(
           "input_tensor_types",
           "Input types of autograd.Function.backward (including only tensor inputs)."
-          "This attribute is mostly used for input checks for better robustnes.",
+          "This attribute is mostly used for input checks for better robustness.",
           AttributeProto::INTS,
           false)
       .Attr(
@@ -3998,6 +4019,11 @@ Return true if all elements are true and false otherwise.
           "A string inidicating autograd.Function.backward outputs's type."
           "value 'c' - non-tensor output; value 'd' - tensor output.",
           AttributeProto::STRING)
+      .Attr(
+          "comment",
+          "comment only for debugging purposes.",
+          AttributeProto::STRING,
+          false)
       .TypeConstraint(
           "T",
           OpSchema::all_tensor_types(),
@@ -4022,11 +4048,11 @@ Return true if all elements are true and false otherwise.
         // For details, see how we interpret it in PythonOpGrad implementation.
         for (auto i = 1; i < input_count; ++i) {
           const auto inferred_input_type = ctx.getInputType(i);
-          ORT_ENFORCE(inferred_input_type, "PythonOpGrad's ", i, "th input type is missing.");
+          ORT_ENFORCE(inferred_input_type, "PythonOpGrad's ", i, "-th input type is missing.");
           ORT_ENFORCE(inferred_input_type->value_case() == TypeProto::kTensorType,
-                      "PythonOpGrad's ", i, "th input type must be a tensor.");
+                      "PythonOpGrad's ", i, "-th input type must be a tensor.");
           ORT_ENFORCE(inferred_input_type->tensor_type().elem_type() == input_tensor_types_proto->ints().at(i - 1),
-                      "PythonOpGrad's ", i, "th input type must be ", input_tensor_types_proto->ints().at(i - 1));
+                      "PythonOpGrad's ", i, "-th input type must be ", input_tensor_types_proto->ints().at(i - 1));
         }
 
         // Load expected output types.
@@ -4035,20 +4061,36 @@ Return true if all elements are true and false otherwise.
                     "PythonOpGrad's output list and \"output_tensor_types\" attribute should have the same length.");
         // This is a required field.
         ORT_ENFORCE(output_tensor_types_proto, "PythonOpGrad's must have \"output_tensor_types\" attribute.");
-        // Set inferred output types.
-        size_t rank_count = 0;
-        for (auto i = 0; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
-          updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i)));
-          const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks");
-          ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
-          for (int64_t j = 0; j < output_tensor_ranks->ints().at(i); ++j) {
-            std::stringstream ss;
-            ss << "PythonOpGrad_unknown_rank_" << rank_count++;
-            rank_only_shape.add_dim()->set_dim_param(ss.str());
-          }
 
-          // Assign symbolic shape.
-          updateOutputShape(ctx, i, rank_only_shape);
+        std::string func_name = getAttribute(ctx, "name", "");
+        if (func_name == "_InspectActivation" || func_name == "_IncrementStep") {
+          // PythonOpGrad with name attribute being "_InspectActivation" or "_IncrementStep" will behave exactly
+          // the same as a normal PythonOpGrad when execution. The only difference is that:
+          // 1). those ops having the same number of tensor inputs and tensor outputs;
+          // 2). and the i-th output tensor's shape is same as i-th input tensor's shape.
+          ORT_ENFORCE(ctx.getNumOutputs() == ctx.getNumInputs() - 1);  // inputs contains one extra context input
+          for (size_t i = 0; i < static_cast<size_t>(ctx.getNumOutputs()); ++i) {
+            size_t input_idx = i + static_cast<size_t>(1);
+            propagateElemTypeFromInputToOutput(ctx, input_idx, i);
+            propagateShapeFromInputToOutput(ctx, input_idx, i);
+          }
+        } else {
+          // Set inferred output types.
+          size_t rank_count = 0;
+          const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks")->ints();
+          for (auto i = 0; i < static_cast<int64_t>(ctx.getNumOutputs()); ++i) {
+            updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i)));
+
+            ONNX_NAMESPACE::TensorShapeProto rank_only_shape;
+            for (int64_t j = 0; j < output_tensor_ranks.at(i); ++j) {
+              std::stringstream ss;
+              ss << "PythonOpGrad_unknown_rank_" << rank_count++;
+              rank_only_shape.add_dim()->set_dim_param(ss.str());
+            }
+
+            // Assign symbolic shape.
+            updateOutputShape(ctx, i, rank_only_shape);
+          }
         }
       });
 
diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py
index 6787c8603f..73fcf6b2d1 100644
--- a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py
+++ b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py
@@ -118,6 +118,7 @@ def _export_pt_1_10(g, n, *args, **kwargs):
         input_pointer_scalar_positions = []
 
         tensor_args = []
+        debug_comment = ""
         # Encode inputs to autograd.Function.
         for i, arg, call_type in zip(range(len(args)), args, cconv):
             if call_type == "d":
@@ -155,6 +156,11 @@ def _export_pt_1_10(g, n, *args, **kwargs):
                             ORTModuleONNXModelException, Exception(f"Unknown argument type found: {type(arg)}.")
                         )
                 else:
+                    if name == "_InspectActivation" and isinstance(arg, str):
+                        # _InspectActivation is a special case where the first argument is a string
+                        # that is used to determine the activation name to be inspected.
+                        debug_comment += arg
+
                     # All other inputs are accessed via "pointers".
                     input_pointer_scalar_positions.append(i)
                     input_pointer_scalars.append(id(arg))
@@ -187,6 +193,7 @@ def _export_pt_1_10(g, n, *args, **kwargs):
             "output_tensor_types_i": output_tensor_types,
             "output_tensor_ranks_i": output_tensor_ranks,
             "training_mode_i": 1 if training_mode else 0,
+            "comment_s": debug_comment,
         }
 
         if len(input_int_scalars) > 0:
diff --git a/orttraining/orttraining/python/training/utils/hooks/__init__.py b/orttraining/orttraining/python/training/utils/hooks/__init__.py
new file mode 100644
index 0000000000..0944db1067
--- /dev/null
+++ b/orttraining/orttraining/python/training/utils/hooks/__init__.py
@@ -0,0 +1,9 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+
+from ._subscriber_manager import SubscriberManager
+
+from ._statistics_subscriber import StatisticsSubscriber
diff --git a/orttraining/orttraining/python/training/utils/hooks/_statistics_subscriber.py b/orttraining/orttraining/python/training/utils/hooks/_statistics_subscriber.py
new file mode 100644
index 0000000000..d5ecef6030
--- /dev/null
+++ b/orttraining/orttraining/python/training/utils/hooks/_statistics_subscriber.py
@@ -0,0 +1,113 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from pathlib import Path
+from typing import Union
+
+import os
+import warnings
+import shutil
+import torch
+
+from ._subscriber_base import SubscriberBase
+
+
+class StatisticsSubscriber(SubscriberBase):
+    """
+    This subscriber is used to dump the activation statistics into files.
+
+    Each activation will be summarized into 1 or 2 files, depending on whether it is used in the backward pass.
+    > In the forward pass, summarize the tensor's statistics and write to a file.
+    > In the backward pass, summarize the tensor's gradient statistics and write it into another file.
+    So for each run step, there will be many files.
+
+    Currently, the statistics mainly include:
+    > Number of inf/nan values.
+    > Common statistics: tensor shape, data type, total element size, min/max/mean/std of the tensor elements.
+    > Number of zero/negative/positive elements.
+    > Sampled elements (taking the first 128 elements for example).
+    > To be extended...
+
+    `merge_activation_summary.py` can be used to merge the files into one file per training step.
+    """
+
+    def __init__(
+        self,
+        output_dir: str,
+        start_step: Union[None, int] = None,
+        end_step: Union[None, int] = None,
+        override_output_dir: bool = False,
+    ):
+        """
+        Steps in [start_step, end_step) will run subscriber actions.
+
+        Args:
+            output_dir: the directory in all activation statistics files will be stored.
+            start_step: the first step that runs subscriber actions.
+            end_step: the end step (exclusively) that runs subscriber actions.
+            override_output_dir: whether `output_dir` can be overridden if it already exists.
+        """
+        super().__init__(start_step=start_step, end_step=end_step)
+        self._output_dir = output_dir
+        if os.path.exists(self._output_dir):
+            if override_output_dir:
+
+                warnings.warn(f"Output directory {self._output_dir} already exists, overriding it.")
+                shutil.rmtree(self._output_dir)
+            else:
+                raise FileExistsError(
+                    f"Output directory {self._output_dir} already exists. "
+                    "Set override_output_dir=True for StatisticsSubscriber if this is the intention."
+                )
+
+    def module_post_forward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
+        output_file_path = os.path.join(f"{self._output_dir}", f"step_{step}")
+        return self._summarize_activations(activation, depth, name, output_file_path, True)
+
+    def module_pre_backward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
+        output_file_path = os.path.join(f"{self._output_dir}", f"step_{step}")
+        return self._summarize_activations(activation, depth, name, output_file_path, False)
+
+    def _summarize_activations(self, tensor: torch.Tensor, depth: int, name: str, step_folder: str, is_forward: bool):
+        display_name = name + " forward run" if is_forward is True else name + " backward run"
+        output_file_name = name + "_forward" if is_forward is True else name + "_backward"
+
+        if tensor is None or not isinstance(tensor, torch.Tensor):
+            print(f"{display_name} not a torch tensor, value: {tensor}")
+            return
+
+        step_path = Path(step_folder)
+        if not step_path.exists():
+            step_path.mkdir(parents=True, exist_ok=False)
+        order_file_path = step_path / "order.txt"
+        tensor_file_path = step_path / output_file_name
+
+        # This is to try the best effort to align the count of numbers per line for easier comparison in diff views,
+        # though it does not always guarantee to do this way.
+        torch.set_printoptions(precision=6, linewidth=128)
+
+        flatten_array = tensor.flatten()
+        zero_tensor = torch.tensor(0, dtype=flatten_array.dtype, device=flatten_array.device)
+        num_nan = torch.isnan(flatten_array).sum()
+        num_inf = torch.isinf(flatten_array).sum()
+
+        with order_file_path.open(mode="a", encoding="utf-8") as f:
+            f.write(f"{output_file_name}\n")
+
+        with tensor_file_path.open(mode="w", encoding="utf-8") as f:
+            f.write(
+                f"{'>'*depth + display_name} shape: {tensor.shape} dtype: {tensor.dtype} size: {flatten_array.size()} \n"
+                f"min: {flatten_array.min()} max: {flatten_array.max()}, mean: {flatten_array.mean()}, "
+                f"std: {flatten_array.std()} \n"
+                f"nan: {num_nan}, inf: {num_inf}\n"
+            )
+            f.write(f"samples(top 128): {flatten_array[:128]}\n")
+
+            f.write(
+                f"neg: {torch.less(flatten_array, zero_tensor).to(torch.int64).sum()}, "
+                f"pos: {torch.greater(flatten_array, zero_tensor).to(torch.int64).sum()}, "
+                f"zero: {torch.eq(flatten_array, zero_tensor).to(torch.int64).sum()},\n"
+            )
+            f.write(f"{'='*16}\n")
diff --git a/orttraining/orttraining/python/training/utils/hooks/_subscriber_base.py b/orttraining/orttraining/python/training/utils/hooks/_subscriber_base.py
new file mode 100644
index 0000000000..55a55d4ce3
--- /dev/null
+++ b/orttraining/orttraining/python/training/utils/hooks/_subscriber_base.py
@@ -0,0 +1,73 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+
+from typing import Union
+
+import sys
+import torch
+
+
+class SubscriberBase:
+    """
+    Base class for all module hook subscribers.
+    Currently, the hook here only means post-forward hook and pre-backward hook.
+
+    A module hook subscriber is a class that implements the `module_post_forward_impl` and `module_pre_backward_impl`
+    function.
+    > The post_forward hook is called after the activation is generated in the forward path.
+    > The pre_backward hook is called before the activation gradient is computed.
+
+    The post_forward path:
+        Module_A generates activation tensor_a --> Post forward hook (calling subscribers' forward one by one) -->
+        Module_B generates activation tensor_b --> ...
+
+    The pre_backward path:
+        Module_B backward run, tensor_b's gradient is computed as tensor_b_grad -->
+        Pre-backward hook (calling subscribers' backward one by one) -->
+        Module_A backward run, tensor_a's gradient is computed as tensor_a_grad
+
+    Be noted: the "Pre"/"Post" is described from the perspective of Module_A.
+    """
+
+    def __init__(self, start_step: Union[None, int], end_step: Union[None, int]):
+        """
+        Steps in [start_step, end_step) will run the subscriber's actions, and other steps will skip.
+        If start_step is None, 0 is given; if end_step is None, sys.maxsize is given.
+        """
+        self._start_step: int = start_step if start_step is not None else 0
+        self._end_step: int = end_step if end_step is not None else sys.maxsize
+
+    def module_post_forward(self, activation: torch.Tensor, depth: int, name: str, step: int):
+        """
+        This function will be run after the torch Module forward is completed.
+
+        Args:
+            activation: Tensor to be inspected.
+            depth: The indent level of the torch Module generating `activation`.
+            name: The unique name for the `activation`.
+            step: Current execution step.
+        """
+        if self._start_step <= step < self._end_step:
+            self.module_post_forward_impl(activation, depth, name, step)
+
+    def module_pre_backward(self, activation: torch.Tensor, depth: int, name: str, step: int):
+        """
+        This function will be run before the torch Module backward run.
+
+        Args:
+            activation: Tensor to be inspected.
+            depth: The indent level of the torch Module generating `activation`.
+            name: The unique name for the `activation`.
+            step: Current execution step.
+        """
+        if self._start_step <= step < self._end_step:
+            self.module_pre_backward_impl(activation, depth, name, step)
+
+    def module_post_forward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
+        raise NotImplementedError()
+
+    def module_pre_backward_impl(self, activation: torch.Tensor, depth: int, name: str, step: int):
+        raise NotImplementedError()
diff --git a/orttraining/orttraining/python/training/utils/hooks/_subscriber_manager.py b/orttraining/orttraining/python/training/utils/hooks/_subscriber_manager.py
new file mode 100644
index 0000000000..3312182334
--- /dev/null
+++ b/orttraining/orttraining/python/training/utils/hooks/_subscriber_manager.py
@@ -0,0 +1,300 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+
+from collections import abc
+from typing import Callable, List, Union
+
+import torch
+
+from onnxruntime.training.ortmodule import ORTModule
+
+from ._subscriber_base import SubscriberBase
+
+
+class _RuntimeStates:
+    """
+    A data struct holding states for runtime context. Tho kinds of states are included:
+    > Global states that are one-time collected during model hook registration. A global execution step is
+      also initialized to reflect how many steps have been executed, it will get updated after each step
+      completes its forward path.
+    > Intra-execution step states, initialized and cleaned up intended only for current execution step.
+      Usually, it carries intermediate information during the model execution.
+    """
+
+    class _GlobalStates:
+        def __init__(self):
+            # Used to track current execution step, e.g. how many forward/backward path is called.
+            self.execution_step = 0
+            # Used to store the depth of each module, which indicate the indentation level of the module.
+            self.module_index_to_depth = {}
+            # Used to store the unique id of each sequential activation.
+            self.module_to_module_index = {}
+
+            self.subscribers = set()
+
+    class _ExecutionStepStates:
+        def __init__(self):
+            # Used to store the activation tensor names, if already handled, then skipped.
+            # Need to clear after each step.
+            self.observed_activation_names = {}
+
+    def __init__(self):
+        self.global_states = _RuntimeStates._GlobalStates()
+        self.reset_step_states()
+
+    def reset_step_states(self):
+        self.execution_step_states = _RuntimeStates._ExecutionStepStates()
+
+
+class _InspectActivation(torch.autograd.Function):
+    """
+    This class is used to run the subscriber's forward and backward functions.
+    """
+
+    @staticmethod
+    def forward(ctx, activation_name: str, module_idx: int, run_ctx: _RuntimeStates, input_tensor):
+        """
+        Make sure there is a same number of `tensor` type inputs and outputs.
+        This is enforced by ORT's PythonOp's schema check.
+        """
+        depth = run_ctx.global_states.module_index_to_depth[module_idx]
+
+        input_tensor_copied = None
+        if input_tensor is None or not isinstance(input_tensor, torch.Tensor):
+            input_tensor_copied = input_tensor
+        else:
+            input_tensor_copied = input_tensor.detach().clone()
+
+        ctx.current_step = run_ctx.global_states.execution_step
+        ctx.name = activation_name
+        ctx.id = module_idx
+        ctx.depth = depth
+        ctx.subscribers = run_ctx.global_states.subscribers
+
+        # Run subscribers sequentially.
+        for subscriber in run_ctx.global_states.subscribers:
+            subscriber.module_post_forward(input_tensor_copied, depth, activation_name, ctx.current_step)
+
+        return input_tensor.detach() if input_tensor is not None else None
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        val = None
+        if grad_output is None or not isinstance(grad_output, torch.Tensor):
+            val = grad_output
+        else:
+            val = grad_output.detach().clone()
+
+        for subscriber in ctx.subscribers:
+            subscriber.module_pre_backward(val, ctx.depth, ctx.name, ctx.current_step)
+
+        return None, None, None, grad_output.detach() if grad_output is not None else None
+
+
+class _IncrementStep(torch.autograd.Function):
+    """
+    This class is used to manage the global execution step, e.g.
+    global step increment by one, once a full forward path is completed and the state clear.
+
+    This autograd Function is registered as a post-forward hook to the root module. So once the root
+    module's forward path is completed, this backward function will be called immediately, triggering
+    global step increment and state clear.
+    """
+
+    @staticmethod
+    def forward(ctx, run_ctx, input_tensor):
+        """
+        Make sure there is a same number of `tensor` inputs and outputs.
+        This is enforced by ORT's PythonOp's schema check.
+        """
+        ctx.current_step = run_ctx.global_states.execution_step
+        ctx.run_ctx = run_ctx
+
+        # We cannot do the step incremental here. Imagine the outside-most module has multiple outputs,
+        # we need to increase the step only at the very last output handling.
+        # We avoid the complexity to probe the last output handling, and instead, we assume once
+        # the very first backward of the outside-most module is called, then the forward pass MUST be completed.
+
+        # Be noted: it is not safe to register _IncrementStep only for one of the outputs of the outside-most module,
+        # because we are not sure which output branch is executed earlier, for example.
+        #                                   OuterMostModuleOutputs
+        #                                 /                         \
+        #  OuterMostModuleOutputs_0_0th_output             OuterMostModuleOutputs_0_1th_output
+        #                    |                                            |
+        #         PythonOp(_InspectActivation)                  PythonOp(_InspectActivation)
+        #                    |                                            |
+        #          PythonOp(_IncrementStep)                           graph output
+        #                    |
+        #                graph output
+        # The PythonOp(_InspectActivation) (who relies on global step) after 1th output is possible
+        # to run before or after PythonOp(_IncrementStep), so increasing the step is not safe.
+
+        return input_tensor.detach() if isinstance(input_tensor, torch.Tensor) else input_tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        # In case there are multiple backward calls for multiple outputs of the outside-most module.
+        if ctx.current_step == ctx.run_ctx.global_states.execution_step:
+            if ctx.current_step >= 0:
+                print(f"{'='*6} Completed forward pass for STEP {ctx.current_step} {'='*6}")
+            ctx.run_ctx.global_states.execution_step += 1
+            ctx.run_ctx.reset_step_states()
+
+        return None, grad_output.detach() if isinstance(grad_output, torch.Tensor) else grad_output
+
+
+class SubscriberManager:
+    """
+    This class is used to manage all the subscribers and register the post-forward hook to the root module.
+    `subscribe()` is used to register a list of subscribers.
+
+    Currently, the hook handled here is post forward hook for nn.Module. The hook is registered for all nn.Modules
+    recursively. Each hook inserts a PythonOp for every tensor output generated by the corresponding module.
+    Each subscriber implementation is called in the PythonOp's forward function, and backward function.
+
+    There is one special handling for global step increment and state clear. A post-forward hook is registered
+    for the outside-most module, which is the root module. In that hook, _IncrementStep is called, which will
+    increase the step by 1 once the very first time its backward is called (check _IncrementStep for details).
+    """
+
+    def __init__(self):
+        self._run_ctx: _RuntimeStates = _RuntimeStates()
+
+    def subscribe(self, module: Union[torch.nn.Module, ORTModule], subscribers: List[SubscriberBase]):
+        """
+        The API is called externally to register hooks that are implicitly defined by subscribers.
+        Each time all global states will be cleaned up once called.
+        """
+        if not isinstance(module, torch.nn.Module):
+            raise ValueError("module must be a torch.nn.Module instance")
+
+        self._reset_all_states()
+
+        if isinstance(module, ORTModule):
+            module = module.module
+
+        for subscriber in subscribers:
+            if not isinstance(subscriber, SubscriberBase):
+                raise ValueError("subscriber must be a SubscriberBase instance")
+            self._run_ctx.global_states.subscribers.add(subscriber)
+
+        self._initialize(module)
+
+    def _reset_all_states(self):
+        self._run_ctx = _RuntimeStates()
+
+    def _initialize(self, module: torch.nn.Module):
+        """
+        Register hooks for the specified module.
+        """
+        if len(self._run_ctx.global_states.subscribers) == 0:
+            raise RuntimeError("No subscribers are registered.")
+
+        next_module_index = [0]
+        # Register post forward hook for every module, inside the hook, we loop every tensor output of the module,
+        # and wrap it with an autograd Function called _InspectActivation (which takes in a tensor and returns the same
+        # tensor). In this way, we keep ORT and PyTorch run have the same boundary to check activation equality.
+        self._register_hooks_recursively(module, 1, next_module_index)
+
+        # Register post forward hook for the outside-most module, then we increase the dump step.
+        # Be noted, if backward is not triggered, the global dump step remains the original number,
+        # which means the subsequent run will override the previous dump files. This indeed happens to imagine ORTModule
+        # firstly export graph (run the forward only), after the gradient graph is built, another forward+backward is
+        # triggered, override the previous dump files.
+        def _post_forward_outmost_module_hook(module, _, module_outputs):
+            def _apply_to_tensors_func(_, outputs):
+                return _IncrementStep.apply(self._run_ctx, outputs)
+
+            return self._apply_function_to_tensors(module, module_outputs, _apply_to_tensors_func)
+
+        module.register_forward_hook(_post_forward_outmost_module_hook)
+
+    def _register_hooks_recursively(self, module: torch.nn.Module, depth: int, next_module_index: List[int]):
+        """
+        Called to register hooks for every `torch.nn.Module`. Due to `Module` can contain child `Module`s,
+        this function is called recursively by passing in `next_module_index` - a list of int to maintain a
+        global incremental unique module id.
+
+        Args:
+            module: torch.nn.Module to register hook.
+            depth: the indent of the module compared with the outside-most Module.
+            next_module_index: list of int, carrying a global unique module index that can be used next.
+        """
+        module_index = next_module_index[0]
+        self._run_ctx.global_states.module_index_to_depth[module_index] = depth
+        self._run_ctx.global_states.module_to_module_index[module] = module_index
+
+        for child in module.children():
+            if isinstance(child, torch.nn.Module) and child not in self._run_ctx.global_states.module_to_module_index:
+                next_module_index[0] += 1
+                self._register_hooks_recursively(child, depth + 1, next_module_index)
+
+        def _post_forward_module_hook(module, _, module_outputs):
+            if module in self._run_ctx.global_states.module_to_module_index and isinstance(module, torch.nn.Module):
+                module_index = self._run_ctx.global_states.module_to_module_index[module]
+
+                def _apply_to_tensors_func(index, activation_tensor):
+                    name = f"{module.__class__.__name__}_{module_index}_{index}th_output"
+                    if name not in self._run_ctx.execution_step_states.observed_activation_names:
+                        self._run_ctx.execution_step_states.observed_activation_names[name] = True
+                        return _InspectActivation.apply(name, module_index, self._run_ctx, activation_tensor)
+
+                    return activation_tensor
+
+                return self._apply_function_to_tensors(module, module_outputs, _apply_to_tensors_func)
+            return module_outputs
+
+        module.register_forward_hook(_post_forward_module_hook)
+
+    def _is_builtin_type(self, obj):
+        # https://stackoverflow.com/a/17795199
+        return obj.__class__.__module__ in ["__builtin__", "builtins"]
+
+    def _apply_function_to_tensors(self, module: torch.nn.Module, data, func: Callable):
+        """
+        Apply func to all tensors in the given object.
+
+        Args:
+            module: the module that generates the tensors.
+            data: the object that contains activation tensors.
+            func: the function to apply to the tensors.
+        """
+        tensor_output_idx: List[int] = [0]
+
+        def _apply_to_tensors_by_flatten(
+            module: torch.nn.Module,
+            index_for_tensor_output: List[int],
+            outputs,
+            func: Callable,
+        ):
+            if isinstance(outputs, abc.Sequence):
+                touched_outputs = []
+                for output in outputs:
+                    touched_output = _apply_to_tensors_by_flatten(module, index_for_tensor_output, output, func)
+                    touched_outputs.append(touched_output)
+                return outputs.__class__(touched_outputs)
+
+            if isinstance(outputs, abc.Mapping):
+                # apply inplace to avoid recreating dict inherited objects
+                for key in outputs:
+                    outputs[key] = _apply_to_tensors_by_flatten(
+                        module,
+                        index_for_tensor_output,
+                        outputs[key],
+                        func,
+                    )
+                return outputs
+
+            if isinstance(outputs, torch.Tensor):
+                cur_id = index_for_tensor_output[0]
+                index_for_tensor_output[0] += 1
+                return func(cur_id, outputs)
+
+            if not self._is_builtin_type(outputs):
+                raise RuntimeError(f"Unknown type {type(outputs)}")
+            return outputs
+
+        return _apply_to_tensors_by_flatten(module, tensor_output_idx, data, func)
diff --git a/orttraining/orttraining/python/training/utils/hooks/merge_activation_summary.py b/orttraining/orttraining/python/training/utils/hooks/merge_activation_summary.py
new file mode 100644
index 0000000000..ba174712d4
--- /dev/null
+++ b/orttraining/orttraining/python/training/utils/hooks/merge_activation_summary.py
@@ -0,0 +1,156 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+
+"""
+This merges convergence debugging per activation summary files into per-step summary files.
+
+Both PyTorch and ORT run result directories are needed, check `StatisticsSubscriber` usage from
+[ORTModule_Convergence_Notes](docs/ORTModule_Convergence_Notes.md) for how to export the results.
+
+Since ORT and PyTorch run the compute in different orders (e.g. different typological orders sort
+implementations), when we generate a per-step summary, we want the summary to be comparable between
+ORT and PyTorch run. So during the merge, the same typological order is used.
+
+Example:
+    python merge_activation_summary.py --pt_dir pt_out --ort_dir ort_out --output_dir /tmp/output
+
+"""
+
+from pathlib import Path
+
+import argparse
+import logging
+import os
+import shutil
+
+logger = logging.getLogger(__name__)
+
+
+def generate_summaries_per_step(args):
+    pt_dir = args.pt_dir
+    ort_dir = args.ort_dir
+    output_path = Path(args.output_dir)
+
+    output_path.mkdir(parents=True, exist_ok=False)
+
+    # We should use the order.txt generated by PyTorch run, which means, we follow the PyTorch typological order to compare
+    # activation results. Here we assume to get the order.txt from pt_dir/step_0/order.txt
+    topo_order_file_path = Path(f"{pt_dir}/step_0/order.txt")
+
+    src_ort_path = Path(ort_dir)
+    src_pt_path = Path(pt_dir)
+
+    merge_ort_path = output_path / "merge_ort"
+    merge_pt_path = output_path / "merge_pt"
+
+    def generate_summary_per_step(topo_order_file_path: Path, dump_src_path: Path, merge_dest_path: Path):
+        logger.warning(
+            "Start generating summary per step for [%s] following typological order in [%s]",
+            dump_src_path.as_posix(),
+            topo_order_file_path.as_posix(),
+        )
+        with topo_order_file_path.open(mode="r", encoding="utf-8") as order_file:
+            tensor_name_in_order = order_file.readlines()
+
+        if merge_dest_path.exists():
+            shutil.rmtree(merge_dest_path.as_posix())
+        merge_dest_path.mkdir(parents=True, exist_ok=False)
+
+        for dump_step_path in dump_src_path.iterdir():
+            if dump_step_path.is_dir():
+                step_name = dump_step_path.name
+                merge_filename_for_sub_dir = merge_dest_path / f"{step_name}_.txt"
+                # Open merge_filename_for_sub_dir in write mode
+                with merge_filename_for_sub_dir.open(mode="w", encoding="utf-8") as outfile:
+                    for filename in tensor_name_in_order:
+                        filename = filename.rstrip("\n")
+                        full_filename = dump_step_path / filename
+                        if not full_filename.exists():
+                            # Be noted that some tensor handled in PyTorch might be missing in ORT graph
+                            # (if the activation is not used by others, which is pruned during export)
+                            logger.warning("tensor %s not exist", full_filename)
+                            continue
+
+                        with full_filename.open(mode="r", encoding="utf-8") as infile:
+                            outfile.write(infile.read())
+
+                        outfile.write("\n")
+
+        logger.warning(
+            "Finish generating summary per step for [%s] following typological order in [%s], merged files are in [%s]",
+            dump_src_path.as_posix(),
+            topo_order_file_path.as_posix(),
+            merge_dest_path.as_posix(),
+        )
+
+    generate_summary_per_step(topo_order_file_path, src_pt_path, merge_pt_path)
+    generate_summary_per_step(topo_order_file_path, src_ort_path, merge_ort_path)
+
+
+def parse_arguments():
+    """Parse arguments
+    Merge per-step summary from dumped files.
+
+    Data are collected with `StatisticsSubscriber` for ORT and PyTorch runs.
+    Required parameters include
+    > dump root folders for ORT and PyTorch.
+    > output folder.
+
+    Returns:
+        Namespace: arguments
+    """
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--pt_dir",
+        required=True,
+        type=str,
+        help="Root of input directory of PyTorch run result of activation dump.",
+    )
+
+    parser.add_argument(
+        "--ort_dir",
+        required=True,
+        type=str,
+        help="Root of input directory of ORTModule run result of activation dump.",
+    )
+
+    parser.add_argument(
+        "--output_dir",
+        required=True,
+        type=str,
+        help="Root of output directory for generated PyTorch/ORTModule per-step summaries.",
+    )
+
+    parser.add_argument(
+        "--overwrite",
+        required=False,
+        action="store_true",
+        help="Overwrite exists output files.",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_arguments()
+
+    if os.path.exists(args.output_dir):
+        if args.overwrite:
+            logger.warning("Output directory %s already exists, overwriting it.", args.output_dir)
+            shutil.rmtree(args.output_dir)
+        else:
+            raise FileExistsError(
+                f"Output directory {args.output_dir} already exists. " "Enable --overwrite to allow overwriting it."
+            )
+
+    logger.info("Arguments: %s", str(args))
+    generate_summaries_per_step(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/orttraining/orttraining/test/python/orttraining_ortmodule_tests.py b/orttraining/orttraining/test/python/orttraining_ortmodule_tests.py
index 51c18f070c..901015c09d 100644
--- a/orttraining/orttraining/test/python/orttraining_ortmodule_tests.py
+++ b/orttraining/orttraining/test/python/orttraining_ortmodule_tests.py
@@ -142,6 +142,14 @@ def run_data_sampler_tests(cwd, log):
     run_subprocess(command, cwd=cwd, log=log).check_returncode()
 
 
+def run_hooks_tests(cwd, log):
+    log.debug("Running: Data hooks tests")
+
+    command = [sys.executable, "-m", "pytest", "-sv", "orttraining_test_hooks.py"]
+
+    run_subprocess(command, cwd=cwd, log=log).check_returncode()
+
+
 def run_pytorch_export_contrib_ops_tests(cwd, log):
     log.debug("Running: PyTorch Export Contrib Ops Tests")
 
@@ -192,6 +200,8 @@ def main():
 
     run_data_sampler_tests(cwd, log)
 
+    run_hooks_tests(cwd, log)
+
     run_experimental_gradient_graph_tests(cwd, log)
 
     # TODO(bmeswani): Enable this test once it can run with latest pytorch
diff --git a/orttraining/orttraining/test/python/orttraining_test_hooks.py b/orttraining/orttraining/test/python/orttraining_test_hooks.py
new file mode 100644
index 0000000000..4a889d07d6
--- /dev/null
+++ b/orttraining/orttraining/test/python/orttraining_test_hooks.py
@@ -0,0 +1,138 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import tempfile
+import pytest
+import torch
+
+from onnxruntime.training.ortmodule import ORTModule
+from onnxruntime.training.utils.hooks import SubscriberManager, StatisticsSubscriber
+
+
+class NeuralNetSingleOutput(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super().__init__()
+
+        self.fc1 = torch.nn.Linear(input_size, hidden_size)
+        self.relu = torch.nn.ReLU()
+        self.fc2 = torch.nn.Linear(hidden_size, num_classes)
+
+    def forward(self, input1, input2):
+        model_input = input1 + input2
+        out = self.fc1(model_input)
+        out = self.relu(out)
+        out = self.fc2(out)
+        return out
+
+
+class NeuralNetMultipleOutputs(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super().__init__()
+
+        self.fc1 = torch.nn.Linear(input_size, hidden_size)
+        self.relu = torch.nn.ReLU()
+        self.fc2 = torch.nn.Linear(hidden_size, num_classes)
+
+    def forward(self, input1, input2):
+        model_input = input1 + input2
+        out = self.fc1(model_input)
+        out = self.relu(out)
+        out = self.fc2(out)
+        return out, model_input
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("backend", ["torch", "ortmodule"])
+def test_statistic_subscriber_single_output(device, backend):
+    input_size = 8
+    hidden_size = 16
+    num_classes = 32
+    model = NeuralNetSingleOutput(input_size, hidden_size, num_classes)
+    model.to(device)
+    model.train()
+
+    with tempfile.TemporaryDirectory() as temporary_dir:
+        output_dir_path = os.path.join(temporary_dir, f"{backend}_out")
+        sub_manager = SubscriberManager()
+        sub_manager.subscribe(model, [StatisticsSubscriber(output_dir_path, override_output_dir=True)])
+
+        if backend == "ortmodule":
+            model = ORTModule(model)
+
+        batch_size = 4
+        input1_tensor = torch.randn(batch_size, input_size, device=device)
+        input2_tensor = torch.randn(batch_size, input_size, device=device)
+        for _ in range(5):
+            y = model(input1_tensor, input2_tensor)
+            y.sum().backward()
+
+        assert os.path.exists(output_dir_path)
+
+        expected_files = [
+            "order.txt",
+            "Linear_1_0th_output_forward",
+            "Linear_1_0th_output_backward",
+            "NeuralNetSingleOutput_0_0th_output_forward",
+            "NeuralNetSingleOutput_0_0th_output_backward",
+            "ReLU_2_0th_output_forward",
+            "ReLU_2_0th_output_backward",
+            "Linear_3_0th_output_forward",
+            "Linear_3_0th_output_backward",
+        ]
+
+        for i in range(5):
+            step_dir = os.path.join(output_dir_path, f"step_{i}")
+            for file in expected_files:
+                assert os.path.exists(os.path.join(step_dir, file))
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("backend", ["torch", "ortmodule"])
+def test_statistic_subscriber_multiple_outputs(device, backend):
+    input_size = 8
+    hidden_size = 16
+    num_classes = 32
+    model = NeuralNetMultipleOutputs(input_size, hidden_size, num_classes)
+    model.to(device)
+    model.train()
+
+    with tempfile.TemporaryDirectory() as temporary_dir:
+        output_dir_path = os.path.join(temporary_dir, f"{backend}_out")
+
+        sub_manager = SubscriberManager()
+        sub_manager.subscribe(model, [StatisticsSubscriber(output_dir_path, override_output_dir=True)])
+
+        if backend == "ortmodule":
+            model = ORTModule(model)
+
+        batch_size = 4
+        input1_tensor = torch.randn(batch_size, input_size, device=device).requires_grad_(True)
+        input2_tensor = torch.randn(batch_size, input_size, device=device)
+        for _ in range(5):
+            y_output1, y_output2 = model(input1_tensor, input2_tensor)
+            y = y_output1.sum() + y_output2.sum()
+            y.backward()
+
+        assert os.path.exists(output_dir_path)
+
+        expected_files = [
+            "order.txt",
+            "NeuralNetMultipleOutputs_0_0th_output_forward",
+            "NeuralNetMultipleOutputs_0_0th_output_backward",
+            "NeuralNetMultipleOutputs_0_1th_output_forward",
+            "NeuralNetMultipleOutputs_0_1th_output_backward",
+            "Linear_1_0th_output_forward",
+            "Linear_1_0th_output_backward",
+            "ReLU_2_0th_output_forward",
+            "ReLU_2_0th_output_backward",
+            "Linear_3_0th_output_forward",
+            "Linear_3_0th_output_backward",
+        ]
+
+        for i in range(5):
+            step_dir = os.path.join(output_dir_path, f"step_{i}")
+
+            assert len(os.listdir(step_dir)) == len(expected_files)
+            for file in expected_files:
+                assert os.path.exists(os.path.join(step_dir, file))
diff --git a/setup.py b/setup.py
index 8a90b95e1f..7e384b22d6 100644
--- a/setup.py
+++ b/setup.py
@@ -545,6 +545,7 @@ if enable_training or enable_training_apis:
                 "onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator",
                 "onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.fused_ops",
                 "onnxruntime.training.utils.data",
+                "onnxruntime.training.utils.hooks",
             ]
         )