From a7d62907743387003e0c341e8a22b232f6898a70 Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Thu, 4 Aug 2022 22:28:28 -0700 Subject: [PATCH] CUDA kernel for ClipGradNorm for TensorSeq gradients (#12412) --- onnxruntime/core/framework/utils.cc | 5 + .../providers/cuda/multi_tensor/common.cuh | 7 +- .../core/graph/training_op_defs.cc | 30 +++++ .../test/gradient/optimizer_ops_test.cc | 1 + .../cuda/optimizer/clip_grad_norm_test.cc | 62 ++++++++++ .../cuda/cuda_training_kernels.cc | 4 + .../clip_grad_norm/clip_grad_norm.cc | 107 ++++++++++++++++++ .../optimizer/clip_grad_norm/clip_grad_norm.h | 29 +++++ .../clip_grad_norm/clip_grad_norm_impl.cu | 64 +++++++++++ .../clip_grad_norm/clip_grad_norm_impl.h | 23 ++++ 10 files changed, 329 insertions(+), 3 deletions(-) create mode 100644 orttraining/orttraining/test/training_ops/cuda/optimizer/clip_grad_norm_test.cc create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.cc create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.h create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.cu create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.h diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 3cfd52f63b..8059de6c84 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -453,6 +453,11 @@ static void FinalizeFeedFetchCopyInfo(FeedsFetchesManager& feeds_fetches_manager if (fetch.IsAllocated()) { if (fetch.IsTensor()) { fetch_alloc_info[i] = &fetch.Get().Location(); + } else if (fetch.IsTensorSequence()) { + const auto& tensor_seq = fetch.Get(); + if (tensor_seq.Size() != std::size_t{0}) { + fetch_alloc_info[i] = &tensor_seq.Get(0).Location(); + } } else if (fetch.IsSparseTensor()) { #if !defined(DISABLE_SPARSE_TENSORS) fetch_alloc_info[i] = &fetch.Get().Location(); diff --git a/onnxruntime/core/providers/cuda/multi_tensor/common.cuh b/onnxruntime/core/providers/cuda/multi_tensor/common.cuh index 79cc6b7e1b..85e7e218bb 100644 --- a/onnxruntime/core/providers/cuda/multi_tensor/common.cuh +++ b/onnxruntime/core/providers/cuda/multi_tensor/common.cuh @@ -10,6 +10,7 @@ #include #include "core/common/common.h" +#include "gsl/gsl-lite.hpp" namespace onnxruntime { namespace cuda { @@ -77,8 +78,8 @@ template & tensor_sizes, - std::vector>& grouped_tensor_pointers, + gsl::span tensor_sizes, + gsl::span> grouped_tensor_pointers, TMultiTensorFunctor multipleTensorKernel, TFunctorParams&&... kernelParams) { // Check if 32-bit integer is enough. @@ -147,4 +148,4 @@ void launch_multi_tensor_functor( } } // namespace cuda -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 59a15d6391..9cd9df0f7b 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -1387,6 +1387,36 @@ void RegisterTrainingOpSchemas() { } }); + ONNX_CONTRIB_OPERATOR_SCHEMA(InplaceClipGradNorm) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc( + "InplaceClipGradNorm operator, taking multiple gradients as inputs (seq). " + "InplaceClipGradNorm should be used in conjunction with optimizers that accept seq " + "gradients as input, since this op takes a sequence of tensors as input and outputs a sequence of tensors " + "there by avoiding the need for SequenceConstruct (and making any unnecessary copy)." + "Please note that the gradient clipping happens inplace.") + .Input(0, "gradients", "Sequence of gradients computed in this iteration.", "S_GRAD") + .Output(0, "clipped_gradients", "Gradients after being clipped as per given inputs and attributes.", "S_GRAD") + .Attr( + "max_norm", + "Coefficient of previously accumulated gradient in running average.", + AttributeProto::FLOAT, + 1.0f) + .Attr( + "norm_type", + "Type of normalization to perform during execution of clip grad norm." + "Currently, the only norm supported is the frobenius norm (which is also the default).", + AttributeProto::STRING, + std::string("fro")) + .TypeConstraint( + "S_GRAD", + {"seq(tensor(float16))", "seq(tensor(float))", "seq(tensor(double))"}, + "Constrain gradients' types.") + .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { + propagateShapeAndTypeFromFirstInput(ctx); + }); + ONNX_CONTRIB_OPERATOR_SCHEMA_ELSEWHERE(LambOptimizer, RegisterLambOpSchema); ONNX_CONTRIB_OPERATOR_SCHEMA(InPlaceAccumulator) diff --git a/orttraining/orttraining/test/gradient/optimizer_ops_test.cc b/orttraining/orttraining/test/gradient/optimizer_ops_test.cc index 35fd4bc7d9..c100730aac 100644 --- a/orttraining/orttraining/test/gradient/optimizer_ops_test.cc +++ b/orttraining/orttraining/test/gradient/optimizer_ops_test.cc @@ -5,6 +5,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" namespace onnxruntime { namespace test { diff --git a/orttraining/orttraining/test/training_ops/cuda/optimizer/clip_grad_norm_test.cc b/orttraining/orttraining/test/training_ops/cuda/optimizer/clip_grad_norm_test.cc new file mode 100644 index 0000000000..65b0b8d0bf --- /dev/null +++ b/orttraining/orttraining/test/training_ops/cuda/optimizer/clip_grad_norm_test.cc @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { + +#ifdef USE_CUDA + +TEST(OptimizerTest, InplaceClipGradNorm) { + OpTester test("InplaceClipGradNorm", 1, onnxruntime::kMSDomain); + + SeqTensors gradients_input; + gradients_input.AddTensor({3}, {1.f, 2.f, 3.f}); + gradients_input.AddTensor({4}, {4.f, 5.f, 6.f, 7.f}); + gradients_input.AddTensor({5}, {8.f, 9.f, 10.f, 11.f, 12.f}); + + test.AddSeqInput("gradients", gradients_input); + + test.AddAttribute("max_norm", 12.f); + + SeqTensors clipped_gradients; + clipped_gradients.AddTensor({3}, {0.4707f, 0.9414f, 1.4120f}); + clipped_gradients.AddTensor({4}, {1.8827f, 2.3534f, 2.8241f, 3.2948f}); + 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(OptimizerTest, InplaceClipGradNormNoClipping) { + OpTester test("InplaceClipGradNorm", 1, onnxruntime::kMSDomain); + + SeqTensors gradients_input; + gradients_input.AddTensor({3}, {1.f, 2.f, 3.f}); + gradients_input.AddTensor({4}, {4.f, 5.f, 6.f, 7.f}); + gradients_input.AddTensor({5}, {8.f, 9.f, 10.f, 11.f, 12.f}); + + test.AddSeqInput("gradients", gradients_input); + + test.AddAttribute("max_norm", 100.f); + + SeqTensors clipped_gradients; + clipped_gradients.AddTensor({3}, {1.f, 2.f, 3.f}); + clipped_gradients.AddTensor({4}, {4.f, 5.f, 6.f, 7.f}); + clipped_gradients.AddTensor({5}, {8.f, 9.f, 10.f, 11.f, 12.f}); + test.AddSeqOutput("clipped_gradients", clipped_gradients); + + std::vector> providers; + providers.emplace_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +#endif + +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index e2f3998f6d..cf5a7cf7c1 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -209,6 +209,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, ReduceAllL2); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, ReduceAllL2); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, InplaceClipGradNorm); + #if defined(ORT_USE_NCCL) || defined(USE_MPI) // P2P communication operators. class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Send); @@ -434,6 +436,8 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + // P2P communication operators. #if defined(ORT_USE_NCCL) || defined(USE_MPI) BuildKernelCreateInfo, 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 new file mode 100644 index 0000000000..064c1b2486 --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.cc @@ -0,0 +1,107 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.h" +#include "orttraining/training_ops/cuda/reduction/reduction_all_impl.h" +#include "orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.h" + +namespace onnxruntime { +namespace cuda { + +namespace { + +constexpr int ChunkSize = 2048 * 32; +constexpr float Epsilon = 0.000001f; + +void GetGroupedTensors(const TensorSeq* gradients, InlinedVector* tensor_sizes, + InlinedVector>* grouped_tensor_pointers) { + for (size_t i = 0; i < gradients->Size(); ++i) { + (*tensor_sizes)[i] = static_cast(gradients->Get(i).Shape().Size()); + (*grouped_tensor_pointers)[i] = {const_cast(gradients->Get(i).Data())}; + } +} + +Status GetL2Norm(cudaStream_t stream, InlinedVector& tensor_sizes, + InlinedVector>& grouped_tensor_pointers, float** l2_norm) { + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(*l2_norm, 0, sizeof(float), stream)); + MultiTensorReduceL2 multi_tensor_reduce_l2_functor; + launch_multi_tensor_functor>( + stream, ChunkSize, tensor_sizes, grouped_tensor_pointers, multi_tensor_reduce_l2_functor, *l2_norm); + + ScalarSqrt(stream, *l2_norm, *l2_norm); + + return Status::OK(); +} + +Status PopulateOutput(cudaStream_t stream, AllocatorPtr alloc, const TensorSeq* gradients, + 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) { + return Status::OK(); + } + + (*clipped_gradients)->SetType(gradients->DataType()); + (*clipped_gradients)->Reserve(gradients->Size()); + for (size_t gradient_idx = 0; gradient_idx < gradients->Size(); ++gradient_idx) { + const Tensor& source_tensor = gradients->Get(gradient_idx); + std::unique_ptr target_tensor = Tensor::Create(source_tensor.DataType(), + source_tensor.Shape(), alloc); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(target_tensor->MutableDataRaw(), + source_tensor.DataRaw(), + source_tensor.SizeInBytes(), + cudaMemcpyDeviceToDevice, stream)); + (*clipped_gradients)->Add(std::move(*target_tensor)); // Add will check for type consistency + } + + return Status::OK(); +} + +} // namespace + +ONNX_OPERATOR_KERNEL_EX( + InplaceClipGradNorm, + kMSDomain, + 1, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .Alias(0, 0) // Return updated gradients in-place + // Note that the allocation planner may or may not plan for the output + // buffer to be the same as the input buffer. + .TypeConstraint("S_GRAD", DataTypeImpl::AllFixedSizeSequenceTensorTypes()), + InplaceClipGradNorm); + +Status InplaceClipGradNorm::ComputeInternal(OpKernelContext* ctx) const { + // Prepare the inputs + const TensorSeq* gradients = ctx->Input(0); + InlinedVector tensor_sizes(gradients->Size()); + // Need to use InlinedVector> until the signature of launch_multi_tensor_functor + // is updated so that nested InlinedVector can be passed in. + InlinedVector> grouped_tensor_pointers(gradients->Size()); + GetGroupedTensors(gradients, &tensor_sizes, &grouped_tensor_pointers); + + AllocatorPtr alloc; + ORT_ENFORCE(ctx->GetTempSpaceAllocator(&alloc).IsOK(), "InplaceClipGradNorm: Unable to get an allocator."); + + // Get frobenius norm for the grouped inputs + float* total_norm = reinterpret_cast(alloc->Alloc(sizeof(float))); + ORT_RETURN_IF_ERROR(GetL2Norm(Stream(), tensor_sizes, grouped_tensor_pointers, &total_norm)); + + // Perform gradient clipping + ClipGradNormFunctor clip_grad_functor; + launch_multi_tensor_functor( + Stream(), ChunkSize, tensor_sizes, grouped_tensor_pointers, clip_grad_functor, total_norm, + Epsilon, max_norm_); + + // Populate the output sequence tensors. + TensorSeq* clipped_gradients = ctx->Output(0); + ORT_RETURN_IF_ERROR(PopulateOutput(Stream(), alloc, gradients, &clipped_gradients)); + + return Status::OK(); +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.h b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.h new file mode 100644 index 0000000000..df352ab65a --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.h @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include "core/common/common.h" +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace cuda { + +class InplaceClipGradNorm final : public CudaKernel { + public: + InplaceClipGradNorm(const OpKernelInfo& info) : CudaKernel(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 ComputeInternal(OpKernelContext* context) const override; + + private: + float max_norm_; + std::string norm_type_; +}; + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.cu b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.cu new file mode 100644 index 0000000000..e6ddfcaa9b --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.cu @@ -0,0 +1,64 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include "core/providers/cuda/cuda_common.h" +#include "orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.h" + +namespace onnxruntime { +namespace cuda { + +template +__global__ void ClipGradNorm( + ChunkGroup chunks, + const float* total_norm, + const float epsilon, + const float max_norm) { + const int tensor_idx = chunks.block_index_to_tensor_group_index[blockIdx.x]; + const int tensor_size = chunks.tensor_sizes[tensor_idx]; + + const int chunk_start_idx = chunks.block_index_to_chunk_start_index[blockIdx.x]; + // chunk_size is chunks.chunk_size if the loaded chunk is full. Otherwise (this + // chunk is the last one in the source tensor), the actual size is determined + // by the bound of the source tensor. + const int chunk_size = min(tensor_size, chunk_start_idx + chunks.chunk_size) - chunk_start_idx; + + T* gradients_chunk_ptr = static_cast(chunks.tensor_ptrs[0][tensor_idx]) + chunk_start_idx; + +#pragma unroll(4) + for (int i = threadIdx.x; i < chunk_size; i += blockDim.x) { + float clip_coefficient = max_norm / (*total_norm + epsilon); + gradients_chunk_ptr[i] = static_cast(gradients_chunk_ptr[i]) * + static_cast(fminf(clip_coefficient, 1.0f)); + } +} + +template +void ClipGradNormFunctor::operator()( + cudaStream_t stream, + ChunkGroup chunks, + const float* total_norm, + const float epsilon, + const float max_norm) { + const int num_blocks_per_grid = chunks.chunk_count; + const int num_threads_per_block = ChunkGroup::thread_count_per_block; + + ClipGradNorm<<>>(chunks, total_norm, epsilon, max_norm); +} + +#define SPECIALIZE_CLIPGRADNORM_FUNCTOR(T) \ + template void ClipGradNormFunctor::operator()(cudaStream_t stream, \ + ChunkGroup chunks, \ + const float* total_norm, \ + const float epsilon, \ + const float max_norm); \ + \ + template __global__ void ClipGradNorm(ChunkGroup chunks, \ + const float* total_norm, \ + const float epsilon, \ + const float max_norm); + +SPECIALIZE_CLIPGRADNORM_FUNCTOR(float); + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.h b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.h new file mode 100644 index 0000000000..5d75cbf1e1 --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/multi_tensor/common.cuh" + +namespace onnxruntime { +namespace cuda { + +constexpr int ClipGradNormGroupSize = 1; + +template +struct ClipGradNormFunctor { + void operator()(cudaStream_t stream, + ChunkGroup chunks, + const float* l2_norm, + const float epsilon, + const float max_norm); +}; + +} // namespace cuda +} // namespace onnxruntime