From 8fc3037ff44efad1b96f56585380fc6986c68fd3 Mon Sep 17 00:00:00 2001 From: pengwa Date: Fri, 30 Jun 2023 08:36:06 +0800 Subject: [PATCH] Support SCELossInternal/SCELossInternalGrad run with larger sized input (#16363) ### Support SCELoss/SCELossGrad run with larger sized input #### Motivation and Context: Run bigger batch size for Bloom model. For Bloom560M model, ORT has potential to run bigger batch size from initialally 6 to now 10. SCELoss/SCELossGrad's input size is Bsz X 1023 * 250680. When Bsz is bigger than 8, totoal element count cannot be represented by int32_t, which those kernels are using to passing total elem count. There is silent overflow causing other indirectly exceptions, or wrong mistake without errors. #### Changes in this PR - For SCELossInternal/SCELossGradInternal CUDA kernels, use uint64_t if total element count is bigger than int32::max() to pass all element count and element index for the ops mentioned above. - For SCELossInternal/SCELossGradInternal CPU kernels, - always use uint64_t to pass the element count. - update the Eigen functions involved in the two kernels' implementations, to use `ptrdiff_t` to pass element count instead of original `int`. - Parallelize SCELossInternal/SCELossGradInternal CPU kernels, otherwise, it is super slow when handling so many elements. - Others changed needed: - Add `CompareOrtValueNumerals` to compare two OrtValue with different data types (float or float16), without caller explicitly converting to the lower-precision data types. The comparison is also done in parallel, which reduce the comparsion time for the large UT case from 22s to ~1.6s. - The check of `IsResultCloselyMatch` is buggy for nan/inf cases, so fix the bugs. - The cross entropy tests are running CPU base line with float, then the result is used to compare with float16 results of CUDA runs. But there is precision issue when we check the results. Because the randomized input data is represented in float, CPU use it directly, but CUDA use a float16 version of it, so there is precision diff between the inputs, as the test data count increases, it make the results fail even on 1e-2. The fix is: generate data in float16, convert to float for CPU run, directly use float16 for CUDA runs. When compare the output, cast back CPU float to float16 then compare with CUDA outputs. - `RandomValueGenerator ` for the large size take about ~20second, so `ParallelRandomValueGenerator ` is added to random input in parallel, it takes about <2s for preparing input data. #### Non-goals `SoftmaxCrossEntropyLoss` && `SoftmaxCrossEntropyLossGrad` is not covered in this PR --- .../cuda/cu_inc/elementwise_impl.cuh | 18 +- .../providers/cuda/shared_inc/fast_divmod.h | 32 +- .../providers/rocm/shared_inc/fast_divmod.h | 32 +- onnxruntime/core/util/math.h | 14 +- onnxruntime/core/util/math_cpu.cc | 34 +- .../test/common/tensor_op_test_utils.h | 87 ++++- .../test/providers/provider_test_utils.h | 6 +- onnxruntime/test/util/compare_ortvalue.cc | 126 ++++++- .../test/util/include/test/compare_ortvalue.h | 28 +- .../training_ops/cuda/cross_entropy_test.cc | 314 +++++++++--------- .../training_ops/cpu/loss/cross_entropy.cc | 38 ++- .../training_ops/cpu/loss/cross_entropy.h | 3 +- .../cpu/loss/softmax_cross_entropy_loss.cc | 155 +++++---- .../loss/softmax_cross_entropy_loss_impl.cc | 6 +- .../loss/softmax_cross_entropy_loss_impl.cu | 58 ++-- .../cuda/loss/softmaxcrossentropy_impl.cu | 14 +- .../cuda/quantization/fake_quant_impl.cu | 4 +- 17 files changed, 646 insertions(+), 323 deletions(-) diff --git a/onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh b/onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh index 790a612790..07a65bd252 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh @@ -16,12 +16,12 @@ constexpr int kElementsPerThread = GridDim::maxElementsPerThread; constexpr int kThreadsPerBlock = GridDim::maxThreadsPerBlock; #endif -template -__global__ void ElementwiseKernel(T* output_data, const FuncT functor, CUDA_LONG N) { - CUDA_LONG start = kElementsPerThread * kThreadsPerBlock * blockIdx.x + threadIdx.x; +template +__global__ void ElementwiseKernel(T* output_data, const FuncT functor, TIndex N) { + TIndex start = kElementsPerThread * kThreadsPerBlock * blockIdx.x + threadIdx.x; T value[kElementsPerThread]; - CUDA_LONG id = start; + TIndex id = start; #pragma unroll for (int i = 0; i < kElementsPerThread; ++i) { if (id < N) { @@ -40,12 +40,12 @@ __global__ void ElementwiseKernel(T* output_data, const FuncT functor, CUDA_LONG } } -template -void LaunchElementwiseKernel(cudaStream_t stream, T* output_data, const FuncT& functor, size_t output_size) { +template +void LaunchElementwiseKernel(cudaStream_t stream, T* output_data, const FuncT& functor, TIndex output_size) { if (output_size == 0) return; - CUDA_LONG N = static_cast(output_size); - int blocksPerGrid = CeilDiv(N, kThreadsPerBlock * kElementsPerThread); - ElementwiseKernel<<>>(output_data, functor, N); + TIndex N = output_size; + int blocksPerGrid = static_cast(CeilDiv(N, static_cast(kThreadsPerBlock * kElementsPerThread))); + ElementwiseKernel<<>>(output_data, functor, N); } } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/shared_inc/fast_divmod.h b/onnxruntime/core/providers/cuda/shared_inc/fast_divmod.h index 89a4a0c1d5..74a937791a 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/fast_divmod.h +++ b/onnxruntime/core/providers/cuda/shared_inc/fast_divmod.h @@ -14,12 +14,38 @@ namespace onnxruntime { namespace cuda { +// DivMod is a helper class for integer division and modulo operation. +// There is a fast version for int type and a slow version for other type. +template +struct DivMod { + DivMod(T d = 1) { + d_ = d == 0 ? 1 : d; + ORT_ENFORCE(d_ >= 1 && d_ <= std::numeric_limits::max()); + } + + __host__ __device__ inline T div(T n) const { + return n / d_; + } + + __host__ __device__ inline T mod(T n) const { + return n % d_; + } + + __host__ __device__ inline void divmod(T n, T& q, T& r) const { + q = div(n); + r = n - q * d_; + } + + T d_; // divisor +}; + // The code below is based on section 4 Unsigned division of paper https://gmplib.org/~tege/divcnst-pldi94.pdf // In current ORT, fast_divmod is used for calculating the position of a element in tensor, // so unsigned integer division from the paper is good enough for ORT. The advantage is that div is very simple, // then GPU compiler can do loop unroll easilly when divmod is called in a loop. -struct fast_divmod { - fast_divmod(int d = 1) { +template <> +struct DivMod { + DivMod(int d = 1) { d_ = d == 0 ? 1 : d; ORT_ENFORCE(d_ >= 1 && d_ <= static_cast(std::numeric_limits::max())); @@ -58,5 +84,7 @@ struct fast_divmod { uint32_t l_; // l_ = ceil(log2(d_)) }; +using fast_divmod = DivMod; // Keep the old name for backward compatibility. + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/shared_inc/fast_divmod.h b/onnxruntime/core/providers/rocm/shared_inc/fast_divmod.h index 0043e4f153..83ca0a443c 100644 --- a/onnxruntime/core/providers/rocm/shared_inc/fast_divmod.h +++ b/onnxruntime/core/providers/rocm/shared_inc/fast_divmod.h @@ -14,12 +14,38 @@ namespace onnxruntime { namespace rocm { +// DivMod is a helper class for integer division and modulo operation. +// There is a fast version for int type and a slow version for other type. +template +struct DivMod { + DivMod(T d = 1) { + d_ = d == 0 ? 1 : d; + ORT_ENFORCE(d_ >= 1 && d_ <= std::numeric_limits::max()); + } + + __host__ __device__ inline T div(T n) const { + return n / d_; + } + + __host__ __device__ inline T mod(T n) const { + return n % d_; + } + + __host__ __device__ inline void divmod(T n, T& q, T& r) const { + q = div(n); + r = n - q * d_; + } + + T d_; // divisor +}; + // The code below is based on section 4 Unsigned division of paper https://gmplib.org/~tege/divcnst-pldi94.pdf // In current ORT, fast_divmod is used for calculating the position of a element in tensor, // so unsigned integer division from the paper is good enough for ORT. The advantage is that div is very simple, // then GPU compiler can do loop unroll easilly when divmod is called in a loop. -struct fast_divmod { - fast_divmod(int d = 1) { +template <> +struct DivMod { + DivMod(int d = 1) { d_ = d == 0 ? 1 : d; ORT_ENFORCE(d_ >= 1 && d_ <= static_cast(std::numeric_limits::max())); @@ -58,5 +84,7 @@ struct fast_divmod { uint32_t l_; // l_ = ceil(log2(d_)) }; +using fast_divmod = DivMod; // Keep the old name for backward compatibility. + } // namespace rocm } // namespace onnxruntime diff --git a/onnxruntime/core/util/math.h b/onnxruntime/core/util/math.h index dbbec2ae0f..c30d841dd8 100644 --- a/onnxruntime/core/util/math.h +++ b/onnxruntime/core/util/math.h @@ -53,15 +53,15 @@ enum StorageOrder { namespace math { template -void Exp(int N, const T* x, T* y, Provider* provider); +void Exp(ptrdiff_t N, const T* x, T* y, Provider* provider); template -void Log(int N, const T* x, T* y, Provider* provider); +void Log(ptrdiff_t N, const T* x, T* y, Provider* provider); template -void Sqr(int N, const T* x, T* y, Provider* provider); +void Sqr(ptrdiff_t N, const T* x, T* y, Provider* provider); #define DECLARE_BINARY_OP(name) \ template \ - void name(int N, const T* a, const T* b, T* y, Provider* provider); \ + void name(ptrdiff_t N, const T* a, const T* b, T* y, Provider* provider); \ template \ void name##ToRow(int M, int N, const T* a, const T* b, T* y, Provider* provider); \ template \ @@ -90,16 +90,16 @@ void RowwiseSum(int N, int D, const T* x, T* y, // Sum of vector x, and writes the result to a single value y. template -void Sum(int N, const T* x, T* y, Provider* provider); +void Sum(ptrdiff_t N, const T* x, T* y, Provider* provider); template -void Scale(int N, float alpha, const T* x, T* y, Provider* provider); +void Scale(ptrdiff_t N, float alpha, const T* x, T* y, Provider* provider); // Different from the Scale function above, if alpha is passed in // as a pointer, we will assume that it lives on the correct execution provider, // for example on GPU. template -void Scale(int N, const float* alpha, const T* x, T* y, Provider* provider); +void Scale(ptrdiff_t N, const float* alpha, const T* x, T* y, Provider* provider); template void MatMul( diff --git a/onnxruntime/core/util/math_cpu.cc b/onnxruntime/core/util/math_cpu.cc index b4554e0f0d..f11fe5a701 100644 --- a/onnxruntime/core/util/math_cpu.cc +++ b/onnxruntime/core/util/math_cpu.cc @@ -250,10 +250,10 @@ template void Gemv(const CBLAS_TRANSPOSE TransA, int M, int SPECIALIZED_AXPY(float) #undef SPECIALIZED_AXPY -#define DELEGATE_SIMPLE_UNARY_FUNCTION(T, Funcname, expr) \ - template <> \ - void Funcname(int N, const T* x, T* y, CPUMathUtil*) { \ - EigenVectorMap(y, N) = ConstEigenVectorMap(x, N).array().expr(); \ +#define DELEGATE_SIMPLE_UNARY_FUNCTION(T, Funcname, expr) \ + template <> \ + void Funcname(ptrdiff_t N, const T* x, T* y, CPUMathUtil*) { \ + EigenVectorMap(y, N) = ConstEigenVectorMap(x, N).array().expr(); \ } DELEGATE_SIMPLE_UNARY_FUNCTION(float, Exp, exp) DELEGATE_SIMPLE_UNARY_FUNCTION(double, Exp, exp) @@ -263,7 +263,7 @@ DELEGATE_SIMPLE_UNARY_FUNCTION(float, Sqr, square) #define EIGEN_SIMPLE_BINARY_FUNCTION(T, Funcname, expr) \ template <> \ - void Funcname(int N, const T* a, const T* b, T* y, CPUMathUtil*) { \ + void Funcname(ptrdiff_t N, const T* a, const T* b, T* y, CPUMathUtil*) { \ EigenVectorMap(y, N) = ConstEigenVectorMap(a, N).array() expr ConstEigenVectorMap(b, N).array(); \ } @@ -906,10 +906,10 @@ SPECIALIZED_ROWWISESUM(int64_t) SPECIALIZED_ROWWISESUM(double) #undef SPECIALIZED_ROWWISESUM -#define SPECIALIZED_SUM(T) \ - template <> \ - void Sum(int N, const T* x, T* y, CPUMathUtil* /* unused */) { \ - *y = ConstEigenVectorMap(x, N).sum(); \ +#define SPECIALIZED_SUM(T) \ + template <> \ + void Sum(ptrdiff_t N, const T* x, T* y, CPUMathUtil* /* unused */) { \ + *y = ConstEigenVectorMap(x, N).sum(); \ } SPECIALIZED_SUM(float); @@ -918,14 +918,14 @@ SPECIALIZED_SUM(int64_t); #undef SPECIALIZED_SUM -#define SPECIALIZED_SCALE(T) \ - template <> \ - void Scale(int n, float alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \ - EigenVectorMap(y, n) = ConstEigenVectorMap(x, n) * alpha; \ - } \ - template <> \ - void Scale(int n, const float* alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \ - EigenVectorMap(y, n) = ConstEigenVectorMap(x, n) * (*alpha); \ +#define SPECIALIZED_SCALE(T) \ + template <> \ + void Scale(ptrdiff_t n, float alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \ + EigenVectorMap(y, n) = ConstEigenVectorMap(x, n) * alpha; \ + } \ + template <> \ + void Scale(ptrdiff_t n, const float* alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \ + EigenVectorMap(y, n) = ConstEigenVectorMap(x, n) * (*alpha); \ } SPECIALIZED_SCALE(float) #undef SPECIALIZED_SCALE diff --git a/onnxruntime/test/common/tensor_op_test_utils.h b/onnxruntime/test/common/tensor_op_test_utils.h index a6369ed0e6..72cee65c51 100644 --- a/onnxruntime/test/common/tensor_op_test_utils.h +++ b/onnxruntime/test/common/tensor_op_test_utils.h @@ -15,6 +15,8 @@ #include "core/common/type_utils.h" #include "core/framework/tensor.h" #include "core/util/math.h" +#include "core/platform/threadpool.h" +#include "core/util/thread_utils.h" #include "test/common/random_generator.h" namespace onnxruntime { @@ -160,21 +162,100 @@ inline void CheckTensor(const Tensor& expected_tensor, const Tensor& output_tens "] did not match run output shape [" + output_tensor.Shape().ToString() + "]"); - ASSERT_TRUE(expected_tensor.DataType() == DataTypeImpl::GetType()) << "Compare with non float number is not supported yet. "; + ASSERT_TRUE(expected_tensor.DataType() == DataTypeImpl::GetType()) + << "Compare with non float number is not supported yet. "; auto expected = expected_tensor.Data(); auto output = output_tensor.Data(); for (auto i = 0; i < expected_tensor.Shape().Size(); ++i) { const auto expected_value = expected[i], actual_value = output[i]; if (std::isnan(expected_value)) { - ASSERT_TRUE(std::isnan(actual_value)) << "value mismatch at index " << i << "; expected is NaN, actual is not NaN"; + ASSERT_TRUE(std::isnan(actual_value)) << "value mismatch at index " << i + << "; expected is NaN, actual is not NaN"; } else if (std::isinf(expected_value)) { ASSERT_EQ(expected_value, actual_value) << "value mismatch at index " << i; } else { double diff = fabs(expected_value - actual_value); - ASSERT_TRUE(diff <= (atol + rtol * fabs(expected_value))) << "value mismatch at index " << i << "; expected: " << expected_value << ", actual: " << actual_value; + ASSERT_TRUE(diff <= (atol + rtol * fabs(expected_value))) << "value mismatch at index " << i << "; expected: " + << expected_value << ", actual: " << actual_value; } } } +class ParallelRandomValueGenerator { + public: + using RandomEngine = std::default_random_engine; + using RandomSeedType = RandomEngine::result_type; + + ParallelRandomValueGenerator(RandomSeedType base_seed) + : base_seed_{base_seed} { + } + + // Random values generated are in the range [min, max). + template + typename std::enable_if< + std::is_floating_point::value, + std::vector>::type + Uniform(gsl::span dims, TFloat min, TFloat max) { + OrtThreadPoolParams to; + to.thread_pool_size = 16; + auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to, + concurrency::ThreadPoolType::INTRA_OP); + static double cost = 1; + RandomSeedType base_seed = base_seed_; + std::vector val(detail::SizeFromDims(dims)); + concurrency::ThreadPool::TryParallelFor( + tp.get(), val.size(), cost, + [&min, &max, &base_seed, &val]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + RandomSeedType seed = base_seed; + auto new_seed = static_cast(base_seed) + begin; + if (new_seed < static_cast(std::numeric_limits::max())) + seed = static_cast(new_seed); + RandomEngine generator{seed}; + std::uniform_real_distribution distribution(min, max); + for (std::ptrdiff_t di = begin; di != end; ++di) { + val[di] = distribution(generator); + } + }); + + return val; + } + + // Random values generated are in the range [min, max). + template + typename std::enable_if< + std::is_same_v, + std::vector>::type + Uniform(gsl::span dims, float min, float max) { + OrtThreadPoolParams to; + to.thread_pool_size = 16; + auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to, + concurrency::ThreadPoolType::INTRA_OP); + static double cost = 1; + RandomSeedType base_seed = base_seed_; + std::vector val(detail::SizeFromDims(dims)); + concurrency::ThreadPool::TryParallelFor( + tp.get(), val.size(), cost, + [&min, &max, &base_seed, &val]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + RandomSeedType seed = base_seed; + auto new_seed = static_cast(base_seed) + begin; + if (new_seed < static_cast(std::numeric_limits::max())) + seed = static_cast(new_seed); + RandomEngine generator{seed}; + std::uniform_real_distribution distribution(min, max); + for (std::ptrdiff_t di = begin; di != end; ++di) { + val[di] = TFloat16(math::floatToHalf(distribution(generator))); + ; + } + }); + + return val; + } + + private: + const RandomSeedType base_seed_; +}; + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index 4318fd0362..87dc5980e9 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -9,13 +9,13 @@ namespace onnxruntime { namespace test { -inline void ConvertFloatToMLFloat16(const float* f_datat, MLFloat16* h_data, int input_size) { +inline void ConvertFloatToMLFloat16(const float* f_datat, MLFloat16* h_data, size_t input_size) { auto in_vector = ConstEigenVectorMap(f_datat, input_size); auto output_vector = EigenVectorMap(static_cast(static_cast(h_data)), input_size); output_vector = in_vector.template cast(); } -inline void ConvertMLFloat16ToFloat(const MLFloat16* h_data, float* f_data, int input_size) { +inline void ConvertMLFloat16ToFloat(const MLFloat16* h_data, float* f_data, size_t input_size) { auto in_vector = ConstEigenVectorMap(static_cast(static_cast(h_data)), input_size); auto output_vector = EigenVectorMap(f_data, input_size); @@ -24,7 +24,7 @@ inline void ConvertMLFloat16ToFloat(const MLFloat16* h_data, float* f_data, int inline std::vector FloatsToMLFloat16s(const std::vector& f) { std::vector m(f.size()); - ConvertFloatToMLFloat16(f.data(), m.data(), static_cast(f.size())); + ConvertFloatToMLFloat16(f.data(), m.data(), f.size()); return m; } diff --git a/onnxruntime/test/util/compare_ortvalue.cc b/onnxruntime/test/util/compare_ortvalue.cc index 78bbd46f5a..501f7b3a42 100644 --- a/onnxruntime/test/util/compare_ortvalue.cc +++ b/onnxruntime/test/util/compare_ortvalue.cc @@ -2,8 +2,11 @@ // Licensed under the MIT License. #include "test/compare_ortvalue.h" + #include #include +#include +#include #ifdef __GNUC__ #pragma GCC diagnostic push @@ -25,11 +28,12 @@ #pragma GCC diagnostic pop #endif -#include "core/graph/onnx_protobuf.h" #include "core/framework/tensorprotoutils.h" #include "core/framework/utils.h" #include "core/framework/TensorSeq.h" +#include "core/graph/onnx_protobuf.h" #include +#include "core/util/math.h" using namespace onnxruntime; @@ -61,13 +65,34 @@ const char* ElementTypeToString(MLDataType type) { return DataTypeImpl::ToString(type); } +/** + * @brief Check if two values are closely matched with given tolerance. + + * Definition of closely match: + * > If any of actual_value and expected_value is nan, actual_value and expected_value must both be nan. + * > If any of actual_value and expected_value is inf, then actual_value and expected_value + * must both be inf with same sign. + * > Otherwise, diff <= tol. + + * @param actual_value The value to be checked. + * @param expected_value The baseline value used to check against. + * @param diff The absolute difference calculated by the caller from actual_value and expected_value. + * @param tol The absolute tolerance. + * @return True when closely matched; False otherwise. + */ template -bool IsResultCloselyMatch(const T& outvalue, const T& expected_value, const double diff, const double tol) { - if (diff > tol) return false; - if (std::isnan(diff) && !(std::isnan(outvalue) && std::isnan(expected_value)) && - !(std::isinf(outvalue) && std::isinf(expected_value))) - return false; - return true; +bool IsResultCloselyMatch(const T& actual_value, const T& expected_value, const double diff, const double tol) { + if (std::isnan(actual_value) || std::isnan(expected_value)) + return std::isnan(actual_value) && std::isnan(expected_value); // not possible both are not nan if diff is nan. + + if (std::isinf(actual_value) || std::isinf(expected_value)) { + if (std::isinf(actual_value) && std::isinf(expected_value)) + return (actual_value > 0 && expected_value > 0) || (actual_value < 0 && expected_value < 0); + else + return false; + } + + return (diff <= tol); } template @@ -456,6 +481,93 @@ static std::pair CompareTensorOrtValueAndTensorType return std::make_pair(COMPARE_RESULT::SUCCESS, ""); } +namespace { +template +float ParseValueToFloat(T data_value) { + return static_cast(data_value); +} + +template <> +float ParseValueToFloat(MLFloat16 data_value) { + return Eigen::half_impl::half_to_float(Eigen::half_impl::__half_raw(data_value.val)); +} + +template <> +float ParseValueToFloat(float data_value) { + // Covert float to half and then convert back to float to simulate rounding to half + return ParseValueToFloat(MLFloat16(math::floatToHalf(data_value))); +} + +template +std::pair CompareFloat16WithFloatResult(const Tensor& outvalue, + const Tensor& expected_value, + double per_sample_tolerance, + double relative_per_sample_tolerance) { + const size_t size1 = static_cast(expected_value.Shape().Size()); + const ExpectT* expected_output = expected_value.Data(); + const RealT* real_output = outvalue.Data(); + + COMPARE_RESULT result = COMPARE_RESULT::SUCCESS; + std::string error_msg = ""; + + OrtThreadPoolParams to; + to.thread_pool_size = 16; + auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP); + + static double cost = 1; + std::once_flag write_flag; + concurrency::ThreadPool::TryParallelFor( + tp.get(), size1, cost, + [&error_msg, &result, &expected_output, &real_output, per_sample_tolerance, + relative_per_sample_tolerance, size1, &write_flag]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t di = begin; di != end && di < static_cast(size1); ++di) { + float expected = ParseValueToFloat(expected_output[di]); + float real = ParseValueToFloat(real_output[di]); + const double diff = std::fabs(expected - real); + const double rtol = per_sample_tolerance + relative_per_sample_tolerance * std::fabs(expected); + if (!IsResultCloselyMatch(real, expected, diff, rtol)) { + std::ostringstream oss; + oss << "idx: " << di << ", expected " << expected << ", got " << real + << ", diff: " << diff << ", tol=" << rtol << "\n"; + std::call_once(write_flag, [&error_msg, &result, &oss]() { + error_msg = oss.str(); + result = COMPARE_RESULT::RESULT_DIFFERS; + }); + break; + } + } + }); + + return std::make_pair(result, error_msg); +} + +} // namespace + +std::pair CompareOrtValueNumerals(const OrtValue& real_mlvalue, const OrtValue& expected_mlvalue, + double per_sample_tolerance, + double relative_per_sample_tolerance) { + if (real_mlvalue.IsTensor()) { + const Tensor& outvalue = real_mlvalue.Get(); + const Tensor& expected_tensor = expected_mlvalue.Get(); + if (outvalue.DataType() == DataTypeImpl::GetType() && + expected_tensor.DataType() == DataTypeImpl::GetType()) { + return CompareFloat16WithFloatResult(outvalue, + expected_tensor, + per_sample_tolerance, + relative_per_sample_tolerance); + } else if (outvalue.DataType() == DataTypeImpl::GetType() && + expected_tensor.DataType() == DataTypeImpl::GetType()) { + return CompareFloat16WithFloatResult(outvalue, + expected_tensor, + per_sample_tolerance, + relative_per_sample_tolerance); + } + } + + return std::make_pair(COMPARE_RESULT::NOT_SUPPORT, "Unsupported compare with CompareOrtValueNumerals."); +} + std::pair VerifyValueInfo(const ONNX_NAMESPACE::ValueInfoProto& v, const OrtValue* val_ptr) { Ort::ConstValue o{val_ptr}; if (!v.has_type()) return std::make_pair(COMPARE_RESULT::SUCCESS, ""); diff --git a/onnxruntime/test/util/include/test/compare_ortvalue.h b/onnxruntime/test/util/include/test/compare_ortvalue.h index 24b74b9002..545df706c9 100644 --- a/onnxruntime/test/util/include/test/compare_ortvalue.h +++ b/onnxruntime/test/util/include/test/compare_ortvalue.h @@ -16,15 +16,33 @@ struct Value; } namespace onnxruntime { -enum class COMPARE_RESULT { SUCCESS, - RESULT_DIFFERS, - TYPE_MISMATCH, - SHAPE_MISMATCH, - NOT_SUPPORT }; + +enum class COMPARE_RESULT { + SUCCESS, + RESULT_DIFFERS, + TYPE_MISMATCH, + SHAPE_MISMATCH, + NOT_SUPPORT, +}; + std::pair CompareOrtValue(const OrtValue& real, const OrtValue& expected, double per_sample_tolerance, double relative_per_sample_tolerance, bool post_processing); +// Compare two OrtValue numerically equal or not. The difference with CompareOrtValue is that this function +// will only check the numerical values of the OrtValue, and ignore the type, shape, etc. +// +// For example, if some tests run CPU EP baseline with float, and run CUDA test with float16, we could check +// them without converting the float16 to float. Be noted: we will try convert the float value to float16 +// then back to float, to simulate rounding to half. +// +// If not equal, only return one single the differed value pair (to avoid multiple thread +// append into the same error_message string stream). +std::pair CompareOrtValueNumerals(const OrtValue& real, + const OrtValue& expected, + double per_sample_tolerance, + double relative_per_sample_tolerance); + // verify if the 'value' matches the 'expected' ValueInfoProto. 'value' is a model output std::pair VerifyValueInfo(const ONNX_NAMESPACE::ValueInfoProto& expected, const OrtValue* value); diff --git a/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc b/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc index ae752dfe29..a80e07a295 100644 --- a/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc +++ b/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include "test/compare_ortvalue.h" #include "test/util/include/default_providers.h" @@ -283,14 +284,18 @@ TEST(CrossEntropyTest, SparseSoftmaxCrossEntropyGrad_LargeSizeTensor) { TestSparseSoftmaxCrossEntropyGrad(dY_dims, log_prob_dims, index_dims, dX_dims, "sum"); } -static void PrepareSCELossTestData(const std::vector* X_dims, - const std::vector* index_dims, const std::vector* weight_dims, - const std::int64_t ignore_index, - std::vector& X_data, std::vector& index_data, - std::vector& weight_data) { +template +void PrepareSCELossTestData(const std::vector* X_dims, + const std::vector* index_dims, const std::vector* weight_dims, + const std::int64_t ignore_index, + std::vector& X_data, + std::vector& index_data, + std::vector& weight_data) { // create rand inputs + ParallelRandomValueGenerator prandom{2333}; + X_data = prandom.Uniform(*X_dims, -200.0f, 200.0f); + RandomValueGenerator random{2333}; - X_data = random.Uniform(*X_dims, -200.0f, 200.0f); index_data = random.Uniform(*index_dims, 0, (*X_dims)[1]); // Add one data point that has ignore_index. if (index_data.size() > 0) { @@ -298,7 +303,7 @@ static void PrepareSCELossTestData(const std::vector* X_dims, } if (weight_dims) { - weight_data = random.Uniform(*weight_dims, 0.0f, 1.0f); + weight_data = prandom.Uniform(*weight_dims, 0.0f, 1.0f); } } @@ -316,9 +321,9 @@ static std::vector RunSCELossWithEP(const char* op, const std::vector* weight_dims, const std::vector* Y_dims, const std::vector* log_prob_dims, - std::vector& X_data, + std::vector& X_data, std::vector& index_data, - std::vector& weight_data) { + std::vector& weight_data) { /** * OpTester's atol/rtol check is too strict for our testing cases. Imagine expected value is 4.7683704451628728e-07, * real value is 0, even we set rtol=1e-1, atol = 1e-4. The check still fail. @@ -330,11 +335,11 @@ static std::vector RunSCELossWithEP(const char* op, * So here we disable OpTester's check by default, and do the check externally. */ bool need_verify_outputs = false; - std::vector> expected_values; + std::vector> expected_values; // Still need feed the output data even we don't want to verify it. - expected_values.push_back(FillZeros(*Y_dims)); + expected_values.push_back(FillZeros(*Y_dims)); if (log_prob_dims) { - expected_values.push_back(FillZeros(*log_prob_dims)); + expected_values.push_back(FillZeros(*log_prob_dims)); } OpTester test(op, opset_version, domain, need_verify_outputs /*verify_output*/); @@ -349,46 +354,21 @@ static std::vector RunSCELossWithEP(const char* op, test.AddAttribute("ignore_index", ignore_index); } - if (std::is_same::value) { - std::vector X_data_half(X_data.size()); - ConvertFloatToMLFloat16(X_data.data(), X_data_half.data(), static_cast(X_data.size())); - test.AddInput("X", *X_dims, X_data_half); - } else { - test.AddInput("X", *X_dims, X_data); - } + test.AddInput("X", *X_dims, X_data); test.AddInput("index", *index_dims, index_data); if (weight_dims) { - if (std::is_same::value) { - std::vector weight_data_half(weight_data.size()); - ConvertFloatToMLFloat16(weight_data.data(), weight_data_half.data(), static_cast(weight_data.size())); - test.AddInput("weight", *weight_dims, weight_data_half); - } else { - test.AddInput("weight", *weight_dims, weight_data); - } + test.AddInput("weight", *weight_dims, weight_data); } if (is_internal_op && ignore_index != -1) { test.AddInput("ignore_index", {}, &ignore_index, 1); } - if (std::is_same::value) { - std::vector output_half(expected_values[0].size()); - ConvertFloatToMLFloat16(expected_values[0].data(), output_half.data(), static_cast(expected_values[0].size())); - test.AddOutput("output", *Y_dims, output_half); - - if (log_prob_dims) { - std::vector log_prob_half(expected_values[1].size()); - ConvertFloatToMLFloat16(expected_values[1].data(), log_prob_half.data(), static_cast(expected_values[1].size())); - test.AddOutput("log_prob", *log_prob_dims, log_prob_half); - } - - } else { - test.AddOutput("output", *Y_dims, expected_values[0]); - if (log_prob_dims) { - test.AddOutput("log_prob", *log_prob_dims, expected_values[1]); - } + test.AddOutput("output", *Y_dims, expected_values[0]); + if (log_prob_dims) { + test.AddOutput("log_prob", *log_prob_dims, expected_values[1]); } std::vector> eps; @@ -407,23 +387,38 @@ static void TestSCELoss(const char* op, int opset_version, ASSERT_TRUE((std::is_same::value || std::is_same::value)); ASSERT_TRUE((std::is_same::value || std::is_same::value)); - std::vector X_data, weight_data; + std::vector X_data, weight_data; std::vector index_data; - PrepareSCELossTestData(X_dims, index_dims, weight_dims, ignore_index, X_data, index_data, weight_data); + PrepareSCELossTestData(X_dims, index_dims, weight_dims, ignore_index, X_data, index_data, weight_data); // Run on CPU using float input and output // (because the CPU implementation doesn't support variant input output types.) // Be noted, no result comparison is done here. - std::vector cpu_fetches = RunSCELossWithEP( - op, opset_version, domain, - []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, - reduction, ignore_index, error_tolerance, - X_dims, index_dims, weight_dims, - Y_dims, log_prob_dims, - X_data, index_data, weight_data); + std::vector cpu_fetches; + if (std::is_same::value) { + std::vector X_data_temp(X_data.size()); + std::vector weight_data_temp(weight_data.size()); + ConvertMLFloat16ToFloat(reinterpret_cast(X_data.data()), X_data_temp.data(), X_data.size()); + ConvertMLFloat16ToFloat(reinterpret_cast(weight_data.data()), weight_data_temp.data(), weight_data.size()); + cpu_fetches = RunSCELossWithEP( + op, opset_version, domain, + []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, + reduction, ignore_index, error_tolerance, + X_dims, index_dims, weight_dims, + Y_dims, log_prob_dims, + X_data_temp, index_data, weight_data_temp); + } else { + cpu_fetches = RunSCELossWithEP( + op, opset_version, domain, + []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, + reduction, ignore_index, error_tolerance, + X_dims, index_dims, weight_dims, + Y_dims, log_prob_dims, + X_data, index_data, weight_data); + } // Run on CUDA. - // Be noted, no result comparison is done here because OpTester's check is too strick for our test cases. + // Be noted, no result comparison is done here because OpTester's check is too strict for our test cases. // Check more details in comment of RunSCELossWithEP. std::vector target_fetches = RunSCELossWithEP( op, opset_version, domain, @@ -442,22 +437,10 @@ static void TestSCELoss(const char* op, int opset_version, // Compare ASSERT_EQ(cpu_fetches.size(), target_fetches.size()); for (size_t i = 0; i < cpu_fetches.size(); i++) { - if (std::is_same::value) { - auto y_data_size = cpu_fetches[i].Get().Shape().Size(); - std::vector cpu_temp_buffer; - cpu_temp_buffer.resize(y_data_size); - const float* y_buffer = cpu_fetches[i].Get().Data(); - std::copy(y_buffer, y_buffer + y_data_size, cpu_temp_buffer.begin()); - - std::vector ret_half(cpu_temp_buffer.size()); - ConvertFloatToMLFloat16(cpu_temp_buffer.data(), ret_half.data(), static_cast(cpu_temp_buffer.size())); - - OrtValue target; - test::CreateInputOrtValueOnCPU(cpu_fetches[i].Get().Shape().GetDims(), ret_half, &target); - auto ret = CompareOrtValue(target_fetches[i], target, error_tolerance /*per_sample_tolerance*/, - error_tolerance /*relative_per_sample_tolerance*/, false); + if (!std::is_same::value) { // baseline output is float. + auto ret = CompareOrtValueNumerals(target_fetches[i], cpu_fetches[i], error_tolerance /*per_sample_tolerance*/, + error_tolerance /*relative_per_sample_tolerance*/); EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second; - } else { auto ret = CompareOrtValue(target_fetches[i], cpu_fetches[i], error_tolerance /*per_sample_tolerance*/, error_tolerance /*relative_per_sample_tolerance*/, false); @@ -655,6 +638,29 @@ TEST(CrossEntropyTest, DISABLED_SoftmaxCrossEntropyLoss_LargeSizeTensor) { TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); } +#ifndef _WIN32 +// Disable the large size tests for Windows because it is too slow, running on Linux would be enough. +// This test requires lots of memory, currently, it can run with 16GB V100 GPU. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternal_LargeSizeTensorUInt64Index) { + // The element count is bigger than the upper limit of int32_t. + constexpr int64_t bsz = 419431; + constexpr int64_t vocab_size = 5120; + std::vector X_dims{bsz, vocab_size}; + std::vector index_dims{bsz}; + std::vector weight_dims{vocab_size}; + std::vector Y_dims{}; + std::vector Y_dims_none{bsz}; + std::vector log_prob_dims{bsz, vocab_size}; + + constexpr std::int64_t ignore_index = -1; + constexpr double error_tolerance = 1e-3; + // Only test reduce mean because it's too costly to run more tests. + TestSCELoss("SoftmaxCrossEntropyLossInternal", 1, onnxruntime::kMSDomain, + &X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean", + ignore_index, error_tolerance); +} +#endif + static void TestSoftmaxCrossEntropyLossGrad(const std::vector& dY_dims, const std::vector& log_prob_dims, const std::vector& index_dims, @@ -678,11 +684,11 @@ static void TestSoftmaxCrossEntropyLossGrad(const std::vector& dY_dims, } if (test_fp16) { std::vector dY_data_half(dY_data.size()); - ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), static_cast(dY_data.size())); + ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), dY_data.size()); test.AddInput("dY", dY_dims, dY_data_half); std::vector log_prob_data_half(log_prob_data.size()); - ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), static_cast(log_prob_data.size())); + ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), log_prob_data.size()); test.AddInput("log_prob", log_prob_dims, log_prob_data_half); test.AddInput("index", index_dims, index_data); @@ -773,26 +779,28 @@ TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LargeSizeTensor_half) { TestSoftmaxCrossEntropyLossGrad({2, 30528}, log_prob_dims, index_dims, dX_dims, "none", -1, true, 5e-2); } -static void PrepareSCELossInternalGradTestData( +template +void PrepareSCELossInternalGradTestData( const std::vector& dY_dims, const std::vector& log_prob_dims, const std::vector& index_dims, const std::vector& weight_dims, const std::vector& dX_dims, const std::int64_t ignore_index, bool has_bias, - std::vector& dY_data, std::vector& log_prob_data, - std::vector& index_data, std::vector& weight_data, - std::vector& bias_data) { + std::vector& dY_data, std::vector& log_prob_data, + std::vector& index_data, std::vector& weight_data, + std::vector& bias_data) { // Create rand inputs RandomValueGenerator random{2333}; - dY_data = random.Uniform(dY_dims, -10.0f, 10.0f); - log_prob_data = random.Uniform(log_prob_dims, -10.0f, 10.0f); + ParallelRandomValueGenerator prandom{2333}; + dY_data = prandom.Uniform(dY_dims, -10.0f, 10.0f); + log_prob_data = prandom.Uniform(log_prob_dims, -10.0f, 10.0f); index_data = random.Uniform(index_dims, 0, dX_dims[1]); // Add one data point that has ignore_index. if ((ignore_index != -1) && (index_data.size() > 0)) { index_data[0] = ignore_index; } - weight_data = random.Uniform(weight_dims, 0.0f, 1.0f); + weight_data = prandom.Uniform(weight_dims, 0.0f, 1.0f); if (has_bias) { - bias_data = random.Uniform(dX_dims, 0.0f, 1.0f); + bias_data = prandom.Uniform(dX_dims, 0.0f, 1.0f); } } @@ -809,11 +817,11 @@ static std::vector RunSCELossInternalGradWithEP( const std::vector& index_dims, const std::vector& weight_dims, const std::vector& dX_dims, - std::vector& dY_data, - std::vector& log_prob_data, + std::vector& dY_data, + std::vector& log_prob_data, std::vector& index_data, - std::vector& weight_data, - std::vector& bias_data) { + std::vector& weight_data, + std::vector& bias_data) { /** * OpTester's atol/rtol check is too strict for our testing cases. Imagine expected value is 4.7683704451628728e-07, * real value is 0, even we set rtol=1e-1, atol = 1e-4. The check still fail with the following check: @@ -825,7 +833,7 @@ static std::vector RunSCELossInternalGradWithEP( * So here we disable OpTester's check by default, and do the check externally. */ bool need_verify_outputs = false; - std::vector expected_value = FillZeros(dX_dims); + std::vector expected_value = FillZeros(dX_dims); ORT_ENFORCE((std::is_same::value || std::is_same::value)); ORT_ENFORCE((std::is_same::value || std::is_same::value)); @@ -840,56 +848,21 @@ static std::vector RunSCELossInternalGradWithEP( test->AddAttribute("output_type", static_cast(utils::ToTensorProtoElementType())); } - if (std::is_same::value) { - std::vector dY_data_half(dY_data.size()); - ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), static_cast(dY_data.size())); - test->AddInput("dY", dY_dims, dY_data_half); - - std::vector log_prob_data_half(log_prob_data.size()); - ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), static_cast(log_prob_data.size())); - test->AddInput("log_prob", log_prob_dims, log_prob_data_half); - - test->AddInput("index", index_dims, index_data); - - std::vector weight_data_half(weight_data.size()); - ConvertFloatToMLFloat16(weight_data.data(), weight_data_half.data(), static_cast(weight_data.size())); - test->AddInput("weight", weight_dims, weight_data_half); - - if (ignore_index != -1 || has_bias) { - test->AddInput("ignore_index", {}, &ignore_index, 1); - } - } else { - test->AddInput("dY", dY_dims, dY_data); - test->AddInput("log_prob", log_prob_dims, log_prob_data); - test->AddInput("index", index_dims, index_data); - test->AddInput("weight", weight_dims, weight_data); - if (ignore_index != -1 || has_bias) { - test->AddInput("ignore_index", {}, &ignore_index, 1); - } + test->AddInput("dY", dY_dims, dY_data); + test->AddInput("log_prob", log_prob_dims, log_prob_data); + test->AddInput("index", index_dims, index_data); + test->AddInput("weight", weight_dims, weight_data); + if (ignore_index != -1 || has_bias) { + test->AddInput("ignore_index", {}, &ignore_index, 1); } - if (std::is_same::value) { - // Be noted, bias should be aligned with output's data type. - if (has_bias) { - std::vector bias_data_half(bias_data.size()); - ConvertFloatToMLFloat16(bias_data.data(), bias_data_half.data(), static_cast(bias_data.size())); - test->AddInput("bias", dX_dims, bias_data_half); - } - - std::vector expected_data_half(expected_value.size()); - ConvertFloatToMLFloat16(expected_value.data(), expected_data_half.data(), static_cast(expected_value.size())); - test->AddOutput("dX", dX_dims, expected_data_half, false /*sort_output*/, - error_tolerance /*rel_error*/, error_tolerance /*abs_error*/); - - } else { - if (has_bias) { - test->AddInput("bias", dX_dims, bias_data); - } - - test->AddOutput("dX", dX_dims, expected_value, false /*sort_output*/, - error_tolerance /*rel_error*/, error_tolerance /*abs_error*/); + if (has_bias) { + test->AddInput("bias", dX_dims, bias_data); } + test->AddOutput("dX", dX_dims, expected_value, false /*sort_output*/, + error_tolerance /*rel_error*/, error_tolerance /*abs_error*/); + std::vector> eps; eps.emplace_back(ep_creator()); test->Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); @@ -906,20 +879,56 @@ static void TestSoftmaxCrossEntropyLossInternalGrad(const std::vector& const std::int64_t ignore_index = -1, const double error_tolerance = 1e-4, const bool has_bias = false) { - std::vector dY_data, log_prob_data, weight_data, bias_data; + std::vector dY_data, log_prob_data, weight_data; + std::vector bias_data; std::vector index_data; - PrepareSCELossInternalGradTestData(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, ignore_index, has_bias, - dY_data, log_prob_data, index_data, weight_data, bias_data); + PrepareSCELossInternalGradTestData(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, ignore_index, + has_bias, dY_data, log_prob_data, index_data, weight_data, bias_data); + std::vector cpu_fetches; // Run on CPU using float input and output // (because the CPU implementation doesn't support variant input output types.) // Be noted, no result comparison is done here. - std::vector cpu_fetches = - RunSCELossInternalGradWithEP( + if constexpr (std::is_same::value) { + std::vector dY_data_temp(dY_data.size()); + std::vector log_prob_data_temp(log_prob_data.size()); + std::vector weight_data_temp(weight_data.size()); + ConvertMLFloat16ToFloat(reinterpret_cast(dY_data.data()), dY_data_temp.data(), dY_data.size()); + ConvertMLFloat16ToFloat(reinterpret_cast(log_prob_data.data()), log_prob_data_temp.data(), + log_prob_data.size()); + ConvertMLFloat16ToFloat(reinterpret_cast(weight_data.data()), weight_data_temp.data(), weight_data.size()); + if constexpr (std::is_same::value) { + std::vector bias_data_temp(bias_data.size()); + ConvertMLFloat16ToFloat(reinterpret_cast(bias_data.data()), bias_data_temp.data(), bias_data.size()); + cpu_fetches = RunSCELossInternalGradWithEP( + []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, + reduction, ignore_index, error_tolerance, has_bias, + dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + dY_data_temp, log_prob_data_temp, index_data, weight_data_temp, bias_data_temp); + } else { + cpu_fetches = RunSCELossInternalGradWithEP( + []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, + reduction, ignore_index, error_tolerance, has_bias, + dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + dY_data_temp, log_prob_data_temp, index_data, weight_data_temp, bias_data); + } + } else { + if constexpr (std::is_same::value) { + std::vector bias_data_temp(bias_data.size()); + ConvertMLFloat16ToFloat(reinterpret_cast(bias_data.data()), bias_data_temp.data(), bias_data.size()); + cpu_fetches = RunSCELossInternalGradWithEP( + []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, + reduction, ignore_index, error_tolerance, has_bias, + dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + dY_data, log_prob_data, index_data, weight_data, bias_data_temp); + } else { + cpu_fetches = RunSCELossInternalGradWithEP( []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, reduction, ignore_index, error_tolerance, has_bias, dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, dY_data, log_prob_data, index_data, weight_data, bias_data); + } + } // Run on CUDA and compare results with cpu results. std::vector target_fetches = @@ -938,22 +947,10 @@ static void TestSoftmaxCrossEntropyLossInternalGrad(const std::vector& // Compare ASSERT_EQ(cpu_fetches.size(), target_fetches.size()); for (size_t i = 0; i < cpu_fetches.size(); i++) { - if (std::is_same::value) { - auto y_data_size = cpu_fetches[i].Get().Shape().Size(); - std::vector cpu_temp_buffer; - cpu_temp_buffer.resize(y_data_size); - const float* y_buffer = cpu_fetches[i].Get().Data(); - std::copy(y_buffer, y_buffer + y_data_size, cpu_temp_buffer.begin()); - - std::vector ret_half(cpu_temp_buffer.size()); - ConvertFloatToMLFloat16(cpu_temp_buffer.data(), ret_half.data(), static_cast(cpu_temp_buffer.size())); - - OrtValue target; - test::CreateInputOrtValueOnCPU(cpu_fetches[i].Get().Shape().GetDims(), ret_half, &target); - auto ret = CompareOrtValue(target_fetches[i], target, error_tolerance /*per_sample_tolerance*/, - error_tolerance /*relative_per_sample_tolerance*/, false); + if (!std::is_same::value) { + auto ret = CompareOrtValueNumerals(target_fetches[i], cpu_fetches[i], error_tolerance /*per_sample_tolerance*/, + error_tolerance /*relative_per_sample_tolerance*/); EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second; - } else { auto ret = CompareOrtValue(target_fetches[i], cpu_fetches[i], error_tolerance /*per_sample_tolerance*/, error_tolerance /*relative_per_sample_tolerance*/, false); @@ -1073,5 +1070,26 @@ TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternalGrad_TinySizeTensorFloatIn "none", 0, 5e-2, true); } +#ifndef _WIN32 +// Disable the large size tests for Windows because it is too slow, running on Linux would be enough. +// This test requires lots of memory, currently, it can run with 16GB V100 GPU. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternalGrad_LargeSizeTensorUInt64Index) { + // The element count is bigger than the upper limit of int32_t. + constexpr int64_t bsz = 419431; + constexpr int64_t vocab_size = 5120; + std::vector dY_dims{}; + std::vector log_prob_dims{bsz, vocab_size}; + std::vector index_dims{bsz}; + std::vector weight_dims{vocab_size}; + std::vector dX_dims{bsz, vocab_size}; + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", -1, 1e-3, false /*has_bias*/); + + // This test did not test against reduce-sum because the absolute value after computing is pretty big, sometimes + // around 65535, CPU baseline result is smaller than 65535, but CUDA result is a little bigger than it, generating + // inf, which is not a good test case. +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.cc b/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.cc index e2cdf238a0..b4b137db20 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.cc +++ b/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.cc @@ -6,7 +6,6 @@ #include "core/util/math_cpuonly.h" #include "core/providers/common.h" #include -#include "core/util/math.h" #include "core/providers/cpu/math/matmul_helper.h" #include "core/common/gsl.h" @@ -14,38 +13,38 @@ namespace onnxruntime { namespace contrib { template -void ComputeShareSoftmaxCrossEntropyCPU(const int n, - const int d, - const int nd, +void ComputeShareSoftmaxCrossEntropyCPU(const int nd, + const int c, + const Eigen::Index nd_c, const T* logit_data, T* shifted_logit, T* log_prob_data) { // Find the max in each batch, resulting in a tensor of shape [batch] // logit_max = max(logit_data) - std::vector logit_max(n); - math::RowwiseMax(n, d, logit_data, logit_max.data(), nullptr); + std::vector logit_max(nd); + math::RowwiseMax(nd, c, logit_data, logit_max.data(), nullptr); // Subtract the max in batch b from every element in batch b. // Broadcasts along the batch dimension. // shifted_logit = logit_data - logit_max - gsl::copy(gsl::make_span(logit_data, nd), gsl::make_span(shifted_logit, nd)); - math::SubToCol(n, d, logit_max.data(), shifted_logit, nullptr); + gsl::copy(gsl::make_span(logit_data, nd_c), gsl::make_span(shifted_logit, nd_c)); + math::SubToCol(nd, c, logit_max.data(), shifted_logit, nullptr); // exp_shifted_logit = exp(shifted_logit) - math::Exp(nd, shifted_logit, log_prob_data, nullptr); + math::Exp(nd_c, shifted_logit, log_prob_data, nullptr); // sum_exp = sum_{class} (exp_shifted_logit) float* sum_exp = logit_max.data(); - math::RowwiseSum(n, d, log_prob_data, sum_exp, nullptr); + math::RowwiseSum(nd, c, log_prob_data, sum_exp, nullptr); // log_sum_exp = log(sum_exp) float* log_sum_exp = sum_exp; - math::Log(n, sum_exp, log_sum_exp, nullptr); + math::Log(nd, sum_exp, log_sum_exp, nullptr); // log_prob = shifted_logit - log(sum_exp) // the subtraction broadcasts along the batch dimension - gsl::copy(gsl::make_span(shifted_logit, nd), gsl::make_span(log_prob_data, nd)); - math::SubToCol(n, d, log_sum_exp, log_prob_data, nullptr); + gsl::copy(gsl::make_span(shifted_logit, nd_c), gsl::make_span(log_prob_data, nd_c)); + math::SubToCol(nd, c, log_sum_exp, log_prob_data, nullptr); } ONNX_OPERATOR_KERNEL_EX( @@ -70,7 +69,7 @@ Status SoftmaxCrossEntropy::Compute(OpKernelContext* context) const { int64_t D = logit_shape[logit_shape.NumDimensions() - 1]; const int n = gsl::narrow_cast(N); const int d = gsl::narrow_cast(D); - const int nd = gsl::narrow_cast(N * D); + const uint64_t nd = N * D; Tensor* loss = context->Output(0, TensorShape({})); Tensor* log_prob = context->Output(1, logit_shape); @@ -85,7 +84,10 @@ Status SoftmaxCrossEntropy::Compute(OpKernelContext* context) const { // probability = exp(shifted_logit) / sum(exp(shifted_logit)) // where shifted_logit = logit - max_logit // along classes - ComputeShareSoftmaxCrossEntropyCPU(n, d, nd, logit_data, + + ORT_ENFORCE(nd <= static_cast(std::numeric_limits::max())); + ComputeShareSoftmaxCrossEntropyCPU(n, d, static_cast(nd), + logit_data, shifted_logit.data(), log_prob_data); @@ -180,7 +182,7 @@ Status SparseSoftmaxCrossEntropy::Compute(OpKernelContext* context) const { int64_t D = logit_shape[logit_shape.NumDimensions() - 1]; const int n = gsl::narrow_cast(N); const int d = gsl::narrow_cast(D); - const int nd = gsl::narrow_cast(N * D); + const uint64_t nd = N * D; Tensor* loss = context->Output(0, TensorShape({})); Tensor* log_prob = context->Output(1, logit_shape); @@ -192,7 +194,9 @@ Status SparseSoftmaxCrossEntropy::Compute(OpKernelContext* context) const { // computation begins here std::vector shifted_logit(nd); - ComputeShareSoftmaxCrossEntropyCPU(n, d, nd, logit_data, + + ORT_ENFORCE(nd <= static_cast(std::numeric_limits::max())); + ComputeShareSoftmaxCrossEntropyCPU(n, d, static_cast(nd), logit_data, shifted_logit.data(), log_prob_data); diff --git a/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.h b/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.h index dd0d5750e0..4a180fa66e 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.h +++ b/orttraining/orttraining/training_ops/cpu/loss/cross_entropy.h @@ -5,6 +5,7 @@ #include "core/framework/op_kernel.h" #include "orttraining/training_ops/cpu/loss/reduction_type.h" +#include "core/util/math_cpuonly.h" namespace onnxruntime { namespace contrib { @@ -24,7 +25,7 @@ class LossBase : public OpKernel { template void ComputeShareSoftmaxCrossEntropyCPU(const int n, const int d, - const int nd, + const Eigen::Index nd, const T* logit_data, T* shifted_logit, T* log_prob_data); diff --git a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc index c1b9865b53..224e13f574 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc +++ b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc @@ -5,7 +5,6 @@ #include "core/util/math_cpuonly.h" #include "core/providers/common.h" #include -#include "core/util/math.h" #include "core/providers/cpu/math/matmul_helper.h" #include "core/providers/cpu/tensor/transpose.h" #include "core/providers/cpu/controlflow/scan_utils.h" @@ -113,8 +112,8 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const VerifyLogitWeightAndLabelShape(logit_shape, label_shape, p_weight ? &p_weight->Shape() : nullptr); // N_D = N * D1 * D2...D*K - int64_t N_D; - int64_t C; + int64_t N_D = 0; + int64_t C = 0; GetNDCFromLogitAndLabelShape(logit_shape, label_shape, N_D, C); const T1* logit_data = logit.template Data(); OrtValue transpose_output; @@ -133,7 +132,7 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const const int n_d = gsl::narrow_cast(N_D); const int c = gsl::narrow_cast(C); - const int n_d_c = gsl::narrow_cast(N_D * C); + const uint64_t n_d_c = N_D * C; Tensor* loss = context->Output(0, reduction_ == ReductionType::NONE ? TensorShape(label.Shape()) : TensorShape({})); T1* log_prob_data; std::vector log_prob_data_buffer(0); @@ -150,7 +149,9 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const const T2* label_data = label.template Data(); T1* loss_data = loss->template MutableData(); std::vector shifted_logit(n_d_c); - ComputeShareSoftmaxCrossEntropyCPU(n_d, c, n_d_c, logit_data, shifted_logit.data(), log_prob_data); + ORT_ENFORCE(n_d_c <= static_cast(std::numeric_limits::max())); + ComputeShareSoftmaxCrossEntropyCPU(n_d, c, static_cast(n_d_c), logit_data, shifted_logit.data(), + log_prob_data); std::vector loss_sample_buffer(0); T1* loss_sample; if (reduction_ == ReductionType::NONE) { @@ -260,17 +261,15 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co VerifyLogitWeightAndLabelShape(probability_shape, label_shape, p_weight ? &p_weight->Shape() : nullptr); // N_D = N * D1 * D2...D*K - int64_t N_D; - int64_t C; + int64_t N_D = 0; + int64_t C = 0; GetNDCFromLogitAndLabelShape(probability_shape, label_shape, N_D, C); - const int n_d = gsl::narrow_cast(N_D); - const int c = gsl::narrow_cast(C); const T1* dY_data = dY.template Data(); const T1* log_prob_data = log_prob.template Data(); const T2* label_data = label.template Data(); Tensor* d_logit = context->Output(0, probability_shape); T1* d_logit_data = d_logit->template MutableData(); - std::memset(d_logit_data, 0, sizeof(T1) * n_d); + std::memset(d_logit_data, 0, sizeof(T1) * N_D); OrtValue transpose_output; TensorShapeVector new_shape; std::vector permutations; @@ -285,90 +284,107 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co log_prob_data = (*transpose_output.GetMutable()).template Data(); } - // REVIEW(codemzs): Use parallel for below. + static constexpr double cost = 1.0; + auto* tp = context->GetOperatorThreadPool(); if (p_weight) { const Tensor& weight = *p_weight; const T1* weight_data = weight.template Data(); if (reduction_ == ReductionType::NONE) { - for (int i = 0; i < n_d; i++) { - T2 label_sample = label_data[i]; - T1 weight_smaple = weight_data[label_sample] * dY_data[i]; - for (int j = 0; j < c; j++) { - int index = i * c + j; - if (ignore_index == label_sample) { - d_logit_data[index] = 0; - } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == j)) * weight_smaple; - } - } - } - + concurrency::ThreadPool::TryParallelFor( + tp, N_D * C, cost, + [&label_data, &weight_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_data]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t index = begin; index != end; ++index) { + int64_t row = index / C; + int64_t col = index % C; + T2 label_sample = label_data[row]; + T1 weight_smaple = weight_data[label_sample] * dY_data[row]; + if (ignore_index == label_sample) { + d_logit_data[index] = 0; + } else { + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_smaple; + } + } + }); } else { T1 dY_scaled = *dY_data; if (reduction_ == ReductionType::MEAN) { - T1 sum_weight = (T1)0; - for (int i = 0; i < n_d; i++) { + double sum_weight = (double)0; + for (int64_t i = 0; i < N_D; i++) { if (ignore_index != label_data[i]) { sum_weight += weight_data[label_data[i]]; } } if (sum_weight != 0) { - dY_scaled = *dY_data / sum_weight; + dY_scaled = static_cast(*dY_data / sum_weight); } } - for (int i = 0; i < n_d; i++) { - T2 label_sample = label_data[i]; - T1 weight_smaple = weight_data[label_sample] * dY_scaled; - for (int j = 0; j < c; j++) { - int index = i * c + j; - if (ignore_index == label_sample) { - d_logit_data[index] = 0; - } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == j)) * weight_smaple; - } - } - } + concurrency::ThreadPool::TryParallelFor( + tp, N_D * C, cost, + [&label_data, &weight_data, dY_scaled, &d_logit_data, &log_prob_data, ignore_index, C]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t index = begin; index != end; ++index) { + int64_t row = index / C; + int64_t col = index % C; + T2 label_sample = label_data[row]; + T1 weight_smaple = weight_data[label_sample] * dY_scaled; + if (ignore_index == label_sample) { + d_logit_data[index] = 0; + } else { + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_smaple; + } + } + }); } } else { if (reduction_ == ReductionType::NONE) { - for (int i = 0; i < n_d; i++) { - T2 label_sample = label_data[i]; - for (int j = 0; j < c; j++) { - int index = i * c + j; - if (ignore_index == label_sample) { - d_logit_data[index] = 0; - } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == j)) * dY_data[i]; - } - } - } + concurrency::ThreadPool::TryParallelFor( + tp, N_D * C, cost, + [&label_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_data]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t index = begin; index != end; ++index) { + int64_t row = index / C; + int64_t col = index % C; + T2 label_sample = label_data[row]; + if (ignore_index == label_sample) { + d_logit_data[index] = 0; + } else { + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * dY_data[row]; + } + } + }); + } else { T1 dY_scaled = *dY_data; - int unignored_sample_count = 0; - for (int i = 0; i < n_d; i++) { + uint64_t unignored_sample_count = 0; + for (int64_t i = 0; i < N_D; i++) { if (ignore_index != label_data[i]) { unignored_sample_count += 1; } } if ((reduction_ == ReductionType::MEAN) && (unignored_sample_count != 0)) { - dY_scaled = *dY_data / unignored_sample_count; + dY_scaled = static_cast(*dY_data / unignored_sample_count); } - for (int i = 0; i < n_d; i++) { - T2 label_sample = label_data[i]; - for (int j = 0; j < c; j++) { - int index = i * c + j; - if (ignore_index == label_sample) { - d_logit_data[index] = 0; - } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == j)) * dY_scaled; - } - } - } + concurrency::ThreadPool::TryParallelFor( + tp, N_D * C, cost, + [&label_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_scaled]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t index = begin; index != end; ++index) { + int64_t row = index / C; + int64_t col = index % C; + T2 label_sample = label_data[row]; + if (ignore_index == label_sample) { + d_logit_data[index] = 0; + } else { + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * dY_scaled; + } + } + }); } } @@ -391,9 +407,14 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co if (p_bias) { ORT_ENFORCE(probability_shape.Size() == p_bias->Shape().Size()); const T1* bias_data = p_bias->Data(); - for (size_t i = 0; i < static_cast(probability_shape.Size()); ++i) { - d_logit_data[i] += bias_data[i]; - } + concurrency::ThreadPool::TryParallelFor( + tp, probability_shape.Size(), cost, + [&d_logit_data, &bias_data]( + std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t index = begin; index != end; ++index) { + d_logit_data[index] += bias_data[index]; + } + }); } return Status::OK(); diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc index d0dec4fae6..b3501e62b5 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc @@ -214,7 +214,7 @@ Status SoftmaxCrossEntropyLossGrad::ComputeInternal(OpKernelCon p_weight ? &p_weight->Shape() : nullptr); // N_D = N * D1 * D2...Dk - int64_t N_D, C; + int64_t N_D = 0, C = 0; onnxruntime::contrib::GetNDCFromLogitAndLabelShape(probability_shape, label_shape, N_D, C); Tensor* d_logit = ctx->Output(0, probability_shape); const T* dY_data = dY.Data(); @@ -295,8 +295,8 @@ Status SoftmaxCrossEntropyLossGrad::ComputeInternal(OpKernelCon onnxruntime::contrib::GetPermutationAndShape(false, logit_shape, new_shape, permutations); transpose_output.GetMutable()->Reshape(d_logit->Shape()); d_logit->Reshape(logit_shape); - ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, *d_logit, - *transpose_output.GetMutable())); + ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, + *d_logit, *transpose_output.GetMutable())); auto* transposed_data = (*transpose_output.GetMutable()).template Data(); CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(d_logit_data, transposed_data, sizeof(TOut) * probability_shape.Size(), cudaMemcpyDeviceToDevice, Stream(ctx))); diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu index e2afa69ec1..ec575a39e2 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu @@ -7,12 +7,12 @@ namespace onnxruntime { namespace cuda { -template +template struct OpSoftmaxCrossEntropyWeights { OpSoftmaxCrossEntropyWeights(const TLabel* label_data, const T* weight_data, TLabel C, TLabel ignore_index) : label_data_(label_data), weight_data_(weight_data), C_(C), ignore_index_(ignore_index) {} - __device__ __inline__ TOut operator()(CUDA_LONG idx) const { + __device__ __inline__ TOut operator()(TIndex idx) const { if (label_data_[idx] != ignore_index_) { if (IsWeighted) { CUDA_KERNEL_ASSERT(label_data_[idx] >= 0 && label_data_[idx] < C_); @@ -33,13 +33,13 @@ template void ComputeSoftmaxCrossEntropyWeightsImpl(cudaStream_t stream, const TLabel* label, const T* weight, size_t count, size_t label_depth, int64_t ignore_index, TOut* weight_data_nd) { if (weight) { - OpSoftmaxCrossEntropyWeights op(label, weight, static_cast(label_depth), - static_cast(ignore_index)); - LaunchElementwiseKernel(stream, weight_data_nd, op, count); + typedef OpSoftmaxCrossEntropyWeights OP_Type; + OP_Type op(label, weight, static_cast(label_depth), static_cast(ignore_index)); + LaunchElementwiseKernel(stream, weight_data_nd, op, count); } else { - OpSoftmaxCrossEntropyWeights op(label, nullptr, static_cast(label_depth), - static_cast(ignore_index)); - LaunchElementwiseKernel(stream, weight_data_nd, op, count); + typedef OpSoftmaxCrossEntropyWeights OP_Type; + OP_Type op(label, nullptr, static_cast(label_depth), static_cast(ignore_index)); + LaunchElementwiseKernel(stream, weight_data_nd, op, count); } } @@ -57,7 +57,7 @@ INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(BFloat16, int64_t, BFloat16); #undef INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL -template +template struct OpWeightedSoftmaxCrossEntropyLoss { OpWeightedSoftmaxCrossEntropyLoss(const T* log_prob_data, const TLabel* label_data, const T* weight_data, const TAcc* normalize_factor_data, TLabel C, TLabel ignore_index) @@ -68,11 +68,12 @@ struct OpWeightedSoftmaxCrossEntropyLoss { C_(C), ignore_index_(ignore_index) {} - __device__ __inline__ T operator()(CUDA_LONG idx) const { + __device__ __inline__ T operator()(TIndex idx) const { if (label_data_[idx] != ignore_index_) { CUDA_KERNEL_ASSERT(label_data_[idx] >= 0 && label_data_[idx] < C_); - return static_cast(static_cast(-log_prob_data_[idx * C_ + label_data_[idx]] * weight_data_[idx]) / - (*normalize_factor_data_)); + T ret = static_cast(static_cast(-log_prob_data_[idx * C_ + label_data_[idx]] * weight_data_[idx]) / + (*normalize_factor_data_)); + return ret; } return T(0.f); } @@ -89,12 +90,14 @@ template void SoftmaxCrossEntropyLossImpl(cudaStream_t stream, const T* log_prob, const TLabel* label, const T* weight, const TAcc* normalize_factor, size_t count, size_t label_depth, int64_t ignore_index, T* output_data) { - OpWeightedSoftmaxCrossEntropyLoss op(log_prob, label, weight, normalize_factor, - static_cast(label_depth), static_cast(ignore_index)); - LaunchElementwiseKernel(stream, output_data, op, count); + typedef OpWeightedSoftmaxCrossEntropyLoss OP_Type; + OP_Type op(log_prob, label, weight, normalize_factor, + static_cast(label_depth), + static_cast(ignore_index)); + LaunchElementwiseKernel(stream, output_data, op, count); } -template +template struct OpWeightedSoftmaxCrossEntropyLossGrad { OpWeightedSoftmaxCrossEntropyLossGrad(const T* dY_data, const T* log_prob_data, const TLabel* label_data, const T* weight_data, const TAcc* normalize_factor_data, const TOut* bias_data, @@ -106,15 +109,15 @@ struct OpWeightedSoftmaxCrossEntropyLossGrad { normalize_factor_data_(normalize_factor_data), bias_data_(bias_data), C_(C) { - C_fdm_ = fast_divmod(static_cast(C)); + C_fdm_ = DivMod(static_cast(C)); } - __device__ __inline__ TOut operator()(CUDA_LONG idx) const { + __device__ __inline__ TOut operator()(TIndex idx) const { // normalize_factor is sum of labels' weights. Because zero sum implies all weights are 0, the loss function should // be constant 0 and its corresponding gradient should be 0 as well. TAcc result = TAcc(0.f); if (*normalize_factor_data_ != TAcc(0.f)) { - int row, d; + TIndex row, d; C_fdm_.divmod(idx, row, d); CUDA_KERNEL_ASSERT(weight_data_[row] == T(0.f) || (label_data_[row] >= 0 && label_data_[row] < C_)); result = static_cast((IsReductionNone ? dY_data_[row] : *dY_data_) * weight_data_[row]) * @@ -131,17 +134,24 @@ struct OpWeightedSoftmaxCrossEntropyLossGrad { const TAcc* normalize_factor_data_; const TOut* bias_data_; TLabel C_; - fast_divmod C_fdm_; + DivMod C_fdm_; }; template void SoftmaxCrossEntropyLossGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const TLabel* label, const T* weight, const TAcc* normalize_factor, const TOut* bias_data, size_t count, size_t label_depth, bool reduction_none, TOut* output_data) { -#define LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL(is_reduction_none, has_bias) \ - OpWeightedSoftmaxCrossEntropyLossGrad op( \ - dY, log_prob, label, weight, normalize_factor, bias_data, static_cast(label_depth)); \ - LaunchElementwiseKernel(stream, output_data, op, count * label_depth) +#define LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL(is_reduction_none, has_bias) \ + uint64_t total_count = count * label_depth; \ + if (total_count <= static_cast(std::numeric_limits::max())) { \ + typedef OpWeightedSoftmaxCrossEntropyLossGrad OP_Type; \ + OP_Type op(dY, log_prob, label, weight, normalize_factor, bias_data, static_cast(label_depth)); \ + LaunchElementwiseKernel(stream, output_data, op, static_cast(total_count)); \ + } else { \ + typedef OpWeightedSoftmaxCrossEntropyLossGrad OP_Type; \ + OP_Type op(dY, log_prob, label, weight, normalize_factor, bias_data, static_cast(label_depth)); \ + LaunchElementwiseKernel(stream, output_data, op, total_count); \ + } if (reduction_none) { if (bias_data) { LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL(true, true); diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu b/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu index 9f58073256..c78a9f788d 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu @@ -25,7 +25,7 @@ template void SoftMaxCrossEntropyImpl(cudaStream_t stream, const T* log_prob, const T* label, size_t normalize_factor, T* output_data, size_t count) { OpSoftmaxCrossEntropy op(log_prob, label, static_cast(normalize_factor)); - LaunchElementwiseKernel(stream, output_data, op, count); + LaunchElementwiseKernel(stream, output_data, op, static_cast(count)); } template void SoftMaxCrossEntropyImpl(cudaStream_t stream, const float* log_prob, const float* label, @@ -53,7 +53,7 @@ template void SoftMaxCrossEntropyGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const T* label, size_t normalize_factor, T* output_data, size_t count) { OpSoftmaxCrossEntropyGrad op(dY, log_prob, label, static_cast(normalize_factor)); - LaunchElementwiseKernel(stream, output_data, op, count); + LaunchElementwiseKernel(stream, output_data, op, static_cast(count)); } template void SoftMaxCrossEntropyGradImpl(cudaStream_t stream, const float* dY, const float* log_prob, @@ -92,11 +92,11 @@ void SparseSoftmaxCrossEntropyImpl(cudaStream_t stream, const T* log_prob, const if (weight) { OpSparseSoftmaxCrossEntropy op(log_prob, label, weight, normalize_factor, static_cast(label_depth)); - LaunchElementwiseKernel(stream, output_data, op, count); + LaunchElementwiseKernel(stream, output_data, op, static_cast(count)); } else { OpSparseSoftmaxCrossEntropy op(log_prob, label, nullptr, normalize_factor, static_cast(label_depth)); - LaunchElementwiseKernel(stream, output_data, op, count); + LaunchElementwiseKernel(stream, output_data, op, static_cast(count)); } } @@ -136,11 +136,13 @@ void SparseSoftmaxCrossEntropyGradImpl(cudaStream_t stream, const T* dY, const T if (weight) { OpSparseSoftmaxCrossEntropyGrad op(dY, log_prob, label, weight, normalize_factor, fast_divmod(static_cast(label_depth))); - LaunchElementwiseKernel(stream, output_data, op, count * label_depth); + LaunchElementwiseKernel(stream, output_data, op, + static_cast(count * label_depth)); } else { OpSparseSoftmaxCrossEntropyGrad op(dY, log_prob, label, nullptr, normalize_factor, fast_divmod(static_cast(label_depth))); - LaunchElementwiseKernel(stream, output_data, op, count * label_depth); + LaunchElementwiseKernel(stream, output_data, op, + static_cast(count * label_depth)); } } diff --git a/orttraining/orttraining/training_ops/cuda/quantization/fake_quant_impl.cu b/orttraining/orttraining/training_ops/cuda/quantization/fake_quant_impl.cu index d3d479664b..46178dca69 100644 --- a/orttraining/orttraining/training_ops/cuda/quantization/fake_quant_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/quantization/fake_quant_impl.cu @@ -98,8 +98,8 @@ template void FakeQuantGradImpl(cudaStream_t stream, const int64_t num_elements, const T* dY_data, const bool* gradient_mask_data, T* dX_data) { FakeQuantGradFunctor fake_quant_grad_functor(dY_data, gradient_mask_data); - LaunchElementwiseKernel( - stream, dX_data, fake_quant_grad_functor, num_elements); + LaunchElementwiseKernel( + stream, dX_data, fake_quant_grad_functor, static_cast(num_elements)); } #define SPECIALIZED_FAKEQUANTGRAD_IMPL(T) \