From f63e28c92fcb5b610a2eb11cc2a32f27e5b19e66 Mon Sep 17 00:00:00 2001 From: Ashwini Khade Date: Wed, 15 Jun 2022 11:13:35 -0700 Subject: [PATCH] C API version 0.001 (#11758) * C API version 0.001 * fix linker issues * fixes for save checkpoint api * plus fixes based on tests * plus test_runner and other changes * Plus cosmetic updates * remove unnecessary headers * plus some updates * plus more changes Co-authored-by: Ashwini Khade --- cmake/onnxruntime_session.cmake | 11 +- cmake/onnxruntime_training.cmake | 8 - .../core/session/onnxruntime_c_api.h | 28 ++- onnxruntime/core/session/onnxruntime_c_api.cc | 16 ++ onnxruntime/core/session/ort_apis.h | 33 +++ .../common/synthetic_data_loader.cc | 48 +++- .../common/synthetic_data_loader.h | 6 +- .../test/training_api/core/checkpoint_test.cc | 11 +- .../test/training_api/trainer/trainer.cc | 138 +++++++---- .../training_api/include/checkpoint.h | 21 +- .../orttraining/training_api/include/module.h | 10 +- .../include/onnxruntime_training_c_api.h | 41 ++++ .../training_api/include/optimizer.h | 2 +- .../training_api/include/training_session.h | 55 +++++ .../orttraining/training_api/include/utils.h | 2 +- .../orttraining/training_api/module.cc | 84 ++++--- .../onnxruntime_training_c_api.cc | 218 ++++++++++++++++++ .../orttraining/training_api/optimizer.cc | 4 +- .../training_api/training_session.cc | 69 ++++++ 19 files changed, 667 insertions(+), 138 deletions(-) create mode 100644 orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h create mode 100644 orttraining/orttraining/training_api/include/training_session.h create mode 100644 orttraining/orttraining/training_api/onnxruntime_training_c_api.cc create mode 100644 orttraining/orttraining/training_api/training_session.cc diff --git a/cmake/onnxruntime_session.cmake b/cmake/onnxruntime_session.cmake index 7656c4e2cd..83129616c8 100644 --- a/cmake/onnxruntime_session.cmake +++ b/cmake/onnxruntime_session.cmake @@ -7,6 +7,15 @@ file(GLOB onnxruntime_session_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/core/session/*.cc" ) +if (onnxruntime_ENABLE_TRAINING_ON_DEVICE) + file(GLOB_RECURSE on_device_training_api_srcs CONFIGURE_DEPENDS + "${ORTTRAINING_SOURCE_DIR}/training_api/*.cc" + ) + + list(APPEND onnxruntime_session_srcs ${on_device_training_api_srcs}) +endif() + + if (onnxruntime_MINIMAL_BUILD) set(onnxruntime_session_src_exclude "${ONNXRUNTIME_ROOT}/core/session/provider_bridge_ort.cc" @@ -48,7 +57,7 @@ if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS) endif() if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) - onnxruntime_add_include_to_target(onnxruntime_session Python::Module) + onnxruntime_add_include_to_target(onnxruntime_session Python::Module) endif() if (NOT onnxruntime_BUILD_SHARED_LIB) diff --git a/cmake/onnxruntime_training.cmake b/cmake/onnxruntime_training.cmake index f028475728..3b93ff85c9 100644 --- a/cmake/onnxruntime_training.cmake +++ b/cmake/onnxruntime_training.cmake @@ -17,14 +17,6 @@ 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 diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 9c0080a467..fd3b09067e 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -270,6 +270,11 @@ ORT_RUNTIME_CLASS(CUDAProviderOptionsV2); ORT_RUNTIME_CLASS(Op); ORT_RUNTIME_CLASS(OpAttr); +#ifdef ENABLE_TRAINING_ON_DEVICE +ORT_RUNTIME_CLASS(TrainingSession); +ORT_RUNTIME_CLASS(CheckpointState); +#endif + #ifdef _WIN32 typedef _Return_type_success_(return == 0) OrtStatus* OrtStatusPtr; #else @@ -3343,13 +3348,13 @@ struct OrtApi { _In_reads_(input_len) const OrtValue* const* initializers, size_t initializers_num); /** \brief: Create attribute of onnxruntime operator - * + * * \param[in] name of the attribute * \param[in] data of the attribute * \param[in] data length * \param[in] data type * \param[out] attribute that has been created, which must be released by OrtApi::ReleaseOpAttr - * + * * \since Version 1.12. */ ORT_API2_STATUS(CreateOpAttr, @@ -3362,14 +3367,14 @@ struct OrtApi { /* \brief: Release op attribute * * \param[in] attribute created by OrtApi::CreateOpAttr - * + * * \since Version 1.12. */ ORT_CLASS_RELEASE(OpAttr); /** \brief: Create onnxruntime native operator - * - * \param[in] kernel info + * + * \param[in] kernel info * \param[in] operator name * \param[in] operator domain * \param[in] operator opset @@ -3379,7 +3384,7 @@ struct OrtApi { * \param[in] attributes used to initialize the operator * \param[in] number of the attributes * \param[out] operator that has been created - * + * * \since Version 1.12. */ ORT_API2_STATUS(CreateOp, @@ -3396,14 +3401,14 @@ struct OrtApi { /** \brief: Invoke the operator created by OrtApi::CreateOp * The inputs must follow the order as specified in onnx specification - * + * * \param[in] kernel context * \param[in] operator that has been created * \param[in] inputs * \param[in] number of inputs * \param[in] outputs * \param[in] number of outputs - * + * * \since Version 1.12. */ ORT_API2_STATUS(InvokeOp, @@ -3417,10 +3422,15 @@ struct OrtApi { /* \brief: Release an onnxruntime operator * * \param[in] operator created by OrtApi::CreateOp - * + * * \since Version 1.12. */ ORT_CLASS_RELEASE(Op); + +#ifdef ENABLE_TRAINING_ON_DEVICE + // defines c apis for on device training scenarios + #include "../../../orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h" +#endif }; /* diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index cf114077a9..a411991769 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -2527,6 +2527,22 @@ static constexpr OrtApi ort_api_1_to_12 = { &OrtApis::CreateOp, &OrtApis::InvokeOp, &OrtApis::ReleaseOp, +#ifdef ENABLE_TRAINING_ON_DEVICE + // Experimental for on-device training. Always keep at the bottom. + &OrtApis::LoadCheckpoint, + &OrtApis::SaveCheckpoint, + &OrtApis::CreateTrainingSession, + &OrtApis::InitializeTrainingSession, + &OrtApis::TrainingSessionGetTrainModeOutputCount, + &OrtApis::TrainingSessionGetEvalModeOutputCount, + &OrtApis::ResetGrad, + &OrtApis::TrainStep, + &OrtApis::EvalStep, + &OrtApis::OptimizerStep, + &OrtApis::ReleaseTrainingSession, + &OrtApis::ReleaseCheckpointState, +#endif + }; // Asserts to do a some checks to ensure older Versions of the OrtApi never change (will detect an addition or deletion but not if they cancel out each other) diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h index a56b85ee5b..f6d14e7b44 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/onnxruntime/core/session/ort_apis.h @@ -375,4 +375,37 @@ ORT_API_STATUS_IMPL(InvokeOp, ORT_API(void, ReleaseOp, _Frees_ptr_opt_ OrtOp* op); +#ifdef ENABLE_TRAINING_ON_DEVICE +ORT_API_STATUS_IMPL(CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options, + _Inout_ OrtCheckpointState* checkpoint_state, _Outptr_ OrtTrainingSession** out); + +ORT_API_STATUS_IMPL(InitializeTrainingSession, _Inout_ OrtTrainingSession* session, + _In_ const ORTCHAR_T* train_model_path, _In_ const ORTCHAR_T* eval_model_path, + _In_ const ORTCHAR_T* optimizer_model_path); + +ORT_API(void, ReleaseTrainingSession, _Frees_ptr_opt_ OrtTrainingSession* session); + +ORT_API_STATUS_IMPL(TrainingSessionGetTrainModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out); + +ORT_API_STATUS_IMPL(TrainingSessionGetEvalModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out); + +ORT_API_STATUS_IMPL(ResetGrad, _Inout_ OrtTrainingSession* session); + +ORT_API_STATUS_IMPL(TrainStep, _Inout_ OrtTrainingSession* session, _In_opt_ const OrtRunOptions* run_options, + size_t inputs_len, _In_reads_(input_len) const OrtValue* const* inputs, + size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + +ORT_API_STATUS_IMPL(EvalStep, _Inout_ OrtTrainingSession* session, _In_opt_ const OrtRunOptions* run_options, + size_t inputs_len, _In_reads_(input_len) const OrtValue* const* inputs, + size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + +ORT_API_STATUS_IMPL(OptimizerStep, _Inout_ OrtTrainingSession* session, _In_opt_ const OrtRunOptions* run_options); + +ORT_API_STATUS_IMPL(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Outptr_ OrtCheckpointState** checkpoint_state); + +ORT_API_STATUS_IMPL(SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Inout_ OrtTrainingSession* session, + bool save_optimizer_state); + +ORT_API(void, ReleaseCheckpointState, _Frees_ptr_opt_ OrtCheckpointState* session); +#endif } // namespace OrtApis diff --git a/orttraining/orttraining/test/training_api/common/synthetic_data_loader.cc b/orttraining/orttraining/test/training_api/common/synthetic_data_loader.cc index 510c155f97..2d71b26478 100644 --- a/orttraining/orttraining/test/training_api/common/synthetic_data_loader.cc +++ b/orttraining/orttraining/test/training_api/common/synthetic_data_loader.cc @@ -53,7 +53,20 @@ void SyntheticSampleBatch::AddFloatInput(const std::vector& shape) { RandomFloats(data_vector_.back()->GetData()); } -bool SyntheticDataLoader::GetNextSampleBatch(std::vector& batches) { +#define ORT_RETURN_ON_ERROR(expr) \ + do { \ + OrtStatus* onnx_status = (expr); \ + if (onnx_status != NULL) { \ + auto code = ort_api->GetErrorCode(onnx_status); \ + const char* msg = ort_api->GetErrorMessage(onnx_status); \ + ort_api->ReleaseStatus(onnx_status); \ + printf("Run failed with error code :%d\n", code); \ + printf("Error message :%s\n", msg); \ + return false; \ + } \ + } while (0); + +bool SyntheticDataLoader::GetNextSampleBatch(std::vector& batches) { if (sample_batch_iter_index_ >= num_of_sample_batches) { return false; } @@ -62,31 +75,44 @@ bool SyntheticDataLoader::GetNextSampleBatch(std::vector& batches) { auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); auto& sample = sample_batch_collections_[sample_batch_iter_index_]; + const auto* ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION); for (size_t i = 0; i < sample->NumOfInput(); ++i) { auto input_ptr = sample->GetInputAtIndex(i); auto shape_vector = input_ptr->ShapeVector(); - // Be noted: the created Ort::Value won't clean the raw data after its lifetime ended. + // Be noted: the created OrtValue won't clean the raw data after its lifetime ended. auto ptr_flt = dynamic_cast*>(input_ptr); if (ptr_flt) { - batches.push_back(Ort::Value::CreateTensor( - memory_info, input_ptr->GetData().data(), - input_ptr->NumOfBytesPerSample(), shape_vector.data(), shape_vector.size())); + OrtValue* value = NULL; + ORT_RETURN_ON_ERROR(ort_api->CreateTensorWithDataAsOrtValue(memory_info, + input_ptr->GetData().data(), (input_ptr->NumOfBytesPerSample() * sizeof(float)), + shape_vector.data(), shape_vector.size(), + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, + &value)); + batches.emplace_back(value); continue; } auto ptr_int = dynamic_cast*>(input_ptr); if (ptr_int) { - batches.push_back(Ort::Value::CreateTensor( - memory_info, input_ptr->GetData().data(), - input_ptr->NumOfBytesPerSample(), shape_vector.data(), shape_vector.size())); + OrtValue* value = NULL; + ORT_RETURN_ON_ERROR(ort_api->CreateTensorWithDataAsOrtValue(memory_info, + input_ptr->GetData().data(), (input_ptr->NumOfBytesPerSample() * sizeof(int64_t)), + shape_vector.data(), shape_vector.size(), + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, + &value)); + batches.emplace_back(value); continue; } auto ptr_int32 = dynamic_cast*>(input_ptr); if (ptr_int32) { - batches.push_back(Ort::Value::CreateTensor( - memory_info, input_ptr->GetData().data(), - input_ptr->NumOfBytesPerSample(), shape_vector.data(), shape_vector.size())); + OrtValue* value = nullptr; + ORT_RETURN_ON_ERROR(ort_api->CreateTensorWithDataAsOrtValue(memory_info, + input_ptr->GetData().data(), (input_ptr->NumOfBytesPerSample() * sizeof(int32_t)), + shape_vector.data(), shape_vector.size(), + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32, + &value)); + batches.emplace_back(value); continue; } diff --git a/orttraining/orttraining/test/training_api/common/synthetic_data_loader.h b/orttraining/orttraining/test/training_api/common/synthetic_data_loader.h index 058f59e6ba..23990b46a9 100644 --- a/orttraining/orttraining/test/training_api/common/synthetic_data_loader.h +++ b/orttraining/orttraining/test/training_api/common/synthetic_data_loader.h @@ -95,7 +95,7 @@ struct SyntheticDataLoader { num_of_sample_batches += 1; } - bool GetNextSampleBatch(std::vector& batches); + bool GetNextSampleBatch(std::vector& batches); size_t NumOfSampleBatches() { return num_of_sample_batches; @@ -106,9 +106,9 @@ struct SyntheticDataLoader { } private: - // Be noted: all raw data MUST remain during the training, because all Ort::Value created as session inputs + // Be noted: all raw data MUST remain during the training, because all OrtValue created as session inputs // did not explicitly copy the data in. - // And also, the created Ort::Value also won't clean the raw data pointer. The raw data should be removed when + // And also, the created OrtValue also won't clean the raw data pointer. The raw data should be removed when // the life time of this struct ends. std::vector> sample_batch_collections_; int64_t sample_batch_count_; diff --git a/orttraining/orttraining/test/training_api/core/checkpoint_test.cc b/orttraining/orttraining/test/training_api/core/checkpoint_test.cc index 8d917402cd..9fc5edc57f 100644 --- a/orttraining/orttraining/test/training_api/core/checkpoint_test.cc +++ b/orttraining/orttraining/test/training_api/core/checkpoint_test.cc @@ -181,15 +181,14 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CPU) { sample->AddFloatInput(fc2_bias_shape); data_loader.AddSyntheticSampleBatch(std::move(sample)); - std::vector all_weights_values; + std::vector all_weights_values; data_loader.GetNextSampleBatch(all_weights_values); ASSERT_EQ(all_weights_values.size(), 4); - Ort::Value* data_ptr = all_weights_values.data(); NameMLValMap name_to_ort_value{ - {"fc1.weight", **reinterpret_cast<::OrtValue**>(data_ptr)}, - {"fc1.bias", **reinterpret_cast<::OrtValue**>(data_ptr + 1)}, - {"fc2.weight", **reinterpret_cast<::OrtValue**>(data_ptr + 2)}, - {"fc2.bias", **reinterpret_cast<::OrtValue**>(data_ptr + 3)}, + {"fc1.weight", *all_weights_values[0]}, + {"fc1.bias", *all_weights_values[1]}, + {"fc2.weight", *all_weights_values[2]}, + {"fc2.bias", *all_weights_values[3]}, }; // Module/Optimizer creation and trainable parameter name definitions. diff --git a/orttraining/orttraining/test/training_api/trainer/trainer.cc b/orttraining/orttraining/test/training_api/trainer/trainer.cc index 8dfc1e1a54..82290226e2 100644 --- a/orttraining/orttraining/test/training_api/trainer/trainer.cc +++ b/orttraining/orttraining/test/training_api/trainer/trainer.cc @@ -1,8 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include -#include "orttraining/training_api/include/interfaces.h" +#include +#include "orttraining/training_api/include/utils.h" #include "cxxopts.hpp" #include "../common/synthetic_data_loader.h" @@ -16,6 +16,8 @@ using namespace onnxruntime::training::api; using namespace std; +const OrtApi* g_ort_api = nullptr; + struct TestRunnerParameters { // Models configs. std::string model_training_graph_path; @@ -44,6 +46,19 @@ void EnforceCheck(bool run_ret, std::string err_msg) { } } +#define ORT_RETURN_ON_ERROR(expr) \ + do { \ + OrtStatus* onnx_status = (expr); \ + if (onnx_status != NULL) { \ + auto code = g_ort_api->GetErrorCode(onnx_status); \ + const char* msg = g_ort_api->GetErrorMessage(onnx_status); \ + g_ort_api->ReleaseStatus(onnx_status); \ + printf("Run failed with error code :%d\n", code); \ + printf("Error message :%s\n", msg); \ + return -1; \ + } \ + } while (0); + bool ParseArguments(int argc, char* argv[], TestRunnerParameters& params) { cxxopts::Options options("Training API Test", "Main Program to test training C++ APIs."); // clang-format off @@ -152,38 +167,47 @@ void InitSyntheticDataLoader( } } -void RunTraining(const TestRunnerParameters& params) { - const auto& api = Ort::GetApi(); +int RunTraining(const TestRunnerParameters& params) { + g_ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION); - CheckpointState state; - // TODO: update using public API's calling pattern, e.g. api.LoadCheckpoint(). - EnforceCheck(LoadCheckpoint(params.checkpoint_to_load_path, state).IsOK(), "Failed to load checkpoint"); + // Create Env + OrtEnv* env; + // TODO: enable global threadpool + OrtThreadingOptions* threading_options = nullptr; + ORT_RETURN_ON_ERROR(g_ort_api->CreateThreadingOptions(&threading_options)); + ORT_RETURN_ON_ERROR(g_ort_api->CreateEnvWithGlobalThreadPools( + ORT_LOGGING_LEVEL_VERBOSE, "log", threading_options, &env)); + g_ort_api->ReleaseThreadingOptions(threading_options); - OrtSessionOptions* session_options = nullptr; - EnforceCheck(api.CreateSessionOptions(&session_options) == nullptr, "Failed to create session options."); + // Load Checkpoint State + OrtCheckpointState* checkpoint_state; + ORT_RETURN_ON_ERROR(g_ort_api->LoadCheckpoint(params.checkpoint_to_load_path.c_str(), &checkpoint_state)); + + // Create TrainingSession + OrtSessionOptions* soptions; + ORT_RETURN_ON_ERROR(g_ort_api->CreateSessionOptions(&soptions)); #ifdef USE_CUDA OrtCUDAProviderOptionsV2* cuda_options = nullptr; - EnforceCheck(api.CreateCUDAProviderOptions(&cuda_options) == nullptr, "Failed to create cuda provider options"); - EnforceCheck(api.SessionOptionsAppendExecutionProvider_CUDA_V2(session_options, cuda_options) == nullptr, - "Failed to append cuda ep."); + ORT_RETURN_ON_ERROR(g_ort_api->CreateCUDAProviderOptions(&cuda_options)); + ORT_RETURN_ON_ERROR(g_ort_api->SessionOptionsAppendExecutionProvider_CUDA_V2(soptions, cuda_options)); #endif - OrtEnv* env = nullptr; - EnforceCheck(api.CreateEnv(ORT_LOGGING_LEVEL_WARNING, "e2e_test_runner", &env) == nullptr, "Failed to create env"); - - // TODO: update using public API's calling pattern, e.g. api.CreateModule(). - Ort::OrtModule module(env, session_options, - params.model_training_graph_path, - state.module_checkpoint_state.named_parameters, - params.model_evaluation_graph_path); + OrtTrainingSession* session; + ORT_RETURN_ON_ERROR(g_ort_api->CreateTrainingSession(env, soptions, checkpoint_state, &session)); + // Initialize Training Session bool do_eval = params.model_evaluation_graph_path.has_value(); + ORT_RETURN_ON_ERROR(g_ort_api->InitializeTrainingSession(session, params.model_training_graph_path.c_str(), + do_eval ? params.model_evaluation_graph_path.value().c_str() : nullptr, + params.optimizer_training_graph_path.size() > 0 ? params.optimizer_training_graph_path.c_str() : nullptr)); - // TODO: update using public API's calling pattern, e.g. api.CreateOptimizer(). - Ort::OrtOptimizer optimizer(env, session_options, - params.optimizer_training_graph_path, - module.NamedParameters()); + size_t train_mode_output_count, eval_mode_output_count = 0; + ORT_RETURN_ON_ERROR(g_ort_api->TrainingSessionGetTrainModeOutputCount(session, &train_mode_output_count)); + + if (do_eval) { + ORT_RETURN_ON_ERROR(g_ort_api->TrainingSessionGetEvalModeOutputCount(session, &eval_mode_output_count)); + } int64_t sample_batch_count_per_epoch = 4; if (sample_batch_count_per_epoch < params.train_batch_size || sample_batch_count_per_epoch % params.train_batch_size != 0) { @@ -194,12 +218,12 @@ void RunTraining(const TestRunnerParameters& params) { onnxruntime::training::test::training_api::SyntheticDataLoader data_loader; InitSyntheticDataLoader(data_loader, params, num_of_batches_per_epoch); - int64_t total_step_count = params.num_train_epochs * num_of_batches_per_epoch; - int64_t warmup_step_count = total_step_count / 3; - - // TODO: update using public API's calling pattern, e.g. api.CreateLinearLRScheduler(). - Ort::OrtLinearLRScheduler scheduler = Ort::OrtLinearLRScheduler(optimizer, warmup_step_count, total_step_count); + // TODO: Add C API for LRScheduler + //int64_t total_step_count = params.num_train_epochs * num_of_batches_per_epoch; + //int64_t warmup_step_count = total_step_count / 3; + //Ort::OrtLinearLRScheduler scheduler = Ort::OrtLinearLRScheduler(optimizer, warmup_step_count, total_step_count); + std::cout << "Initialization completed. Now starting training loop." << std::endl; const int64_t stabilized_perf_start_step = 0; double stabilized_total_end_to_end_time{0}; auto end_to_end_start = std::chrono::high_resolution_clock::now(); @@ -210,7 +234,7 @@ void RunTraining(const TestRunnerParameters& params) { end_to_end_start = std::chrono::high_resolution_clock::now(); } - std::vector inputs; + std::vector inputs; data_loader.GetNextSampleBatch(inputs); #if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE) @@ -220,14 +244,15 @@ void RunTraining(const TestRunnerParameters& params) { train_step_range.Begin(); #endif - std::vector fetches; - EnforceCheck(module.TrainStep(inputs, fetches), "Failed during module.TrainStep."); - + std::vector fetches(train_mode_output_count); + ORT_RETURN_ON_ERROR(g_ort_api->TrainStep(session, nullptr, + inputs.size(), (const OrtValue* const*)inputs.data(), + train_mode_output_count, fetches.data())); #if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE) train_step_range.End(); #endif - float loss = *(fetches[0].GetTensorMutableData()); + float loss = onnxruntime::training::api::utils::GetValue(*fetches[0]); std::cout << "Batch # : " << batch_idx << " Loss: " << loss << std::endl; if ((batch_idx + 1) % params.gradient_accumulation_steps == 0) { @@ -238,14 +263,14 @@ void RunTraining(const TestRunnerParameters& params) { onnxruntime::profile::Color::Blue); opt_step_range.Begin(); #endif - EnforceCheck(optimizer.Step(), "Failed during optimizer.Step()."); + ORT_RETURN_ON_ERROR(g_ort_api->OptimizerStep(session, nullptr)); #if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE) opt_step_range.End(); #endif // Update learning rate. - EnforceCheck(scheduler.Step(), "Failed during shceduler.Step()"); + //EnforceCheck(scheduler.Step(), "Failed during shceduler.Step()"); #if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE) onnxruntime::profile::NvtxRangeCreator resetgrad_range( @@ -254,7 +279,7 @@ void RunTraining(const TestRunnerParameters& params) { resetgrad_range.Begin(); #endif - EnforceCheck(module.ResetGrad(), "Failed during module.ResetGrad()."); + ORT_RETURN_ON_ERROR(g_ort_api->ResetGrad(session)); #if defined(USE_CUDA) && defined(ENABLE_NVTX_PROFILE) resetgrad_range.End(); @@ -262,41 +287,55 @@ void RunTraining(const TestRunnerParameters& params) { } if (do_eval && (batch_idx + 1) % params.eval_interval == 0) { - std::vector eval_results; - EnforceCheck(module.EvalStep(inputs, eval_results), "Failed during Module.EvalStep()."); + std::vector eval_results(eval_mode_output_count); + ORT_RETURN_ON_ERROR(g_ort_api->EvalStep(session, nullptr, + inputs.size(), (const OrtValue* const*)inputs.data(), + train_mode_output_count, eval_results.data())); } if ((batch_idx + 1) % params.checkpoint_interval == 0) { // Save trained weights - CheckpointState state_to_save; - EnforceCheck(module.GetStateDict(state_to_save.module_checkpoint_state), "Failed to load module states."); - EnforceCheck(optimizer.GetStateDict(state_to_save.optimizer_checkpoint_state), "Failed to load optimizer states."); - state_to_save.property_bag.AddProperty(std::string("epoch"), epoch); std::string ckpt_file = params.output_dir + "/ckpt_" + params.model_name + std::to_string(batch_idx); + ORT_RETURN_ON_ERROR(g_ort_api->SaveCheckpoint(ckpt_file.c_str(), session, true)); - // TODO: update using public API's calling pattern, e.g. api.SaveCheckpoint(). - EnforceCheck(SaveCheckpoint(state_to_save, ckpt_file).IsOK(), "Failed to save checkpoint."); + // TODO: enable adding more properties to checkpoint + // state_to_save.property_bag.AddProperty(std::string("epoch"), epoch); } batch_idx++; + + // release input ortvalues + for (size_t i = 0; i < inputs.size(); i++) { + g_ort_api->ReleaseValue(inputs[i]); + } + + // TODO(askhade): release output values. Needs changes from Aishwarya's PR. } data_loader.ResetIterateIndex(); } + // Save trained weights + std::string ckpt_file = params.output_dir + "/ckpt_" + params.model_name; + ORT_RETURN_ON_ERROR(g_ort_api->SaveCheckpoint(ckpt_file.c_str(), session, true)); + auto end = std::chrono::high_resolution_clock::now(); std::chrono::duration duration_seconds = end - end_to_end_start; stabilized_total_end_to_end_time = duration_seconds.count(); std::cout << "Training completed - end to end latency: " << stabilized_total_end_to_end_time << "(s)" << std::endl; - api.ReleaseEnv(env); + // Delete all the ptrs + g_ort_api->ReleaseTrainingSession(session); #ifdef USE_CUDA // Finally, don't forget to release the provider options - api.ReleaseCUDAProviderOptions(cuda_options); + g_ort_api->ReleaseCUDAProviderOptions(cuda_options); #endif + g_ort_api->ReleaseSessionOptions(soptions); + g_ort_api->ReleaseCheckpointState(checkpoint_state); + g_ort_api->ReleaseEnv(env); - api.ReleaseSessionOptions(session_options); + return 0; } int main(int argc, char* argv[]) { @@ -309,6 +348,5 @@ int main(int argc, char* argv[]) { EnforceCheck(ParseArguments(argc, argv, params), "Parse arguments failed."); // Start training session - RunTraining(params); - return 0; + return RunTraining(params); } diff --git a/orttraining/orttraining/training_api/include/checkpoint.h b/orttraining/orttraining/training_api/include/checkpoint.h index 4b9efa28d4..e85585bfa7 100644 --- a/orttraining/orttraining/training_api/include/checkpoint.h +++ b/orttraining/orttraining/training_api/include/checkpoint.h @@ -49,6 +49,17 @@ struct CheckpointState { PropertyBag property_bag; }; +/** + * @brief Save training states as ORT checkpoint. + * + * @param state parameter/optimizer and other user defined training states. + * @param checkpoint_path folder where checkpoint is saved. + * @return Status + * TODO: change state to const ref + */ +Status SaveCheckpoint(CheckpointState& state, + const PathString& checkpoint_path); + /** * @brief Save ONNX initializers as ORT checkpoint. * @@ -61,16 +72,6 @@ Status SaveCheckpoint(const std::vector& trainable_ const std::vector& non_trainable_tensor_protos, const PathString& checkpoint_path); -/** - * @brief Save training states as ORT checkpoint. - * - * @param state parameter/optimizer and other user defined training states. - * @param checkpoint_path folder where checkpoint is saved. - * @return Status - */ -Status SaveCheckpoint(CheckpointState& state, - const PathString& checkpoint_path); - /** * @brief Load training states from ORT checkpoint. * diff --git a/orttraining/orttraining/training_api/include/module.h b/orttraining/orttraining/training_api/include/module.h index 2c973041b6..b6be7714c6 100644 --- a/orttraining/orttraining/training_api/include/module.h +++ b/orttraining/orttraining/training_api/include/module.h @@ -58,7 +58,7 @@ struct Module { // 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>& named_parameters, + const std::unordered_map>& named_parameters, const onnxruntime::SessionOptions& session_options, const Environment& env, const std::optional& eval_model_path_or_bytes = std::nullopt); @@ -84,10 +84,15 @@ struct Module { // Return the states of the module as a map. Status GetStateDict(ModuleCheckpointState& module_checkpoint_states); + // Returns the output count for training graph + size_t GetTrainModeOutputCount() const noexcept; + + // Returns the output count for eval graph + size_t GetEvalModeOutputCount() const noexcept; + private: std::unique_ptr train_sess_{nullptr}; std::unique_ptr eval_sess_{nullptr}; - std::unordered_map> named_parameters_; std::vector train_input_names_; std::vector train_output_names_; std::vector eval_input_names_; @@ -95,6 +100,7 @@ struct Module { std::vector weights_; std::vector gradients_; bool accumulate_gradient_ = true; + const std::unordered_map>& named_parameters_; }; } // namespace api diff --git a/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h b/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h new file mode 100644 index 0000000000..b40d88fe0f --- /dev/null +++ b/orttraining/orttraining/training_api/include/onnxruntime_training_c_api.h @@ -0,0 +1,41 @@ +// This file contains c apis for on device training +// This file should never be included standalone +// It is included from within core/session/onnxruntime_c_api.h when +// on device training is enabled +// These apis can be moved to core/session/onnxruntime_c_api.h once they stabilize + +// DO NOT UNCOMMENT +//#include "core/session/onnxruntime_c_api.h" + +ORT_API2_STATUS(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Outptr_ OrtCheckpointState** checkpoint_state); + + +ORT_API2_STATUS(SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Inout_ OrtTrainingSession* session, + bool save_optimizer_state); + +ORT_API2_STATUS(CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options, + _Inout_ OrtCheckpointState* checkpoint_state, _Outptr_ OrtTrainingSession** out); + +ORT_API2_STATUS(InitializeTrainingSession, _Inout_ OrtTrainingSession* session, + _In_ const ORTCHAR_T* train_model_path, _In_ const ORTCHAR_T* eval_model_path, + _In_ const ORTCHAR_T* optimizer_model_path); + +ORT_API2_STATUS(TrainingSessionGetTrainModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out); + +ORT_API2_STATUS(TrainingSessionGetEvalModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out); + +ORT_API2_STATUS(ResetGrad, _Inout_ OrtTrainingSession* session); + +ORT_API2_STATUS(TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options, + size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, + size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + +ORT_API2_STATUS(EvalStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options, + size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, + size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs); + +ORT_API2_STATUS(OptimizerStep, _Inout_ OrtTrainingSession* sess, + _In_opt_ const OrtRunOptions* run_options); + +ORT_CLASS_RELEASE(TrainingSession); +ORT_CLASS_RELEASE(CheckpointState); diff --git a/orttraining/orttraining/training_api/include/optimizer.h b/orttraining/orttraining/training_api/include/optimizer.h index c60f8f20ce..56ae0a2ccc 100644 --- a/orttraining/orttraining/training_api/include/optimizer.h +++ b/orttraining/orttraining/training_api/include/optimizer.h @@ -86,7 +86,7 @@ struct Optimizer { // TODO: load this info from checkpoint OptimizerType optimizer_type_ = OptimizerType::AdamW; std::unique_ptr optim_sess_; - std::unordered_map> named_parameters_; + const std::unordered_map>& named_parameters_; GroupOptimizerState optimizer_state_; std::vector input_names_; std::vector output_names_; diff --git a/orttraining/orttraining/training_api/include/training_session.h b/orttraining/orttraining/training_api/include/training_session.h new file mode 100644 index 0000000000..36a66baa4b --- /dev/null +++ b/orttraining/orttraining/training_api/include/training_session.h @@ -0,0 +1,55 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/common/common.h" +#include "module.h" +#include "optimizer.h" +#include "checkpoint.h" + +namespace onnxruntime { +namespace training { +namespace api { +using namespace common; + +// Wrapper on top of module and optimizer classes and is the only class exposed via capis +class TrainingSession { + public: + TrainingSession(const Environment& session_env, + const SessionOptions& session_options, + const std::unordered_map>& parameters); + + Status Initialize(const std::string& train_model_uri, + const std::optional& eval_model_uri, + const std::optional& optim_model_uri); + + size_t GetTrainModeOutputCount() const noexcept; + + size_t GetEvalModeOutputCount() const noexcept; + + Status TrainStep(const RunOptions& run_options, + const std::vector& inputs, + std::vector& fetches); + + Status EvalStep(const RunOptions& run_options, + const std::vector& inputs, + std::vector& fetches); + + Status ResetGrad(); + + Status OptimizerStep(const RunOptions& run_options); + + Status CreateCheckpointState(CheckpointState& chkpt_state, bool save_optimizer_state); + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TrainingSession); + + const Environment& environment_; + SessionOptions session_options_; + const std::unordered_map> named_parameters_; + std::unique_ptr module_; + std::unique_ptr optimizer_; +}; +} // namespace api +} // namespace training +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_api/include/utils.h b/orttraining/orttraining/training_api/include/utils.h index 8d44885c8e..9d00b05a00 100644 --- a/orttraining/orttraining/training_api/include/utils.h +++ b/orttraining/orttraining/training_api/include/utils.h @@ -29,7 +29,7 @@ Status OrtValueLike(const SessionState& sess_state, const OrtValue& input_val, O // Create OrtValue from a single value of type T template -void WarpInOrtValue(T value, +void WrapInOrtValue(T value, OrtValue* p_ortvalue, AllocatorPtr alloc = nullptr) { static CPUExecutionProviderInfo info; diff --git a/orttraining/orttraining/training_api/module.cc b/orttraining/orttraining/training_api/module.cc index d7e2c8015c..8084c71a41 100644 --- a/orttraining/orttraining/training_api/module.cc +++ b/orttraining/orttraining/training_api/module.cc @@ -52,24 +52,27 @@ Status Parameter::ResetGrad() { } Module::Module(const std::string& train_model_path_or_bytes, - std::unordered_map>& named_parameters, + const std::unordered_map>& named_parameters, const onnxruntime::SessionOptions& session_options, const Environment& env, - const std::optional& eval_model_path_or_bytes) { + const std::optional& eval_model_path_or_bytes) : named_parameters_{named_parameters} { + // Create session for training model train_sess_ = std::make_unique(session_options, env); ORT_THROW_IF_ERROR(train_sess_->Load(train_model_path_or_bytes)); ORT_THROW_IF_ERROR(train_sess_->Initialize()); + // Extract model input and output names utils::GetGraphInputOutputNames(train_sess_, train_input_names_, train_output_names_); - auto& train_sess_state = train_sess_->GetSessionState(); - std::vector param_input_names, grad_input_names, user_input_names, reset_grad_name; + // Reorder the extracted input names in the following order: + // user inputs, weights, gradients, reset_grad + std::vector user_input_names, param_input_names, grad_input_names, reset_grad_name; std::string param_name; std::unordered_map param_name_to_grad_input_index_map; for (const auto& input_name : train_input_names_) { - auto it = named_parameters.find(input_name); - if (it != named_parameters.end()) { + auto it = named_parameters_.find(input_name); + if (it != named_parameters_.end()) { param_input_names.emplace_back(input_name); } else if (input_name == ACCUMULATE_GRAD_CONTROL_INPUT_NAME) { reset_grad_name.emplace_back(input_name); @@ -89,12 +92,12 @@ Module::Module(const std::string& train_model_path_or_bytes, train_input_names_.insert(train_input_names_.end(), reset_grad_name.begin(), reset_grad_name.end()); // Loop each parameter, allocate it's memory based on user specified device. + auto& train_sess_state = train_sess_->GetSessionState(); for (auto& param_name : param_input_names) { - ORT_ENFORCE(named_parameters.find(param_name) != named_parameters.end()); - OrtValue& source_ortvalue = named_parameters[param_name]->Data(); - ORT_ENFORCE(source_ortvalue.IsTensor()); - const Tensor& source_tensor = source_ortvalue.Get(); + auto params_iter = named_parameters_.find(param_name); + ORT_ENFORCE(params_iter != named_parameters_.end()); + // Retrieve the target device for "param_name" std::vector node_info_vec; ORT_THROW_IF_ERROR(train_sess_state.GetInputNodeInfo(param_name, node_info_vec)); const auto& node_info = node_info_vec.front(); @@ -103,39 +106,44 @@ Module::Module(const std::string& train_model_path_or_bytes, ORT_ENFORCE(target_device == *(it->device), "Inconsistent device requirements found for input: ", param_name); } - // Create parameter value copy with corresponding device user sets the session on. - // We did not re-use the data even CPU tensor is needed. // TODO(pengwa): consider whether we should alloc contiguous buffer for parameters or gradients. - OrtValue target_ortvalue; - auto allocator = train_sess_state.GetAllocator(target_device); - ORT_ENFORCE(allocator != nullptr); + // Copy ortvalue buffer from CPU to target_device for this "param_name" (based on graph partitioning) + // Only copies data if target device is not the same as the current device the buffer is placed on - Tensor::InitOrtValue(source_tensor.DataType(), - source_tensor.Shape(), - allocator, target_ortvalue); - Tensor* target_tensor_ptr = target_ortvalue.GetMutable(); - ORT_THROW_IF_ERROR(train_sess_state.GetDataTransferMgr().CopyTensor(source_tensor, *target_tensor_ptr)); + OrtValue& param_data = params_iter->second->Data(); + ORT_ENFORCE(param_data.IsTensor()); + const Tensor& param_data_tensor = param_data.Get(); + // If the source device type is already same as target device skip copy + if (param_data_tensor.Location().device.Type() != target_device.Type()) { + // TODO: move this outside of the for loop? + auto target_allocator = train_sess_state.GetAllocator(target_device); + ORT_ENFORCE(target_allocator != nullptr); - auto param_share_ptr = - std::make_shared(param_name, target_ortvalue, named_parameters[param_name]->RequiresGrad()); - named_parameters_.insert({param_name, param_share_ptr}); - weights_.push_back(param_share_ptr->Data()); + // Create a new tensor on the target_device and switch the source_ortvalue to point to this new tensor + auto target_tensor = std::make_unique(param_data_tensor.DataType(), param_data_tensor.Shape(), target_allocator); + ORT_THROW_IF_ERROR(train_sess_state.GetDataTransferMgr().CopyTensor(param_data_tensor, *target_tensor.get())); + auto ml_tensor_type = DataTypeImpl::GetType(); + // TODO test the original buffer is released. + param_data.Init(target_tensor.release(), ml_tensor_type, ml_tensor_type->GetDeleteFunc()); + } + + weights_.push_back(param_data); // Create gradient buffer when parameter requires gradient. - if (param_share_ptr->RequiresGrad()) { + if (params_iter->second->RequiresGrad()) { // Create gradient accumulation buffer. auto it = param_name_to_grad_input_index_map.find(param_name); ORT_ENFORCE(it != param_name_to_grad_input_index_map.end(), "Gradient buffer input not providered for param: ", param_name); const size_t grad_input_index = it->second; - auto& param_grad_buffer_name = grad_input_names[grad_input_index]; + auto& param_grad_name = grad_input_names[grad_input_index]; // TODO: don't pre-allocate the gradient buffer. // Gradient usually stays on the same device of its parameter. - OrtValue param_grad_buffer_ortvalue; - ORT_THROW_IF_ERROR(utils::OrtValueLike(train_sess_state, target_ortvalue, param_grad_buffer_ortvalue)); - ORT_THROW_IF_ERROR(param_share_ptr->SetGrad(param_grad_buffer_name, param_grad_buffer_ortvalue)); - gradients_[grad_input_index] = param_share_ptr->Gradient(); + OrtValue param_grad; + ORT_THROW_IF_ERROR(utils::OrtValueLike(train_sess_state, param_data, param_grad)); + ORT_THROW_IF_ERROR(params_iter->second->SetGrad(param_grad_name, param_grad)); + gradients_[grad_input_index] = params_iter->second->Gradient(); } } @@ -146,9 +154,9 @@ Module::Module(const std::string& train_model_path_or_bytes, utils::GetGraphInputOutputNames(eval_sess_, eval_input_names_, eval_output_names_); // Eval model validation - // We are making certain assumptions: Like the order in which parameters - // occur will be same between train and eval graphs, - // and all the weights present in both graphs match. + // We are making certain assumptions: Like the order in which parameters occur will be same between train and eval + // graphs, and all the weights present in both graphs match. + // TODO: Add the checks instead of making assumptions?? std::vector eval_user_input_names, param_input_names; for (const auto& input_name : eval_input_names_) { if (named_parameters_.find(input_name) != named_parameters_.end()) { @@ -168,6 +176,14 @@ Module::Module(const std::string& train_model_path_or_bytes, } } +size_t Module::GetTrainModeOutputCount() const noexcept { + return train_output_names_.size(); +} + +size_t Module::GetEvalModeOutputCount() const noexcept { + return eval_output_names_.size(); +} + std::vector> Module::Parameters() const { std::vector> params; for (auto& it : named_parameters_) { @@ -187,7 +203,7 @@ Status Module::TrainStep(const std::vector& inputs, std::vector(accumulate_gradient_, &do_update_input); + utils::WrapInOrtValue(accumulate_gradient_, &do_update_input); feeds.push_back(do_update_input); // TODO: need to filter out the grads from the output ortvalues diff --git a/orttraining/orttraining/training_api/onnxruntime_training_c_api.cc b/orttraining/orttraining/training_api/onnxruntime_training_c_api.cc new file mode 100644 index 0000000000..9ec36cb68e --- /dev/null +++ b/orttraining/orttraining/training_api/onnxruntime_training_c_api.cc @@ -0,0 +1,218 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/framework/error_code_helper.h" +#include "core/framework/ort_value.h" +#include "core/session/ort_apis.h" +#include "core/session/ort_env.h" +#include "orttraining/training_api/include/checkpoint.h" +#include "orttraining/training_api/include/training_session.h" +#include "core/session/abi_session_options_impl.h" + +ORT_API_STATUS_IMPL(OrtApis::CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options, + _Inout_ OrtCheckpointState* checkpoint_state, _Outptr_ OrtTrainingSession** out) { + API_IMPL_BEGIN + std::unique_ptr train_sess; + auto chkpt_state = reinterpret_cast(checkpoint_state); + OrtStatus* status = nullptr; + *out = nullptr; + + ORT_TRY { + train_sess = std::make_unique( + env->GetEnvironment(), + options == nullptr ? onnxruntime::SessionOptions() : options->value, + chkpt_state->module_checkpoint_state.named_parameters); + + *out = reinterpret_cast(train_sess.release()); + } + ORT_CATCH(const std::exception& e) { + ORT_HANDLE_EXCEPTION([&]() { + status = OrtApis::CreateStatus(ORT_FAIL, e.what()); + }); + } + + return status; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::InitializeTrainingSession, _Inout_ OrtTrainingSession* session, + _In_ const ORTCHAR_T* train_model_path, _In_ const ORTCHAR_T* eval_model_path, + _In_ const ORTCHAR_T* optimizer_model_path) { + API_IMPL_BEGIN + + auto train_sess = reinterpret_cast(session); + ORT_API_RETURN_IF_STATUS_NOT_OK(train_sess->Initialize(train_model_path, + eval_model_path ? std::optional{eval_model_path} + : std::nullopt, + optimizer_model_path ? std::optional{optimizer_model_path} + : std::nullopt)); + + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::TrainingSessionGetTrainModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out) { + API_IMPL_BEGIN + auto session = reinterpret_cast(sess); + *out = session->GetTrainModeOutputCount(); + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::TrainingSessionGetEvalModeOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out) { + API_IMPL_BEGIN + auto session = reinterpret_cast(sess); + *out = session->GetEvalModeOutputCount(); + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::ResetGrad, _Inout_ OrtTrainingSession* session) { + API_IMPL_BEGIN + auto train_session = reinterpret_cast(session); + ORT_API_RETURN_IF_STATUS_NOT_OK(train_session->ResetGrad()); + + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options, + size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, + size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs) { + API_IMPL_BEGIN + auto session = reinterpret_cast(sess); + constexpr int queue_id = 0; + + std::vector feeds(inputs_len); + + for (size_t i = 0; i != inputs_len; ++i) { + auto& ort_value = feeds[i] = *reinterpret_cast(inputs[i]); + if (ort_value.Fence()) { + ort_value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id); + } + } + + // Create output feed + std::vector fetches(outputs_len); + for (size_t i = 0; i != outputs_len; ++i) { + if (outputs[i] != nullptr) { + ::OrtValue& value = *(outputs[i]); + if (value.Fence()) + value.Fence()->BeforeUsingAsOutput(onnxruntime::kCpuExecutionProvider, queue_id); + fetches[i] = value; + } + } + Status status; + if (run_options == nullptr) { + OrtRunOptions op; + status = session->TrainStep(op, feeds, fetches); + } else { + status = session->TrainStep(*run_options, feeds, fetches); + } + + if (!status.IsOK()) + return onnxruntime::ToOrtStatus(status); + for (size_t i = 0; i != outputs_len; ++i) { + ::OrtValue& value = fetches[i]; + if (value.Fence()) + value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id); + if (outputs[i] == nullptr) { + outputs[i] = new OrtValue(value); + } + } + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::EvalStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options, + size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs, + size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs) { + API_IMPL_BEGIN + auto session = reinterpret_cast(sess); + constexpr int queue_id = 0; + + std::vector feeds(inputs_len); + + for (size_t i = 0; i != inputs_len; ++i) { + auto& ort_value = feeds[i] = *reinterpret_cast(inputs[i]); + + if (ort_value.Fence()) ort_value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id); + } + + // Create output feed + std::vector fetches(outputs_len); + for (size_t i = 0; i != outputs_len; ++i) { + if (outputs[i] != nullptr) { + ::OrtValue& value = *(outputs[i]); + if (value.Fence()) + value.Fence()->BeforeUsingAsOutput(onnxruntime::kCpuExecutionProvider, queue_id); + fetches[i] = value; + } + } + Status status; + if (run_options == nullptr) { + OrtRunOptions op; + status = session->EvalStep(op, feeds, fetches); + } else { + status = session->EvalStep(*run_options, feeds, fetches); + } + + if (!status.IsOK()) + return onnxruntime::ToOrtStatus(status); + for (size_t i = 0; i != outputs_len; ++i) { + ::OrtValue& value = fetches[i]; + if (value.Fence()) + value.Fence()->BeforeUsingAsInput(onnxruntime::kCpuExecutionProvider, queue_id); + if (outputs[i] == nullptr) { + outputs[i] = new OrtValue(value); + } + } + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::OptimizerStep, _Inout_ OrtTrainingSession* sess, + _In_opt_ const OrtRunOptions* run_options) { + API_IMPL_BEGIN + auto session = reinterpret_cast(sess); + if (run_options == nullptr) { + OrtRunOptions op; + ORT_API_RETURN_IF_STATUS_NOT_OK(session->OptimizerStep(op)); + } else { + ORT_API_RETURN_IF_STATUS_NOT_OK(session->OptimizerStep(*run_options)); + } + + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Outptr_ OrtCheckpointState** checkpoint_state) { + API_IMPL_BEGIN + *checkpoint_state = nullptr; + auto chkpt_state = std::make_unique(); + ORT_API_RETURN_IF_STATUS_NOT_OK(onnxruntime::training::api::LoadCheckpoint(checkpoint_path, *chkpt_state)); + *checkpoint_state = reinterpret_cast(chkpt_state.release()); + + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::SaveCheckpoint, _In_ const ORTCHAR_T* checkpoint_path, _Inout_ OrtTrainingSession* sess, + bool save_optimizer_state) { + API_IMPL_BEGIN + auto session = reinterpret_cast(sess); + onnxruntime::training::api::CheckpointState chkpt_state; + ORT_API_RETURN_IF_STATUS_NOT_OK(session->CreateCheckpointState(chkpt_state, save_optimizer_state)); + ORT_API_RETURN_IF_STATUS_NOT_OK(onnxruntime::training::api::SaveCheckpoint(chkpt_state, checkpoint_path)); + + return nullptr; + API_IMPL_END +} + +ORT_API(void, OrtApis::ReleaseTrainingSession, _Frees_ptr_opt_ OrtTrainingSession* session) { + delete reinterpret_cast(session); +} + +ORT_API(void, OrtApis::ReleaseCheckpointState, _Frees_ptr_opt_ OrtCheckpointState* checkpoint_state) { + delete reinterpret_cast(checkpoint_state); +} diff --git a/orttraining/orttraining/training_api/optimizer.cc b/orttraining/orttraining/training_api/optimizer.cc index 1dd2923bf2..e2e988b5c4 100644 --- a/orttraining/orttraining/training_api/optimizer.cc +++ b/orttraining/orttraining/training_api/optimizer.cc @@ -106,8 +106,8 @@ Optimizer::Optimizer(const std::string& optim_path_or_bytes, Status Optimizer::Step() { OrtValue learning_rate_input, step_input; - utils::WarpInOrtValue(optimizer_state_.learning_rate, &learning_rate_input); - utils::WarpInOrtValue(optimizer_state_.step, &step_input); + utils::WrapInOrtValue(optimizer_state_.learning_rate, &learning_rate_input); + utils::WrapInOrtValue(optimizer_state_.step, &step_input); std::vector feeds({learning_rate_input, step_input}); feeds.insert(feeds.end(), inputs_.begin(), inputs_.end()); diff --git a/orttraining/orttraining/training_api/training_session.cc b/orttraining/orttraining/training_api/training_session.cc new file mode 100644 index 0000000000..020a708c29 --- /dev/null +++ b/orttraining/orttraining/training_api/training_session.cc @@ -0,0 +1,69 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "orttraining/training_api/include/training_session.h" + +namespace onnxruntime { +namespace training { +namespace api { + +TrainingSession::TrainingSession(const Environment& session_env, + const SessionOptions& session_options, + const std::unordered_map>& parameters) + : environment_(session_env), + session_options_{session_options}, + named_parameters_{parameters} {} + +Status TrainingSession::Initialize(const std::string& train_model_uri, const std::optional& eval_model_uri, + const std::optional& optim_model_uri) { + module_ = std::move(std::make_unique(train_model_uri, named_parameters_, session_options_, + environment_, eval_model_uri)); + + if (optim_model_uri.has_value()) { + optimizer_ = std::move(std::make_unique(optim_model_uri.value(), named_parameters_, + session_options_, environment_)); + } + + return Status::OK(); +} + +size_t TrainingSession::GetTrainModeOutputCount() const noexcept { + return module_->GetTrainModeOutputCount(); +} + +size_t TrainingSession::GetEvalModeOutputCount() const noexcept { + return module_->GetEvalModeOutputCount(); +} + +Status TrainingSession::TrainStep(const RunOptions&, + const std::vector& inputs, + std::vector& fetches) { + return module_->TrainStep(inputs, fetches); +} + +Status TrainingSession::EvalStep(const RunOptions&, + const std::vector& inputs, + std::vector& fetches) { + return module_->EvalStep(inputs, fetches); +} + +Status TrainingSession::ResetGrad() { + return module_->ResetGrad(); +} + +Status TrainingSession::OptimizerStep(const RunOptions&) { + return optimizer_->Step(); +} + +Status TrainingSession::CreateCheckpointState(CheckpointState& chkpt_state, bool save_optimizer_state) { + ORT_RETURN_IF_ERROR(module_->GetStateDict(chkpt_state.module_checkpoint_state)); + if (save_optimizer_state) { + ORT_RETURN_IF_ERROR(optimizer_->GetStateDict(chkpt_state.optimizer_checkpoint_state)); + } + + return Status::OK(); +} + +} // namespace api +} // namespace training +} // namespace onnxruntime