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

Support SCELossInternal/SCELossInternalGrad run with larger sized input (#16363)

Browse source 
### 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
This commit is contained in:
pengwa 2023-06-30 08:36:06 +08:00 committed by GitHub
parent 322237f482
commit 8fc3037ff4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
17 changed files with 646 additions and 323 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh
View file

@ -16,12 +16,12 @@ constexpr int kElementsPerThread = GridDim::maxElementsPerThread;
constexpr int kThreadsPerBlock = GridDim::maxThreadsPerBlock;
#endif
template <typename T, typename FuncT>
__global__ void ElementwiseKernel(T* output_data, const FuncT functor, CUDA_LONG N) {
CUDA_LONG start = kElementsPerThread * kThreadsPerBlock * blockIdx.x + threadIdx.x;
template <typename T, typename FuncT, typename TIndex>
__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 <typename T, typename FuncT>
void LaunchElementwiseKernel(cudaStream_t stream, T* output_data, const FuncT& functor, size_t output_size) {
template <typename T, typename FuncT, typename TIndex>
void LaunchElementwiseKernel(cudaStream_t stream, T* output_data, const FuncT& functor, TIndex output_size) {
if (output_size == 0) return;
CUDA_LONG N = static_cast<CUDA_LONG>(output_size);
int blocksPerGrid = CeilDiv(N, kThreadsPerBlock * kElementsPerThread);
ElementwiseKernel<T, FuncT><<<blocksPerGrid, kThreadsPerBlock, 0, stream>>>(output_data, functor, N);
TIndex N = output_size;
int blocksPerGrid = static_cast<int>(CeilDiv(N, static_cast<TIndex>(kThreadsPerBlock * kElementsPerThread)));
ElementwiseKernel<T, FuncT, TIndex><<<blocksPerGrid, kThreadsPerBlock, 0, stream>>>(output_data, functor, N);
}
} // namespace cuda

32
onnxruntime/core/providers/cuda/shared_inc/fast_divmod.h
View file

@ -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 <typename T>
struct DivMod {
DivMod(T d = 1) {
d_ = d == 0 ? 1 : d;
ORT_ENFORCE(d_ >= 1 && d_ <= std::numeric_limits<T>::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<int> {
DivMod(int d = 1) {
d_ = d == 0 ? 1 : d;
ORT_ENFORCE(d_ >= 1 && d_ <= static_cast<uint32_t>(std::numeric_limits<int>::max()));
@ -58,5 +84,7 @@ struct fast_divmod {
uint32_t l_; // l_ = ceil(log2(d_))
};
using fast_divmod = DivMod<int>; // Keep the old name for backward compatibility.
} // namespace cuda
} // namespace onnxruntime

32
onnxruntime/core/providers/rocm/shared_inc/fast_divmod.h
View file

@ -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 <typename T>
struct DivMod {
DivMod(T d = 1) {
d_ = d == 0 ? 1 : d;
ORT_ENFORCE(d_ >= 1 && d_ <= std::numeric_limits<T>::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<int> {
DivMod(int d = 1) {
d_ = d == 0 ? 1 : d;
ORT_ENFORCE(d_ >= 1 && d_ <= static_cast<uint32_t>(std::numeric_limits<int>::max()));
@ -58,5 +84,7 @@ struct fast_divmod {
uint32_t l_; // l_ = ceil(log2(d_))
};
using fast_divmod = DivMod<int>; // Keep the old name for backward compatibility.
} // namespace rocm
} // namespace onnxruntime

14
onnxruntime/core/util/math.h
View file

@ -53,15 +53,15 @@ enum StorageOrder {
namespace math {
template <typename T, class Provider>
void Exp(int N, const T* x, T* y, Provider* provider);
void Exp(ptrdiff_t N, const T* x, T* y, Provider* provider);
template <typename T, class Provider>
void Log(int N, const T* x, T* y, Provider* provider);
void Log(ptrdiff_t N, const T* x, T* y, Provider* provider);
template <typename T, class Provider>
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 <typename T, class Provider> \
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 <typename T, class Provider> \
void name##ToRow(int M, int N, const T* a, const T* b, T* y, Provider* provider); \
template <typename T, class Provider> \
@ -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 <typename T, class Provider>
void Sum(int N, const T* x, T* y, Provider* provider);
void Sum(ptrdiff_t N, const T* x, T* y, Provider* provider);
template <typename T, class Provider>
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 <typename T, class Provider>
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 <typename T>
void MatMul(

34
onnxruntime/core/util/math_cpu.cc
View file

@ -250,10 +250,10 @@ template void Gemv<double, CPUMathUtil>(const CBLAS_TRANSPOSE TransA, int M, int
SPECIALIZED_AXPY(float)
#undef SPECIALIZED_AXPY
#define DELEGATE_SIMPLE_UNARY_FUNCTION(T, Funcname, expr) \
template <> \
void Funcname<T, CPUMathUtil>(int N, const T* x, T* y, CPUMathUtil*) { \
EigenVectorMap<T>(y, N) = ConstEigenVectorMap<T>(x, N).array().expr(); \
#define DELEGATE_SIMPLE_UNARY_FUNCTION(T, Funcname, expr) \
template <> \
void Funcname<T, CPUMathUtil>(ptrdiff_t N, const T* x, T* y, CPUMathUtil*) { \
EigenVectorMap<T>(y, N) = ConstEigenVectorMap<T>(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<T, CPUMathUtil>(int N, const T* a, const T* b, T* y, CPUMathUtil*) { \
void Funcname<T, CPUMathUtil>(ptrdiff_t N, const T* a, const T* b, T* y, CPUMathUtil*) { \
EigenVectorMap<T>(y, N) = ConstEigenVectorMap<T>(a, N).array() expr ConstEigenVectorMap<T>(b, N).array(); \
}
@ -906,10 +906,10 @@ SPECIALIZED_ROWWISESUM(int64_t)
SPECIALIZED_ROWWISESUM(double)
#undef SPECIALIZED_ROWWISESUM
#define SPECIALIZED_SUM(T) \
template <> \
void Sum<T, CPUMathUtil>(int N, const T* x, T* y, CPUMathUtil* /* unused */) { \
*y = ConstEigenVectorMap<T>(x, N).sum(); \
#define SPECIALIZED_SUM(T) \
template <> \
void Sum<T, CPUMathUtil>(ptrdiff_t N, const T* x, T* y, CPUMathUtil* /* unused */) { \
*y = ConstEigenVectorMap<T>(x, N).sum(); \
}
SPECIALIZED_SUM(float);
@ -918,14 +918,14 @@ SPECIALIZED_SUM(int64_t);
#undef SPECIALIZED_SUM
#define SPECIALIZED_SCALE(T) \
template <> \
void Scale<T, CPUMathUtil>(int n, float alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \
EigenVectorMap<T>(y, n) = ConstEigenVectorMap<T>(x, n) * alpha; \
} \
template <> \
void Scale<T, CPUMathUtil>(int n, const float* alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \
EigenVectorMap<T>(y, n) = ConstEigenVectorMap<T>(x, n) * (*alpha); \
#define SPECIALIZED_SCALE(T) \
template <> \
void Scale<T, CPUMathUtil>(ptrdiff_t n, float alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \
EigenVectorMap<T>(y, n) = ConstEigenVectorMap<T>(x, n) * alpha; \
} \
template <> \
void Scale<T, CPUMathUtil>(ptrdiff_t n, const float* alpha, const T* x, T* y, CPUMathUtil* /*provider*/) { \
EigenVectorMap<T>(y, n) = ConstEigenVectorMap<T>(x, n) * (*alpha); \
}
SPECIALIZED_SCALE(float)
#undef SPECIALIZED_SCALE

87
onnxruntime/test/common/tensor_op_test_utils.h
View file

@ -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<float>()) << "Compare with non float number is not supported yet. ";
ASSERT_TRUE(expected_tensor.DataType() == DataTypeImpl::GetType<float>())
<< "Compare with non float number is not supported yet. ";
auto expected = expected_tensor.Data<float>();
auto output = output_tensor.Data<float>();
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 TFloat>
typename std::enable_if<
std::is_floating_point<TFloat>::value,
std::vector<TFloat>>::type
Uniform(gsl::span<const int64_t> 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<TFloat> 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<std::ptrdiff_t>(base_seed) + begin;
if (new_seed < static_cast<std::ptrdiff_t>(std::numeric_limits<RandomSeedType>::max()))
seed = static_cast<RandomSeedType>(new_seed);
RandomEngine generator{seed};
std::uniform_real_distribution<TFloat> 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 TFloat16>
typename std::enable_if<
std::is_same_v<TFloat16, MLFloat16>,
std::vector<TFloat16>>::type
Uniform(gsl::span<const int64_t> 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<TFloat16> 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<std::ptrdiff_t>(base_seed) + begin;
if (new_seed < static_cast<std::ptrdiff_t>(std::numeric_limits<RandomSeedType>::max()))
seed = static_cast<RandomSeedType>(new_seed);
RandomEngine generator{seed};
std::uniform_real_distribution<float> 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

6
onnxruntime/test/providers/provider_test_utils.h
View file

@ -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<float>(f_datat, input_size);
auto output_vector = EigenVectorMap<Eigen::half>(static_cast<Eigen::half*>(static_cast<void*>(h_data)), input_size);
output_vector = in_vector.template cast<Eigen::half>();
}
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<Eigen::half>(static_cast<const Eigen::half*>(static_cast<const void*>(h_data)), input_size);
auto output_vector = EigenVectorMap<float>(f_data, input_size);
@ -24,7 +24,7 @@ inline void ConvertMLFloat16ToFloat(const MLFloat16* h_data, float* f_data, int
inline std::vector<MLFloat16> FloatsToMLFloat16s(const std::vector<float>& f) {
std::vector<MLFloat16> m(f.size());
ConvertFloatToMLFloat16(f.data(), m.data(), static_cast<int>(f.size()));
ConvertFloatToMLFloat16(f.data(), m.data(), f.size());
return m;
}

126
onnxruntime/test/util/compare_ortvalue.cc
View file

@ -2,8 +2,11 @@
// Licensed under the MIT License.
#include "test/compare_ortvalue.h"
#include <cmath>
#include <sstream>
#include <mutex>
#include <thread>
#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 <core/session/onnxruntime_cxx_api.h>
#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 <typename T>
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 <typename FLOAT_TYPE>
@ -456,6 +481,93 @@ static std::pair<COMPARE_RESULT, std::string> CompareTensorOrtValueAndTensorType
return std::make_pair(COMPARE_RESULT::SUCCESS, "");
}
namespace {
template <typename T>
float ParseValueToFloat(T data_value) {
return static_cast<float>(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 <typename RealT, typename ExpectT>
std::pair<COMPARE_RESULT, std::string> CompareFloat16WithFloatResult(const Tensor& outvalue,
const Tensor& expected_value,
double per_sample_tolerance,
double relative_per_sample_tolerance) {
const size_t size1 = static_cast<size_t>(expected_value.Shape().Size());
const ExpectT* expected_output = expected_value.Data<ExpectT>();
const RealT* real_output = outvalue.Data<RealT>();
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<std::ptrdiff_t>(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<float>(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<COMPARE_RESULT, std::string> 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<Tensor>();
const Tensor& expected_tensor = expected_mlvalue.Get<Tensor>();
if (outvalue.DataType() == DataTypeImpl::GetType<float>() &&
expected_tensor.DataType() == DataTypeImpl::GetType<MLFloat16>()) {
return CompareFloat16WithFloatResult<float, MLFloat16>(outvalue,
expected_tensor,
per_sample_tolerance,
relative_per_sample_tolerance);
} else if (outvalue.DataType() == DataTypeImpl::GetType<MLFloat16>() &&
expected_tensor.DataType() == DataTypeImpl::GetType<float>()) {
return CompareFloat16WithFloatResult<MLFloat16, float>(outvalue,
expected_tensor,
per_sample_tolerance,
relative_per_sample_tolerance);
}
}
return std::make_pair(COMPARE_RESULT::NOT_SUPPORT, "Unsupported compare with CompareOrtValueNumerals.");
}
std::pair<COMPARE_RESULT, std::string> 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, "");

28
onnxruntime/test/util/include/test/compare_ortvalue.h
View file

@ -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<COMPARE_RESULT, std::string> 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<COMPARE_RESULT, std::string> 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<COMPARE_RESULT, std::string> VerifyValueInfo(const ONNX_NAMESPACE::ValueInfoProto& expected,
const OrtValue* value);

314
orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc
View file

@ -1,6 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <iostream>
#include <memory>
#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<int64_t>* X_dims,
const std::vector<int64_t>* index_dims, const std::vector<int64_t>* weight_dims,
const std::int64_t ignore_index,
std::vector<float>& X_data, std::vector<int64_t>& index_data,
std::vector<float>& weight_data) {
template <typename FloatT>
void PrepareSCELossTestData(const std::vector<int64_t>* X_dims,
const std::vector<int64_t>* index_dims, const std::vector<int64_t>* weight_dims,
const std::int64_t ignore_index,
std::vector<FloatT>& X_data,
std::vector<int64_t>& index_data,
std::vector<FloatT>& weight_data) {
// create rand inputs
ParallelRandomValueGenerator prandom{2333};
X_data = prandom.Uniform<FloatT>(*X_dims, -200.0f, 200.0f);
RandomValueGenerator random{2333};
X_data = random.Uniform<float>(*X_dims, -200.0f, 200.0f);
index_data = random.Uniform<int64_t>(*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<int64_t>* X_dims,
}
if (weight_dims) {
weight_data = random.Uniform<float>(*weight_dims, 0.0f, 1.0f);
weight_data = prandom.Uniform<FloatT>(*weight_dims, 0.0f, 1.0f);
}
}
@ -316,9 +321,9 @@ static std::vector<OrtValue> RunSCELossWithEP(const char* op,
const std::vector<int64_t>* weight_dims,
const std::vector<int64_t>* Y_dims,
const std::vector<int64_t>* log_prob_dims,
std::vector<float>& X_data,
std::vector<T>& X_data,
std::vector<int64_t>& index_data,
std::vector<float>& weight_data) {
std::vector<T>& 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<OrtValue> 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<std::vector<float>> expected_values;
std::vector<std::vector<TOut>> expected_values;
// Still need feed the output data even we don't want to verify it.
expected_values.push_back(FillZeros<float>(*Y_dims));
expected_values.push_back(FillZeros<TOut>(*Y_dims));
if (log_prob_dims) {
expected_values.push_back(FillZeros<float>(*log_prob_dims));
expected_values.push_back(FillZeros<TOut>(*log_prob_dims));
}
OpTester test(op, opset_version, domain, need_verify_outputs /*verify_output*/);
@ -349,46 +354,21 @@ static std::vector<OrtValue> RunSCELossWithEP(const char* op,
test.AddAttribute("ignore_index", ignore_index);
}
if (std::is_same<T, MLFloat16>::value) {
std::vector<MLFloat16> X_data_half(X_data.size());
ConvertFloatToMLFloat16(X_data.data(), X_data_half.data(), static_cast<int>(X_data.size()));
test.AddInput<MLFloat16>("X", *X_dims, X_data_half);
} else {
test.AddInput<float>("X", *X_dims, X_data);
}
test.AddInput<T>("X", *X_dims, X_data);
test.AddInput<int64_t>("index", *index_dims, index_data);
if (weight_dims) {
if (std::is_same<T, MLFloat16>::value) {
std::vector<MLFloat16> weight_data_half(weight_data.size());
ConvertFloatToMLFloat16(weight_data.data(), weight_data_half.data(), static_cast<int>(weight_data.size()));
test.AddInput<MLFloat16>("weight", *weight_dims, weight_data_half);
} else {
test.AddInput<float>("weight", *weight_dims, weight_data);
}
test.AddInput<T>("weight", *weight_dims, weight_data);
}
if (is_internal_op && ignore_index != -1) {
test.AddInput<int64_t>("ignore_index", {}, &ignore_index, 1);
}
if (std::is_same<TOut, MLFloat16>::value) {
std::vector<MLFloat16> output_half(expected_values[0].size());
ConvertFloatToMLFloat16(expected_values[0].data(), output_half.data(), static_cast<int>(expected_values[0].size()));
test.AddOutput<MLFloat16>("output", *Y_dims, output_half);
if (log_prob_dims) {
std::vector<MLFloat16> log_prob_half(expected_values[1].size());
ConvertFloatToMLFloat16(expected_values[1].data(), log_prob_half.data(), static_cast<int>(expected_values[1].size()));
test.AddOutput<MLFloat16>("log_prob", *log_prob_dims, log_prob_half);
}
} else {
test.AddOutput<float>("output", *Y_dims, expected_values[0]);
if (log_prob_dims) {
test.AddOutput<float>("log_prob", *log_prob_dims, expected_values[1]);
}
test.AddOutput<TOut>("output", *Y_dims, expected_values[0]);
if (log_prob_dims) {
test.AddOutput<TOut>("log_prob", *log_prob_dims, expected_values[1]);
}
std::vector<std::unique_ptr<IExecutionProvider>> eps;
@ -407,23 +387,38 @@ static void TestSCELoss(const char* op, int opset_version,
ASSERT_TRUE((std::is_same<T, MLFloat16>::value || std::is_same<T, float>::value));
ASSERT_TRUE((std::is_same<TOut, MLFloat16>::value || std::is_same<TOut, float>::value));
std::vector<float> X_data, weight_data;
std::vector<T> X_data, weight_data;
std::vector<int64_t> index_data;
PrepareSCELossTestData(X_dims, index_dims, weight_dims, ignore_index, X_data, index_data, weight_data);
PrepareSCELossTestData<T>(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<OrtValue> cpu_fetches = RunSCELossWithEP<float, float>(
op, opset_version, domain,
[]() -> std::unique_ptr<IExecutionProvider> { 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<OrtValue> cpu_fetches;
if (std::is_same<T, MLFloat16>::value) {
std::vector<float> X_data_temp(X_data.size());
std::vector<float> weight_data_temp(weight_data.size());
ConvertMLFloat16ToFloat(reinterpret_cast<MLFloat16*>(X_data.data()), X_data_temp.data(), X_data.size());
ConvertMLFloat16ToFloat(reinterpret_cast<MLFloat16*>(weight_data.data()), weight_data_temp.data(), weight_data.size());
cpu_fetches = RunSCELossWithEP<float, float>(
op, opset_version, domain,
[]() -> std::unique_ptr<IExecutionProvider> { 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<T, float>(
op, opset_version, domain,
[]() -> std::unique_ptr<IExecutionProvider> { 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<OrtValue> target_fetches = RunSCELossWithEP<T, TOut>(
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<TOut, MLFloat16>::value) {
auto y_data_size = cpu_fetches[i].Get<Tensor>().Shape().Size();
std::vector<float> cpu_temp_buffer;
cpu_temp_buffer.resize(y_data_size);
const float* y_buffer = cpu_fetches[i].Get<Tensor>().Data<float>();
std::copy(y_buffer, y_buffer + y_data_size, cpu_temp_buffer.begin());
std::vector<MLFloat16> ret_half(cpu_temp_buffer.size());
ConvertFloatToMLFloat16(cpu_temp_buffer.data(), ret_half.data(), static_cast<int>(cpu_temp_buffer.size()));
OrtValue target;
test::CreateInputOrtValueOnCPU<MLFloat16>(cpu_fetches[i].Get<Tensor>().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<TOut, float>::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<float, float>(&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<int64_t> X_dims{bsz, vocab_size};
std::vector<int64_t> index_dims{bsz};
std::vector<int64_t> weight_dims{vocab_size};
std::vector<int64_t> Y_dims{};
std::vector<int64_t> Y_dims_none{bsz};
std::vector<int64_t> 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<MLFloat16, MLFloat16>("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<int64_t>& dY_dims,
const std::vector<int64_t>& log_prob_dims,
const std::vector<int64_t>& index_dims,
@ -678,11 +684,11 @@ static void TestSoftmaxCrossEntropyLossGrad(const std::vector<int64_t>& dY_dims,
}
if (test_fp16) {
std::vector<MLFloat16> dY_data_half(dY_data.size());
ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), static_cast<int>(dY_data.size()));
ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), dY_data.size());
test.AddInput<MLFloat16>("dY", dY_dims, dY_data_half);
std::vector<MLFloat16> log_prob_data_half(log_prob_data.size());
ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), static_cast<int>(log_prob_data.size()));
ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), log_prob_data.size());
test.AddInput<MLFloat16>("log_prob", log_prob_dims, log_prob_data_half);
test.AddInput<int64_t>("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 <typename T, typename TOut>
void PrepareSCELossInternalGradTestData(
const std::vector<int64_t>& dY_dims, const std::vector<int64_t>& log_prob_dims,
const std::vector<int64_t>& index_dims, const std::vector<int64_t>& weight_dims,
const std::vector<int64_t>& dX_dims, const std::int64_t ignore_index, bool has_bias,
std::vector<float>& dY_data, std::vector<float>& log_prob_data,
std::vector<int64_t>& index_data, std::vector<float>& weight_data,
std::vector<float>& bias_data) {
std::vector<T>& dY_data, std::vector<T>& log_prob_data,
std::vector<int64_t>& index_data, std::vector<T>& weight_data,
std::vector<TOut>& bias_data) {
// Create rand inputs
RandomValueGenerator random{2333};
dY_data = random.Uniform<float>(dY_dims, -10.0f, 10.0f);
log_prob_data = random.Uniform<float>(log_prob_dims, -10.0f, 10.0f);
ParallelRandomValueGenerator prandom{2333};
dY_data = prandom.Uniform<T>(dY_dims, -10.0f, 10.0f);
log_prob_data = prandom.Uniform<T>(log_prob_dims, -10.0f, 10.0f);
index_data = random.Uniform<int64_t>(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<float>(weight_dims, 0.0f, 1.0f);
weight_data = prandom.Uniform<T>(weight_dims, 0.0f, 1.0f);
if (has_bias) {
bias_data = random.Uniform<float>(dX_dims, 0.0f, 1.0f);
bias_data = prandom.Uniform<TOut>(dX_dims, 0.0f, 1.0f);
}
}
@ -809,11 +817,11 @@ static std::vector<OrtValue> RunSCELossInternalGradWithEP(
const std::vector<int64_t>& index_dims,
const std::vector<int64_t>& weight_dims,
const std::vector<int64_t>& dX_dims,
std::vector<float>& dY_data,
std::vector<float>& log_prob_data,
std::vector<T>& dY_data,
std::vector<T>& log_prob_data,
std::vector<int64_t>& index_data,
std::vector<float>& weight_data,
std::vector<float>& bias_data) {
std::vector<T>& weight_data,
std::vector<TOut>& 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<OrtValue> RunSCELossInternalGradWithEP(
* So here we disable OpTester's check by default, and do the check externally.
*/
bool need_verify_outputs = false;
std::vector<float> expected_value = FillZeros<float>(dX_dims);
std::vector<TOut> expected_value = FillZeros<TOut>(dX_dims);
ORT_ENFORCE((std::is_same<T, MLFloat16>::value || std::is_same<T, float>::value));
ORT_ENFORCE((std::is_same<TOut, MLFloat16>::value || std::is_same<TOut, float>::value));
@ -840,56 +848,21 @@ static std::vector<OrtValue> RunSCELossInternalGradWithEP(
test->AddAttribute("output_type", static_cast<int64_t>(utils::ToTensorProtoElementType<TOut>()));
}
if (std::is_same<T, MLFloat16>::value) {
std::vector<MLFloat16> dY_data_half(dY_data.size());
ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), static_cast<int>(dY_data.size()));
test->AddInput<MLFloat16>("dY", dY_dims, dY_data_half);
std::vector<MLFloat16> log_prob_data_half(log_prob_data.size());
ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), static_cast<int>(log_prob_data.size()));
test->AddInput<MLFloat16>("log_prob", log_prob_dims, log_prob_data_half);
test->AddInput<int64_t>("index", index_dims, index_data);
std::vector<MLFloat16> weight_data_half(weight_data.size());
ConvertFloatToMLFloat16(weight_data.data(), weight_data_half.data(), static_cast<int>(weight_data.size()));
test->AddInput<MLFloat16>("weight", weight_dims, weight_data_half);
if (ignore_index != -1 || has_bias) {
test->AddInput<int64_t>("ignore_index", {}, &ignore_index, 1);
}
} else {
test->AddInput<float>("dY", dY_dims, dY_data);
test->AddInput<float>("log_prob", log_prob_dims, log_prob_data);
test->AddInput<int64_t>("index", index_dims, index_data);
test->AddInput<float>("weight", weight_dims, weight_data);
if (ignore_index != -1 || has_bias) {
test->AddInput<int64_t>("ignore_index", {}, &ignore_index, 1);
}
test->AddInput<T>("dY", dY_dims, dY_data);
test->AddInput<T>("log_prob", log_prob_dims, log_prob_data);
test->AddInput<int64_t>("index", index_dims, index_data);
test->AddInput<T>("weight", weight_dims, weight_data);
if (ignore_index != -1 || has_bias) {
test->AddInput<int64_t>("ignore_index", {}, &ignore_index, 1);
}
if (std::is_same<TOut, MLFloat16>::value) {
// Be noted, bias should be aligned with output's data type.
if (has_bias) {
std::vector<MLFloat16> bias_data_half(bias_data.size());
ConvertFloatToMLFloat16(bias_data.data(), bias_data_half.data(), static_cast<int>(bias_data.size()));
test->AddInput<MLFloat16>("bias", dX_dims, bias_data_half);
}
std::vector<MLFloat16> expected_data_half(expected_value.size());
ConvertFloatToMLFloat16(expected_value.data(), expected_data_half.data(), static_cast<int>(expected_value.size()));
test->AddOutput<MLFloat16>("dX", dX_dims, expected_data_half, false /*sort_output*/,
error_tolerance /*rel_error*/, error_tolerance /*abs_error*/);
} else {
if (has_bias) {
test->AddInput<float>("bias", dX_dims, bias_data);
}
test->AddOutput<float>("dX", dX_dims, expected_value, false /*sort_output*/,
error_tolerance /*rel_error*/, error_tolerance /*abs_error*/);
if (has_bias) {
test->AddInput<TOut>("bias", dX_dims, bias_data);
}
test->AddOutput<TOut>("dX", dX_dims, expected_value, false /*sort_output*/,
error_tolerance /*rel_error*/, error_tolerance /*abs_error*/);
std::vector<std::unique_ptr<IExecutionProvider>> eps;
eps.emplace_back(ep_creator());
test->Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps);
@ -906,20 +879,56 @@ static void TestSoftmaxCrossEntropyLossInternalGrad(const std::vector<int64_t>&
const std::int64_t ignore_index = -1,
const double error_tolerance = 1e-4,
const bool has_bias = false) {
std::vector<float> dY_data, log_prob_data, weight_data, bias_data;
std::vector<T> dY_data, log_prob_data, weight_data;
std::vector<TOut> bias_data;
std::vector<int64_t> 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<T, TOut>(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<OrtValue> 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<OrtValue> cpu_fetches =
RunSCELossInternalGradWithEP<float, float>(
if constexpr (std::is_same<T, MLFloat16>::value) {
std::vector<float> dY_data_temp(dY_data.size());
std::vector<float> log_prob_data_temp(log_prob_data.size());
std::vector<float> weight_data_temp(weight_data.size());
ConvertMLFloat16ToFloat(reinterpret_cast<MLFloat16*>(dY_data.data()), dY_data_temp.data(), dY_data.size());
ConvertMLFloat16ToFloat(reinterpret_cast<MLFloat16*>(log_prob_data.data()), log_prob_data_temp.data(),
log_prob_data.size());
ConvertMLFloat16ToFloat(reinterpret_cast<MLFloat16*>(weight_data.data()), weight_data_temp.data(), weight_data.size());
if constexpr (std::is_same<TOut, MLFloat16>::value) {
std::vector<float> bias_data_temp(bias_data.size());
ConvertMLFloat16ToFloat(reinterpret_cast<MLFloat16*>(bias_data.data()), bias_data_temp.data(), bias_data.size());
cpu_fetches = RunSCELossInternalGradWithEP<float, float>(
[]() -> std::unique_ptr<IExecutionProvider> { 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<float, TOut>(
[]() -> std::unique_ptr<IExecutionProvider> { 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<TOut, MLFloat16>::value) {
std::vector<float> bias_data_temp(bias_data.size());
ConvertMLFloat16ToFloat(reinterpret_cast<MLFloat16*>(bias_data.data()), bias_data_temp.data(), bias_data.size());
cpu_fetches = RunSCELossInternalGradWithEP<T, float>(
[]() -> std::unique_ptr<IExecutionProvider> { 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<T, TOut>(
[]() -> std::unique_ptr<IExecutionProvider> { 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<OrtValue> target_fetches =
@ -938,22 +947,10 @@ static void TestSoftmaxCrossEntropyLossInternalGrad(const std::vector<int64_t>&
// Compare
ASSERT_EQ(cpu_fetches.size(), target_fetches.size());
for (size_t i = 0; i < cpu_fetches.size(); i++) {
if (std::is_same<TOut, MLFloat16>::value) {
auto y_data_size = cpu_fetches[i].Get<Tensor>().Shape().Size();
std::vector<float> cpu_temp_buffer;
cpu_temp_buffer.resize(y_data_size);
const float* y_buffer = cpu_fetches[i].Get<Tensor>().Data<float>();
std::copy(y_buffer, y_buffer + y_data_size, cpu_temp_buffer.begin());
std::vector<MLFloat16> ret_half(cpu_temp_buffer.size());
ConvertFloatToMLFloat16(cpu_temp_buffer.data(), ret_half.data(), static_cast<int>(cpu_temp_buffer.size()));
OrtValue target;
test::CreateInputOrtValueOnCPU<MLFloat16>(cpu_fetches[i].Get<Tensor>().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<TOut, float>::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<int64_t> dY_dims{};
std::vector<int64_t> log_prob_dims{bsz, vocab_size};
std::vector<int64_t> index_dims{bsz};
std::vector<int64_t> weight_dims{vocab_size};
std::vector<int64_t> dX_dims{bsz, vocab_size};
TestSoftmaxCrossEntropyLossInternalGrad<MLFloat16, MLFloat16>(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

38
orttraining/orttraining/training_ops/cpu/loss/cross_entropy.cc
View file

@ -6,7 +6,6 @@
#include "core/util/math_cpuonly.h"
#include "core/providers/common.h"
#include <unsupported/Eigen/SpecialFunctions>
#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 <typename T>
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<T> logit_max(n);
math::RowwiseMax<T, CPUMathUtil>(n, d, logit_data, logit_max.data(), nullptr);
std::vector<T> logit_max(nd);
math::RowwiseMax<T, CPUMathUtil>(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<T, CPUMathUtil>(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<T, CPUMathUtil>(nd, c, logit_max.data(), shifted_logit, nullptr);
// exp_shifted_logit = exp(shifted_logit)
math::Exp<T, CPUMathUtil>(nd, shifted_logit, log_prob_data, nullptr);
math::Exp<T, CPUMathUtil>(nd_c, shifted_logit, log_prob_data, nullptr);
// sum_exp = sum_{class} (exp_shifted_logit)
float* sum_exp = logit_max.data();
math::RowwiseSum<T, CPUMathUtil>(n, d, log_prob_data, sum_exp, nullptr);
math::RowwiseSum<T, CPUMathUtil>(nd, c, log_prob_data, sum_exp, nullptr);
// log_sum_exp = log(sum_exp)
float* log_sum_exp = sum_exp;
math::Log<T, CPUMathUtil>(n, sum_exp, log_sum_exp, nullptr);
math::Log<T, CPUMathUtil>(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<T, CPUMathUtil>(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<T, CPUMathUtil>(nd, c, log_sum_exp, log_prob_data, nullptr);
}
ONNX_OPERATOR_KERNEL_EX(
@ -70,7 +69,7 @@ Status SoftmaxCrossEntropy<T>::Compute(OpKernelContext* context) const {
int64_t D = logit_shape[logit_shape.NumDimensions() - 1];
const int n = gsl::narrow_cast<int>(N);
const int d = gsl::narrow_cast<int>(D);
const int nd = gsl::narrow_cast<int>(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<T>::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<uint64_t>(std::numeric_limits<Eigen::Index>::max()));
ComputeShareSoftmaxCrossEntropyCPU(n, d, static_cast<Eigen::Index>(nd),
logit_data,
shifted_logit.data(),
log_prob_data);
@ -180,7 +182,7 @@ Status SparseSoftmaxCrossEntropy<T>::Compute(OpKernelContext* context) const {
int64_t D = logit_shape[logit_shape.NumDimensions() - 1];
const int n = gsl::narrow_cast<int>(N);
const int d = gsl::narrow_cast<int>(D);
const int nd = gsl::narrow_cast<int>(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<T>::Compute(OpKernelContext* context) const {
// computation begins here
std::vector<float> shifted_logit(nd);
ComputeShareSoftmaxCrossEntropyCPU(n, d, nd, logit_data,
ORT_ENFORCE(nd <= static_cast<uint64_t>(std::numeric_limits<Eigen::Index>::max()));
ComputeShareSoftmaxCrossEntropyCPU(n, d, static_cast<Eigen::Index>(nd), logit_data,
shifted_logit.data(),
log_prob_data);

3
orttraining/orttraining/training_ops/cpu/loss/cross_entropy.h
View file

@ -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 <typename T>
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);

155
orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc
View file

@ -5,7 +5,6 @@
#include "core/util/math_cpuonly.h"
#include "core/providers/common.h"
#include <unsupported/Eigen/SpecialFunctions>
#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<T1, T2>::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<T1>();
OrtValue transpose_output;
@ -133,7 +132,7 @@ Status SoftmaxCrossEntropyLoss<T1, T2>::Compute(OpKernelContext* context) const
const int n_d = gsl::narrow_cast<int>(N_D);
const int c = gsl::narrow_cast<int>(C);
const int n_d_c = gsl::narrow_cast<int>(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<T1> log_prob_data_buffer(0);
@ -150,7 +149,9 @@ Status SoftmaxCrossEntropyLoss<T1, T2>::Compute(OpKernelContext* context) const
const T2* label_data = label.template Data<T2>();
T1* loss_data = loss->template MutableData<T1>();
std::vector<T1> 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<uint64_t>(std::numeric_limits<Eigen::Index>::max()));
ComputeShareSoftmaxCrossEntropyCPU(n_d, c, static_cast<Eigen::Index>(n_d_c), logit_data, shifted_logit.data(),
log_prob_data);
std::vector<T1> loss_sample_buffer(0);
T1* loss_sample;
if (reduction_ == ReductionType::NONE) {
@ -260,17 +261,15 @@ Status SoftmaxCrossEntropyLossGrad<T1, T2>::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<int>(N_D);
const int c = gsl::narrow_cast<int>(C);
const T1* dY_data = dY.template Data<T1>();
const T1* log_prob_data = log_prob.template Data<T1>();
const T2* label_data = label.template Data<T2>();
Tensor* d_logit = context->Output(0, probability_shape);
T1* d_logit_data = d_logit->template MutableData<T1>();
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<size_t> permutations;
@ -285,90 +284,107 @@ Status SoftmaxCrossEntropyLossGrad<T1, T2>::Compute(OpKernelContext* context) co
log_prob_data = (*transpose_output.GetMutable<Tensor>()).template Data<T1>();
}
// 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<T1>();
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<T1>(*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<T1>(*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<T1, T2>::Compute(OpKernelContext* context) co
if (p_bias) {
ORT_ENFORCE(probability_shape.Size() == p_bias->Shape().Size());
const T1* bias_data = p_bias->Data<T1>();
for (size_t i = 0; i < static_cast<size_t>(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();

6
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc
View file

@ -214,7 +214,7 @@ Status SoftmaxCrossEntropyLossGrad<T, TLabel, TOut>::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<T>();
@ -295,8 +295,8 @@ Status SoftmaxCrossEntropyLossGrad<T, TLabel, TOut>::ComputeInternal(OpKernelCon
onnxruntime::contrib::GetPermutationAndShape(false, logit_shape, new_shape, permutations);
transpose_output.GetMutable<Tensor>()->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<Tensor>()));
ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations,
*d_logit, *transpose_output.GetMutable<Tensor>()));
auto* transposed_data = (*transpose_output.GetMutable<Tensor>()).template Data<TOut>();
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(d_logit_data, transposed_data, sizeof(TOut) * probability_shape.Size(),
cudaMemcpyDeviceToDevice, Stream(ctx)));

58
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu
View file

@ -7,12 +7,12 @@
namespace onnxruntime {
namespace cuda {
template <typename T, typename TLabel, typename TOut, bool IsWeighted>
template <typename T, typename TLabel, typename TOut, bool IsWeighted, typename TIndex>
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 <typename T, typename TLabel, typename TOut>
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<T, TLabel, TOut, true> op(label, weight, static_cast<TLabel>(label_depth),
static_cast<TLabel>(ignore_index));
LaunchElementwiseKernel<TOut, decltype(op)>(stream, weight_data_nd, op, count);
typedef OpSoftmaxCrossEntropyWeights<T, TLabel, TOut, true, size_t> OP_Type;
OP_Type op(label, weight, static_cast<TLabel>(label_depth), static_cast<TLabel>(ignore_index));
LaunchElementwiseKernel<TOut, OP_Type, size_t>(stream, weight_data_nd, op, count);
} else {
OpSoftmaxCrossEntropyWeights<T, TLabel, TOut, false> op(label, nullptr, static_cast<TLabel>(label_depth),
static_cast<TLabel>(ignore_index));
LaunchElementwiseKernel<TOut, decltype(op)>(stream, weight_data_nd, op, count);
typedef OpSoftmaxCrossEntropyWeights<T, TLabel, TOut, false, size_t> OP_Type;
OP_Type op(label, nullptr, static_cast<TLabel>(label_depth), static_cast<TLabel>(ignore_index));
LaunchElementwiseKernel<TOut, OP_Type, size_t>(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 <typename T, typename TAcc, typename TLabel>
template <typename T, typename TAcc, typename TLabel, typename TIndex>
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<T>(static_cast<TAcc>(-log_prob_data_[idx * C_ + label_data_[idx]] * weight_data_[idx]) /
(*normalize_factor_data_));
T ret = static_cast<T>(static_cast<TAcc>(-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 <typename T, typename TAcc, typename TLabel>
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<T, TAcc, TLabel> op(log_prob, label, weight, normalize_factor,
static_cast<TLabel>(label_depth), static_cast<TLabel>(ignore_index));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count);
typedef OpWeightedSoftmaxCrossEntropyLoss<T, TAcc, TLabel, size_t> OP_Type;
OP_Type op(log_prob, label, weight, normalize_factor,
static_cast<TLabel>(label_depth),
static_cast<TLabel>(ignore_index));
LaunchElementwiseKernel<T, OP_Type, size_t>(stream, output_data, op, count);
}
template <typename T, typename TAcc, typename TLabel, typename TOut, bool IsReductionNone, bool HasBias>
template <typename T, typename TAcc, typename TLabel, typename TOut, bool IsReductionNone, bool HasBias, typename TIndex>
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<int>(C));
C_fdm_ = DivMod(static_cast<TIndex>(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<TAcc>((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<TIndex> C_fdm_;
};
template <typename T, typename TAcc, typename TLabel, typename TOut>
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<T, TAcc, TLabel, TOut, is_reduction_none, has_bias> op( \
dY, log_prob, label, weight, normalize_factor, bias_data, static_cast<TLabel>(label_depth)); \
LaunchElementwiseKernel<TOut, decltype(op)>(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<uint64_t>(std::numeric_limits<int>::max())) { \
typedef OpWeightedSoftmaxCrossEntropyLossGrad<T, TAcc, TLabel, TOut, is_reduction_none, has_bias, int> OP_Type; \
OP_Type op(dY, log_prob, label, weight, normalize_factor, bias_data, static_cast<TLabel>(label_depth)); \
LaunchElementwiseKernel<TOut, OP_Type, int>(stream, output_data, op, static_cast<int>(total_count)); \
} else { \
typedef OpWeightedSoftmaxCrossEntropyLossGrad<T, TAcc, TLabel, TOut, is_reduction_none, has_bias, uint64_t> OP_Type; \
OP_Type op(dY, log_prob, label, weight, normalize_factor, bias_data, static_cast<TLabel>(label_depth)); \
LaunchElementwiseKernel<TOut, decltype(op), uint64_t>(stream, output_data, op, total_count); \
}
if (reduction_none) {
if (bias_data) {
LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL(true, true);

14
orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu
View file

@ -25,7 +25,7 @@ template <typename T>
void SoftMaxCrossEntropyImpl(cudaStream_t stream, const T* log_prob, const T* label, size_t normalize_factor,
T* output_data, size_t count) {
OpSoftmaxCrossEntropy<T> op(log_prob, label, static_cast<T>(normalize_factor));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count);
LaunchElementwiseKernel<T, decltype(op), CUDA_LONG>(stream, output_data, op, static_cast<CUDA_LONG>(count));
}
template void SoftMaxCrossEntropyImpl(cudaStream_t stream, const float* log_prob, const float* label,
@ -53,7 +53,7 @@ template <typename T>
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<T> op(dY, log_prob, label, static_cast<T>(normalize_factor));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count);
LaunchElementwiseKernel<T, decltype(op), CUDA_LONG>(stream, output_data, op, static_cast<CUDA_LONG>(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<T, Tin, true> op(log_prob, label, weight, normalize_factor,
static_cast<Tin>(label_depth));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count);
LaunchElementwiseKernel<T, decltype(op), CUDA_LONG>(stream, output_data, op, static_cast<CUDA_LONG>(count));
} else {
OpSparseSoftmaxCrossEntropy<T, Tin, false> op(log_prob, label, nullptr, normalize_factor,
static_cast<Tin>(label_depth));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count);
LaunchElementwiseKernel<T, decltype(op), CUDA_LONG>(stream, output_data, op, static_cast<CUDA_LONG>(count));
}
}
@ -136,11 +136,13 @@ void SparseSoftmaxCrossEntropyGradImpl(cudaStream_t stream, const T* dY, const T
if (weight) {
OpSparseSoftmaxCrossEntropyGrad<T, Tin, true> op(dY, log_prob, label, weight, normalize_factor,
fast_divmod(static_cast<int>(label_depth)));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count * label_depth);
LaunchElementwiseKernel<T, decltype(op), CUDA_LONG>(stream, output_data, op,
static_cast<CUDA_LONG>(count * label_depth));
} else {
OpSparseSoftmaxCrossEntropyGrad<T, Tin, false> op(dY, log_prob, label, nullptr, normalize_factor,
fast_divmod(static_cast<int>(label_depth)));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count * label_depth);
LaunchElementwiseKernel<T, decltype(op), CUDA_LONG>(stream, output_data, op,
static_cast<CUDA_LONG>(count * label_depth));
}
}

4
orttraining/orttraining/training_ops/cuda/quantization/fake_quant_impl.cu
View file

@ -98,8 +98,8 @@ template <typename T>
void FakeQuantGradImpl(cudaStream_t stream, const int64_t num_elements, const T* dY_data,
const bool* gradient_mask_data, T* dX_data) {
FakeQuantGradFunctor<T> fake_quant_grad_functor(dY_data, gradient_mask_data);
LaunchElementwiseKernel<T, decltype(fake_quant_grad_functor)>(
stream, dX_data, fake_quant_grad_functor, num_elements);
LaunchElementwiseKernel<T, decltype(fake_quant_grad_functor), CUDA_LONG>(
stream, dX_data, fake_quant_grad_functor, static_cast<CUDA_LONG>(num_elements));
}
#define SPECIALIZED_FAKEQUANTGRAD_IMPL(T) \