diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake
index a73682efbf..cbba796e0b 100644
--- a/cmake/onnxruntime_python.cmake
+++ b/cmake/onnxruntime_python.cmake
@@ -223,6 +223,9 @@ if (onnxruntime_ENABLE_TRAINING)
   file(GLOB onnxruntime_python_optim_srcs CONFIGURE_DEPENDS
     "${ORTTRAINING_SOURCE_DIR}/python/training/optim/*.py"
   )
+  file(GLOB onnxruntime_python_ortmodule_srcs CONFIGURE_DEPENDS
+    "${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/*.py"
+  )
   file(GLOB onnxruntime_python_train_tools_srcs CONFIGURE_DEPENDS
     "${REPO_ROOT}/tools/python/register_custom_ops_pytorch_exporter.py"
   )
@@ -388,6 +391,7 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/amp
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/optim
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_capi_training_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/capi/training/
@@ -400,6 +404,9 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_optim_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/optim/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_ortmodule_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_train_tools_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/
diff --git a/orttraining/orttraining/python/training/__init__.py b/orttraining/orttraining/python/training/__init__.py
index 84d9b95199..93d69ef231 100644
--- a/orttraining/orttraining/python/training/__init__.py
+++ b/orttraining/orttraining/python/training/__init__.py
@@ -2,12 +2,11 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
+
+
 from onnxruntime.capi._pybind_state import TrainingParameters
 from onnxruntime.capi.training.training_session import TrainingSession
 
 from .orttrainer_options import ORTTrainerOptions
 from .orttrainer import ORTTrainer, TrainStepInfo
 from . import amp, checkpoint, optim, model_desc_validation
