From 56bae3b196ad958edc2ae8923e6a554ca30b6cfa Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Fri, 2 Sep 2022 07:47:17 -0700 Subject: [PATCH] Use InplaceClipGradNorm for offline processing for on-device training (#12603) --- .../kernel_def_hashes/training_ops.cpu.json | 4 + .../test/testdata/training_api/adamw.onnx | Bin 925 -> 580 bytes .../python/training/onnxblock/model.py | 1 - .../python/training/onnxblock/optim/optim.py | 58 +++++----- .../test/python/orttraining_test_onnxblock.py | 39 +++---- .../cuda/optimizer/clip_grad_norm_test.cc | 38 +++++-- .../orttraining/training_api/optimizer.cc | 102 +++++++++++------- .../training_ops/cpu/cpu_training_kernels.cc | 4 + .../clip_grad_norm/clip_grad_norm.cc | 84 +++++++++++++++ .../optimizer/clip_grad_norm/clip_grad_norm.h | 30 ++++++ .../clip_grad_norm/clip_grad_norm.cc | 4 +- 11 files changed, 262 insertions(+), 102 deletions(-) create mode 100644 orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.cc create mode 100644 orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.h diff --git a/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json b/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json index 1cc69e3975..44ea17b717 100644 --- a/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json +++ b/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json @@ -274,5 +274,9 @@ [ "InPlaceAccumulatorV2 com.microsoft CPUExecutionProvider", 12968279839987729832 + ], + [ + "InplaceClipGradNorm com.microsoft CPUExecutionProvider", + 10251631611024722504 ] ] diff --git a/onnxruntime/test/testdata/training_api/adamw.onnx b/onnxruntime/test/testdata/training_api/adamw.onnx index 6eb4e83a76dee3f8f19a108e7b74d32570bdbffa..eed8ef5e4883b70a98fedeabcb187530ccd97c66 100644 GIT binary patch delta 89 zcmbQseuQO$97{I~*W!u$<+(W1ixN{ZQ}arSC$liRvs#)N85m3sV3c4nwzROEoWuBm iTNI&YasZPbbCNFCWCKQN9fgc5UtZRb%y3^TO()@wSu@rB5IN*5qm+DGvdGjDHm(A-c4O>ue-ZefD5;t zd+w=!1o$)j1bzVPIOMA!>cd+5=IzXzH)GrAv54bO<519#uHZZjV&-Zg%zuJ^mf*&M z8mFE1lm(M>?M)={p^^dPfT;Xx;@!=7cNs^|vjP|^cA2u6GAd%FWlFUtQk|q4<6-Zl zGdO&j)85Z%Z~o9QKJA|NUN#$dBJiUCm5>n`8?l}wnn6)%mOv>Xl0*vbR~hzsAeE6L zeI~s~L@d?{4pgSpUq{h#E|T~pnzqx#H#L2ahh&0rZ_w!u507%fYeV?(F3~|0d|}eG z*dGr1&FVVbszW);WNP~Ck#%boS~b`qVRB9!&iSib%W74jIA)qS)9+tr_kGx-LP}=r z%O;S&YQQc_R1k{T`8K@jXYU@u-h5;}qvW%*Wj=%B3Kda13aAuHaE+XqB6p$GeQTtI zqoWO95rA!5s0any8ll|f8MeE?=AI+FS%RE*{8wXF2<|^jUZ!|>* providers) { OpTester test("InplaceClipGradNorm", 1, onnxruntime::kMSDomain); SeqTensors gradients_input; @@ -28,12 +28,10 @@ TEST(OptimizerTest, InplaceClipGradNorm) { clipped_gradients.AddTensor({5}, {3.7654f, 4.2361f, 4.7068f, 5.1775f, 5.6481f}); test.AddSeqOutput("clipped_gradients", clipped_gradients); - std::vector> providers; - providers.emplace_back(DefaultCudaExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, providers); } -TEST(OptimizerTest, InplaceClipGradNormNoClipping) { +void InplaceClipGradNormNoClippingTest(std::vector>* providers) { OpTester test("InplaceClipGradNorm", 1, onnxruntime::kMSDomain); SeqTensors gradients_input; @@ -51,9 +49,35 @@ TEST(OptimizerTest, InplaceClipGradNormNoClipping) { clipped_gradients.AddTensor({5}, {8.f, 9.f, 10.f, 11.f, 12.f}); test.AddSeqOutput("clipped_gradients", clipped_gradients); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, providers); +} + +} // namespace + +TEST(OptimizerTest, InplaceClipGradNorm_CPU) { + std::vector> providers; + providers.emplace_back(DefaultCpuExecutionProvider()); + InplaceClipGradNormTest(&providers); +} + +TEST(OptimizerTest, InplaceClipGradNormNoClipping_CPU) { + std::vector> providers; + providers.emplace_back(DefaultCpuExecutionProvider()); + InplaceClipGradNormNoClippingTest(&providers); +} + +#ifdef USE_CUDA + +TEST(OptimizerTest, InplaceClipGradNorm_CUDA) { std::vector> providers; providers.emplace_back(DefaultCudaExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + InplaceClipGradNormTest(&providers); +} + +TEST(OptimizerTest, InplaceClipGradNormNoClipping_CUDA) { + std::vector> providers; + providers.emplace_back(DefaultCudaExecutionProvider()); + InplaceClipGradNormNoClippingTest(&providers); } #endif diff --git a/orttraining/orttraining/training_api/optimizer.cc b/orttraining/orttraining/training_api/optimizer.cc index ddefa6333f..80916f9c3c 100644 --- a/orttraining/orttraining/training_api/optimizer.cc +++ b/orttraining/orttraining/training_api/optimizer.cc @@ -22,11 +22,50 @@ constexpr char GROUP_ZERO_NAME[] = "group0"; // TODO: Conolidate with frontend tooling const std::vector MOMENT_STATE_NAMES{"momentum0", "momentum1"}; -constexpr char LearningRateName[] = "learning_rate"; -constexpr char StepName[] = "step"; -constexpr char ParamsName[] = "params"; -constexpr char FirstOrderMomentsName[] = "first_order_moments"; -constexpr char SecondOrderMomentsName[] = "second_order_moments"; +constexpr std::array AdamWOptimizerInputs = { + "learning_rate", + "step", + "params", + "gradients", + "first_order_moments", + "second_order_moments"}; + +Status GraphInputsAreExpected(gsl::span actual_graph_inputs, + gsl::span expected_graph_inputs) { + const auto stringify = [](const auto& container) { + if (container.empty()) { + return std::string("[]"); + } + std::string container_str("["); + for (const auto& val : container) { + container_str += std::string(val) + ", "; + } + container_str.pop_back(); + container_str.back() = ']'; + + return container_str; + }; + + const auto construct_unexpected_input_status = [&stringify](const auto& actual_inputs, const auto& expected_inputs) { + std::ostringstream error_stream; + error_stream << "Invalid graph inputs." + << "\n\tExpected: " << stringify(expected_inputs) + << "\n\tActual: " << stringify(actual_inputs); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, error_stream.str()); + }; + + if (actual_graph_inputs.size() != expected_graph_inputs.size()) { + return construct_unexpected_input_status(actual_graph_inputs, expected_graph_inputs); + } + + for (size_t input_idx = 0; input_idx < expected_graph_inputs.size(); ++input_idx) { + if (actual_graph_inputs[input_idx] != expected_graph_inputs[input_idx]) { + return construct_unexpected_input_status(actual_graph_inputs, expected_graph_inputs); + } + } + + return Status::OK(); +} } // namespace @@ -48,38 +87,30 @@ Status Optimizer::GenerateMomentumNamedStates() { return Status::OK(); } -// Constructs the ortvalue inputs to be fed to the graph -// at each step +// Constructs the ortvalue inputs to be fed to the graph at each step Status Optimizer::ConstructInputs() { if (optimizer_type_ == OptimizerType::AdamW) { auto& param_named_optimizer_states = optimizer_state_.param_named_optimizer_states; - std::vector params, first_order_moments, second_order_moments; - // TODO: Change to tensor seq implementation once clip grad norm op - // that accepts tensor seq as input for gradients is complete. - std::vector grads; - - // Input names 0-4 are reserved for lr, step, params, first order moments, second order moments - // input names 5 onwards are all the gradient names. - // Collect all the inputs based on the gradient names order. - for (size_t i = 5; i < input_names_.size(); i++) { - std::string param_name; - if (utils::GetParamNameFromGradient(input_names_[i], param_name)) { - const auto named_parameter_it = named_parameters_.find(param_name); - ORT_ENFORCE(named_parameter_it != named_parameters_.end(), - "Unknown param: ", param_name, " for field: ", input_names_[i]); - - // Collect the gradients as ortvalues - grads.push_back(named_parameter_it->second->Gradient()); + std::vector params, grads, first_order_moments, second_order_moments; + // Collect all the non user defined inputs from the named_parameters_. + for (auto& [parameter_name, parameter] : named_parameters_) { + if (parameter->RequiresGrad()) { // Collect parameters and prepare for tensorseq creation - auto* param_tensor = named_parameter_it->second->Data().GetMutable(); + auto* param_tensor = parameter->Data().GetMutable(); params.emplace_back( Tensor(param_tensor->DataType(), param_tensor->Shape(), param_tensor->MutableDataRaw(), param_tensor->Location())); + // Collect gradients and prepare for tensorseq creation + auto* grad_tensor = parameter->Gradient().GetMutable(); + grads.emplace_back( + Tensor(grad_tensor->DataType(), grad_tensor->Shape(), + grad_tensor->MutableDataRaw(), grad_tensor->Location())); + // Collect first order moments and prepare for tensorseq creation - auto* first_order_moment_tensor = param_named_optimizer_states.at(param_name) + auto* first_order_moment_tensor = param_named_optimizer_states.at(parameter_name) .momentum_named_states.at(MOMENT_STATE_NAMES[0]) .GetMutable(); first_order_moments.emplace_back( @@ -87,20 +118,17 @@ Status Optimizer::ConstructInputs() { first_order_moment_tensor->MutableDataRaw(), first_order_moment_tensor->Location())); // Collect second order moments and prepare for tensorseq creation - auto* second_order_moment_tensor = param_named_optimizer_states.at(param_name) + auto* second_order_moment_tensor = param_named_optimizer_states.at(parameter_name) .momentum_named_states.at(MOMENT_STATE_NAMES[1]) .GetMutable(); second_order_moments.emplace_back( Tensor(second_order_moment_tensor->DataType(), second_order_moment_tensor->Shape(), second_order_moment_tensor->MutableDataRaw(), second_order_moment_tensor->Location())); - } else { - ORT_ENFORCE( - false, "This is an invalid graph. Optimizer graph contains unknown user input:", input_names_[i]); } } const auto tensorseq_inserter = [](auto& tensors, auto* inputs) { - ORT_ENFORCE(!tensors.empty(), "Tensors cannot be empty while building a tensor sequence."); + ORT_ENFORCE(!tensors.empty(), "Tensors vector cannot be empty while building a tensor sequence."); auto tensor_seq = std::make_unique(tensors.front().DataType()); tensor_seq->SetElements(std::move(tensors)); @@ -111,13 +139,9 @@ Status Optimizer::ConstructInputs() { // Add the params and moments as tensorseq ortvalues to inputs tensorseq_inserter(params, &inputs_); + tensorseq_inserter(grads, &inputs_); tensorseq_inserter(first_order_moments, &inputs_); tensorseq_inserter(second_order_moments, &inputs_); - - // Add the gradients as ortvalues to inputs - inputs_.insert(inputs_.end(), - std::make_move_iterator(grads.begin()), - std::make_move_iterator(grads.end())); } // Add other optimizer reordering logic here return Status::OK(); @@ -138,13 +162,9 @@ Optimizer::Optimizer(const std::string& optim_path_or_bytes, ORT_THROW_IF_ERROR(optim_sess_->Initialize()); utils::GetGraphInputOutputNames(optim_sess_, input_names_, output_names_); - ORT_ENFORCE(input_names_[0] == LearningRateName); // TODO: make this better - ORT_ENFORCE(input_names_[1] == StepName); // TODO: make this better - ORT_ENFORCE(input_names_[2] == ParamsName); // TODO: make this better if (optimizer_type_ == OptimizerType::AdamW) { - ORT_ENFORCE(input_names_[3] == FirstOrderMomentsName); // TODO: make this better - ORT_ENFORCE(input_names_[4] == SecondOrderMomentsName); // TODO: make this better + ORT_THROW_IF_ERROR(GraphInputsAreExpected(input_names_, AdamWOptimizerInputs)); ORT_THROW_IF_ERROR(GenerateMomentumNamedStates()); } else { diff --git a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc index 082be5eef1..c30851425e 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc @@ -86,6 +86,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float_float, ReduceAllL2); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, InplaceClipGradNorm); + // the kernels within the following ifdef are not included in a build with // --enable_training_ops but without --enable_training #ifdef ENABLE_TRAINING @@ -198,6 +200,8 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, + BuildKernelCreateInfo, + // the kernels within the following ifdef are not included in a build with // --enable_training_ops but without --enable_training #ifdef ENABLE_TRAINING diff --git a/orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.cc b/orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.cc new file mode 100644 index 0000000000..8d441e18d4 --- /dev/null +++ b/orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.cc @@ -0,0 +1,84 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.h" +#include "core/providers/cpu/math/element_wise_ops.h" +#include "core/providers/cpu/tensor/utils.h" +#include "core/providers/cpu/reduction/reduction_ops.h" + +namespace onnxruntime { +namespace contrib { + +namespace { + +constexpr float Epsilon = 0.000001f; + +template +T GetL2Norm(const TensorSeq& gradients) { + T l2_norm = 0; + for (const auto& tensor : gradients) { + l2_norm += + ReduceAggregatorSumSquare(tensor.Shape().Size(), *tensor.Data()).aggall(tensor.Data()); + } + return reduce_sqrt(l2_norm); +} + +template +void ClipGradNorm(T total_norm, T max_norm, TensorSeq& gradients) { + const T clip_coefficient = std::min(max_norm / (total_norm + static_cast(Epsilon)), static_cast(1.0f)); + + for (const auto& grad : gradients) { + auto& tensor = const_cast(grad); + MakeEigenArrayMap(tensor) *= clip_coefficient; + } +} + +Status PopulateOutput(OpKernelContext* ctx, const TensorSeq* gradients, TensorSeq* clipped_gradients) { + if (gradients == clipped_gradients) { + return Status::OK(); + } + + AllocatorPtr alloc; + ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc)); + + clipped_gradients->SetType(gradients->DataType()); + clipped_gradients->Reserve(gradients->Size()); + for (const auto& grad : *gradients) { + Tensor target_tensor(grad.DataType(), grad.Shape(), alloc); + CopyCpuTensor(&grad, &target_tensor); + clipped_gradients->Add(std::move(target_tensor)); // Add will check for type consistency + } + + return Status::OK(); +} + +} // namespace + +ONNX_OPERATOR_KERNEL_EX( + InplaceClipGradNorm, + kMSDomain, + 1, + kCpuExecutionProvider, + (*KernelDefBuilder::Create()) + .Alias(0, 0) /* Return updated gradients in-place */ + .TypeConstraint("S_GRAD", DataTypeImpl::AllFixedSizeSequenceTensorTypes()), + InplaceClipGradNorm); + +template +Status InplaceClipGradNorm::Compute(OpKernelContext* ctx) const { + const TensorSeq* gradients = ctx->Input(0); + + const T total_norm = GetL2Norm(*gradients); + + auto grads = const_cast(gradients); + ClipGradNorm(total_norm, max_norm_, *grads); + + // Populate the output sequence tensors. + TensorSeq* clipped_gradients = ctx->Output(0); + ORT_RETURN_IF_ERROR(PopulateOutput(ctx, gradients, clipped_gradients)); + + return Status::OK(); +} + +} // namespace contrib +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.h b/orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.h new file mode 100644 index 0000000000..d5da3445ed --- /dev/null +++ b/orttraining/orttraining/training_ops/cpu/optimizer/clip_grad_norm/clip_grad_norm.h @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" + +namespace onnxruntime { +namespace contrib { + +template +class InplaceClipGradNorm final : public OpKernel { + public: + InplaceClipGradNorm(const OpKernelInfo& info) + : OpKernel(info) { + info.GetAttrOrDefault("max_norm", &max_norm_, 1.0f); + info.GetAttrOrDefault("norm_type", &norm_type_, std::string("fro")); + ORT_ENFORCE(norm_type_ == "fro", "Given norm type ", norm_type_, " is not supported for InplaceClipGradNorm."); + } + + Status Compute(OpKernelContext* context) const override; + + private: + float max_norm_; + std::string norm_type_; +}; + +} // namespace contrib +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.cc b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.cc index 064c1b2486..9841f2b7f6 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.cc @@ -40,7 +40,7 @@ Status PopulateOutput(cudaStream_t stream, AllocatorPtr alloc, const TensorSeq* TensorSeq** clipped_gradients) { // If the output buffer is the same as the input buffer, the planner has // decided to reuse the buffer. No need to perform a memcpy in that case. - if (const_cast(gradients) == *clipped_gradients) { + if (gradients == *clipped_gradients) { return Status::OK(); } @@ -84,7 +84,7 @@ Status InplaceClipGradNorm::ComputeInternal(OpKernelContext* ctx) const { GetGroupedTensors(gradients, &tensor_sizes, &grouped_tensor_pointers); AllocatorPtr alloc; - ORT_ENFORCE(ctx->GetTempSpaceAllocator(&alloc).IsOK(), "InplaceClipGradNorm: Unable to get an allocator."); + ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc)); // Get frobenius norm for the grouped inputs float* total_norm = reinterpret_cast(alloc->Alloc(sizeof(float)));