diff --git a/orttraining/orttraining/test/training_api/core/checkpoint_test.cc b/orttraining/orttraining/test/training_api/core/checkpoint_test.cc index 8edad98472..58c1a11e8e 100644 --- a/orttraining/orttraining/test/training_api/core/checkpoint_test.cc +++ b/orttraining/orttraining/test/training_api/core/checkpoint_test.cc @@ -191,8 +191,8 @@ TEST(CheckpointApiTest, LoadCheckpointToModel) { } /** - * Create Module with sets of parameters, - * Create Optimizer passing in Module's parameters. + * Create Module passing in checkpoint state, + * Create Optimizer passing in checkpoint state. * Save Optimizer states into ORT checkpoint files, * Then load it into ORT, compare with the initial optimizer states values. */ @@ -206,7 +206,8 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) { auto optim_uri = "testdata/training_api/adamw.onnx"; // Generate randomized weight values using synthetic data generator. - constexpr int64_t fc2_weight_dim_in = 10, fc2_weight_dim_out = 500, fc1_weight_dim_in = 500, fc1_weight_dim_out = 784; + constexpr int64_t fc2_weight_dim_in = 10, fc2_weight_dim_out = 500, + fc1_weight_dim_in = 500, fc1_weight_dim_out = 784; const std::vector fc1_weight_shape{fc1_weight_dim_in, fc1_weight_dim_out}; const std::vector fc1_bias_shape{fc1_weight_dim_in}; const std::vector fc2_weight_shape{fc2_weight_dim_in, fc2_weight_dim_out}; @@ -249,12 +250,6 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) { auto optimizer = std::make_unique(optim_uri, &state, session_option, *env, cuda_provider); - /// Phase 2 - Run Optimizer.GetStateDict and call save checkpoint APIs. - /// And check the result checkpoint files. - - CheckpointState checkpoint_state; - ORT_ENFORCE(optimizer->GetStateDict(checkpoint_state.optimizer_checkpoint_state).IsOK()); - // Remove the temporary directory if it already exists. auto ckpt_test_root_dir = ORT_TSTR("checkpointing_api_test_dir"); if (Env::Default().FolderExists(ckpt_test_root_dir)) { @@ -265,13 +260,14 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) { // Call Save APIs. PathString checkpoint_path{ ConcatPathComponent(tmp_dir.Path(), ORT_TSTR("e2e_ckpt_save_cpu"))}; - ASSERT_STATUS_OK(SaveCheckpoint(checkpoint_state, checkpoint_path, true)); + ASSERT_STATUS_OK(SaveCheckpoint(state, checkpoint_path, true)); // Check the ckpt files in the directory. std::set expected_file_names{ ORT_TSTR("optim_group0_momentum0_tensors.pbseq"), ORT_TSTR("optim_group0_momentum1_tensors.pbseq"), ORT_TSTR("optim_group0_properties.pbseq"), + ORT_TSTR("paramtrain_tensors.pbseq"), }; std::set valid_file_names; @@ -289,27 +285,27 @@ TEST(CheckpointApiTest, SaveOptimizerStateAsCheckpoint_ThenLoad_CUDA) { ASSERT_EQ(expected_file_names, valid_file_names); - /// Phase 3 - Run load checkpoint APIs. + /// Phase 2 - Run load checkpoint APIs. /// Validate the result matches with initial optimizer state values. // Call Load APIs CheckpointState checkpoint_state_to_load; ASSERT_STATUS_OK(LoadCheckpoint(checkpoint_path, checkpoint_state_to_load)); OptimizerCheckpointState optimizer_state = checkpoint_state_to_load.optimizer_checkpoint_state; - std::unordered_map>& + InlinedHashMap>& group_optimizer_states = optimizer_state.group_named_optimizer_states; ASSERT_EQ(group_optimizer_states.size(), 1); ASSERT_EQ(group_optimizer_states.begin()->first, "group0"); - std::unordered_map& + InlinedHashMap& param_named_optimizer_states = group_optimizer_states["group0"]->param_named_optimizer_states; ASSERT_EQ(param_named_optimizer_states.size(), named_parameters.size()); for (auto it = param_named_optimizer_states.begin(); it != param_named_optimizer_states.end(); ++it) { ASSERT_TRUE(named_parameters.find(it->first) != named_parameters.end()); - for (auto& [momentum_name, restored_ort_value] : it->second.momentum_named_states) { + for (auto& [momentum_name, restored_ort_value] : it->second) { ASSERT_TRUE(momentum_name == "momentum0" || momentum_name == "momentum1"); const OrtValue& param_ort_value = name_to_ort_value[it->first]; ASSERT_TRUE(restored_ort_value.IsTensor() && param_ort_value.IsTensor()); diff --git a/orttraining/orttraining/test/training_api/core/training_api_tests.cc b/orttraining/orttraining/test/training_api/core/training_api_tests.cc index e08687eb18..f41b72dba2 100644 --- a/orttraining/orttraining/test/training_api/core/training_api_tests.cc +++ b/orttraining/orttraining/test/training_api/core/training_api_tests.cc @@ -42,7 +42,7 @@ constexpr float INITIAL_LR = 1e-3f; */ Status CreateFakeOptimizerCheckpointStateOnCPU( const std::unordered_map>& named_parameters, - const std::vector& momentum_keys, + const InlinedVector& momentum_keys, OptimizerCheckpointState& optimizer_checkpoint_state) { auto& grouped_optimizer_states = optimizer_checkpoint_state.group_named_optimizer_states; grouped_optimizer_states.insert({"group0", std::make_shared()}); @@ -58,7 +58,7 @@ Status CreateFakeOptimizerCheckpointStateOnCPU( OrtValue param = pair.second->Data(); const auto& param_tensor = param.template Get(); GenerateRandomInput(param_tensor.Shape().GetDims(), param_moment_state); - cur_param_optimizer_states.momentum_named_states.insert({state_name, std::move(param_moment_state)}); + cur_param_optimizer_states.insert({state_name, std::move(param_moment_state)}); } } } @@ -287,39 +287,7 @@ TEST(TrainingApiTest, OptimizerCreatedWithOptimizerCheckpointState) { std::shared_ptr optim = std::make_shared( ToUTF8String(optim_uri), &new_state, session_option, *env, providers); - // After loading state dict, check if optim state is updated to new states. - OptimizerCheckpointState optimizer_states; - ASSERT_STATUS_OK(optim->GetStateDict(optimizer_states)); - - for (auto& p : model->NamedParameters()) { - auto param_name = p.first; - ParameterOptimizerState& param_state = - optimizer_states.group_named_optimizer_states["group0"]->param_named_optimizer_states.at(param_name); - - ParameterOptimizerState& external_param_state = - external_optimizer_checkpoint_state.group_named_optimizer_states["group0"] - ->param_named_optimizer_states.at(param_name); - for (auto& param_p : param_state.momentum_named_states) { - std::vector moment_vec; - if (run_cuda) { - CudaOrtValueToCpuVec(param_state.momentum_named_states.at(param_p.first), moment_vec); - } else { - CpuOrtValueToVec(param_state.momentum_named_states.at(param_p.first), moment_vec); - } - std::vector external_moment_vect; - - if (run_cuda) { - CudaOrtValueToCpuVec(external_param_state.momentum_named_states.at(param_p.first), external_moment_vect); - } else { - CpuOrtValueToVec(external_param_state.momentum_named_states.at(param_p.first), external_moment_vect); - } - - ASSERT_EQ(moment_vec.size(), external_moment_vect.size()); - for (size_t i = 0; i < moment_vec.size(); i++) { - ASSERT_EQ(moment_vec[i], external_moment_vect[i]); - } - } - } + ASSERT_TRUE(optim.get() != nullptr); } } @@ -328,9 +296,9 @@ void TestLRSchduler(const std::basic_string& test_file_name, int64_t total_step_count, int64_t warmup_step_count) { std::vector run_cuda_list{false}; - // #ifdef USE_CUDA - // run_cuda_list.push_back(true); - // #endif +#ifdef USE_CUDA + run_cuda_list.push_back(true); +#endif for (auto run_cuda : run_cuda_list) { std::vector> providers; @@ -392,8 +360,7 @@ void TestLRSchduler(const std::basic_string& test_file_name, optim, warmup_step_count, total_step_count); for (auto it = test_data.begin(); it != test_data.end(); ++it) { - OptimizerCheckpointState optimizer_states; - ASSERT_STATUS_OK(optim->GetStateDict(optimizer_states)); + OptimizerCheckpointState& optimizer_states = state.optimizer_checkpoint_state; auto group_optimizer_state = optimizer_states.group_named_optimizer_states["group0"]; constexpr const float rtol = 1e-4f, atol = 1e-5f; @@ -493,11 +460,10 @@ TEST(TrainingApiTest, OptimStep) { std::string param_name = "fc2.weight"; // before training, check if optim state is initialized to 0 - onnxruntime::training::api::OptimizerCheckpointState optimizer_states; - ASSERT_STATUS_OK(optim->GetStateDict(optimizer_states)); + onnxruntime::training::api::OptimizerCheckpointState& optimizer_states = state.optimizer_checkpoint_state; onnxruntime::training::api::ParameterOptimizerState& param_state = optimizer_states.group_named_optimizer_states["group0"]->param_named_optimizer_states.at(param_name); - OrtValue& moment_1 = param_state.momentum_named_states.at("momentum0"); + OrtValue& moment_1 = param_state.at("momentum0"); std::vector param_vec_before_optimizer_step; CudaOrtValueToCpuVec(model->NamedParameters().at(param_name)->Data(), param_vec_before_optimizer_step); diff --git a/orttraining/orttraining/training_api/checkpoint.cc b/orttraining/orttraining/training_api/checkpoint.cc index 73b24ae49f..df64b521c3 100644 --- a/orttraining/orttraining/training_api/checkpoint.cc +++ b/orttraining/orttraining/training_api/checkpoint.cc @@ -260,7 +260,7 @@ Status OrtSaveOptimizerStatesInternal(OptimizerCheckpointState& optimizer_state, // Firstly indexed by momentum names; Secondly indexed by parameter names. InlinedHashMap> optimizer_state_ort_values; for (const auto& [param_name, param_optimizer_state] : group_optimizer_state_ptr->param_named_optimizer_states) { - for (const auto& [momentum_name, m_state_val] : param_optimizer_state.momentum_named_states) { + for (const auto& [momentum_name, m_state_val] : param_optimizer_state) { if (optimizer_state_ort_values.find(momentum_name) == optimizer_state_ort_values.end()) { std::unordered_map param_name_to_ortvalue{{param_name, m_state_val}}; optimizer_state_ort_values.insert({momentum_name, param_name_to_ortvalue}); @@ -421,7 +421,7 @@ Status OrtLoadOptimizerStatesInternal(const PathString& optimizer_folder_path, } auto& group_optimizer_state = grouped_optimizer_states[group_name]; - std::unordered_map& + InlinedHashMap& param_optimizer_states = group_optimizer_state->param_named_optimizer_states; const PathString& tensor_file_path = GetTensorProtoFilePath(optimizer_folder_path, @@ -437,7 +437,7 @@ Status OrtLoadOptimizerStatesInternal(const PathString& optimizer_folder_path, ParameterOptimizerState param_state; param_optimizer_states.insert({param_name, param_state}); } - param_optimizer_states[param_name].momentum_named_states.insert({momentum_name, std::move(pair.second)}); + param_optimizer_states[param_name].insert({momentum_name, std::move(pair.second)}); } } diff --git a/orttraining/orttraining/training_api/module.cc b/orttraining/orttraining/training_api/module.cc index a5bbd1a6c4..7ba9cb6d6b 100644 --- a/orttraining/orttraining/training_api/module.cc +++ b/orttraining/orttraining/training_api/module.cc @@ -178,12 +178,12 @@ Module::Module(const std::string& train_model_path_or_bytes, state_->module_checkpoint_state.train_session_data_transfer_mgr = &train_sess_->GetDataTransferManager(); // Extract model input and output names - std::vector train_input_names, train_output_names; + InlinedVector train_input_names, train_output_names; utils::GetGraphInputOutputNames(train_sess_, train_input_names, train_output_names); // 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; + InlinedVector user_input_names, param_input_names, grad_input_names, reset_grad_name; std::unordered_map param_name_to_grad_input_index_map; for (const auto& input_name : train_input_names) { @@ -283,7 +283,7 @@ Module::Module(const std::string& train_model_path_or_bytes, // 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, eval_param_input_names; + InlinedVector eval_user_input_names, eval_param_input_names; for (const auto& input_name : eval_input_names_) { if (state_->module_checkpoint_state.named_parameters.find(input_name) != state_->module_checkpoint_state.named_parameters.end()) { diff --git a/orttraining/orttraining/training_api/module.h b/orttraining/orttraining/training_api/module.h index 4cbddc9656..8fc4228f83 100644 --- a/orttraining/orttraining/training_api/module.h +++ b/orttraining/orttraining/training_api/module.h @@ -137,18 +137,22 @@ struct Module { private: std::unique_ptr train_sess_{nullptr}; std::unique_ptr eval_sess_{nullptr}; - std::vector train_input_names_; - std::vector train_output_names_; - std::vector eval_input_names_; - std::vector eval_output_names_; - std::vector weight_names_; - std::vector weights_; - std::vector gradients_; - bool accumulate_gradient_ = false; + + InlinedVector train_input_names_; + InlinedVector train_output_names_; + InlinedVector eval_input_names_; + InlinedVector eval_output_names_; + InlinedVector weight_names_; + + InlinedVector weights_; + InlinedVector gradients_; + CheckpointState* state_; // Non owning pointer to the state. + + bool accumulate_gradient_ = false; std::string eval_model_path_; - size_t train_user_input_count_ = 0U; - size_t eval_user_input_count_ = 0U; + size_t train_user_input_count_{0U}; + size_t eval_user_input_count_{0U}; }; } // namespace api diff --git a/orttraining/orttraining/training_api/optimizer.cc b/orttraining/orttraining/training_api/optimizer.cc index 7e0d21b4e5..b7bc67f0c4 100644 --- a/orttraining/orttraining/training_api/optimizer.cc +++ b/orttraining/orttraining/training_api/optimizer.cc @@ -113,7 +113,7 @@ Status Optimizer::GenerateMomentumNamedStates(OptimizerCheckpointState& optimize OrtValue param_state; ORT_ENFORCE(utils::CreateZeroValuedOrtValueLike(optim_sess_state, pair.second->Data(), param_state).IsOK(), "Error generating moment state for ", pair.first); - cur_param_optimizer_states.momentum_named_states.insert({state_name, std::move(param_state)}); + cur_param_optimizer_states.insert({state_name, std::move(param_state)}); } } } @@ -127,8 +127,8 @@ Status Optimizer::ConstructInputs() { auto& param_named_optimizer_states = optimizer_state_->param_named_optimizer_states; - std::vector params, grads; - std::vector> list_of_momentums; + InlinedVector params, grads; + InlinedVector> list_of_momentums; list_of_momentums.resize(optimizer_algo_ptr_->momentum_keys.size()); // Collect all the non-user-defined inputs from the named_parameters_. @@ -150,7 +150,7 @@ Status Optimizer::ConstructInputs() { for (size_t m_index = 0; m_index < optimizer_algo_ptr_->momentum_keys.size(); ++m_index) { auto* moment_tensor = param_named_optimizer_states.at(parameter_name) - .momentum_named_states.at(optimizer_algo_ptr_->momentum_keys[m_index]) + .at(optimizer_algo_ptr_->momentum_keys[m_index]) .GetMutable(); list_of_momentums[m_index].emplace_back( Tensor(moment_tensor->DataType(), moment_tensor->Shape(), @@ -258,20 +258,6 @@ Status Optimizer::Step() { return Status::OK(); } -Status Optimizer::GetStateDict(OptimizerCheckpointState& optimizer_checkpoint_state) { - auto& grouped_optimizer_states = optimizer_checkpoint_state.group_named_optimizer_states; - - // To support multiple groups, the Optimizer constructor needs to accept information for grouping. - grouped_optimizer_states.insert({GROUP_ZERO_NAME, std::make_shared(*optimizer_state_)}); - - // Pass the optimizer session data transfer manager for data copying when saving. - // An alternative is, we can do copy at this stage. - ORT_RETURN_IF_NOT(optim_sess_, "optimizer session not initialized"); - const DataTransferManager& sess_data_transfer_manager = optim_sess_->GetDataTransferManager(); - optimizer_checkpoint_state.optimizer_session_data_transfer_mgr = &sess_data_transfer_manager; - return Status::OK(); -} - Status Optimizer::LoadStateDict(OptimizerCheckpointState& optimizer_checkpoint_states) { auto group_optimizer_state_it = optimizer_checkpoint_states.group_named_optimizer_states.find(GROUP_ZERO_NAME); @@ -293,8 +279,8 @@ Status Optimizer::LoadStateDict(OptimizerCheckpointState& optimizer_checkpoint_s ORT_ENFORCE(src_exist || !strict_match, "Parameter ", params_iter.first, " not found in the source optimizer checkpoint states."); - std::unordered_map& momentum_named_states = - param_named_optimizer_states.at(params_iter.first).momentum_named_states; + InlinedHashMap& momentum_named_states = + param_named_optimizer_states.at(params_iter.first); OrtValue& param_data = params_iter.second->Data(); ORT_ENFORCE(param_data.IsTensor()); diff --git a/orttraining/orttraining/training_api/optimizer.h b/orttraining/orttraining/training_api/optimizer.h index a6769172e3..ffbe293c88 100644 --- a/orttraining/orttraining/training_api/optimizer.h +++ b/orttraining/orttraining/training_api/optimizer.h @@ -19,9 +19,7 @@ namespace api { * For Adam optimizer, it looks like: * { "moment_0": OrtValue, "moment_1": OrtValue,}. */ -struct ParameterOptimizerState { - std::unordered_map momentum_named_states; -}; +typedef InlinedHashMap ParameterOptimizerState; /** * @brief States belong to one specific group of trainable Parameters. @@ -33,7 +31,7 @@ struct GroupOptimizerState { // Adaptive learning rate as training proceeds. Be noted, learning_rate can be // restored by lr scheduler from given step and initial_lr, though, we still save/load this in checkpoint. float learning_rate{initial_lr}; - std::unordered_map param_named_optimizer_states; + InlinedHashMap param_named_optimizer_states; }; /** @@ -43,16 +41,16 @@ struct GroupOptimizerState { */ struct OptimizerCheckpointState { public: - std::unordered_map> group_named_optimizer_states; + InlinedHashMap> group_named_optimizer_states; const DataTransferManager* optimizer_session_data_transfer_mgr; }; struct OptimizerAlgorithmBase { - OptimizerAlgorithmBase(const std::vector& momentum_keys, - const std::vector& optimizer_states_inputs) + OptimizerAlgorithmBase(const InlinedVector& momentum_keys, + const InlinedVector& optimizer_states_inputs) : momentum_keys(momentum_keys), optimizer_states_inputs(optimizer_states_inputs) {} - std::vector momentum_keys; - std::vector optimizer_states_inputs; + InlinedVector momentum_keys; + InlinedVector optimizer_states_inputs; }; struct AdamWOptimizerAlgorithm : public OptimizerAlgorithmBase { @@ -106,15 +104,6 @@ struct Optimizer { Status Step(); - /** - * @brief Get the current optimizer state. - * - * Be noted the returned optimizer_checkpoint_states will hold new references to - * original momentum states. - * @return Status - */ - Status GetStateDict(OptimizerCheckpointState& optimizer_checkpoint_states); - Status SetLearningRate(float lr) { optimizer_state_->learning_rate = lr; return Status::OK(); @@ -159,11 +148,13 @@ struct Optimizer { std::unique_ptr optimizer_algo_ptr_; std::unique_ptr optim_sess_; - CheckpointState* state_; // Non owning pointer to the state. + + CheckpointState* state_; // Non owning pointer to the state std::shared_ptr optimizer_state_; - std::vector input_names_; - std::vector output_names_; - std::vector inputs_; + + InlinedVector input_names_; + InlinedVector output_names_; + InlinedVector inputs_; int32_t group_count_{0}; }; diff --git a/orttraining/orttraining/training_api/utils.cc b/orttraining/orttraining/training_api/utils.cc index aa2c0173a7..6a00aa229b 100644 --- a/orttraining/orttraining/training_api/utils.cc +++ b/orttraining/orttraining/training_api/utils.cc @@ -19,9 +19,9 @@ namespace utils { const std::vector GRAD_SUFFIX{"_grad.accumulation.buffer", "_grad", "_grad.accumulation.out"}; void GetGraphInputOutputNames(const std::unique_ptr& session_object, - std::vector& input_names, - std::vector& output_names) { - auto get_names = [](const std::vector* node_args, std::vector& names) { + InlinedVector& input_names, + InlinedVector& output_names) { + auto get_names = [](const std::vector* node_args, InlinedVector& names) { ORT_ENFORCE(nullptr != node_args); for (const auto* arg : *node_args) { names.push_back(arg->Name()); diff --git a/orttraining/orttraining/training_api/utils.h b/orttraining/orttraining/training_api/utils.h index 37ae4b6747..e28d667de5 100644 --- a/orttraining/orttraining/training_api/utils.h +++ b/orttraining/orttraining/training_api/utils.h @@ -14,8 +14,8 @@ namespace utils { // Get names of graph inputs and outputs void GetGraphInputOutputNames(const std::unique_ptr& session_object, - std::vector& input_names, - std::vector& output_names); + InlinedVector& input_names, + InlinedVector& output_names); // Fetch the parameter name from suffix: name = param_name+suffix, // returns True if suffix is present in name else False bool GetParamNameFromSuffix(const std::string& name, const std::string& suffix, std::string& param_name);