Module step

cmake/onnxruntime_training.cmake

@@ -17,6 +17,15 @@ file(GLOB_RECURSE onnxruntime_training_srcs
     "${ORTTRAINING_SOURCE_DIR}/core/agent/*.cc"
     )
 
+if (onnxruntime_ENABLE_TRAINING_ON_DEVICE)
+  file(GLOB_RECURSE onnxruntime_training_api_srcs CONFIGURE_DEPENDS
+    "${ORTTRAINING_SOURCE_DIR}/training_api/*.h"
+    "${ORTTRAINING_SOURCE_DIR}/training_api/*.cc"
+  )
+
+  list(APPEND onnxruntime_training_srcs ${onnxruntime_training_api_srcs})
+endif()
+
 # This needs to be built in framework.cmake
 file(GLOB_RECURSE onnxruntime_training_framework_excluded_srcs CONFIGURE_DEPENDS
     "${ORTTRAINING_SOURCE_DIR}/core/framework/torch/*.h"
@@ -245,11 +254,29 @@ if (onnxruntime_BUILD_UNIT_TESTS)
       endif()
     endif()
 
-    onnxruntime_add_include_to_target(onnxruntime_training_api_test_runner onnxruntime_common onnx onnx_proto ${PROTOBUF_LIB} onnxruntime_training flatbuffers)
-    target_include_directories(onnxruntime_training_api_test_runner PUBLIC ${CMAKE_CURRENT_BINARY_DIR} ${ONNXRUNTIME_ROOT} ${ORTTRAINING_ROOT} ${MPI_CXX_INCLUDE_DIRS} ${eigen_INCLUDE_DIRS} ${CXXOPTS} ${extra_includes} ${onnxruntime_graph_header} ${onnxruntime_exec_src_dir} ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR}/onnx)
+    onnxruntime_add_include_to_target(onnxruntime_training_api_test_runner onnxruntime_training 
+      onnxruntime_framework onnxruntime_common onnx onnx_proto ${PROTOBUF_LIB} flatbuffers)
+   
+    target_include_directories(onnxruntime_training_api_test_runner PUBLIC 
+      ${CMAKE_CURRENT_BINARY_DIR} 
+      ${ONNXRUNTIME_ROOT} 
+      ${ORTTRAINING_ROOT} 
+      ${MPI_CXX_INCLUDE_DIRS} 
+      ${eigen_INCLUDE_DIRS} 
+      ${CXXOPTS} 
+      ${extra_includes} 
+      ${onnxruntime_graph_header} 
+      ${onnxruntime_exec_src_dir} 
+      ${CMAKE_CURRENT_BINARY_DIR} 
+      ${CMAKE_CURRENT_BINARY_DIR}/onnx
+    )
 
-    target_link_libraries(onnxruntime_training_api_test_runner PRIVATE onnxruntime_training ${ONNXRUNTIME_LIBS} ${onnxruntime_EXTERNAL_LIBRARIES})
+    target_link_libraries(onnxruntime_training_api_test_runner PRIVATE 
+      onnxruntime_training
+      ${ONNXRUNTIME_LIBS}
+      ${onnxruntime_EXTERNAL_LIBRARIES}
+    )
     set_target_properties(onnxruntime_training_api_test_runner PROPERTIES FOLDER "ONNXRuntimeTest")
   endif()
 
-endif()
+endif()

orttraining/orttraining/test/training_api/test_runner.cc

@@ -12,6 +12,7 @@
 #include "core/providers/cpu/cpu_provider_factory_creator.h"
 #include "orttraining/core/framework/tensorboard/event_writer.h"
 #include "orttraining/training_api/interfaces.h"
+#include "orttraining/training_api/utils.h"
 
 using namespace onnxruntime;
 using namespace onnxruntime::common;

orttraining/orttraining/training_api/interfaces.cc