-from .execution_agent import InferenceAgent, TrainingAgent
-from .ortmodule import ORTModule
-from .runstateinfo import RunStateInfo
diff --git a/orttraining/orttraining/python/training/_ortmodule_inference_manager.py b/orttraining/orttraining/python/training/_ortmodule_inference_manager.py
deleted file mode 100644
index 808677858a..0000000000
--- a/orttraining/orttraining/python/training/_ortmodule_inference_manager.py
+++ /dev/null
@@ -1,85 +0,0 @@
-# -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT License.
-# --------------------------------------------------------------------------
-
-from . import _ortmodule_utils as _utils, _ortmodule_io as _io
-from ._ortmodule_graph_execution_manager import GraphExecutionManager, _run_forward
-
-import copy
-import onnx
-import onnxruntime
-
-import torch
-
-
-class InferenceManager(GraphExecutionManager):
-    """Concrete instance of GraphExecutionManager that is able to manage the inference model
-
-    InferenceManager is resposible for building and running the forward graph of the inference model
-    """
-
-    def __init__(self, model):
-        super().__init__(model)
-        self._export_mode = torch.onnx.TrainingMode.EVAL
-
-    def forward(self, *inputs, **kwargs):
-        '''Forward pass of the inference model
-
-        ONNX model is exported the first time this method is executed.
-        Next, we build an optimized inference graph with module_graph_builder.
-        Finally, we instantiate the ONNX Runtime InferenceSession through the InferenceAgent.
-        '''
-
-        # Exporting module to ONNX for the first time
-        build_graph = self._export_model(*inputs, **kwargs)
-        if build_graph:
-            # If model was exported, then initialize the graph builder
-            self._initialize_graph_builder(training=False)
-
-            # Save the onnx model if the model was exported
-            if self._save_onnx:
-                onnx.save(self._onnx_model, self._save_onnx_prefix + '_exported_inference_model.onnx')
-
-        # Build the inference graph
-        if build_graph:
-            self._build_graph()
-
-        module_device = _utils.get_device_from_module(self._original_module)
-        # The inference session should be created every time
-        # the graph was built or if the device changed between calls to forward
-        create_execution_session = build_graph or self._device != module_device
-        if self._device != module_device:
-            self._device = module_device
-        if create_execution_session:
-            # Create execution session creates the inference_session
-            self._create_execution_agent()
-
-        user_outputs, _ = _run_forward(self._execution_agent,
-                                       self._optimized_onnx_model,
-                                       self._device,
-                                       *_io._combine_input_buffers_initializers(
-                                           self._flattened_module.named_parameters(),
-                                           self._graph_info.user_input_names,
-                                           self._input_info,
-                                           self._flattened_module.named_buffers(),
-                                           inputs,
-                                           kwargs))
-
-        return _io.unflatten_user_output(self._module_output_schema,
-                                         self._graph_info.user_output_names,
-                                         user_outputs)
-
-    def _build_graph(self):
-        """Build an optimized inference graph using the module_graph_builder"""
-
-        super()._build_graph()
-        if self._save_onnx:
-            onnx.save(self._optimized_onnx_model, self._save_onnx_prefix + '_inference.onnx')
-
-    def _create_execution_agent(self):
-        """Creates an InferenceAgent that can run forward graph on an inference model"""
-
-        session_options, providers, provider_options = self._get_session_config()
-        self._execution_agent = onnxruntime.training.InferenceAgent(self._optimized_onnx_model.SerializeToString(),
-                                                                    session_options, providers, provider_options)
diff --git a/orttraining/orttraining/python/training/ortmodule/__init__.py b/orttraining/orttraining/python/training/ortmodule/__init__.py
new file mode 100644
index 0000000000..a9f373d7a2
--- /dev/null
+++ b/orttraining/orttraining/python/training/ortmodule/__init__.py
@@ -0,0 +1,25 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from packaging import version
+
+# All global constant goes here, before ORTModule is imported
+ONNX_OPSET_VERSION = 12
+MINIMUM_TORCH_VERSION_STR = '1.8.1'
+
+from .ortmodule import ORTModule
+
+
+# Verify proper PyTorch is installed before proceding to ONNX Runtime initializetion
+try:
+    import torch
+    torch_version = version.parse(torch.__version__.split('+')[0])
+    minimum_torch_version = version.parse(MINIMUM_TORCH_VERSION_STR)
+    if torch_version < minimum_torch_version:
+        raise RuntimeError(
+            f'ONNXRuntime ORTModule frontend requires PyTorch version greater or equal to {MINIMUM_TORCH_VERSION_STR}, '
+            f'but version {torch.__version__} was found instead.')
+except:
+    raise(f'PyTorch {MINIMUM_TORCH_VERSION_STR} must be installed in order to run ONNXRuntime ORTModule frontend!')
diff --git a/orttraining/orttraining/python/training/execution_agent.py b/orttraining/orttraining/python/training/ortmodule/_execution_agent.py
similarity index 100%
rename from orttraining/orttraining/python/training/execution_agent.py
rename to orttraining/orttraining/python/training/ortmodule/_execution_agent.py
diff --git a/orttraining/orttraining/python/training/_ortmodule_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
similarity index 85%
rename from orttraining/orttraining/python/training/_ortmodule_graph_execution_manager.py
rename to orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
index 5fffa701c6..b3abee5556 100644
--- a/orttraining/orttraining/python/training/_ortmodule_graph_execution_manager.py
+++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
@@ -2,11 +2,11 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
-from . import _ortmodule_utils as _utils, _ortmodule_io as _io
-from . import _ortmodule_logger as _logger
+
+from . import _utils, _io, _logger
+from onnxruntime.training.ortmodule import ONNX_OPSET_VERSION
 
 from onnxruntime.capi import _pybind_state as C
-
 from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
 
 from abc import ABC, abstractmethod
@@ -20,49 +20,15 @@ import warnings
 
 from torch.utils.cpp_extension import ROCM_HOME
 
-ONNX_OPSET_VERSION = 12
 
