Create dedicated build for training api (#14136)

### Description Enable creating dedicated build for on device training. With this PR we can build a lean binary for on device training using flag --enable_training_apis. This binary includes only the essentials like training ops, optimizers etc and NOT features like Aten fallback, strided tensors, gradient builders etc . This binary also removes all the deprecated components like training::TrainingSession and OrtTrainer etc ### Motivation and Context This enables our partners to create a lean binary for on device training.
2026-07-26 19:52:38 +00:00 · 2023-01-10 20:58:04 -08:00 · 2023-01-10 20:58:04 -08:00 · d92c663f28
commit d92c663f28
parent 3a39736a2c
56 changed files with 197 additions and 172 deletions
--- a/cmake/CMakeLists.txt
+++ b/cmake/CMakeLists.txt
@ -142,7 +142,6 @@ option(onnxruntime_FUZZ_TEST "Enable Fuzz testing" OFF)
 option(onnxruntime_ENABLE_NVTX_PROFILE "Enable NVTX profile." OFF)
 option(onnxruntime_ENABLE_MEMORY_PROFILE "Enable memory profile." OFF)
 option(onnxruntime_ENABLE_TRAINING "Enable full training functionality. Includes ORTModule and ORT Training APIs" OFF)
-# For building with onnxruntime_ENABLE_TRAINING_APIS, onnxruntime_ENABLE_TRAINING must be ON.
 option(onnxruntime_ENABLE_TRAINING_APIS "Enable ort training apis." OFF)
 option(onnxruntime_ENABLE_TRAINING_OPS "Include training operators but no training session support." OFF)
 option(onnxruntime_ENABLE_TRAINING_E2E_TESTS "Enable training end-to-end tests." OFF)
@ -221,6 +220,29 @@ if (onnxruntime_ENABLE_TRAINING)
  set(onnxruntime_ENABLE_ATEN ON)
 endif()

+# ENABLE_TRAINING includes all training functionality
+# The following 2 entry points
+# 1. ORTModule
+# 2. ORT Training APIs
+# It includes all the feature additions as well like
+# 1. Python OP
+# 2. Aten Fallback
+# 3. Strided Tensors
+# 4. All training ops including communication and  collectives ops
+# 5. ONNXBlock (Front end for training preparation when using training apis)
+# Some features are only enabled when onnxruntime_ENABLE_PYTHON is ON as they are only relevant
+# when using python env
+if (onnxruntime_ENABLE_TRAINING)
+  set(onnxruntime_ENABLE_TRAINING_OPS ON)
+  set(onnxruntime_ENABLE_TRAINING_APIS ON)
+  set(onnxruntime_ENABLE_TRAINING_TORCH_INTEROP ON)
+  set(onnxruntime_ENABLE_ATEN ON)
+endif()
+
+if (onnxruntime_ENABLE_TRAINING_APIS)
+  set(onnxruntime_ENABLE_TRAINING_OPS ON)
+endif()
+
 if (onnxruntime_USE_CUDA)
  set(onnxruntime_DISABLE_RTTI OFF)
 endif()
@ -1285,6 +1307,7 @@ if (onnxruntime_USE_DML)
 endif()

 if (onnxruntime_ENABLE_TRAINING_APIS)
+  add_compile_definitions(ENABLE_TRAINING_CORE)
  add_compile_definitions(ENABLE_TRAINING_APIS)
 endif()

@ -1301,8 +1324,9 @@ if (onnxruntime_ENABLE_ROCM_PROFILING)
 endif()

 if (onnxruntime_ENABLE_TRAINING)
-  add_compile_definitions(ENABLE_TRAINING)
+  add_compile_definitions(ENABLE_TRAINING_CORE)
  add_compile_definitions(ENABLE_STRIDED_TENSORS)
+  add_compile_definitions(ENABLE_TRAINING)

  if (UNIX)
    if (EXISTS "${onnxruntime_MPI_HOME}")
--- a/cmake/external/onnxruntime_external_deps.cmake
+++ b/cmake/external/onnxruntime_external_deps.cmake
@ -429,7 +429,9 @@ if(onnxruntime_ENABLE_ATEN)
  FetchContent_Populate(dlpack)
 endif()

-if(onnxruntime_ENABLE_TRAINING)
+if(onnxruntime_ENABLE_TRAINING OR (onnxruntime_ENABLE_TRAINING_APIS AND onnxruntime_BUILD_UNIT_TESTS))
+  # Once code under orttraining/orttraining/models dir is removed "onnxruntime_ENABLE_TRAINING" should be removed from
+  # this conditional
  FetchContent_Declare(
    cxxopts
    URL ${DEP_URL_cxxopts}
@ -470,4 +472,3 @@ FILE(TO_NATIVE_PATH ${PROJECT_SOURCE_DIR}  ORT_SOURCE_DIR)
 if (onnxruntime_USE_CLOUD)
    include(triton)
 endif()
-
--- a/cmake/onnxruntime_graph.cmake
+++ b/cmake/onnxruntime_graph.cmake
@ -72,6 +72,7 @@ if (onnxruntime_ENABLE_TRAINING_OPS AND NOT onnxruntime_ENABLE_TRAINING)
      "${ORTTRAINING_SOURCE_DIR}/core/graph/training_op_defs.h"
      )
 endif()
+
 if (onnxruntime_ENABLE_TRAINING)
  file(GLOB_RECURSE orttraining_graph_src CONFIGURE_DEPENDS
      "${ORTTRAINING_SOURCE_DIR}/core/graph/*.h"
--- a/cmake/onnxruntime_optimizer.cmake
+++ b/cmake/onnxruntime_optimizer.cmake
@ -72,7 +72,9 @@ else()
  )
 endif()

-if (onnxruntime_ENABLE_TRAINING)
+if (onnxruntime_ENABLE_TRAINING_APIS)
+  # we need optimizers for both full build as well as training api only build.
+  # Using onnxruntime_ENABLE_TRAINING_APIS since it is always ON in a full training build.
  list(APPEND onnxruntime_optimizer_src_patterns
    "${ORTTRAINING_SOURCE_DIR}/core/optimizer/*.h"
    "${ORTTRAINING_SOURCE_DIR}/core/optimizer/*.cc"
@ -98,7 +100,7 @@ onnxruntime_add_static_library(onnxruntime_optimizer ${onnxruntime_optimizer_src
 install(DIRECTORY ${PROJECT_SOURCE_DIR}/../include/onnxruntime/core/optimizer  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/onnxruntime/core)
 onnxruntime_add_include_to_target(onnxruntime_optimizer onnxruntime_common onnxruntime_framework onnx onnx_proto ${PROTOBUF_LIB} flatbuffers::flatbuffers Boost::mp11 safeint_interface)
 target_include_directories(onnxruntime_optimizer PRIVATE ${ONNXRUNTIME_ROOT})
-if (onnxruntime_ENABLE_TRAINING)
+if (onnxruntime_ENABLE_TRAINING_APIS)
  target_include_directories(onnxruntime_optimizer PRIVATE ${ORTTRAINING_ROOT})
 endif()
 add_dependencies(onnxruntime_optimizer ${onnxruntime_EXTERNAL_DEPENDENCIES})
--- a/cmake/onnxruntime_providers.cmake
+++ b/cmake/onnxruntime_providers.cmake
@ -471,7 +471,9 @@ if (onnxruntime_USE_CUDA)
  onnxruntime_add_include_to_target(onnxruntime_providers_cuda onnxruntime_common onnxruntime_framework onnx onnx_proto ${PROTOBUF_LIB} flatbuffers::flatbuffers)
  if (onnxruntime_ENABLE_TRAINING_OPS)
    onnxruntime_add_include_to_target(onnxruntime_providers_cuda onnxruntime_training)
-    target_link_libraries(onnxruntime_providers_cuda PRIVATE onnxruntime_training)
+    if (onnxruntime_ENABLE_TRAINING)
+      target_link_libraries(onnxruntime_providers_cuda PRIVATE onnxruntime_training)
+    endif()
    if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
      onnxruntime_add_include_to_target(onnxruntime_providers_cuda Python::Module)
    endif()
--- a/cmake/onnxruntime_session.cmake
+++ b/cmake/onnxruntime_session.cmake
@ -10,6 +10,8 @@ file(GLOB onnxruntime_session_srcs CONFIGURE_DEPENDS
 if (onnxruntime_ENABLE_TRAINING_APIS)
  file(GLOB_RECURSE training_api_srcs CONFIGURE_DEPENDS
    "${ORTTRAINING_SOURCE_DIR}/training_api/*.cc"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/checkpoint_common.cc"
+    "${ORTTRAINING_SOURCE_DIR}/core/framework/checkpoint_common.h"
  )

  list(APPEND onnxruntime_session_srcs ${training_api_srcs})
--- a/cmake/onnxruntime_training.cmake
+++ b/cmake/onnxruntime_training.cmake
@ -232,38 +232,4 @@ if (onnxruntime_BUILD_UNIT_TESTS)
  target_link_libraries(onnxruntime_training_gpt2 PRIVATE onnxruntime_training_runner onnxruntime_training ${ONNXRUNTIME_LIBS} ${onnxruntime_EXTERNAL_LIBRARIES})
  set_target_properties(onnxruntime_training_gpt2 PROPERTIES FOLDER "ONNXRuntimeTest")

-  # Training API Tests
-  # Currently disable it by default for internal development usage.
-  if (onnxruntime_ENABLE_TRAINING_APIS)
-    # Only files in the trainer and common folder will be compiled into test trainer.
-    file(GLOB training_api_test_trainer_src
-      "${ORTTRAINING_SOURCE_DIR}/test/training_api/common/*.cc"
-      "${ORTTRAINING_SOURCE_DIR}/test/training_api/common/*.h"
-      "${ORTTRAINING_SOURCE_DIR}/test/training_api/trainer/*.cc"
-      "${ORTTRAINING_SOURCE_DIR}/test/training_api/trainer/*.h"
-    )
-    onnxruntime_add_executable(onnxruntime_test_trainer ${training_api_test_trainer_src})
-
-    onnxruntime_add_include_to_target(onnxruntime_test_trainer onnxruntime_training
-      onnxruntime_framework onnxruntime_common onnx onnx_proto ${PROTOBUF_LIB} flatbuffers::flatbuffers)
-
-    target_include_directories(onnxruntime_test_trainer PRIVATE
-      ${CMAKE_CURRENT_BINARY_DIR}
-      ${ONNXRUNTIME_ROOT}
-      ${ORTTRAINING_ROOT}
-      ${eigen_INCLUDE_DIRS}
-      ${CXXOPTS}
-      ${extra_includes}
-      ${onnxruntime_graph_header}
-      ${onnxruntime_exec_src_dir}
-    )
-
-    target_link_libraries(onnxruntime_test_trainer PRIVATE
-      onnxruntime_training
-      ${ONNXRUNTIME_LIBS}
-      ${onnxruntime_EXTERNAL_LIBRARIES}
-    )
-    set_target_properties(onnxruntime_test_trainer PROPERTIES FOLDER "ONNXRuntimeTest")
-  endif()
-
 endif()
--- a/cmake/onnxruntime_unittests.cmake
+++ b/cmake/onnxruntime_unittests.cmake
@ -328,12 +328,19 @@ if (onnxruntime_USE_CANN)
  list(APPEND onnxruntime_test_providers_src ${onnxruntime_test_providers_cann_src})
 endif()

-if (onnxruntime_ENABLE_TRAINING)
+if (onnxruntime_ENABLE_TRAINING_APIS)
  file(GLOB_RECURSE orttraining_test_trainingops_cpu_src CONFIGURE_DEPENDS
    "${ORTTRAINING_SOURCE_DIR}/test/training_ops/compare_provider_test_utils.cc"
    "${ORTTRAINING_SOURCE_DIR}/test/training_ops/function_op_test_utils.cc"
    "${ORTTRAINING_SOURCE_DIR}/test/training_ops/cpu/*"
    )
+
+  if (NOT onnxruntime_ENABLE_TRAINING)
+    list(REMOVE_ITEM orttraining_test_trainingops_cpu_src
+      "${ORTTRAINING_SOURCE_DIR}/test/training_ops/cpu/tensorboard/summary_op_test.cc"
+      )
+  endif()
+
  list(APPEND onnxruntime_test_providers_src ${orttraining_test_trainingops_cpu_src})

  if (onnxruntime_USE_CUDA OR onnxruntime_USE_ROCM)
@ -684,9 +691,10 @@ set(all_dependencies ${onnxruntime_test_providers_dependencies} )

 if (onnxruntime_ENABLE_TRAINING)
  list(APPEND all_tests ${onnxruntime_test_training_src})
-  if (onnxruntime_ENABLE_TRAINING_APIS)
+endif()
+
+if (onnxruntime_ENABLE_TRAINING_APIS)
    list(APPEND all_tests ${onnxruntime_test_training_api_src})
-  endif()
 endif()

 if (onnxruntime_USE_TVM)
@ -1285,6 +1293,53 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
      set_target_properties(onnxruntime_mlas_test PROPERTIES LINK_FLAGS "-s ALLOW_MEMORY_GROWTH=1")
    endif()
  endif()
+
+  # Training API Tests
+  if (onnxruntime_ENABLE_TRAINING_APIS)
+    # Only files in the trainer and common folder will be compiled into test trainer.
+    file(GLOB training_api_test_trainer_src
+      "${ORTTRAINING_SOURCE_DIR}/test/training_api/common/*.cc"
+      "${ORTTRAINING_SOURCE_DIR}/test/training_api/common/*.h"
+      "${ORTTRAINING_SOURCE_DIR}/test/training_api/trainer/*.cc"
+      "${ORTTRAINING_SOURCE_DIR}/test/training_api/trainer/*.h"
+    )
+    onnxruntime_add_executable(onnxruntime_test_trainer ${training_api_test_trainer_src})
+
+    onnxruntime_add_include_to_target(onnxruntime_test_trainer onnxruntime_session
+      onnxruntime_framework onnxruntime_common onnx onnx_proto ${PROTOBUF_LIB} flatbuffers::flatbuffers)
+
+    set(CXXOPTS ${cxxopts_SOURCE_DIR}/include)
+    target_include_directories(onnxruntime_test_trainer PRIVATE
+      ${CMAKE_CURRENT_BINARY_DIR}
+      ${ONNXRUNTIME_ROOT}
+      ${ORTTRAINING_ROOT}
+      ${eigen_INCLUDE_DIRS}
+      ${CXXOPTS}
+      ${extra_includes}
+      ${onnxruntime_graph_header}
+      ${onnxruntime_exec_src_dir}
+    )
+
+    set(ONNXRUNTIME_TEST_LIBS
+      onnxruntime_session
+      ${onnxruntime_libs}
+      # CUDA is dynamically loaded at runtime
+      onnxruntime_optimizer
+      onnxruntime_providers
+      onnxruntime_util
+      onnxruntime_framework
+      onnxruntime_util
+      onnxruntime_graph
+      ${ONNXRUNTIME_MLAS_LIBS}
+      onnxruntime_common
+      onnxruntime_flatbuffers
+    )
+    target_link_libraries(onnxruntime_test_trainer PRIVATE
+      ${ONNXRUNTIME_TEST_LIBS}
+      ${onnxruntime_EXTERNAL_LIBRARIES}
+    )
+    set_target_properties(onnxruntime_test_trainer PROPERTIES FOLDER "ONNXRuntimeTest")
+  endif()
 endif()

 if (NOT onnxruntime_BUILD_WEBASSEMBLY)
--- a/csharp/src/Microsoft.ML.OnnxRuntime/Training/CheckpointState.shared.cs
+++ b/csharp/src/Microsoft.ML.OnnxRuntime/Training/CheckpointState.shared.cs
@ -33,7 +33,7 @@ namespace Microsoft.ML.OnnxRuntime
            }
            else
            {
-                throw new InvalidOperationException("Training is disabled in the current build. Please build ONNXRuntime from source with the build flags enable_training. \n");
+                throw new InvalidOperationException("Training is disabled in the current build. Please build ONNXRuntime from source with the build flags enable_training_apis. \n");
            }
        }

--- a/csharp/src/Microsoft.ML.OnnxRuntime/Training/TrainingSession.shared.cs
+++ b/csharp/src/Microsoft.ML.OnnxRuntime/Training/TrainingSession.shared.cs
@ -336,7 +336,7 @@ namespace Microsoft.ML.OnnxRuntime
        {
            if (!NativeTrainingMethods.TrainingEnabled())
            {
-                throw new InvalidOperationException("Training is disabled in the current build. Please build ONNXRuntime from source with the build flags enable_training. \n");
+                throw new InvalidOperationException("Training is disabled in the current build. Please build ONNXRuntime from source with the build flags enable_training_apis. \n");
            }
            var options = sessOptions;
            if (sessOptions == null)
--- a/include/onnxruntime/core/framework/run_options.h
+++ b/include/onnxruntime/core/framework/run_options.h
@ -32,6 +32,8 @@ struct OrtRunOptions {
  bool synchronize_execution_providers = true;

 #ifdef ENABLE_TRAINING
+  // Used by onnxruntime::training::TrainingSession. This class is now deprecated.
+  // Delete training_mode when TrainingSession is deleted.
  // Set to 'true' to run in training mode.
  bool training_mode = true;
 #endif
--- a/onnxruntime/core/framework/allocation_planner.cc
+++ b/onnxruntime/core/framework/allocation_planner.cc
@ -297,7 +297,7 @@ class PlannerImpl {
    *is_strided_tensor = false;
 #ifdef ENABLE_TRAINING
    // Inputs of Yields are essentially the outputs for FW partial subgraph
-    // Thses tensors will be pass back to pytorch, thus cannot share the buffer with other tensors
+    // These tensors will be passed back to pytorch, thus cannot share the buffer with other tensors

    // Unhandled corner case:
    // If FW output tensor is consumed by BW graph, and pytorch performs an inplace operation on th returned tensor,
@ -1584,7 +1584,7 @@ class PlannerImpl {
  }
 #endif

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
  bool AllocateInputsContiguously(const Node& node) const {
    const KernelCreateInfo& ci = GetKernelCreateInfo(kernel_create_info_map_, node.Index());
    if (ci.kernel_def == nullptr) {
@ -1961,7 +1961,7 @@ class PlannerImpl {
    }
 #ifdef ENABLE_TRAINING
    // 6. build the node_execution_order_in_training
-    //  the training memory optmization rely on a stable order how kernel get launched to calculate memory pattern
+    //  the training memory optimization rely on a stable order how kernel get launched to calculate memory pattern
    //  so we limit training scenario to run with single stream and single thread mode
    //  the code below will simulate the execution and get the stable execution order
    InlinedVector<int> execution_offsets(num_logic_streams_, -1);
@ -2150,7 +2150,7 @@ Status PlannerImpl::CreatePlan(
  AdjustInplaceLifeIntervals();
 #endif

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
  // Determine allocation order for weights and activations. This needs to be done after ComputeReusePlan.
  ORT_RETURN_IF_ERROR(ComputeAllocationOrder());
 #endif
--- a/onnxruntime/core/framework/execution_frame.h
+++ b/onnxruntime/core/framework/execution_frame.h
@ -57,6 +57,8 @@ class IExecutionFrame {
 #endif

 #ifdef ENABLE_TRAINING
+  // Referenced by PartialGraphExecutionState which is applicable when using ORTModule.
+  // These wont be needed when using ORT Training APIs
  void UpdateFeeds(gsl::span<const int> feed_mlvalue_idxs, gsl::span<const OrtValue> feeds);
  void UpdateFetches(gsl::span<const int> fetch_mlvalue_idxs, gsl::span<const OrtValue> fetches,

--- a/onnxruntime/core/framework/execution_steps.cc
+++ b/onnxruntime/core/framework/execution_steps.cc
@ -57,6 +57,7 @@ Status LaunchKernelStep::Execute(StreamExecutionContext& ctx,
                                 const bool& terminate_flag,
                                 bool& continue_flag) {
 #ifdef ENABLE_TRAINING
+  // legacy code required by ORTTrainer. Should be removed when ORTTrainer is removed
  auto* node_to_execute = ctx.GetNodeToExecute();
  if (node_to_execute && node_to_execute->count(node_index_) == 0) {
    continue_flag = true;
--- a/onnxruntime/core/framework/execution_steps.h
+++ b/onnxruntime/core/framework/execution_steps.h
@ -21,6 +21,7 @@ class BarrierStep : public SequentialExecutionPlan::ExecutionStep {

  std::string ToString() const override;
 #ifdef ENABLE_TRAINING
+  // Only applicable when using PartialExecutor
  bool IsBarrier() const override;
 #endif
 private:
--- a/onnxruntime/core/framework/sequential_execution_plan.h
+++ b/onnxruntime/core/framework/sequential_execution_plan.h
@ -117,7 +117,7 @@ struct SequentialExecutionPlan : public ExecutionPlanBase {
                           bool& continue_flag) = 0;
    virtual std::string ToString() const = 0;
 #ifdef ENABLE_TRAINING
-    // the partial execution mode for training need special handle for barrier
+    // the partial execution mode for training needs special handling for barrier
    virtual bool IsBarrier() const { return false; }
 #endif
  };
--- a/onnxruntime/core/framework/sequential_executor.cc
+++ b/onnxruntime/core/framework/sequential_executor.cc
@ -441,10 +441,12 @@ onnxruntime::Status ExecuteKernel(StreamExecutionContext& ctx,
  } else {
    KernelScope kernel_scope(session_scope, kernel_ctx, *p_kernel);
    ORT_TRY {
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
      if (p_kernel->KernelDef().AllocateInputsContiguously()) {
        ORT_RETURN_IF_ERROR(utils::VerifyInputTensorsAllocatedContiguously(&kernel_ctx));
      }
+#endif
+#ifdef ENABLE_TRAINING
      // Cache lookup. Currently we only cache single-output nodes,
      // to keep memory overhead impact in check. Hence we only look in cache
      // if the current node has one output.
--- a/onnxruntime/core/framework/sequential_executor.h
+++ b/onnxruntime/core/framework/sequential_executor.h
@ -24,8 +24,11 @@ namespace onnxruntime {
 class StreamExecutionContext;
 class DeviceStreamCollection;
 class SessionScope;
+
+#ifdef ENABLE_TRAINING
 using OrtValueCache = InlinedHashMap<std::string, OrtValue>;
 using OrtValueCachePtr = std::shared_ptr<OrtValueCache>;
+#endif

 onnxruntime::Status ExecuteKernel(StreamExecutionContext& ctx,
                                  NodeIndex idx,
--- a/onnxruntime/core/framework/session_state.cc
+++ b/onnxruntime/core/framework/session_state.cc
@ -294,7 +294,7 @@ bool SessionState::IsSparseInitializer(int ort_value_index) const {
 }
 #endif

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 Status SessionState::GetInitializedTensors(
    const std::unordered_set<std::string>& interested_weights,
    bool allow_missing_weights, NameMLValMap& retrieved_weights) const {
@ -1383,11 +1383,11 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
  // sharing a single buffer makes it hard to release individual ones, leading
  // to memory waste.
  //
-  // TODO!! disabling memory pattern tracer increases fragementation, leading to
+  // TODO!! disabling memory pattern tracer increases fragmentation, leading to
  //  out of memory error in some training tests. Need to create kernel first,
-  //  and let the kernel tells us whether the initalizer needs to be traced.
+  //  and let the kernel tells us whether the initializer needs to be traced.
  //
-#if defined(ENABLE_TRAINING)
+#if defined(ENABLE_TRAINING_CORE)
  std::unique_ptr<ITensorAllocator> tensor_allocator(
      ITensorAllocator::Create(enable_mem_pattern_, *p_seq_exec_plan_, *this, weights_buffers_));
 #else
@ -1456,7 +1456,7 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_

  ORT_RETURN_IF_ERROR(CreateKernels(kernel_registry_manager));

-#ifndef ENABLE_TRAINING
+#ifndef ENABLE_TRAINING_CORE
  const auto disable_prepacking =
      session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigDisablePrepacking, "0");

--- a/onnxruntime/core/framework/session_state.h
+++ b/onnxruntime/core/framework/session_state.h
@ -161,7 +161,7 @@ class SessionState {
  bool IsSparseInitializer(int ort_value_index) const;
 #endif

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
  /**
    Get some initialized tensors (weights).
    @param interested_weights The names of the weights to retrieve.
@ -516,6 +516,7 @@ class SessionState {
  PrepackedWeightsContainer* const prepacked_weights_container_{};

 #ifdef ENABLE_TRAINING
+// Needed for ORTTrainer. Should be removed along with ORTTrainer code
 #ifndef DISABLE_ABSEIL
  InlinedHashMap<InlinedVector<int>, InlinedHashSet<NodeIndex>> to_be_executed_nodes_;
 #else
--- a/onnxruntime/core/framework/utils.cc
+++ b/onnxruntime/core/framework/utils.cc
@ -631,11 +631,11 @@ ExecuteGraphImpl(const SessionState& session_state,
  bool is_subgraph = session_state.GetGraphViewer().ParentNode() != nullptr;
  // in following two cases, we execute the workload in main thread:
  // 1. execution mode is sequential.
-  // 2. execute a subgraph. Because in current implmentation, the execute of subgraph is launched through parent kernel.
+  // 2. execute a subgraph. Because in current implementation, the execute of subgraph is launched through parent kernel.
  //    the parent kernel will occupy a thread in thread pool. if we use multiple threads to execute subgraph, it may cause
  //    deadlock when we reach the limitation of thread pool.
  bool single_thread_mode = execution_mode == ExecutionMode::ORT_SEQUENTIAL || is_subgraph;
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
  single_thread_mode = true;
 #endif

@ -1004,7 +1004,7 @@ int32_t ONNXTensorElementDataTypeToProtoTensorType(ONNXTensorElementDataType onn
  }
 }

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 common::Status VerifyInputTensorsAllocatedContiguously(OpKernelContext* context) {
  const Tensor* prev_input = context->Input<Tensor>(0);
  for (int i = 1; i < context->InputCount(); i++) {
--- a/onnxruntime/core/framework/utils.h
+++ b/onnxruntime/core/framework/utils.h
@ -194,7 +194,7 @@ constexpr ONNXTensorElementDataType GetONNXTensorElementDataType<uint64_t>() {

 int32_t ONNXTensorElementDataTypeToProtoTensorType(ONNXTensorElementDataType);

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 common::Status VerifyInputTensorsAllocatedContiguously(OpKernelContext* context);
 #endif

--- a/onnxruntime/core/graph/graph.cc
+++ b/onnxruntime/core/graph/graph.cc
@ -1630,7 +1630,7 @@ Status Graph::BuildConnections(std::unordered_set<std::string>& outer_scope_node
            // (they're internally available to the fused node but removed from the Graph instance).
            // Fusion happens after the model was loaded in full so we know the inputs were valid originally.
            bool check = node.NodeType() != Node::Type::Fused;
-#if defined(ENABLE_TRAINING)
+#if defined(ENABLE_TRAINING_CORE)
            // Only check initial model load for training as graph modifications there also render inputs 'invalid'.
            check = check && num_resolves_ == 0;
 #endif
--- a/onnxruntime/core/optimizer/compute_optimizer/compute_optimizer.cc
+++ b/onnxruntime/core/optimizer/compute_optimizer/compute_optimizer.cc
@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 #include <onnx/defs/attr_proto_util.h>

 #include "core/common/safeint.h"
--- a/onnxruntime/core/optimizer/compute_optimizer/compute_optimizer.h
+++ b/onnxruntime/core/optimizer/compute_optimizer/compute_optimizer.h
@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 #pragma once

 #include "core/optimizer/compute_optimizer/passthrough_actors.h"
--- a/onnxruntime/core/optimizer/compute_optimizer/passthrough_actors.cc
+++ b/onnxruntime/core/optimizer/compute_optimizer/passthrough_actors.cc
@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 #include <onnx/defs/attr_proto_util.h>

 #include "core/common/safeint.h"
--- a/onnxruntime/core/optimizer/compute_optimizer/passthrough_actors.h
+++ b/onnxruntime/core/optimizer/compute_optimizer/passthrough_actors.h
@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 #pragma once

 // Uncomment for debugging
--- a/onnxruntime/core/optimizer/dropout_elimination.cc
+++ b/onnxruntime/core/optimizer/dropout_elimination.cc
@ -26,7 +26,7 @@ bool EliminateDropout::SatisfyCondition(const Graph& graph, const Node& node, co
    return false;
  }

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
  // allow Dropout elimination when:
  //    1. ratio input is an initializer of 0
 	//    2. ratio input is not a graph input, so it cannot be overridden
--- a/onnxruntime/core/optimizer/graph_transformer_utils.cc
+++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc
@ -67,12 +67,14 @@
 #include "core/optimizer/transpose_optimizer/ort_transpose_optimizer.h"
 #include "core/optimizer/unsqueeze_elimination.h"
 #include "core/optimizer/identical_children_consolidation.h"
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 #include "orttraining/core/optimizer/bitmask_dropout_replacement.h"
 #include "orttraining/core/optimizer/bias_softmax_dropout_fusion.h"
-#include "orttraining/core/optimizer/memory_optimizer.h"
 #include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h"
 #endif
+#ifdef ENABLE_TRAINING
+#include "orttraining/core/optimizer/memory_optimizer.h"
+#endif

 #endif  // !defined(ORT_MINIMAL_BUILD)

@ -276,7 +278,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
      transformers.emplace_back(std::make_unique<BiasGeluFusion>(cpu_cuda_dml_rocm_eps));
      transformers.emplace_back(std::make_unique<BiasSoftmaxFusion>(cpu_cuda_rocm_eps));
      transformers.emplace_back(std::make_unique<BiasDropoutFusion>(cuda_rocm_eps));
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
      transformers.emplace_back(std::make_unique<BitmaskDropoutReplacement>(cuda_rocm_eps));
      transformers.emplace_back(std::make_unique<BiasSoftmaxDropoutFusion>(cuda_rocm_eps));
      transformers.emplace_back(std::make_unique<SceLossGradBiasFusion>(cpu_cuda_rocm_eps));
--- a/onnxruntime/core/optimizer/layer_norm_fusion.cc
+++ b/onnxruntime/core/optimizer/layer_norm_fusion.cc
@ -345,7 +345,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
    for (size_t i = 0; i < mul_node.MutableInputDefs().size(); i++) {
      if (graph_utils::NodeArgIsConstant(graph, *(mul_node.MutableInputDefs()[i])) ||
          graph_utils::IsGraphInput(graph, mul_node.MutableInputDefs()[i])) {
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
        if (axes_values.empty() ||
            mul_node.MutableInputDefs()[i]->Shape()->dim_size() == static_cast<int>(axes_values.size())) {
          scale = mul_node.MutableInputDefs()[i];
@ -362,7 +362,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
    for (size_t i = 0; i < last_add_node.MutableInputDefs().size(); i++) {
      if (graph_utils::NodeArgIsConstant(graph, *(last_add_node.MutableInputDefs()[i])) ||
          graph_utils::IsGraphInput(graph, last_add_node.MutableInputDefs()[i])) {
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
        if (axes_values.empty() ||
            last_add_node.MutableInputDefs()[i]->Shape()->dim_size() == static_cast<int>(axes_values.size())) {
          bias = last_add_node.MutableInputDefs()[i];
@ -423,7 +423,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
    // remove all the other nodes.
    graph_utils::FinalizeNodeFusion(graph, nodes_to_remove, layer_norm_node);

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
    // add two extra output defs, so we have 3 output defs that match what gradient builder expected
    layer_norm_node.MutableOutputDefs().push_back(&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("saved_mean"), nullptr));
    layer_norm_node.MutableOutputDefs().push_back(&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("saved_inv_std_var"), nullptr));
@ -594,7 +594,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr
    for (size_t i = 0; i < mul_node.MutableInputDefs().size(); i++) {
      if (graph_utils::NodeArgIsConstant(graph, *(mul_node.MutableInputDefs()[i])) ||
          graph_utils::IsGraphInput(graph, mul_node.MutableInputDefs()[i])) {
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
        if (axes_values.empty() ||
            mul_node.MutableInputDefs()[i]->Shape()->dim_size() == static_cast<int>(axes_values.size())) {
          scale = mul_node.MutableInputDefs()[i];
@ -643,7 +643,7 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr
    // remove all the other nodes.
    graph_utils::FinalizeNodeFusion(graph, nodes_to_remove, layer_norm_node);

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
    // add one extra output def, so we have 2 output defs that match what gradient builder expected
    layer_norm_node.MutableOutputDefs().push_back(
        &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("saved_inv_std_var"), nullptr));
--- a/onnxruntime/core/session/environment.cc
+++ b/onnxruntime/core/session/environment.cc
@ -34,13 +34,14 @@
 #if defined(ENABLE_TRAINING_OPS)
 #include "orttraining/core/graph/training_op_defs.h"
 #endif
+#ifdef ENABLE_TRAINING_CORE
+#include "orttraining/core/optimizer/graph_transformer_registry.h"
+#endif
 #ifdef ENABLE_TRAINING
 #include "orttraining/core/graph/gradient_builder_registry.h"
-#include "orttraining/core/graph/loss_function_registry.h"
 #include "orttraining/core/graph/optimizer_builder.h"
 #include "orttraining/core/graph/optimizer_graph_builder_registry.h"
-#include "orttraining/core/optimizer/graph_transformer_registry.h"
-
+#include "orttraining/core/graph/loss_function_registry.h"
 #endif

 namespace onnxruntime {
@ -264,13 +265,16 @@ Status Environment::Initialize(std::unique_ptr<logging::LoggingManager> logging_
      // preserve this order until <training schemas>: this depends on operatorsetschema registration.
      training::RegisterTrainingOpSchemas();
 #endif
+#ifdef ENABLE_TRAINING_CORE
+      // <training schemas>
+      // This can also be moved inside enable_training. Needs more investigation
+      training::GraphTransformerRegistry::GetInstance().RegisterExternalGraphTransformers();
+#endif
 #ifdef ENABLE_TRAINING
      training::GradientBuilderRegistry::GetInstance().RegisterGradientBuilders();
      training::LossFunctionRegistry::GetInstance().RegisterNonOperatorLossFunctions();
      training::OptimizerBuilderRegistry::GetInstance().RegisterBuilders();
      training::OptimizerGraphBuilderRegistry::GetInstance().RegisterGraphBuilders();
-      // <training schemas>
-      training::GraphTransformerRegistry::GetInstance().RegisterExternalGraphTransformers();
 #endif
    });

--- a/onnxruntime/core/session/onnxruntime_c_api.cc
+++ b/onnxruntime/core/session/onnxruntime_c_api.cc
@ -2299,8 +2299,7 @@ ORT_API(const OrtTrainingApi*, OrtApis::GetTrainingApi, uint32_t version) {
  ORT_UNUSED_PARAMETER(version);
  fprintf(stderr,
          "Training APIs are not supported with this build. Please build onnxruntime "
-          "from source with the build flags enable_training and enable_training_apis to "
-          "retrieve the training APIs.\n");
+          "from source with the build flags enable_training_apis to retrieve the training APIs.\n");

  return nullptr;
 #endif
--- a/onnxruntime/test/contrib_ops/attention_op_test.cc
+++ b/onnxruntime/test/contrib_ops/attention_op_test.cc
@ -3654,7 +3654,7 @@ TEST(AttentionTest, DISABLED_Attention_Mask1D_Fp16_B2_FusedNoPadding) {
  }
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(AttentionTest, SharedPrepackedWeights) {
  int batch_size = 2;
  int sequence_length = 2;
--- a/onnxruntime/test/contrib_ops/qordered_attention_test.cc
+++ b/onnxruntime/test/contrib_ops/qordered_attention_test.cc
@ -6,7 +6,7 @@
 // The "Attention_WithData_ROW_ORDER", "MatMul_COL_16x64x32", "MatMul_COL_16x64x32_perchannel", "MatMul_addC_COL_16x64x32", "MatMul_addC_COL_16x64x32_perchannel", "MatMul_COL_16x64x32_b3_1", "MatMul_addC_COL_16x64x32_b2_1", "MatMul_addC_COL_16x64x32_b2_1_perchannel", "MatMul_addC_broadcastC_COL_16x64x32_b2_1" tests fails in Windows Orttraining build with errors like:
 //"qkv_bias_const_cout_ == 3 && scale_qkv_weight_const_count_ == 3 && qkv_weight_const_count_ == 3 was false. qkv gemm weight and their scales, qkv gemm bias must all be constant!"

-#if defined(USE_CUDA) && !defined(ENABLE_TRAINING)
+#if defined(USE_CUDA) && !defined(ENABLE_TRAINING_CORE)

 #include <cuda.h>

--- a/onnxruntime/test/contrib_ops/qordered_matmul_op_test.cc
+++ b/onnxruntime/test/contrib_ops/qordered_matmul_op_test.cc
@ -5,7 +5,7 @@

 // The "Attention_WithData_ROW_ORDER", "MatMul_COL_16x64x32", "MatMul_COL_16x64x32_perchannel", "MatMul_addC_COL_16x64x32", "MatMul_addC_COL_16x64x32_perchannel", "MatMul_COL_16x64x32_b3_1", "MatMul_addC_COL_16x64x32_b2_1", "MatMul_addC_COL_16x64x32_b2_1_perchannel", "MatMul_addC_broadcastC_COL_16x64x32_b2_1" tests fails in Windows Orttraining build with errors like:
 //"qkv_bias_const_cout_ == 3 && scale_qkv_weight_const_count_ == 3 && qkv_weight_const_count_ == 3 was false. qkv gemm weight and their scales, qkv gemm bias must all be constant!"
-#if defined(USE_CUDA) && !defined(ENABLE_TRAINING)
+#if defined(USE_CUDA) && !defined(ENABLE_TRAINING_CORE)

 #include <cuda.h>

--- a/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc
+++ b/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc
@ -980,7 +980,7 @@ TEST(QAttentionTest, QAttentionPrunedModel) {
                   input_hidden_size);
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(QAttentionTest, SharedPrepackedWeights) {
  int batch_size = 1;
  int sequence_length = 2;
--- a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc
+++ b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc
@ -362,7 +362,7 @@ TEST(DynamicQuantLSTMTest, LargeSize) {
  RunQuantLSTM<uint8_t>(12, 3, 278);
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) {
  OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/);

--- a/onnxruntime/test/framework/session_state_test.cc
+++ b/onnxruntime/test/framework/session_state_test.cc
@ -309,7 +309,7 @@ TEST(SessionStateTest, TestInitializerMemoryAllocatedUsingNonArenaMemory) {

 INSTANTIATE_TEST_SUITE_P(SessionStateTests, SessionStateTestP, testing::ValuesIn(param_list));

-#ifndef ENABLE_TRAINING
+#ifndef ENABLE_TRAINING_CORE
 class PrePackingTestOpKernel : public OpKernel {
 public:
  PrePackingTestOpKernel(const OpKernelInfo& info) : OpKernel(info) {}
--- a/onnxruntime/test/optimizer/graph_transform_test.cc
+++ b/onnxruntime/test/optimizer/graph_transform_test.cc
@ -84,7 +84,7 @@
 #include "test/util/include/inference_session_wrapper.h"
 #include "test/util/include/temp_dir.h"
 #include "test/util/include/test_utils.h"
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 #include "orttraining/core/optimizer/bitmask_dropout_replacement.h"
 #endif

@ -4270,7 +4270,7 @@ TEST_F(GraphTransformationTests, BiasDropoutFusionTest) {
  TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_residual_same_shape_fusion_dim_is_param.onnx", *logger_);
 }

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 static void TestBitmaskDropoutFusion(const PathString& file_path, bool is_bias_dropout, const logging::Logger& logger,
                                     const int add_count, const int dropout_count, const int bitmask_dropout_count,
                                     const int bias_dropout_count, const int bitmask_bias_dropout_count,
@ -4361,7 +4361,7 @@ TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest) {

  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
  ASSERT_TRUE(op_to_count["Cast"] == 0);
  ASSERT_TRUE(op_to_count["LayerNormalization"] == 1);
 #else
@ -4667,7 +4667,7 @@ TEST_F(GraphTransformationTests, SkipLayerNormFusion_Input_Output_Check) {

      // check outputs
      std::vector<NodeArg*>& output_defs = node.MutableOutputDefs();
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
      EXPECT_EQ(node.OutputDefs().size(), 3u) << "SkipLayerNormalization number of outputs does not equal to 3. Got:" << node.OutputDefs().size();
 #else
      EXPECT_EQ(node.OutputDefs().size(), 1u) << "SkipLayerNormalization number of outputs does not equal to 1. Got:" << node.OutputDefs().size();
@ -5515,7 +5515,7 @@ TEST_F(GraphTransformationTests, PropagateCastOpsTests) {
  }
 }

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
 TEST_F(GraphTransformationTests, PropagateCastOpsTests_Gelu) {
  using Strategy = GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy;
  {
--- a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc
+++ b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc
@ -1110,7 +1110,7 @@ TEST(NchwcOptimizerTests, BatchNormalization) {
  // should be skipped if the batch normalization node has the optional training
  // outputs supplied.
  test_case(false);
-#if defined(ENABLE_TRAINING)
+#if defined(ENABLE_TRAINING_CORE)
  test_case(true);
 #endif
 }
--- a/onnxruntime/test/providers/cpu/math/gemm_test.cc
+++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc
@ -750,7 +750,7 @@ TEST(GemmOpTest, GemmWithAlphaOpset11) {
  TestGemmWithAlphaOpset11<double>();
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(GemmOpTest, SharedPrepackedWeights) {
  OpTester test("Gemm");

--- a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc
+++ b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc
@ -368,7 +368,7 @@ TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_GEMM) {
  RunMatMulIntegerU8X8TestBatch(4, 8, 68);
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE // Prepacking is enabled only on non-training builds
 TEST(MatmulIntegerOpTest, SharedPrepackedWeights) {
  OpTester test("MatMulInteger", 10);
  test.AddInput<uint8_t>("T1", {1, 1}, {11});
--- a/onnxruntime/test/providers/cpu/math/matmul_test.cc
+++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc
@ -287,7 +287,7 @@ TEST(MathOpTest, MatMul_bfloat16) {
 }
 #endif

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(MathOpTest, MatMulSharedPrepackedWeights) {
  OpTester test("MatMul");

--- a/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc
+++ b/onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc
@ -492,7 +492,7 @@ struct PrePackTestOp {
  }
 };

-#ifndef ENABLE_TRAINING
+#ifndef ENABLE_TRAINING_CORE
 TEST(QuantizeLinearMatmulOpTest, QLinearMatMulPrePack) {
  auto registry = std::make_shared<CustomRegistry>();
  std::vector<ONNX_NAMESPACE::OpSchema> schemas{PrePackTestOp::OpSchema()};
--- a/onnxruntime/test/providers/cpu/model_tests.cc
+++ b/onnxruntime/test/providers/cpu/model_tests.cc
@ -195,7 +195,7 @@ TEST_P(ModelTest, Run) {
      {"shape_start_negative_1", "type error", {}},
      {"simple_rnn_batchwise", "type error", {}},
      {"mod_float_mixed_sign_example", "fmod attribute must be true for floating point types", {}},
-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_TRAINING_CORE
      {"adagrad", "not a registered function/op", {}},                  // Op not registered.
      {"adagrad_multiple", "not a registered function/op", {}},         // Op not registered.
      {"adam", "not a registered function/op", {}},                     // Op not registered.
@ -627,17 +627,17 @@ TEST_P(ModelTest, Run) {
  }

  // TODO(leca): move the parallel run test list to a config file and load it in GetParameterStrings() to make the load process run only once
-  std::set<std::string> tests_run_parallel = {"test_resnet18v2", 
+  std::set<std::string> tests_run_parallel = {"test_resnet18v2",
      "test_resnet34v2",
      "test_resnet50",
-      "test_resnet50v2", 
-      "test_resnet101v2", 
-      "test_resnet152v2", 
-      "keras_lotus_resnet3D", 
-      "coreml_Resnet50_ImageNet", 
-      "mlperf_mobilenet", 
-      "mlperf_resnet", 
-      "mlperf_ssd_mobilenet_300", 
+      "test_resnet50v2",
+      "test_resnet101v2",
+      "test_resnet152v2",
+      "keras_lotus_resnet3D",
+      "coreml_Resnet50_ImageNet",
+      "mlperf_mobilenet",
+      "mlperf_resnet",
+      "mlperf_ssd_mobilenet_300",
      "mlperf_ssd_resnet34_1200"};
  bool is_single_node = !model_info->GetNodeName().empty();
  std::vector<ExecutionMode> execution_modes = {ExecutionMode::ORT_SEQUENTIAL};
@ -752,7 +752,7 @@ TEST_P(ModelTest, Run) {
                                                  << provider_name << ", test name:" << results->GetName() << ", result: " << res;
        }
        continue;
-      } 
+      }
 #endif // !USE_DNNL
      std::unique_ptr<OrtSessionOptions, decltype(&OrtApis::ReleaseSessionOptions)> rel_ort_session_option(
        ortso, &OrtApis::ReleaseSessionOptions);
--- a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc
+++ b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc
@ -1129,7 +1129,7 @@ TEST(ConvTransposeTest, ConvTranspose_AutoPad_with_non_default_strides) {
                      OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); //Accuracy Mismatch on OpenVINO-EP
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(ConvTransposeTest, SharedPrepackedWeights) {
  OpTester test("ConvTranspose", 11);
  test.AddAttribute("kernel_shape", vector<int64_t>{3, 3});
--- a/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc
+++ b/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc
@ -1521,7 +1521,7 @@ TEST(QLinearConvTest, Conv2D_S8S8_Depthwise_Kernelsize_PerChannel) {
  TestQLinearConv2dDepthwiseKernelsizePerChannel<int8_t, int8_t>();
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(QLinearConvTest, SharedPrepackedWeights) {
  QuantizedTensor X({0.45246148109436035f, 0.15498268604278564f, 0.11199361085891724f, -0.39421093463897705f,
                     0.2626858949661255f, 0.13414543867111206f, -0.27184486389160156f, -0.43028733134269714f,
--- a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc
+++ b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc
@ -1347,7 +1347,7 @@ TEST(LSTMTest, ONNXRuntime_TestLSTMZeroSeqInMiddle) {
                  &sequence_length, use_bias, use_peepholes, 0.0f, false, false);
 }

-#ifndef ENABLE_TRAINING  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
 TEST(LSTMTest, SharedPrepackedWeights) {
  int64_t seq_length = 2;
  int batch_size = 2;
--- a/onnxruntime/test/providers/kernel_compute_test_utils.cc
+++ b/onnxruntime/test/providers/kernel_compute_test_utils.cc
@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_STRIDED_TENSORS

 #include "test/providers/kernel_compute_test_utils.h"

--- a/onnxruntime/test/providers/kernel_compute_test_utils.h
+++ b/onnxruntime/test/providers/kernel_compute_test_utils.h
@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.

-#ifdef ENABLE_TRAINING
+#ifdef ENABLE_STRIDED_TENSORS

 #pragma once

--- a/onnxruntime/test/providers/provider_test_utils.cc
+++ b/onnxruntime/test/providers/provider_test_utils.cc
@ -339,7 +339,7 @@ struct TensorCheck<MLFloat16> {
    const bool has_rel_err = params.relative_error_.has_value();

    float threshold = 0.001f;
-#if defined(USE_TENSORRT) || defined(ENABLE_TRAINING) || defined(USE_CUDA) || defined(USE_ROCM)
+#if defined(USE_TENSORRT) || defined(ENABLE_TRAINING_CORE) || defined(USE_CUDA) || defined(USE_ROCM)
    threshold = 0.005f;
 #elif defined(USE_DML)
    threshold = 0.008f;
@ -396,7 +396,7 @@ struct TensorCheck<BFloat16> {
    /// XXX: May need to adjust threshold as BFloat is coarse
    float abs_threshold = 0.0001f;
    float threshold = 0.001f;
-#if defined(USE_TENSORRT) || defined(ENABLE_TRAINING) || defined(USE_CUDA) || defined(USE_ROCM) || defined(USE_DML) || defined(USE_DNNL)
+#if defined(USE_TENSORRT) || defined(ENABLE_TRAINING_CORE) || defined(USE_CUDA) || defined(USE_ROCM) || defined(USE_DML) || defined(USE_DNNL)
    threshold = 0.05f;  // expect at least 95% close
 #endif

@ -912,7 +912,7 @@ std::vector<OrtValue> OpTester::ExecuteModel(
      size_t idx = 0;
      for (auto& expected_data : output_data_) {
        OrtValue& ort_value = fetches[idx];
-        
+
        if (expected_data.def_.Exists()) {           // optional edges won't exist (so skip them)
          if (!expected_data.data_.IsAllocated()) {  // optional type output (None)
            EXPECT_TRUE(!ort_value.IsAllocated())
@ -1463,6 +1463,7 @@ void OpTester::AddReferenceOutputs(const std::string& model_path, float abs_erro
 }

 #ifdef ENABLE_TRAINING
+// Deprecated code. Remove this when training::TrainingSession is removed.
 template std::vector<OrtValue> OpTester::ExecuteModel<training::TrainingSession>(
    Model& model, training::TrainingSession& session_object,
    ExpectResult expect_result, const std::string& expected_failure_string,
--- a/orttraining/orttraining/core/agent/training_agent.h
+++ b/orttraining/orttraining/core/agent/training_agent.h
@ -12,11 +12,9 @@
 #include "core/session/inference_session.h"

 namespace onnxruntime {
-#ifdef ENABLE_TRAINING
 struct PartialGraphExecutionState;
 typedef InlinedHashMap<std::string, OrtValue> OrtValueCache;
 typedef std::shared_ptr<OrtValueCache> OrtValueCachePtr;
-#endif

 namespace training {

--- a/orttraining/orttraining/test/training_ops/function_op_test_utils.cc
+++ b/orttraining/orttraining/test/training_ops/function_op_test_utils.cc
@ -52,7 +52,11 @@ TwoDArray OpFunctionTester::RunFunctionBodyGraphOnCPU() {
  RunOptions run_options;
  run_options.run_tag = op_;
  run_options.run_log_verbosity_level = 1;
+
+  #ifdef ENABLE_TRAINING
+  // Remove when training::TrainingSession class is removed.
  run_options.training_mode = true;
+  #endif

  std::vector<OrtValue> cpu_fetches;
  status = cpu_session_object.Run(run_options, feeds, output_names, &cpu_fetches);
--- a/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-on-device-training.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-on-device-training.yml
@ -1,48 +0,0 @@
-trigger: none
-
-jobs:
- job: Onnxruntime_Linux_GPU_TrainingAPIs
-
-  timeoutInMinutes: 120
-  pool: 'Onnxruntime-Linux-GPU-NC6sv3'
-
-  steps:
-  - checkout: self
-    clean: true
-    submodules: recursive
-
-  - template: templates/run-docker-build-steps.yml
-    parameters:
-      RunDockerBuildArgs: |
-        -o ubuntu20.04 -d gpu -e \
-        -t onnxruntime_training_apis_tests_image \
-        -x " \
-          --config RelWithDebInfo \
-          --enable_training \
-          --enable_training_apis \
-          --use_cuda --cuda_version=11.6 --cuda_home=/usr/local/cuda-11.6 --cudnn_home=/usr/local/cuda-11.6 \
-          --build_wheel \
-          --skip_tests \
-          " \
-        -u
-      DisplayName: 'Build onnxruntime'
-
-  # Entry point for all ort training api tests
-  - script: |
-      docker run \
-        --gpus all \
-        --shm-size=1024m \
-        --rm \
-        --volume $(Build.SourcesDirectory):/onnxruntime_src \
-        --volume $(Build.BinariesDirectory):/build \
-        onnxruntime_training_apis_tests_image \
-          /build/RelWithDebInfo/launch_test.py \
-            --cwd /build/RelWithDebInfo --cmd_line_with_args "python orttraining_test_ort_apis.py --cwd /build/RelWithDebInfo" \
-    displayName: 'Run ORT Training APIs Tests'
-    condition: succeededOrFailed()
-    timeoutInMinutes: 120
-  - template: templates/component-governance-component-detection-steps.yml
-    parameters:
-      condition: 'succeeded'
-
-  - template: templates/clean-agent-build-directory-step.yml
--- a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda116.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda116.yml
@ -11,7 +11,7 @@ resources:
 stages:
 - template: templates/py-packaging-training-cuda-stage.yml
  parameters:
-    build_py_parameters: --enable_training --update --build --enable_training_apis
+    build_py_parameters: --enable_training --update --build
    torch_version: '1.13.1'
    opset_version: '15'
    cuda_version: '11.6'
@ -24,7 +24,7 @@ stages:

 - template: templates/py-packaging-training-cuda-stage.yml
  parameters:
-    build_py_parameters: --enable_training --update --build --enable_training_apis
+    build_py_parameters: --enable_training --update --build
    torch_version: '1.13.1'
    opset_version: '15'
    cuda_version: '11.6'
@ -37,7 +37,7 @@ stages:

 - template: templates/py-packaging-stage.yml
  parameters:
-    build_py_parameters: --enable_training --enable_training_apis
+    build_py_parameters: --enable_training
    enable_linux_gpu: false
    enable_ubuntu_cpu: false
    enable_linux_cpu: false
--- a/tools/ci_build/github/azure-pipelines/win-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/win-ci-pipeline.yml
@ -177,7 +177,7 @@ stages:
        BuildConfig: 'RelWithDebInfo'
        EnvSetupScript: setup_env.bat
        buildArch: x64
-        additionalBuildFlags: --enable_training --enable_training_apis
+        additionalBuildFlags: --enable_training_apis
        msbuildPlatform: x64
        isX86: false
        job_name_suffix: ort_training_apis_x64_release