@@ -0,0 +1,78 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#if defined(ENABLE_TRAINING) && defined(ENABLE_TRAINING_ON_DEVICE)
+#include "orttraining/training_api/interfaces.h"
+#include "core/graph/model.h"
+
+namespace onnxruntime {
+namespace training {
+namespace api_test {
+
+static std::unique_ptr<Environment> env;
+
+void GetGraphInputOutputNames(const Graph& graph,
+                              std::vector<std::string> input_names,
+                              std::vector<std::string> output_names) {
+  auto inputs = graph.GetInputs();
+  auto outputs = graph.GetOutputs();
+
+  auto get_names = [&](const std::vector<const NodeArg*>& node_args, std::vector<std::string>& names) {
+    for (const auto* arg : node_args) {
+      names.push_back(arg->Name());
+    }
+  };
+
+  get_names(inputs, input_names);
+  get_names(outputs, output_names);
+}
+
+Module::Module(const std::string& train_model_path_or_bytes,
+               std::map<std::string, std::shared_ptr<Parameter>>& parameters,
+               const std::optional<std::string>& eval_model_path_or_bytes) {
+  parameters_ = std::move(parameters);
+
+  for (auto it = parameters_.begin(); it != parameters_.end(); it++) {
+    ORT_ENFORCE(it->first == it->second->name());
+    weights_.push_back(it->second->data());
+    gradients_.push_back(it->second->gradient());
+  }
+
+  auto so = onnxruntime::SessionOptions();
+  std::string default_logger_id{"Default"};
+
+  ORT_THROW_IF_ERROR(Environment::Create(std::make_unique<logging::LoggingManager>(std::make_unique<logging::CLogSink>(),
+                                                                                   logging::Severity::kWARNING,
+                                                                                   false,
+                                                                                   logging::LoggingManager::InstanceType::Default,
+                                                                                   &default_logger_id),
+                                         env));
+
+  train_sess_ = std::make_unique<onnxruntime::InferenceSession>(so, *env, train_model_path_or_bytes);
+  if (eval_model_path_or_bytes.has_value()) {
+    eval_sess_ = std::make_unique<onnxruntime::InferenceSession>(so, *env, eval_model_path_or_bytes.value());
+  }
+  
+  std::shared_ptr<onnxruntime::Model> model;
+  ORT_THROW_IF_ERROR(onnxruntime::Model::Load(train_model_path_or_bytes, model, nullptr, env->GetLoggingManager()->DefaultLogger()));
+  GetGraphInputOutputNames(model->MainGraph(), input_names_, output_names_);
+}
+
+Status Module::TrainStep(const std::vector<OrtValue>& inputs, std::vector<OrtValue>& outputs) {
+  ORT_NOT_IMPLEMENTED("Not implemented.");
+  std::vector<OrtValue> feeds{inputs};
+  feeds.insert(feeds.end(), weights_.begin(), weights_.end());
+
+  std::vector<OrtValue> fetches{outputs};
+  fetches.insert(fetches.end(), gradients_.begin(), gradients_.end());
+
+  auto status = train_sess_->Run(RunOptions(), input_names_, feeds, output_names_, &fetches);
+
+  return status;
+}
+
+}  // namespace api_test
+}  // namespace training
+}  // namespace onnxruntime
+
+#endif

orttraining/orttraining/training_api/interfaces.h

@@ -2,7 +2,12 @@
 // Licensed under the MIT License.
 
 #if defined(ENABLE_TRAINING) && defined(ENABLE_TRAINING_ON_DEVICE)