-
-def _run_forward(execution_session, onnx_model, device, *inputs, **kwargs):
-    """Runs the forward graph on execution_session with given model inputs and device"""
-
-    # Assert that the input and model device match
-    _utils._check_same_device(device, "Input argument to forward", *inputs)
-
-    # TODO: Try to reuse the output buffers as some of the output tensors are same sizes,
-    #   especially the backward graph outputs.
-    # REVIEW(codemzs): Consolidate Training Agent with InferenceAgent on C++ side to not
-    # have the need for passing IOBinding.
-    if isinstance(execution_session, onnxruntime.training.InferenceAgent):
-        io_binding = execution_session.io_binding()
-        run_options = C.RunOptions()
-
-        # Use IO binding
-        _utils._create_iobinding(io_binding, inputs, onnx_model, device)
-
-        # Run and return module outputs.
-        ort_output = execution_session.run_forward(io_binding, run_options)
-        forward_outputs, run_id = ort_output.ortvalues, ort_output.run_id
-        user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output._ortvalue) for forward_output in forward_outputs)
-        state = None
-    else:
-        state = C.PartialGraphExecutionState()
-        forward_inputs = C.OrtValueVector()
-        for input in inputs:
-            forward_inputs.append(_utils._ortvalue_from_torch_tensor(input))
-
-        forward_outputs = C.OrtValueVector()
-        # Run and return module outputs.
-        execution_session.run_forward(forward_inputs, forward_outputs, state)
-        user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output) for forward_output in forward_outputs)
-
-    # Assert that the outputs and model device match
-    _utils._check_same_device(device, "Output argument from forward", *user_outputs)
-
-    output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
-    run_info = onnxruntime.training.RunStateInfo(state, output_info)
-    # Return user outputs and forward run information
-    return user_outputs, run_info
+class RunStateInfo(object):
+    def __init__(self, state, output_info):
+        """
+        :param state: State of partial run that contains intermediate tensors needed to resume the run later.
+        :param output_info: Output info.
+        """
+        self.state = state
+        self.output_info = output_info
 
 class GraphExecutionManager(ABC):
     def __init__(self, module):
@@ -145,6 +111,26 @@ class GraphExecutionManager(ABC):
             self._torch_alloc = self._torch_gpu_allocator.gpu_caching_allocator_raw_alloc_address()
             self._torch_free = self._torch_gpu_allocator.gpu_caching_allocator_raw_delete_address()
 
