From 2cdc2be57e3db63cc03ba61cd490f003297981e5 Mon Sep 17 00:00:00 2001 From: Aishwarya Bhandare Date: Mon, 25 Apr 2022 20:37:32 +0000 Subject: [PATCH] Module step --- cmake/onnxruntime_training.cmake | 35 +++++- .../test/training_api/test_runner.cc | 1 + .../orttraining/training_api/interfaces.cc | 78 +++++++++++++ .../orttraining/training_api/interfaces.h | 108 +++++++++--------- orttraining/orttraining/training_api/utils.cc | 40 +++++++ orttraining/orttraining/training_api/utils.h | 46 ++++++++ tools/ci_build/build.py | 3 + 7 files changed, 255 insertions(+), 56 deletions(-) create mode 100644 orttraining/orttraining/training_api/interfaces.cc create mode 100644 orttraining/orttraining/training_api/utils.cc create mode 100644 orttraining/orttraining/training_api/utils.h diff --git a/cmake/onnxruntime_training.cmake b/cmake/onnxruntime_training.cmake index 097484a41f..2bd49d01c8 100644 --- a/cmake/onnxruntime_training.cmake +++ b/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() \ No newline at end of file diff --git a/orttraining/orttraining/test/training_api/test_runner.cc b/orttraining/orttraining/test/training_api/test_runner.cc index 868efb75f8..6bea413cd0 100644 --- a/orttraining/orttraining/test/training_api/test_runner.cc +++ b/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; diff --git a/orttraining/orttraining/training_api/interfaces.cc b/orttraining/orttraining/training_api/interfaces.cc new file mode 100644 index 0000000000..816ac7f132 --- /dev/null +++ b/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 env; + +void GetGraphInputOutputNames(const Graph& graph, + std::vector input_names, + std::vector output_names) { + auto inputs = graph.GetInputs(); + auto outputs = graph.GetOutputs(); + + auto get_names = [&](const std::vector& node_args, std::vector& 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>& parameters, + const std::optional& 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(std::make_unique(), + logging::Severity::kWARNING, + false, + logging::LoggingManager::InstanceType::Default, + &default_logger_id), + env)); + + train_sess_ = std::make_unique(so, *env, train_model_path_or_bytes); + if (eval_model_path_or_bytes.has_value()) { + eval_sess_ = std::make_unique(so, *env, eval_model_path_or_bytes.value()); + } + + std::shared_ptr 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& inputs, std::vector& outputs) { + ORT_NOT_IMPLEMENTED("Not implemented."); + std::vector feeds{inputs}; + feeds.insert(feeds.end(), weights_.begin(), weights_.end()); + + std::vector 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 diff --git a/orttraining/orttraining/training_api/interfaces.h b/orttraining/orttraining/training_api/interfaces.h index 1ae67accff..bfe004aadd 100644 --- a/orttraining/orttraining/training_api/interfaces.h +++ b/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>& /*parameters*/, - const std::optional& /*eval_model_path_or_bytes*/) { - ORT_NOT_IMPLEMENTED("Not implemented."); - } + Module(const std::string& train_model_path_or_bytes, + std::map>& parameters, + const std::optional& eval_model_path_or_bytes); // Return the trainable/nontrainable parameters std::vector> parameters() const { - return parameters_; - } - std::unordered_map> named_parameters() const { ORT_NOT_IMPLEMENTED("Not implemented."); - return {}; + // return parameters_; + } + std::map> 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& /*inputs*/, std::vector& /*outputs*/) { - ORT_NOT_IMPLEMENTED("Not implemented."); - return Status::OK(); - } + Status TrainStep(const std::vector& /*inputs*/, std::vector& /*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>& /*module_state_dict*/) { + Status GetStateDict(const std::map>& /*module_state_dict*/) { ORT_NOT_IMPLEMENTED("Not implemented."); return Status::OK(); } @@ -85,7 +85,11 @@ class Module { private: std::unique_ptr train_sess_; std::unique_ptr eval_sess_; - std::vector> parameters_; + std::map> parameters_; + std::vector input_names_; + std::vector output_names_; + std::vector weights_; + std::vector 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>& /*parameters*/) { + std::map>& /*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> named_parameters; - OptimizerState optimizer_states; - std::unordered_map named_properties; -}; +// struct CheckpointStates { +// CheckpointStates() { +// ORT_NOT_IMPLEMENTED("Not implemented."); +// } +// std::map> named_parameters; +// OptimizerState optimizer_states; +// std::unordered_map 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 diff --git a/orttraining/orttraining/training_api/utils.cc b/orttraining/orttraining/training_api/utils.cc new file mode 100644 index 0000000000..f3fefa7c08 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/orttraining/orttraining/training_api/utils.h b/orttraining/orttraining/training_api/utils.h new file mode 100644 index 0000000000..9cc408a318 --- /dev/null +++ b/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> named_parameters; + OptimizerState optimizer_states; + std::unordered_map 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 \ No newline at end of file diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 432f2d0917..dcf8913925 100644 --- a/tools/ci_build/build.py +++ b/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"),