saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

CUDA kernel for ClipGradNorm for TensorSeq gradients (#12412)

Browse source
This commit is contained in:
Baiju Meswani 2022-08-04 22:28:28 -07:00 committed by GitHub
parent 3e1b0ac4b3
commit a7d6290774
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 329 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
onnxruntime/core/framework/utils.cc
View file

@ -453,6 +453,11 @@ static void FinalizeFeedFetchCopyInfo(FeedsFetchesManager& feeds_fetches_manager
if (fetch.IsAllocated()) {
if (fetch.IsTensor()) {
fetch_alloc_info[i] = &fetch.Get<Tensor>().Location();
} else if (fetch.IsTensorSequence()) {
const auto& tensor_seq = fetch.Get<TensorSeq>();
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<SparseTensor>().Location();

7
onnxruntime/core/providers/cuda/multi_tensor/common.cuh
View file

@ -10,6 +10,7 @@
#include <vector>
#include "core/common/common.h"
#include "gsl/gsl-lite.hpp"
namespace onnxruntime {
namespace cuda {
@ -77,8 +78,8 @@ template <int TensorGroupSize, typename TMultiTensorFunctor, typename... TFuncto
void launch_multi_tensor_functor(
cudaStream_t stream,
const int chunk_size,
std::vector<int>& tensor_sizes,
std::vector<std::vector<void*>>& grouped_tensor_pointers,
gsl::span<int> tensor_sizes,
gsl::span<std::vector<void*>> 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
} // namespace onnxruntime

30
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -1387,6 +1387,36 @@ void RegisterTrainingOpSchemas() {
}
});
ONNX_CONTRIB_OPERATOR_SCHEMA(InplaceClipGradNorm)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetDoc(
"InplaceClipGradNorm operator, taking multiple gradients as inputs (seq<tensor>). "
"InplaceClipGradNorm should be used in conjunction with optimizers that accept seq<tensor> "
"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)

1
orttraining/orttraining/test/gradient/optimizer_ops_test.cc
View file

@ -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 {

62
orttraining/orttraining/test/training_ops/cuda/optimizer/clip_grad_norm_test.cc Normal file
View file

@ -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<float> 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<float>("gradients", gradients_input);
test.AddAttribute("max_norm", 12.f);
SeqTensors<float> 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<float>("clipped_gradients", clipped_gradients);
std::vector<std::unique_ptr<IExecutionProvider>> providers;
providers.emplace_back(DefaultCudaExecutionProvider());
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
}
TEST(OptimizerTest, InplaceClipGradNormNoClipping) {
OpTester test("InplaceClipGradNorm", 1, onnxruntime::kMSDomain);
SeqTensors<float> 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<float>("gradients", gradients_input);
test.AddAttribute("max_norm", 100.f);
SeqTensors<float> 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<float>("clipped_gradients", clipped_gradients);
std::vector<std::unique_ptr<IExecutionProvider>> providers;
providers.emplace_back(DefaultCudaExecutionProvider());
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
}
#endif
} // namespace test
} // namespace onnxruntime

4
orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc
View file

@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, ReduceAllL2)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, ReduceAllL2)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, InplaceClipGradNorm)>,
// P2P communication operators.
#if defined(ORT_USE_NCCL) || defined(USE_MPI)
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Send)>,

107
orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.cc Normal file
View file

@ -0,0 +1,107 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <memory>
#include <utility>
#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<int>* tensor_sizes,
InlinedVector<std::vector<void*>>* grouped_tensor_pointers) {
for (size_t i = 0; i < gradients->Size(); ++i) {
(*tensor_sizes)[i] = static_cast<int>(gradients->Get(i).Shape().Size());
(*grouped_tensor_pointers)[i] = {const_cast<float*>(gradients->Get(i).Data<float>())};
}
}
Status GetL2Norm(cudaStream_t stream, InlinedVector<int>& tensor_sizes,
InlinedVector<std::vector<void*>>& grouped_tensor_pointers, float** l2_norm) {
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(*l2_norm, 0, sizeof(float), stream));
MultiTensorReduceL2<float, float> multi_tensor_reduce_l2_functor;
launch_multi_tensor_functor<ClipGradNormGroupSize, MultiTensorReduceL2<float, float>>(
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<TensorSeq*>(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<Tensor> 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<TensorSeq>(0);
InlinedVector<int> tensor_sizes(gradients->Size());
// Need to use InlinedVector<std::vector<void*>> until the signature of launch_multi_tensor_functor
// is updated so that nested InlinedVector can be passed in.
InlinedVector<std::vector<void*>> 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<float*>(alloc->Alloc(sizeof(float)));
ORT_RETURN_IF_ERROR(GetL2Norm(Stream(), tensor_sizes, grouped_tensor_pointers, &total_norm));
// Perform gradient clipping
ClipGradNormFunctor<float> clip_grad_functor;
launch_multi_tensor_functor<ClipGradNormGroupSize, decltype(clip_grad_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<TensorSeq>(0);
ORT_RETURN_IF_ERROR(PopulateOutput(Stream(), alloc, gradients, &clipped_gradients));
return Status::OK();
}
} // namespace cuda
} // namespace onnxruntime

29
orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm.h Normal file
View file

@ -0,0 +1,29 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>
#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

64
orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.cu Normal file
View file

@ -0,0 +1,64 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <algorithm>
#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 <typename T>
__global__ void ClipGradNorm(
ChunkGroup<ClipGradNormGroupSize> 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<T*>(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<T>(gradients_chunk_ptr[i]) *
static_cast<T>(fminf(clip_coefficient, 1.0f));
}
}
template <typename T>
void ClipGradNormFunctor<T>::operator()(
cudaStream_t stream,
ChunkGroup<ClipGradNormGroupSize> 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<ClipGradNormGroupSize>::thread_count_per_block;
ClipGradNorm<T><<<num_blocks_per_grid, num_threads_per_block, 0, stream>>>(chunks, total_norm, epsilon, max_norm);
}
#define SPECIALIZE_CLIPGRADNORM_FUNCTOR(T) \
template void ClipGradNormFunctor<T>::operator()(cudaStream_t stream, \
ChunkGroup<ClipGradNormGroupSize> chunks, \
const float* total_norm, \
const float epsilon, \
const float max_norm); \
\
template __global__ void ClipGradNorm<T>(ChunkGroup<ClipGradNormGroupSize> chunks, \
const float* total_norm, \
const float epsilon, \
const float max_norm);
SPECIALIZE_CLIPGRADNORM_FUNCTOR(float);
} // namespace cuda
} // namespace onnxruntime

23
orttraining/orttraining/training_ops/cuda/optimizer/clip_grad_norm/clip_grad_norm_impl.h Normal file
View file

@ -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 <typename T>
struct ClipGradNormFunctor {
void operator()(cudaStream_t stream,
ChunkGroup<ClipGradNormGroupSize> chunks,
const float* l2_norm,
const float epsilon,
const float max_norm);
};
} // namespace cuda
} // namespace onnxruntime