+    @staticmethod
+    def execution_session_run_forward(execution_session, onnx_model, device, *inputs):
+        """Runs the forward pass on `execution_session` with given `onnx_model`, `device` and `inputs`
+
+        This is a helper that can be called by the actual `GraphExecutionManager.forward` method
+
+        Args:
+            execution_session (InferenceAgent or InferenceAgent): Agent which runs either inference or train
+            onnx_model (onnx.ModelProto): ONNX model
+            device (torch.device): PyTorch device
+            inputs: (torch.Tensor or a container of): User input
+
+        Returns:
+            Returns a tuple (user_outputs, run_info):
+            user_outputs: The model output (either torch.Tensor or a container of torch.Tensor)
+            run_info: A RunStateInfo which contains extra information about the execution of the graph
+        """
+
+        raise NotImplemented
+
     @abstractmethod
     def forward(self):
         """Executes the forward method for ORTModule
diff --git a/orttraining/orttraining/python/training/_ortmodule_graph_execution_manager_factory.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py
similarity index 83%
rename from orttraining/orttraining/python/training/_ortmodule_graph_execution_manager_factory.py
rename to orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py
index 0eae122071..6d61b586a6 100644
--- a/orttraining/orttraining/python/training/_ortmodule_graph_execution_manager_factory.py
+++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager_factory.py
@@ -3,8 +3,8 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
-from ._ortmodule_training_manager import TrainingManager
-from ._ortmodule_inference_manager import InferenceManager
+from ._training_manager import TrainingManager
+from ._inference_manager import InferenceManager
 
 
 class GraphExecutionManagerFactory(object):
diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py
new file mode 100644
index 0000000000..1f9aea92f2
--- /dev/null
+++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py
@@ -0,0 +1,115 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from . import _utils, _io
+from ._graph_execution_manager import GraphExecutionManager, RunStateInfo
+from ._execution_agent import InferenceAgent
+
+from onnxruntime.capi import _pybind_state as C
+import onnx
+import torch
+
+
+class InferenceManager(GraphExecutionManager):
+    """Concrete instance of GraphExecutionManager that is able to manage the inference model
+
+    InferenceManager is resposible for building and running the forward graph of the inference model
+    """
+
+    def __init__(self, model):
+        super().__init__(model)
+        self._export_mode = torch.onnx.TrainingMode.EVAL
+
+    @staticmethod
+    def execution_session_run_forward(execution_session, onnx_model, device, *inputs):
+        """Runs the forward graph on execution_session with given model inputs and device"""
+
+        # Assert that the input and model device match
+        _utils._check_same_device(device, "Input argument to forward", *inputs)
+
+        # TODO: Try to reuse the output buffers as some of the output tensors are same sizes,
+        #   especially the backward graph outputs.
+        # REVIEW(codemzs): Consolidate Training Agent with InferenceAgent on C++ side to not
+        # have the need for passing IOBinding.
+        io_binding = execution_session.io_binding()
+        run_options = C.RunOptions()
+
+        # Use IO binding
+        _utils._create_iobinding(io_binding, inputs, onnx_model, device)
+
+        # Run and return module outputs.
+        ort_output = execution_session.run_forward(io_binding, run_options)
+        forward_outputs, run_id = ort_output.ortvalues, ort_output.run_id
+        user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output._ortvalue) for forward_output in forward_outputs)
+        state = None
+
+        # Assert that the outputs and model device match
+        _utils._check_same_device(device, "Output argument from forward", *user_outputs)
+
+        output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
+        run_info = RunStateInfo(state, output_info)
+        # Return user outputs and forward run information
+        return user_outputs, run_info
+
+    def forward(self, *inputs, **kwargs):
+        '''Forward pass of the inference model
+
+        ONNX model is exported the first time this method is executed.
+        Next, we build an optimized inference graph with module_graph_builder.
+        Finally, we instantiate the ONNX Runtime InferenceSession through the InferenceAgent.
+        '''
+
+        # Exporting module to ONNX for the first time
+        build_graph = self._export_model(*inputs, **kwargs)
+        if build_graph:
+            # If model was exported, then initialize the graph builder
+            self._initialize_graph_builder(training=False)
+
+            # Save the onnx model if the model was exported
+            if self._save_onnx:
+                onnx.save(self._onnx_model, self._save_onnx_prefix + '_exported_inference_model.onnx')
+
+        # Build the inference graph
+        if build_graph:
+            self._build_graph()
+
+        module_device = _utils.get_device_from_module(self._original_module)
+        # The inference session should be created every time
+        # the graph was built or if the device changed between calls to forward
+        create_execution_session = build_graph or self._device != module_device
+        if self._device != module_device:
+            self._device = module_device
+        if create_execution_session:
+            # Create execution session creates the inference_session
+            self._create_execution_agent()
+
+        user_outputs, _ = InferenceManager.execution_session_run_forward(self._execution_agent,
+                                                                         self._optimized_onnx_model,
+                                                                         self._device,
+                                                                         *_io._combine_input_buffers_initializers(
+                                                                             self._flattened_module.named_parameters(),
+                                                                             self._graph_info.user_input_names,
+                                                                             self._input_info,
+                                                                             self._flattened_module.named_buffers(),
+                                                                             inputs,
+                                                                             kwargs))
+
+        return _io.unflatten_user_output(self._module_output_schema,
+                                         self._graph_info.user_output_names,
+                                         user_outputs)
+
+    def _build_graph(self):
+        """Build an optimized inference graph using the module_graph_builder"""
+
+        super()._build_graph()
+        if self._save_onnx:
+            onnx.save(self._optimized_onnx_model, self._save_onnx_prefix + '_inference.onnx')
+
+    def _create_execution_agent(self):
+        """Creates an InferenceAgent that can run forward graph on an inference model"""
+
+        session_options, providers, provider_options = self._get_session_config()
+        self._execution_agent = InferenceAgent(self._optimized_onnx_model.SerializeToString(),
+                                               session_options, providers, provider_options)
diff --git a/orttraining/orttraining/python/training/_ortmodule_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py
similarity index 100%
rename from orttraining/orttraining/python/training/_ortmodule_io.py
rename to orttraining/orttraining/python/training/ortmodule/_io.py
diff --git a/orttraining/orttraining/python/training/_ortmodule_logger.py b/orttraining/orttraining/python/training/ortmodule/_logger.py
similarity index 100%
rename from orttraining/orttraining/python/training/_ortmodule_logger.py
rename to orttraining/orttraining/python/training/ortmodule/_logger.py
diff --git a/orttraining/orttraining/python/training/_ortmodule_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py
similarity index 77%
rename from orttraining/orttraining/python/training/_ortmodule_training_manager.py
rename to orttraining/orttraining/python/training/ortmodule/_training_manager.py
index 4e650d34c4..f0cedd6aa9 100644
--- a/orttraining/orttraining/python/training/_ortmodule_training_manager.py
+++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py
@@ -3,14 +3,14 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
-from . import _ortmodule_utils as _utils, _ortmodule_io as _io
-from ._ortmodule_graph_execution_manager import GraphExecutionManager, _run_forward
+from . import _utils, _io
+from ._graph_execution_manager import GraphExecutionManager, RunStateInfo
+from ._execution_agent import TrainingAgent
+
 from onnxruntime.capi import _pybind_state as C
-from . import _utils as _utils_ort
 from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_type
 
 import onnx
-import onnxruntime
 import torch
 
 
@@ -24,6 +24,35 @@ class TrainingManager(GraphExecutionManager):
         super().__init__(model)
         self._export_mode = torch.onnx.TrainingMode.TRAINING
 
+    @staticmethod
+    def execution_session_run_forward(execution_session, onnx_model, device, *inputs):
+        """Runs the forward graph on execution_session with given model inputs and device"""
+
+        # Assert that the input and model device match
+        _utils._check_same_device(device, "Input argument to forward", *inputs)
+
+        # TODO: Try to reuse the output buffers as some of the output tensors are same sizes,
+        #   especially the backward graph outputs.
+        # REVIEW(codemzs): Consolidate Training Agent with InferenceAgent on C++ side to not
+        # have the need for passing IOBinding.
+        state = C.PartialGraphExecutionState()
+        forward_inputs = C.OrtValueVector()
+        for input in inputs:
+            forward_inputs.append(_utils._ortvalue_from_torch_tensor(input))
+
+        forward_outputs = C.OrtValueVector()
+        # Run and return module outputs.
+        execution_session.run_forward(forward_inputs, forward_outputs, state)
+        user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output) for forward_output in forward_outputs)
+
+        # Assert that the outputs and model device match
+        _utils._check_same_device(device, "Output argument from forward", *user_outputs)
+
+        output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
+        run_info = RunStateInfo(state, output_info)
+        # Return user outputs and forward run information
+        return user_outputs, run_info
+
     def forward(self, *inputs, **kwargs):
         '''Forward pass starts here and continues at `_ORTModuleFunction.forward`
 
@@ -65,24 +94,24 @@ class TrainingManager(GraphExecutionManager):
             gradient implementation for self.forward_graph.'''
 
             @staticmethod
-            def forward(ctx, *inputs, **kwargs):
+            def forward(ctx, *inputs):
                 '''Performs forward pass based on user input and PyTorch initializer
 
                 Autograd Function's apply() doesn't support keyword arguments,
                 so `*inputs` has all the arguments - keyword arguments converted
-                to positional by the caller.
+                to positional/keywords during `TrainingManager.forward`.
 
                 Module outputs are returned to the user
                 '''
 
-                user_outputs, ctx.run_info = _run_forward(self._execution_agent,
-                                                          self._optimized_onnx_model,
-                                                          self._device,
-                                                          *inputs,
-                                                          **kwargs)
+                user_outputs, ctx.run_info = TrainingManager.execution_session_run_forward(self._execution_agent,
+                                                                                           self._optimized_onnx_model,
+                                                                                           self._device,
+                                                                                           *inputs)
 
-                # Disable materializing grads then None object will not be converted to a tensor filled with zeros prior to calling backward.
-                # Also save shape, device and type info to ctx for materializing tensor in backward if output grad is None.
+                # Disable materializing grads then None object will not be
+                # converted to a tensor filled with zeros prior to calling backward.
+                # Save shape, device and type info to ctx for materializing tensor in backward if output grad is None.
                 ctx.set_materialize_grads(False)
                 return user_outputs
 
@@ -179,19 +208,24 @@ class TrainingManager(GraphExecutionManager):
         fw_outputs_device_info = []
         for idx in range(len(self._graph_info.user_output_names)):
             fw_outputs_device_info.append(C.OrtDevice(get_ort_device_type(self._device.type),
-            C.OrtDevice.default_memory(), _utils_ort.get_device_index(self._device)))
+            C.OrtDevice.default_memory(), _utils.get_device_index(self._device)))
 
         bw_fetches_names = [output.name for output in self._optimized_onnx_model.graph.output]
         bw_outputs_device_info = []
         for idx in range(len(bw_fetches_names)):
             bw_outputs_device_info.append(C.OrtDevice(get_ort_device_type(self._device.type),
-            C.OrtDevice.default_memory(), _utils_ort.get_device_index(self._device)))
+            C.OrtDevice.default_memory(), _utils.get_device_index(self._device)))
 
-        self._execution_agent = onnxruntime.training.TrainingAgent(self._optimized_onnx_model.SerializeToString(),
-                                                                    fw_feed_names, self._graph_info.user_output_names,
-                                                                    fw_outputs_device_info, self._graph_info.module_output_gradient_name,
-                                                                    bw_fetches_names, bw_outputs_device_info,
-                                                                    session_options, providers, provider_options)
+        self._execution_agent = TrainingAgent(self._optimized_onnx_model.SerializeToString(),
+                                              fw_feed_names,
+                                              self._graph_info.user_output_names,
+                                              fw_outputs_device_info,
+                                              self._graph_info.module_output_gradient_name,
+                                              bw_fetches_names,
+                                              bw_outputs_device_info,
+                                              session_options,
+                                              providers,
+                                              provider_options)
 
     def _reinitialize_graph_builder(self, input_info):
         """Return true if the module graph builder was reinitialized"""
diff --git a/orttraining/orttraining/python/training/_ortmodule_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py
similarity index 80%
rename from orttraining/orttraining/python/training/_ortmodule_utils.py
rename to orttraining/orttraining/python/training/ortmodule/_utils.py
index 2e47babb51..255adeb1cc 100644
--- a/orttraining/orttraining/python/training/_ortmodule_utils.py
+++ b/orttraining/orttraining/python/training/ortmodule/_utils.py
@@ -3,8 +3,6 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
-from . import _utils
-
 from onnxruntime.capi.onnxruntime_inference_collection import OrtValue
 from onnxruntime.capi import _pybind_state as C
 
@@ -60,6 +58,32 @@ def _check_same_device(device, argument_str, *args):
                     f"{argument_str} found on device {arg_device}, but expected it to be on module device {device}.")
 
 
+def get_device_index(device):
+    if isinstance(device, str):
+        # could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0
+        device = torch.device(device)
+    elif isinstance(device, int):
+        return device
+    return 0 if device.index is None else device.index
+
+
+def get_device_str(device):
+    if isinstance(device, str):
+        # could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0
+        if device.find(':') == -1:
+            device += ':' + str(torch.cuda.current_device())
+    elif isinstance(device, int):
+        device = 'cuda:' + str(device)
+    elif isinstance(device, torch.device):
+        if device.index is None:
+            device = device.type + ':' + str(torch.cuda.current_device())
+        else:
+            device = device.type + ':' + str(device.index)
+    else:
+        raise RuntimeError('Unsupported device type')
+    return device
+
+
 def get_device_from_module(module):
     '''Returns the first device found in the `module`'s parameters or None'''
     device = None
@@ -79,4 +103,4 @@ def _create_iobinding(io_binding, inputs, model, device):
         io_binding.bind_ortvalue_input(value_info.name, OrtValue(_ortvalue_from_torch_tensor(inputs[idx])))
 
     for value_info in model.graph.output:
-        io_binding.bind_output(value_info.name, device.type, device_id=_utils.get_device_index(device))
+        io_binding.bind_output(value_info.name, device.type, device_id=get_device_index(device))
diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule/ortmodule.py
similarity index 98%
rename from orttraining/orttraining/python/training/ortmodule.py
rename to orttraining/orttraining/python/training/ortmodule/ortmodule.py
index 0a91f97bc8..42f523d789 100644
--- a/orttraining/orttraining/python/training/ortmodule.py
+++ b/orttraining/orttraining/python/training/ortmodule/ortmodule.py
@@ -3,8 +3,8 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
-from . import _ortmodule_io as _io
-from ._ortmodule_graph_execution_manager_factory import GraphExecutionManagerFactory
+from . import _io
+from ._graph_execution_manager_factory import GraphExecutionManagerFactory
 
 from onnxruntime.training import register_custom_ops_pytorch_exporter
 
diff --git a/orttraining/orttraining/python/training/runstateinfo.py b/orttraining/orttraining/python/training/runstateinfo.py
deleted file mode 100644
index 27a9b99640..0000000000
--- a/orttraining/orttraining/python/training/runstateinfo.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT License.
-# --------------------------------------------------------------------------
-
-class RunStateInfo(object):
-   def __init__(self, state, output_info):
-      """
-      :param state: State of partial run that contains intermediate tensors needed to resume the run later.
-      :param output_info: Output info.
-      """
-      self.state = state
-      self.output_info = output_info
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
index 683ef920c0..99ac87c7f0 100644
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
@@ -16,7 +16,7 @@ from collections import OrderedDict
 from collections import namedtuple
 from inspect import signature
 
-from onnxruntime.training import _utils, ORTModule
+from onnxruntime.training.ortmodule import ORTModule, _utils
 import _test_helpers
 
 # Import autocasting libs
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py
index 02c7980971..f63391288f 100644
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py
@@ -17,7 +17,7 @@ import datetime
 
 
 import onnxruntime
-from onnxruntime.training import ORTModule
+from onnxruntime.training.ortmodule import ORTModule
 
 def train(model, optimizer, scheduler, train_dataloader, epoch, device, args):
     # ========================================
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_pipeline_parallel.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_pipeline_parallel.py
index 0b4b14fade..d61863867f 100755
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_pipeline_parallel.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_pipeline_parallel.py
@@ -5,8 +5,7 @@ import deepspeed
 from deepspeed.pipe import PipelineModule, LayerSpec
 from deepspeed.utils import RepeatingLoader
 
-import onnxruntime
-from onnxruntime.training import ORTModule
+from onnxruntime.training.ortmodule import ORTModule, _utils
 
 import argparse
 
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py
index 52f6dda670..eab0339c9b 100644
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_deepspeed_zero_stage_1.py
@@ -10,14 +10,13 @@ $ deepspeed orttraining_test_ortmodule_deepspeed_zero_stage_1.py \
 """
 import argparse
 import logging
