Enable a single build with optimized inference and on device training (#14241)

### Description
Right now prepacking code is not compiled when training is enabled. Our
partners want a single build of ort which can do both optimized
inference + training on device. This PR enables prepacking code in a
training build and controls whether it is enabled or not using already
existing session option - kOrtSessionOptionsConfigDisablePrepacking

For Inference scenarios - prepacking will be turned on by default and
this behavior remains the same after this PR too.
For training scenarios - prepacking will be disabled by default and if
user explicitly enables it then an error will be thrown.



### Motivation and Context
Enable both optimized inference as well as on device training in a
single build. For on device training use flag --enable_training_apis.

onnxruntime/core/framework/sequential_executor.cc

@@ -441,12 +441,15 @@ onnxruntime::Status ExecuteKernel(StreamExecutionContext& ctx,
   } else {
     KernelScope kernel_scope(session_scope, kernel_ctx, *p_kernel);
     ORT_TRY {
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
+      // AllocateInputsContiguously - is only required for NCCL kernels
+      // can be moved under USE_NCCL
       if (p_kernel->KernelDef().AllocateInputsContiguously()) {
         ORT_RETURN_IF_ERROR(utils::VerifyInputTensorsAllocatedContiguously(&kernel_ctx));
       }
-#endif
-#ifdef ENABLE_TRAINING
+
+      // This is most probably deprecated code and is causing unnecessary complexity.
+      // Can be removed.
       // 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.

onnxruntime/core/framework/session_state.cc

@@ -294,7 +294,7 @@ bool SessionState::IsSparseInitializer(int ort_value_index) const {
 }
 #endif
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 Status SessionState::GetInitializedTensors(
     const std::unordered_set<std::string>& interested_weights,
     bool allow_missing_weights, NameMLValMap& retrieved_weights) const {
@@ -1387,9 +1387,13 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
   SetMemoryProfiler(mem_profiler.release());
 #endif
 
+  // Note: For Training Prepacking should be always disabled.
+  // For inference it is enabled by default, but users can choose to disable it via session options.
+  const bool disable_prepacking =
+      session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigDisablePrepacking, "0") == "1";
   // Memory pattern tracer allocates all initializers on a single contiguous
   // buffer. This has the effect of reducing memory fragmentation.
-  // Further more, NCCL kernels require initializers to be allocated
+  // Further more, in training scenarios NCCL kernels require initializers to be allocated
   // contiguously.
   //
   // In inferencing scenarios, however, we often want to pre-process and then
@@ -1401,13 +1405,12 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
   //  out of memory error in some training tests. Need to create kernel first,
   //  and let the kernel tells us whether the initializer needs to be traced.
   //
-#if defined(ENABLE_TRAINING_CORE)
-  std::unique_ptr<ITensorAllocator> tensor_allocator(
-      ITensorAllocator::Create(enable_mem_pattern_, *p_seq_exec_plan_, *this, weights_buffers_));
-#else
-  std::unique_ptr<ITensorAllocator> tensor_allocator(
-      ITensorAllocator::Create(false, *p_seq_exec_plan_, *this, weights_buffers_));
-#endif
+  std::unique_ptr<ITensorAllocator> tensor_allocator = nullptr;
+  if (disable_prepacking) {
+    tensor_allocator = ITensorAllocator::Create(enable_mem_pattern_, *p_seq_exec_plan_, *this, weights_buffers_);
+  } else {
+    tensor_allocator = ITensorAllocator::Create(false, *p_seq_exec_plan_, *this, weights_buffers_);
+  }
 
   const auto& initializer_allocation_order = p_seq_exec_plan_->initializer_allocation_order;
 
@@ -1470,15 +1473,10 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
 
   ORT_RETURN_IF_ERROR(CreateKernels(kernel_registry_manager));
 
-#ifndef ENABLE_TRAINING_CORE
-  const auto disable_prepacking =
-      session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigDisablePrepacking, "0");
-
-  if (disable_prepacking != "1") {
+  if (!disable_prepacking) {
     ORT_RETURN_IF_ERROR(PrepackConstantInitializedTensors(constant_initializers_use_count,
                                                           session_options.initializers_to_share_map));
   }
-#endif
 
   ORT_RETURN_IF_ERROR(
       session_state_utils::SaveInputOutputNamesToNodeMapping(*graph_viewer_, *this, valid_outer_scope_node_args));

onnxruntime/core/framework/session_state.h

@@ -161,7 +161,8 @@ class SessionState {
   bool IsSparseInitializer(int ort_value_index) const;
 #endif
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
+// This is referenced in training::TrainingSession. Should be removed when this class is removed.
   /**
     Get some initialized tensors (weights).
     @param interested_weights The names of the weights to retrieve.

onnxruntime/core/framework/utils.cc

@@ -635,9 +635,6 @@ ExecuteGraphImpl(const SessionState& session_state,
   //    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_CORE
-  single_thread_mode = true;
-#endif
 
   // see if we can skip copies due to the types of execution providers available
   if (device_copy_checks.status == DeviceCopyCheck::NoCopy) {
@@ -1004,7 +1001,8 @@ int32_t ONNXTensorElementDataTypeToProtoTensorType(ONNXTensorElementDataType onn
   }
 }
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
+// Needed only when NCCL kernels are enabled.
 common::Status VerifyInputTensorsAllocatedContiguously(OpKernelContext* context) {
   const Tensor* prev_input = context->Input<Tensor>(0);
   for (int i = 1; i < context->InputCount(); i++) {

onnxruntime/core/session/environment.cc

@@ -34,14 +34,12 @@
 #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/optimizer_builder.h"
 #include "orttraining/core/graph/optimizer_graph_builder_registry.h"
 #include "orttraining/core/graph/loss_function_registry.h"
+#include "orttraining/core/optimizer/graph_transformer_registry.h"
 #endif
 
 namespace onnxruntime {
@@ -265,16 +263,15 @@ 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>
+      // This was added for a partner team and is most probably no longer in use.
+      // Can be removed once we have the confirmation.
+      training::GraphTransformerRegistry::GetInstance().RegisterExternalGraphTransformers();
 #endif
     });
 

onnxruntime/test/contrib_ops/attention_op_test.cc

@@ -3654,7 +3654,8 @@ TEST(AttentionTest, DISABLED_Attention_Mask1D_Fp16_B2_FusedNoPadding) {
   }
 }
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(AttentionTest, SharedPrepackedWeights) {
   int batch_size = 2;
   int sequence_length = 2;

onnxruntime/test/contrib_ops/quantize_attention_op_test.cc

@@ -980,7 +980,8 @@ TEST(QAttentionTest, QAttentionPrunedModel) {
                    input_hidden_size);
 }
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(QAttentionTest, SharedPrepackedWeights) {
   int batch_size = 1;
   int sequence_length = 2;

onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc

@@ -362,7 +362,8 @@ TEST(DynamicQuantLSTMTest, LargeSize) {
   RunQuantLSTM<uint8_t>(12, 3, 278);
 }
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) {
   OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/);
 

onnxruntime/test/framework/inference_session_test.cc

@@ -2091,6 +2091,11 @@ TEST(InferenceSessionTests, TestArenaShrinkageAfterRun) {
   arena_cfg.arena_extend_strategy = 1;  // kSameAsRequested
 
   SessionOptions so;
+#ifdef ENABLE_TRAINING
+  // Disable weight prepacking
+  // Without this assert for alloc_stats.num_arena_extensions will fail.
+  so.config_options.configurations["session.disable_prepacking"] = "1";
+#endif
   InferenceSession session_object{so, GetEnvironment()};
   OrtCUDAProviderOptions provider_options{};
   provider_options.default_memory_arena_cfg = &arena_cfg;

onnxruntime/test/providers/cpu/math/gemm_test.cc

@@ -750,7 +750,8 @@ TEST(GemmOpTest, GemmWithAlphaOpset11) {
   TestGemmWithAlphaOpset11<double>();
 }
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in training builds so no need to test the feature in a training build.
 TEST(GemmOpTest, SharedPrepackedWeights) {
   OpTester test("Gemm");
 

onnxruntime/test/providers/cpu/math/matmul_integer_test.cc

@@ -368,7 +368,8 @@ TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_GEMM) {
   RunMatMulIntegerU8X8TestBatch(4, 8, 68);
 }
 
-#ifndef ENABLE_TRAINING_CORE // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(MatmulIntegerOpTest, SharedPrepackedWeights) {
   OpTester test("MatMulInteger", 10);
   test.AddInput<uint8_t>("T1", {1, 1}, {11});

onnxruntime/test/providers/cpu/math/matmul_test.cc

@@ -287,7 +287,8 @@ TEST(MathOpTest, MatMul_bfloat16) {
 }
 #endif
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(MathOpTest, MatMulSharedPrepackedWeights) {
   OpTester test("MatMul");
 

onnxruntime/test/providers/cpu/math/quantize_linear_matmul_test.cc

@@ -492,7 +492,8 @@ struct PrePackTestOp {
   }
 };
 
-#ifndef ENABLE_TRAINING_CORE
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(QuantizeLinearMatmulOpTest, QLinearMatMulPrePack) {
   auto registry = std::make_shared<CustomRegistry>();
   std::vector<ONNX_NAMESPACE::OpSchema> schemas{PrePackTestOp::OpSchema()};

onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc

@@ -1129,7 +1129,8 @@ TEST(ConvTransposeTest, ConvTranspose_AutoPad_with_non_default_strides) {
                       OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); //Accuracy Mismatch on OpenVINO-EP
 }
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(ConvTransposeTest, SharedPrepackedWeights) {
   OpTester test("ConvTranspose", 11);
   test.AddAttribute("kernel_shape", vector<int64_t>{3, 3});

onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc

@@ -1521,7 +1521,8 @@ TEST(QLinearConvTest, Conv2D_S8S8_Depthwise_Kernelsize_PerChannel) {
   TestQLinearConv2dDepthwiseKernelsizePerChannel<int8_t, int8_t>();
 }
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(QLinearConvTest, SharedPrepackedWeights) {
   QuantizedTensor X({0.45246148109436035f, 0.15498268604278564f, 0.11199361085891724f, -0.39421093463897705f,
                      0.2626858949661255f, 0.13414543867111206f, -0.27184486389160156f, -0.43028733134269714f,

onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc

@@ -1347,7 +1347,8 @@ TEST(LSTMTest, ONNXRuntime_TestLSTMZeroSeqInMiddle) {
                   &sequence_length, use_bias, use_peepholes, 0.0f, false, false);
 }
 
-#ifndef ENABLE_TRAINING_CORE  // Prepacking is enabled only on non-training builds
+#ifndef ENABLE_TRAINING
+// Prepacking is disabled in full training build so no need to test the feature in a training build.
 TEST(LSTMTest, SharedPrepackedWeights) {
   int64_t seq_length = 2;
   int batch_size = 2;

orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py

@@ -307,6 +307,8 @@ class GraphExecutionManager(GraphExecutionInterface):
         probe_level = ortmodule._defined_from_envvar("ORTMODULE_MEMORY_OPT_PROBE_RECOMPUTE_LEVEL", "1", warn=True)
         session_options.add_session_config_entry("optimization.enable_memory_optimizer", alleviation_config)
         session_options.add_session_config_entry("optimization.enable_memory_probe_recompute_level", probe_level)
+        # Disable weight prepacking
+        session_options.add_session_config_entry("session.disable_prepacking", "1")
 
         if self._debug_options.save_onnx_models.save:
             session_options.optimized_model_filepath = os.path.join(

orttraining/orttraining/training_api/module.cc

@@ -6,6 +6,7 @@
 #include "core/framework/execution_provider.h"
 #include "core/session/inference_session.h"
 #include "core/session/environment.h"
+#include "core/session/onnxruntime_session_options_config_keys.h"
 
 #include "orttraining/training_api/include/module.h"
 #include "orttraining/training_api/include/utils.h"
@@ -148,6 +149,17 @@ Module::Module(const std::string& train_model_path_or_bytes,
                const std::vector<std::shared_ptr<IExecutionProvider>>& providers,
                const std::optional<std::string>& eval_model_path_or_bytes)
     : named_parameters_{named_parameters} {
+
+  // Enforce weight prepacking is disabled
+  // If user explicitly enabled weight prepacking then return error.
+  // Default value is enabled. Therefore, explicitly disable it if the value is not set by user.
+  std::string disable_prepacking = "";
+  if (session_options.config_options.TryGetConfigEntry(kOrtSessionOptionsConfigDisablePrepacking, disable_prepacking)) {
+    ORT_ENFORCE(disable_prepacking == "1", "Prepacking is not supported for training scenarios.");
+  } else {
+    const_cast<SessionOptions&>(session_options).config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "1";
+  }
+
   train_sess_ = std::make_unique<onnxruntime::InferenceSession>(session_options, env);
   ORT_THROW_IF_ERROR(train_sess_->Load(train_model_path_or_bytes));
   for (const auto& provider : providers) {