-
+#pragma once
+#include "core/common/logging/logging.h"
+#include "core/common/logging/sinks/clog_sink.h"
+#include "core/providers/cpu/cpu_execution_provider.h"
+#include "core/session/environment.h"
+#include "core/session/inference_session.h"
 namespace onnxruntime {
 namespace training {
 namespace api_test {
@@ -10,8 +15,8 @@ namespace api_test {
 class Parameter {
  public:
   // create parameter
-  Parameter(std::string /*name*/, const OrtValue& /*data*/) {
-    ORT_NOT_IMPLEMENTED("Not implemented.");
+  Parameter(std::string& name, const OrtValue& data) : name_(name), data_(data) {
+    ORT_ENFORCE(data.IsAllocated());
   }
 
   // Return the mutable data
@@ -33,8 +38,8 @@ class Parameter {
   }
   // need to set grad but not public api
  private:
-  OrtValue data_;
   std::string name_;
+  OrtValue data_;
 
   OrtValue gradient_;
   std::string gradient_name_;
@@ -48,26 +53,21 @@ class Module {
  public:
   // Initialize a module from an ORT inference session with loaded
   // training ONNX model and load parameters
-  Module(const std::string& /*train_model_path_or_bytes*/,
-         std::unordered_map<std::string, std::shared_ptr<Parameter>>& /*parameters*/,
-         const std::optional<std::string>& /*eval_model_path_or_bytes*/) {
-    ORT_NOT_IMPLEMENTED("Not implemented.");
-  }
+  Module(const std::string& train_model_path_or_bytes,
+         std::map<std::string, std::shared_ptr<Parameter>>& parameters,
+         const std::optional<std::string>& eval_model_path_or_bytes);
 
   // Return the trainable/nontrainable parameters
   std::vector<std::shared_ptr<Parameter>> parameters() const {
-    return parameters_;
-  }
-  std::unordered_map<std::string, std::shared_ptr<Parameter>> named_parameters() const {
     ORT_NOT_IMPLEMENTED("Not implemented.");
-    return {};
+    // return parameters_;
+  }
+  std::map<std::string, std::shared_ptr<Parameter>> named_parameters() const {
+    return parameters_;
   }
 
   // Train Step – does forward and backward computation. The outputs will be the forward’s outputs. Gradients will be accumulated within the Parameter object
-  Status TrainStep(const std::vector<OrtValue>& /*inputs*/, std::vector<OrtValue>& /*outputs*/) {
-    ORT_NOT_IMPLEMENTED("Not implemented.");
-    return Status::OK();
-  }
+  Status TrainStep(const std::vector<OrtValue>& /*inputs*/, std::vector<OrtValue>& /*outputs*/);
 
   // Eval Step – does forward computation. This will use a separate inference session
   // and take in a separate inference graph, while sharing the parameters
@@ -77,7 +77,7 @@ class Module {
   }
 
   // Return the states of the module as a map.
-  Status GetStateDict(const std::unordered_map<std::string, std::shared_ptr<Parameter>>& /*module_state_dict*/) {
+  Status GetStateDict(const std::map<std::string, std::shared_ptr<Parameter>>& /*module_state_dict*/) {
     ORT_NOT_IMPLEMENTED("Not implemented.");
     return Status::OK();
   }
@@ -85,7 +85,11 @@ class Module {
  private:
   std::unique_ptr<onnxruntime::InferenceSession> train_sess_;
   std::unique_ptr<onnxruntime::InferenceSession> eval_sess_;
-  std::vector<std::shared_ptr<Parameter>> parameters_;
+  std::map<std::string, std::shared_ptr<Parameter>> parameters_;
+  std::vector<std::string> input_names_;
+  std::vector<std::string> output_names_;
+  std::vector<OrtValue> weights_;
+  std::vector<OrtValue> gradients_;
 };
 
 // Internal state
@@ -112,7 +116,7 @@ class Optimizer {
   // training ONNX model For each parameter, initialize the OptimizerState based
   // on the graph input’s ValueInfoProto if the parameter doesn’t have it already.
   Optimizer(const std::string& /*optim_path_or_bytes*/,
-            std::unordered_map<std::string, std::shared_ptr<Parameter>>& /*parameters*/) {
+            std::map<std::string, std::shared_ptr<Parameter>>& /*parameters*/) {
     ORT_NOT_IMPLEMENTED("Not implemented.");
   }
 
@@ -187,43 +191,43 @@ class LinearScheduler : public LearningRateScheduler {
   int64_t total_iters_;
 };
 
-namespace utils {
+// namespace utils {
 
-struct CheckpointProperty {
-  int value;
-  // Support primitive types like int, float, string leveraging type trait.
-};
+// struct CheckpointProperty {
+//   int value;
+//   // Support primitive types like int, float, string leveraging type trait.
+// };
 
-struct CheckpointStates {
-  CheckpointStates() {
-    ORT_NOT_IMPLEMENTED("Not implemented.");
-  }
-  std::unordered_map<std::string, std::shared_ptr<Parameter>> named_parameters;
-  OptimizerState optimizer_states;
-  std::unordered_map<std::string, CheckpointProperty> named_properties;
-};
+// struct CheckpointStates {
+//   CheckpointStates() {
+//     ORT_NOT_IMPLEMENTED("Not implemented.");
+//   }
+//   std::map<std::string, std::shared_ptr<Parameter>> named_parameters;
+//   OptimizerState optimizer_states;
+//   std::unordered_map<std::string, CheckpointProperty> named_properties;
+// };
 
-// Save properties into a checkpoint property file (with postfix .prop).
-Status Ort_Save(CheckpointStates& /*state_dicts*/, const PathString& /*checkpoint_path*/) {
-  ORT_NOT_IMPLEMENTED("Not implemented.");
-  return Status::OK();
-}
+// // Save properties into a checkpoint property file (with postfix .prop).
+// Status Ort_Save(CheckpointStates& /*state_dicts*/, const PathString& /*checkpoint_path*/) {
+//   ORT_NOT_IMPLEMENTED("Not implemented.");
+//   return Status::OK();
+// }
 
-// Load properties file having postfix being '.prop'.
-Status Ort_Load(const PathString& /*checkpoint_path*/, CheckpointStates& /*state_dicts*/) {
-  ORT_NOT_IMPLEMENTED("Not implemented.");
-  return Status::OK();
-}
+// // Load properties file having postfix being '.prop'.
+// Status Ort_Load(const PathString& /*checkpoint_path*/, CheckpointStates& /*state_dicts*/) {
+//   ORT_NOT_IMPLEMENTED("Not implemented.");
+//   return Status::OK();
+// }
 
-/*
-  module.train_sess.RegisterExecutionProvider(provider);
-  module.eval_sess.RegisterExecutionProvider(provider);
-  optimizer.optim_sess.RegisterExecutionProvider(provider);
-*/
-void SetExecutionProvider(const Module& /*module*/, const Optimizer& /*optimizer*/, IExecutionProvider* /*provider*/) {
-  ORT_NOT_IMPLEMENTED("Not implemented.");
-}
-}  // namespace utils
+// /*
+//   module.train_sess.RegisterExecutionProvider(provider);
+//   module.eval_sess.RegisterExecutionProvider(provider);
+//   optimizer.optim_sess.RegisterExecutionProvider(provider);
+// */
+// void SetExecutionProvider(const Module& /*module*/, const Optimizer& /*optimizer*/, IExecutionProvider* /*provider*/) {
+//   ORT_NOT_IMPLEMENTED("Not implemented.");
+// }
+// }  // namespace utils
 
 }  // namespace api_test
 }  // namespace training

orttraining/orttraining/training_api/utils.cc

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#if defined(ENABLE_TRAINING) && defined(ENABLE_TRAINING_ON_DEVICE)
+
+#include "core/session/inference_session.h"
+#include "orttraining/training_api/utils.h"
+
+namespace onnxruntime {
+namespace training {
+namespace api_test {
+namespace utils {
+
+// Save properties into a checkpoint property file (with postfix .prop).
+Status Ort_Save(CheckpointStates& /*state_dicts*/, const PathString& /*checkpoint_path*/) {
+  ORT_NOT_IMPLEMENTED("Not implemented.");
+  return Status::OK();
+}
+
+// Load properties file having postfix being '.prop'.
+Status Ort_Load(const PathString& /*checkpoint_path*/, CheckpointStates& /*state_dicts*/) {
+  ORT_NOT_IMPLEMENTED("Not implemented.");
+  return Status::OK();
+}
+
+/*
+  module.train_sess.RegisterExecutionProvider(provider);
+  module.eval_sess.RegisterExecutionProvider(provider);
+  optimizer.optim_sess.RegisterExecutionProvider(provider);
+*/
+void SetExecutionProvider(const Module& /*module*/, const Optimizer& /*optimizer*/, IExecutionProvider* /*provider*/) {
+  ORT_NOT_IMPLEMENTED("Not implemented.");
+}
+
+}  // namespace utils
+}  // namespace api_test
+}  // namespace training
+}  // namespace onnxruntime
+
+#endif

orttraining/orttraining/training_api/utils.h

@@ -0,0 +1,46 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#if defined(ENABLE_TRAINING) && defined(ENABLE_TRAINING_ON_DEVICE)
+
+#pragma once
+#include "core/session/inference_session.h"
+#include "orttraining/training_api/interfaces.h"
+
+namespace onnxruntime {
+namespace training {
+namespace api_test {
+namespace utils {
+struct CheckpointProperty {
+  int value;
+  // Support primitive types like int, float, string leveraging type trait.
+};
+
+struct CheckpointStates {
+  CheckpointStates() {
+    ORT_NOT_IMPLEMENTED("Not implemented.");
+  }
+  std::map<std::string, std::shared_ptr<Parameter>> named_parameters;
+  OptimizerState optimizer_states;
+  std::unordered_map<std::string, CheckpointProperty> named_properties;
+};
+
+// Save properties into a checkpoint property file (with postfix .prop).
+Status Ort_Save(CheckpointStates& /*state_dicts*/, const PathString& /*checkpoint_path*/);
+
+// Load properties file having postfix being '.prop'.
+Status Ort_Load(const PathString& /*checkpoint_path*/, CheckpointStates& /*state_dicts*/);
+
+/*
+  module.train_sess.RegisterExecutionProvider(provider);
+  module.eval_sess.RegisterExecutionProvider(provider);
+  optimizer.optim_sess.RegisterExecutionProvider(provider);
+*/
+void SetExecutionProvider(const Module& /*module*/, const Optimizer& /*optimizer*/, IExecutionProvider* /*provider*/);
+
+}  // namespace utils
+}  // namespace api_test
+}  // namespace training
+}  // namespace onnxruntime
+
+#endif

tools/ci_build/build.py

@@ -176,6 +176,8 @@ def parse_arguments():
         "--enable_training_ops", action='store_true', help="Enable training ops in inference graph.")
     parser.add_argument(
         "--enable_training_torch_interop", action='store_true', help="Enable training kernels interop with torch.")
+    parser.add_argument(
+        "--enable_training_on_device", action='store_true', help="Enable training on device API.")
     parser.add_argument(
         "--disable_nccl", action='store_true', help="Disable Nccl.")
     parser.add_argument(
@@ -836,6 +838,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
         "-Donnxruntime_ENABLE_TRAINING=" + ("ON" if args.enable_training else "OFF"),
         "-Donnxruntime_ENABLE_TRAINING_OPS=" + ("ON" if args.enable_training_ops else "OFF"),
         "-Donnxruntime_ENABLE_TRAINING_TORCH_INTEROP=" + ("ON" if args.enable_training_torch_interop else "OFF"),
+        "-Donnxruntime_ENABLE_TRAINING_ON_DEVICE=" + ("ON" if args.enable_training_on_device else "OFF"),
         # Enable advanced computations such as AVX for some traininig related ops.
         "-Donnxruntime_ENABLE_CPU_FP16_OPS=" + ("ON" if args.enable_training else "OFF"),
         "-Donnxruntime_USE_NCCL=" + ("OFF" if args.disable_nccl else "ON"),