-import os
 import torch
 import time
 from torchvision import datasets, transforms
 import torch.distributed as dist
 
 import onnxruntime
-from onnxruntime.training import ORTModule
+from onnxruntime.training.ortmodule import ORTModule
 
 import deepspeed
 
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py
index dc1c0642a3..2ba77a6a5a 100755
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fairscale_sharded_optimizer.py
@@ -9,8 +9,7 @@ from torchvision import datasets, transforms
 import time
 from torch.nn.parallel import DistributedDataParallel as DDP
 import os
-import onnxruntime
-from onnxruntime.training import ORTModule
+from onnxruntime.training.ortmodule import ORTModule
 import numpy as np
 
 # Usage : 
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py
index 8fa062f78e..a806c67c32 100644
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_poc.py
@@ -5,7 +5,7 @@ import time
 from torchvision import datasets, transforms
 
 import onnxruntime
-from onnxruntime.training import ORTModule
+from onnxruntime.training.ortmodule import ORTModule
 
 
 class NeuralNet(torch.nn.Module):
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_torch_lightning_basic.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_torch_lightning_basic.py
index 5e8c2f747b..7c8e1e5a55 100644
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_torch_lightning_basic.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_torch_lightning_basic.py
@@ -10,7 +10,7 @@ from torch.utils.data import DataLoader
 import pytorch_lightning as pl
 
 import onnxruntime
-from onnxruntime.training import ORTModule
+from onnxruntime.training.ortmodule import ORTModule
 
 
 class LitAutoEncoder(pl.LightningModule):