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

SCELoss(SCELossGrad) support half(float) input float(half) output (#13972)

Browse source 
### Description

A follow up change for
https://github.com/microsoft/onnxruntime/pull/13616.

SoftmaxCrossEntropyLossInternal/SoftmaxCrossEntropyLossInternalGrad
support different type for input and output.

Add SCELoss(SCELossGrad) support half(float) input float(half) output

### Test Note

#### Add tests for variant input and output types. To add such tests,
have to refactor existing testing code for sce loss and scelossinternal
gradient.

Originally, 

FP32 input and output, the CPU kernels, runs with CPU kernels the
baseline, CUDA/RCOM then runs with same data, user CompareTester to
compare with CPU run results.

FP16 input and output, the CPU kernels (did not have half kernels), runs
with Cast_to_float->CPU kernel->cast_to_half as the baseline, CUDA/RCOM
then runs with same data but using Half implementation, user
CompareTester to compare with CPU run results.

Now, we want the support run different input and output types. The
proposed change here is, to run CPU kernels always with float input and
output as baseline (because CPU only have float type kernels impl), this
step is the very first thing for every test.

Then, we run CUDA/RCOM kernels using half_input_half_output,
float_input_float_output, half_input_float_output,
float_input_half_output if there is corresponding kernel registered.

Afterwards, compare the CUDA/ROCM run results with CPU float baselines. 

Be noted, there is one thing that deserved a special note:
CompareOpTester's result compare can be loose than OpTester's.
Roughly speaking: the former tolerant diff <= atol +
rtol*expected_value, while the later one telerant diff < atol && diff <
rtol*expected_value. When the expected value is super small in many
cases of our tests cases, the former one can pass but the later one
fails. So the refactoring also move the check outside of OpTester,
explicitly check the values using the way CompareOPTester did (to align
the previous behaviour).

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
pengwa 2023-02-28 18:02:08 +08:00 committed by GitHub
parent 08699c8052
commit 79aa0acdd0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
23 changed files with 1193 additions and 639 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/core/providers/cpu/cpu_provider_shared.h
View file

@ -277,7 +277,9 @@ inline void wait_event_in_tensor(const Tensor& event_id_tensor) { return g_host_
// From aten_op.h
inline bool IsATenOperatorExecutorInitialized() { return g_host_cpu.contrib__IsATenOperatorExecutorInitialized(); }
inline Status ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span<const int64_t>& axes, bool keepdims) { return g_host_cpu.contrib__ExecuteReduceSumATen(p_ctx, axes, keepdims); }
inline Status ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span<const int64_t>& axes, bool keepdims) {
return g_host_cpu.contrib__ExecuteReduceSumATen(p_ctx, axes, keepdims);
}
} // namespace contrib
#endif // ENABLE_TRAINING
#endif // USE_CUDA || USE_ROCM

47
onnxruntime/core/providers/cuda/math/softmax.cc
View file

@ -11,37 +11,44 @@
namespace onnxruntime {
namespace cuda {
template <typename T, bool is_log_softmax>
template <typename T, typename TOut, bool is_log_softmax>
Status SoftMaxComputeHelper(
cudaStream_t stream,
const T* X,
const TensorShape& input_shape,
T* Y,
TOut* Y,
int64_t axis) {
typedef typename ToCudaType<T>::MappedType CudaT;
typedef typename ToCudaType<T>::MappedType CudaT_IN;
typedef typename ToCudaType<TOut>::MappedType CudaT_OUT;
typedef typename ToCudaType<T>::MappedType CudaT_ACCUM;
int64_t N = input_shape.SizeToDimension(axis);
int64_t D = input_shape.SizeFromDimension(axis);
auto Y_data = reinterpret_cast<CudaT*>(Y);
auto X_data = reinterpret_cast<const CudaT*>(X);
auto Y_data = reinterpret_cast<CudaT_OUT*>(Y);
auto X_data = reinterpret_cast<const CudaT_IN*>(X);
if (D <= 1024 && D * sizeof(T) <= 4096) {
return dispatch_warpwise_softmax_forward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>(
return dispatch_warpwise_softmax_forward<
CudaT_IN, CudaT_OUT, AccumulationType_t<CudaT_ACCUM>, is_log_softmax>(
stream, Y_data, X_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(N));
}
return dispatch_blockwise_softmax_forward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>(
return dispatch_blockwise_softmax_forward<CudaT_IN, CudaT_OUT, AccumulationType_t<CudaT_ACCUM>, is_log_softmax>(
stream, Y_data, X_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D),
gsl::narrow_cast<int>(N));
}
#define SPECIALIZED_SOFTMAX_HELPER_IMPL(T) \
template Status SoftMaxComputeHelper<T, false>(cudaStream_t stream, const T* input, const TensorShape& shape, T* Y, int64_t axis); \
template Status SoftMaxComputeHelper<T, true>(cudaStream_t stream, const T* input, const TensorShape& shape, T* Y, int64_t axis);
#define SPECIALIZED_SOFTMAX_HELPER_IMPL(T, TOut) \
template Status SoftMaxComputeHelper<T, TOut, false>(cudaStream_t stream, const T* input, \
const TensorShape& shape, TOut* Y, int64_t axis); \
template Status SoftMaxComputeHelper<T, TOut, true>(cudaStream_t stream, const T* input, \
const TensorShape& shape, TOut* Y, int64_t axis);
SPECIALIZED_SOFTMAX_HELPER_IMPL(float)
SPECIALIZED_SOFTMAX_HELPER_IMPL(double)
SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16)
SPECIALIZED_SOFTMAX_HELPER_IMPL(BFloat16)
SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16, float)
SPECIALIZED_SOFTMAX_HELPER_IMPL(float, float)
SPECIALIZED_SOFTMAX_HELPER_IMPL(double, double)
SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16, MLFloat16)
SPECIALIZED_SOFTMAX_HELPER_IMPL(BFloat16, BFloat16)
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
@ -170,13 +177,13 @@ Status Softmax<T>::ComputeInternal(OpKernelContext* ctx) const {
Status status;
if (log_softmax_) {
status = SoftMaxComputeHelper<T, true>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis));
status = SoftMaxComputeHelper<T, T, true>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis));
} else {
status = SoftMaxComputeHelper<T, false>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis));
status = SoftMaxComputeHelper<T, T, false>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis));
}
if (!status.IsOK())

4
onnxruntime/core/providers/cuda/math/softmax.h
View file

@ -9,12 +9,12 @@
namespace onnxruntime {
namespace cuda {
template <typename T, bool is_log_softmax>
template <typename T, typename TOut, bool is_log_softmax>
Status SoftMaxComputeHelper(
cudaStream_t stream,
const T* input,
const TensorShape& shape,
T* Y,
TOut* Y,
int64_t axis);
template <typename input_t, typename output_t, typename acc_t, bool is_log_softmax>

58
onnxruntime/core/providers/cuda/math/softmax_impl.cu
View file

@ -1,18 +1,18 @@
/**
* Copyright (c) 2016-present, Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
* Copyright (c) 2016-present, Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/* Modifications Copyright (c) Microsoft. */
@ -29,7 +29,8 @@ namespace onnxruntime {
namespace cuda {
template <typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
Status dispatch_warpwise_softmax_forward(cudaStream_t stream, output_t* dst, const input_t* src, int softmax_elements, int softmax_elements_stride, int batch_count) {
Status dispatch_warpwise_softmax_forward(cudaStream_t stream, output_t* dst, const input_t* src, int softmax_elements,
int softmax_elements_stride, int batch_count) {
if (softmax_elements == 0) {
return Status::OK();
} else {
@ -102,16 +103,23 @@ Status dispatch_warpwise_softmax_forward(cudaStream_t stream, output_t* dst, con
return CUDA_CALL(cudaGetLastError());
}
#define SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(input_t, output_t, acc_t) \
template Status dispatch_warpwise_softmax_forward<input_t, output_t, acc_t, false>(cudaStream_t stream, output_t * dst, \
const input_t* src, int softmax_elements, \
int softmax_elements_stride, int batch_count); \
template Status dispatch_warpwise_softmax_forward<input_t, output_t, acc_t, true>(cudaStream_t stream, output_t * dst, \
const input_t* src, int softmax_elements, \
int softmax_elements_stride, int batch_count);
#define SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(input_t, output_t, acc_t) \
template Status dispatch_warpwise_softmax_forward<input_t, output_t, acc_t, false>(cudaStream_t stream, \
output_t * dst, \
const input_t* src, \
int softmax_elements, \
int softmax_elements_stride, \
int batch_count); \
template Status dispatch_warpwise_softmax_forward<input_t, output_t, acc_t, true>(cudaStream_t stream, \
output_t * dst, \
const input_t* src, \
int softmax_elements, \
int softmax_elements_stride, \
int batch_count);
SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(float, float, float)
SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(half, half, float)
SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(half, float, float)
SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(double, double, double)
SPECIALIZED_WRAPWISE_SOFTMAX_IMPL(BFloat16, BFloat16, float)
@ -143,12 +151,8 @@ Status dispatch_blockwise_softmax_forward(cudaStream_t stream, output_t* output,
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(float, float, float)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(half, half, float)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(half, float, float)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(double, double, double)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(BFloat16, BFloat16, float)
#ifndef DISABLE_CONTRIB_OPS
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(half, float, float) // used by BeamSearch op
#endif
} // namespace cuda
} // namespace onnxruntime

51
onnxruntime/core/providers/rocm/math/softmax.cc
View file

@ -11,42 +11,45 @@
namespace onnxruntime {
namespace rocm {
template <typename T, bool IsLogSoftmax>
template <typename T, typename TOut, bool IsLogSoftmax>
Status SoftMaxComputeHelper(
hipStream_t stream,
const T* X,
const TensorShape& input_shape,
T* Y,
TOut* Y,
int64_t axis,
RocmTuningContext* tuning_ctx) {
typedef typename ToHipType<T>::MappedType HipT;
typedef typename ToHipType<T>::MappedType HipT_IN;
typedef typename ToHipType<TOut>::MappedType HipT_OUT;
typedef typename ToHipType<T>::MappedType HipT_ACCUM;
int64_t N = input_shape.SizeToDimension(axis);
int64_t D = input_shape.SizeFromDimension(axis);
auto Y_data = reinterpret_cast<HipT*>(Y);
auto X_data = reinterpret_cast<const HipT*>(X);
auto Y_data = reinterpret_cast<HipT_OUT*>(Y);
auto X_data = reinterpret_cast<const HipT_IN*>(X);
if (D <= 1024 && D * sizeof(T) <= 4096) {
return dispatch_warpwise_softmax_forward<HipT, HipT, AccumulationType_t<HipT>, IsLogSoftmax>(
return dispatch_warpwise_softmax_forward<HipT_IN, HipT_OUT, AccumulationType_t<HipT_ACCUM>, IsLogSoftmax>(
stream, Y_data, X_data, gsl::narrow_cast<int>(D),
gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(N), tuning_ctx);
}
return dispatch_blockwise_softmax_forward<HipT, HipT, AccumulationType_t<HipT>, IsLogSoftmax>(
return dispatch_blockwise_softmax_forward<HipT_IN, HipT_OUT, AccumulationType_t<HipT_ACCUM>, IsLogSoftmax>(
stream, Y_data, X_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D),
gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(N), tuning_ctx);
}
#define SPECIALIZED_SOFTMAX_HELPER_IMPL(T) \
template Status SoftMaxComputeHelper<T, false>(hipStream_t stream, const T* input, const TensorShape& shape, T* Y, \
int64_t axis, RocmTuningContext* tuning_ctx); \
template Status SoftMaxComputeHelper<T, true>(hipStream_t stream, const T* input, const TensorShape& shape, T* Y, \
int64_t axis, RocmTuningContext* tuning_ctx);
#define SPECIALIZED_SOFTMAX_HELPER_IMPL(T, TOut) \
template Status SoftMaxComputeHelper<T, TOut, false>(hipStream_t stream, const T* input, const TensorShape& shape, TOut* Y, \
int64_t axis, RocmTuningContext* tuning_ctx); \
template Status SoftMaxComputeHelper<T, TOut, true>(hipStream_t stream, const T* input, const TensorShape& shape, TOut* Y, \
int64_t axis, RocmTuningContext* tuning_ctx);
SPECIALIZED_SOFTMAX_HELPER_IMPL(float)
SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16, float)
SPECIALIZED_SOFTMAX_HELPER_IMPL(float, float)
// MIOpen double data type not supported
// SPECIALIZED_SOFTMAX_HELPER_IMPL(double)
SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16)
SPECIALIZED_SOFTMAX_HELPER_IMPL(BFloat16)
// SPECIALIZED_SOFTMAX_HELPER_IMPL(double, double)
SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16, MLFloat16)
SPECIALIZED_SOFTMAX_HELPER_IMPL(BFloat16, BFloat16)
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
@ -175,15 +178,15 @@ Status Softmax<T>::ComputeInternal(OpKernelContext* ctx) const {
Status status;
if (log_softmax_) {
status = SoftMaxComputeHelper<T, true>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis),
GetTuningContext());
status = SoftMaxComputeHelper<T, T, true>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis),
GetTuningContext());
} else {
status = SoftMaxComputeHelper<T, false>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis),
GetTuningContext());
status = SoftMaxComputeHelper<T, T, false>(Stream(ctx), X_data, *compute_input_shape, Y_data,
is_transpose_required ? static_cast<int64_t>(rank) - 1
: static_cast<int64_t>(axis),
GetTuningContext());
}
if (!status.IsOK())

4
onnxruntime/core/providers/rocm/math/softmax.h
View file

@ -11,12 +11,12 @@ namespace rocm {
using tunable::RocmTuningContext;
template <typename T, bool IsLogSoftmax>
template <typename T, typename TOut, bool IsLogSoftmax>
Status SoftMaxComputeHelper(
hipStream_t stream,
const T* input,
const TensorShape& shape,
T* Y,
TOut* Y,
int64_t axis,
RocmTuningContext* tuning_ctx = nullptr);

7
onnxruntime/core/providers/rocm/math/softmax_impl.cu
View file

@ -52,6 +52,7 @@ Status dispatch_warpwise_softmax_forward(hipStream_t stream, OutputT* dst, const
SPECIALIZED_SOFTMAX_IMPL(float, float, float)
SPECIALIZED_SOFTMAX_IMPL(half, half, float)
SPECIALIZED_SOFTMAX_IMPL(half, float, float)
SPECIALIZED_SOFTMAX_IMPL(double, double, double)
SPECIALIZED_SOFTMAX_IMPL(BFloat16, BFloat16, float)
@ -79,12 +80,8 @@ Status dispatch_blockwise_softmax_forward(hipStream_t stream, OutputT* output,
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(float, float, float)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(half, half, float)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(half, float, float)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(double, double, double)
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(BFloat16, BFloat16, float)
#ifndef DISABLE_CONTRIB_OPS
SPECIALIZED_BLOCKWISE_SOFTMAX_IMPL(half, float, float) // used by BeamSearch op
#endif
} // namespace rocm
} // namespace onnxruntime

6
onnxruntime/python/tools/symbolic_shape_infer.py
View file

@ -1830,6 +1830,12 @@ class SymbolicShapeInference:
def _infer_SoftmaxCrossEntropyLoss(self, node):
vi = self.known_vi_[node.output[0]]
elem_type = self.known_vi_[node.input[0]].type.tensor_type.elem_type
# If output type is explicit specified in attribute, we use it as output tensor type.
specified_output_type = get_attribute(node, "output_type", None)
if specified_output_type is not None:
elem_type = specified_output_type
vi.type.tensor_type.elem_type = elem_type
vi.type.tensor_type.shape.CopyFrom(onnx.TensorShapeProto())

25
onnxruntime/test/util/include/test_utils.h
View file

@ -5,6 +5,8 @@
#include "core/framework/framework_common.h"
#include "core/framework/execution_provider.h"
#include "core/framework/ort_value.h"
#include "core/providers/cpu/cpu_execution_provider.h"
#include <memory>
#include <string>
@ -62,5 +64,28 @@ void RunAndVerifyOutputsWithEP(const std::string& model_data,
void CheckShapeEquality(const ONNX_NAMESPACE::TensorShapeProto* shape1,
const ONNX_NAMESPACE::TensorShapeProto* shape2);
// Create OrtValue on CPU copying from provided inputs.
template <typename T>
void CreateInputOrtValueOnCPU(gsl::span<const int64_t> dims, const std::vector<T>& value,
OrtValue* p_ortvalue, AllocatorPtr alloc = nullptr) {
static CPUExecutionProviderInfo info;
static CPUExecutionProvider cpu_provider(info);
static AllocatorPtr cpu_allocator = cpu_provider.GetAllocator(OrtMemTypeDefault);
TensorShape shape(dims);
assert(shape.Size() == static_cast<int64_t>(value.size()));
auto element_type = DataTypeImpl::GetType<T>();
auto allocator = alloc ? alloc : cpu_allocator;
auto p_tensor = std::make_unique<Tensor>(element_type, shape, allocator);
if (value.size() > 0 && !alloc) { // using CPU allocator
memcpy(p_tensor->MutableDataRaw(), value.data(), p_tensor->SizeInBytes());
}
p_ortvalue->Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
}
} // namespace test
} // namespace onnxruntime

13
orttraining/orttraining/core/graph/gradient_builder.cc
View file

@ -1233,8 +1233,19 @@ IMPLEMENT_GRADIENT_BUILDER(GetSoftmaxCrossEntropyLossInternalGradient) {
for (size_t i = 1; i < input_size; i++) {
input_arg_def.emplace_back(I(i));
}
auto src_attrs = SrcNodeAttributes();
std::vector<AttributeProto> attrs;
for (auto& attr : src_attrs) {
if (attr.first == "output_type") {
attrs.push_back(MakeAttribute("output_type", static_cast<int64_t>(IElemType(0))));
continue;
}
attrs.push_back(attr.second);
}
return std::vector<NodeDef>{
NodeDef(OpDef{"SoftmaxCrossEntropyLossInternalGrad", kMSDomain, 1}, input_arg_def, {GI(0)}, SrcNodeAttributes())};
NodeDef(OpDef{"SoftmaxCrossEntropyLossInternalGrad", kMSDomain, 1}, input_arg_def, {GI(0)}, attrs)};
}
IMPLEMENT_GRADIENT_BUILDER(GetGlobalAveragePoolGradient) {

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

@ -149,6 +149,44 @@ TensorProto ToDimensionOneTensor(T value) {
return t;
}
struct InputOutputAdaptorInfo {
bool need_adapt_input = false;
int64_t input_target_elem_type{-1};
bool need_adapt_output = false;
int64_t output_target_elem_type{-1};
};
void HandleDifferedInputOutputDataType(const int64_t input_elem_type,
const int64_t output_elem_type,
InputOutputAdaptorInfo& adaptor_info) {
if (input_elem_type == output_elem_type) {
return;
}
static std::unordered_map<::ONNX_NAMESPACE::TensorProto_DataType, int> bytes_for_elem_type = {
{ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16, 2},
{ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_BFLOAT16, 2},
{ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT, 4},
{ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_DOUBLE, 8},
};
// Use a larger type for computation if the input and output types are different.
bool use_input_elem_type_for_compute =
bytes_for_elem_type[static_cast<::ONNX_NAMESPACE::TensorProto_DataType>(input_elem_type)] >=
bytes_for_elem_type[static_cast<::ONNX_NAMESPACE::TensorProto_DataType>(output_elem_type)];
if (use_input_elem_type_for_compute) {
// Compute in input type and cast to output type before return result.
adaptor_info.need_adapt_output = true;
adaptor_info.output_target_elem_type = output_elem_type;
} else {
// Cast input to output_elem_type, and compute in output_elem_type, return result.
adaptor_info.need_adapt_input = true;
adaptor_info.input_target_elem_type = output_elem_type;
}
}
bool SCELossInternalFunBuilder(
const FunctionBodyBuildContext& ctx,
const OpSchema& schema,
@ -156,32 +194,75 @@ bool SCELossInternalFunBuilder(
bool hasWeight = ctx.hasInput(2);
bool hasIgnoreIndex = ctx.hasInput(3);
InputOutputAdaptorInfo adaptor_info;
// Handle the adaptor only when output_type is specified in attribute.
auto output_type_attr = ctx.getAttribute("output_type");
if (output_type_attr != nullptr) {
const TypeProto* first_input_type_proto = ctx.getInputType(0);
auto output_elem_type = output_type_attr->i();
if (first_input_type_proto != nullptr) {
HandleDifferedInputOutputDataType(first_input_type_proto->tensor_type().elem_type(),
output_elem_type,
adaptor_info);
} else {
// If the input type is not specified, we add input cast to make sure type check successful.
adaptor_info.need_adapt_input = true;
adaptor_info.input_target_elem_type = output_elem_type;
}
}
FunctionBuilder builder(functionProto);
if (adaptor_info.need_adapt_input) {
builder.Add("scores_casted = Cast(scores)", "to", adaptor_info.input_target_elem_type);
if (hasWeight) {
builder.Add("weights_casted = Cast(weights)", "to", adaptor_info.input_target_elem_type);
}
} else {
builder.Add("scores_casted = Identity (scores)");
if (hasWeight) {
builder.Add("weights_casted = Identity (weights)");
}
}
builder
.Const("Shape3D", std::vector<int64_t>({0, 0, -1}))
.Add(R"(
X_NCD = Reshape (scores, Shape3D)
X_NCD = Reshape (scores_casted, Shape3D)
X_NDC = Transpose <perm = [0, 2, 1]> (X_NCD)
X_LogSM = LogSoftmax <axis = 2> (X_NDC)
X_LogSM_NCD = Transpose <perm = [0, 2, 1]> (X_LogSM)
X_shape = Shape (scores)
X_shape = Shape (scores_casted)
X_Log = Reshape (X_LogSM_NCD, X_shape)
)");
if (ctx.hasOutput(1)) {
builder.Add("log_prob = Identity (X_Log)");
builder.Add("intermediate_log_prob = Identity (X_Log)");
}
if (hasWeight)
if (hasIgnoreIndex)
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, weights, ignore_index)");
builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, weights_casted, ignore_index)");
else
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, weights)");
builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, weights_casted)");
else if (hasIgnoreIndex)
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, , ignore_index)");
builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, , ignore_index)");
else
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels)");
builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels)");
if (adaptor_info.need_adapt_output) {
builder.Add("output = Cast(intermediate_output)", "to", adaptor_info.output_target_elem_type);
if (ctx.hasOutput(1)) {
builder.Add("log_prob = Cast(intermediate_log_prob)", "to", adaptor_info.output_target_elem_type);
}
} else {
builder.Add("output = Identity (intermediate_output)");
if (ctx.hasOutput(1)) {
builder.Add("log_prob = Identity(intermediate_log_prob)");
}
}
schema.BuildFunction(functionProto);
return true;
@ -198,6 +279,24 @@ bool SCELossGradFunBuilder(bool ignore_index_as_attr, const FunctionBodyBuildCon
ignore_index_as_attr ? (ctx.getAttribute("ignore_index") != nullptr) : ctx.hasInput(4);
bool has_weight = ctx.hasInput(3);
InputOutputAdaptorInfo adaptor_info;
// Handle the adaptor only when output_type is specified in attribute.
auto output_type_attr = ctx.getAttribute("output_type");
if (output_type_attr != nullptr) {
const TypeProto* first_input_type_proto = ctx.getInputType(0);
auto output_elem_type = output_type_attr->i();
if (first_input_type_proto != nullptr) {
HandleDifferedInputOutputDataType(first_input_type_proto->tensor_type().elem_type(),
output_elem_type,
adaptor_info);
} else {
// If the input type is not specified, we add input cast to make sure type check successful.
adaptor_info.need_adapt_input = true;
adaptor_info.input_target_elem_type = output_elem_type;
}
}
FunctionBuilder builder(functionProto);
// Inputs:
@ -206,6 +305,28 @@ bool SCELossGradFunBuilder(bool ignore_index_as_attr, const FunctionBodyBuildCon
// weight : [C]
// label : [B, d1, d2, ...]
if (adaptor_info.need_adapt_input) {
builder.Add("dY_casted = Cast(dY)", "to", adaptor_info.input_target_elem_type);
builder.Add("log_prob_casted = Cast(log_prob)", "to", adaptor_info.input_target_elem_type);
if (has_weight) {
builder.Add("weight_casted = Cast(weight)", "to", adaptor_info.input_target_elem_type);
}
if (ctx.hasInput(5)) {
builder.Add("bias_casted = Cast(bias)", "to", adaptor_info.input_target_elem_type);
}
} else {
builder.Add("dY_casted = Identity (dY)");
builder.Add("log_prob_casted = Identity (log_prob)");
if (has_weight) {
builder.Add("weight_casted = Identity (weight)");
}
if (ctx.hasInput(5)) {
builder.Add("bias_casted = Identity (bias)");
}
}
// We decompose the forward propagation into two steps, for doing the backward prop.
// Step 1: loss = Neg(Logsoftmax(prediction-for-true-label))
// Step 2: y = Reduce (loss), adjusting for weights and ignore_index
@ -231,55 +352,55 @@ bool SCELossGradFunBuilder(bool ignore_index_as_attr, const FunctionBodyBuildCon
// adj_label_BD is used so we can safely index into tensor-dimensions of size [C]
builder.Add(R"(
adj_label_BD = Where (ignored_BD, zero_label, label)
weight_BD = Gather (weight, adj_label_BD)
zero_weight = CastLike (zero_int64, weight)
weight_BD = Gather (weight_casted, adj_label_BD)
zero_weight = CastLike (zero_int64, weight_casted)
adj_weight_BD = Where (ignored_BD, zero_weight, weight_BD)
)");
if (mean_reduction) {
builder.Add(R"(
sum_weights = ReduceSum <keepdims = 0> (adj_weight_BD)
grad = Div (adj_weight_BD, sum_weights)
d_loss = Mul (grad, dY)
d_loss = Mul (grad, dY_casted)
)");
} else {
builder.Add("d_loss = Mul (adj_weight_BD, dY)");
builder.Add("d_loss = Mul (adj_weight_BD, dY_casted)");
}
} else {
builder.Add(R"(
not_ignored_BD = Not (ignored_BD)
adj_weight_BD = CastLike (not_ignored_BD, dY)
adj_weight_BD = CastLike (not_ignored_BD, dY_casted)
)");
if (mean_reduction) {
builder.Add(R"(
sum_weights = ReduceSum <keepdims = 0> (adj_weight_BD)
grad = Div (adj_weight_BD, sum_weights)
d_loss = Mul (grad, dY)
d_loss = Mul (grad, dY_casted)
)");
} else {
builder.Add("d_loss = Mul (adj_weight_BD, dY)");
builder.Add("d_loss = Mul (adj_weight_BD, dY_casted)");
}
}
} else {
if (has_weight) {
builder.Add("elt_weight = Gather (weight, label)");
builder.Add("elt_weight = Gather (weight_casted, label)");
if (mean_reduction) {
// backward-prop for y = ReduceSum (loss * elt_weight) / ReduceSum(elt_weight)
builder.Add(R"(
sum_weights = ReduceSum <keepdims = 0> (elt_weight)
grad = Div (elt_weight, sum_weights)
d_loss = Mul(grad, dY)
d_loss = Mul(grad, dY_casted)
)");
} else {
// common backward-prop for y = ReduceSum(loss * elt_weight) and y = loss * elt_weight
builder.Add("d_loss = Mul(elt_weight, dY)");
builder.Add("d_loss = Mul(elt_weight, dY_casted)");
}
} else {
if (mean_reduction) {
// backward-prop for y = ReduceSum (loss) / Size(label)
builder.Add(R"(
count = Size(label)
count_T = CastLike (count, dY)
d_div = Div (dY, count_T)
count_T = CastLike (count, dY_casted)
d_div = Div (dY_casted, count_T)
BD = Shape (label)
d_loss = Expand (d_div, BD)
)");
@ -287,7 +408,7 @@ bool SCELossGradFunBuilder(bool ignore_index_as_attr, const FunctionBodyBuildCon
// common backward-prop for y = ReduceSum(loss) and y = loss
builder.Add(R"(
BD = Shape (label)
d_loss = Expand (dY, BD)
d_loss = Expand (dY_casted, BD)
)");
}
}
@ -300,8 +421,8 @@ bool SCELossGradFunBuilder(bool ignore_index_as_attr, const FunctionBodyBuildCon
d_loss_B1Dopt = Unsqueeze (d_loss, axes1)
reshape_arg = Constant < value = int64[3] {0, 0, -1} > ()
d_loss_B1D = Reshape (d_loss_B1Dopt, reshape_arg)
orig_shape = Shape (log_prob)
log_prob_BCD = Reshape (log_prob, reshape_arg)
orig_shape = Shape (log_prob_casted)
log_prob_BCD = Reshape (log_prob_casted, reshape_arg)
prob_BCD = Exp (log_prob_BCD)
)");
@ -341,15 +462,21 @@ bool SCELossGradFunBuilder(bool ignore_index_as_attr, const FunctionBodyBuildCon
if (ctx.hasInput(5)) {
builder.Add(R"(
d_logits_without_bias = Reshape (d_logits_BCD, orig_shape)
bias_shaped = Reshape (bias, orig_shape)
d_logits = Add(d_logits_without_bias, bias_shaped)
bias_shaped = Reshape (bias_casted, orig_shape)
intermediate_d_logits = Add(d_logits_without_bias, bias_shaped)
)");
} else {
builder.Add(R"(
d_logits = Reshape (d_logits_BCD, orig_shape)
intermediate_d_logits = Reshape (d_logits_BCD, orig_shape)
)");
}
if (adaptor_info.need_adapt_output) {
builder.Add("d_logits = Cast(intermediate_d_logits)", "to", adaptor_info.output_target_elem_type);
} else {
builder.Add("d_logits = Identity (intermediate_d_logits)");
}
schema.BuildFunction(functionProto);
return true;
};
@ -3890,6 +4017,11 @@ Return true if all elements are true and false otherwise.
.SetDomain(kMSDomain)
.SinceVersion(1)
.Attr("reduction", reduction_doc, AttributeProto::STRING, std::string("mean"))
.Attr("output_type",
"(Optional) The data type for the output tensor. "
"If not provided, output tensor has the same type as input tensor."
"Strictly must be one of the types from DataType enum in TensorProto",
AttributeProto::INT, OPTIONAL_VALUE)
.Input(0, "scores",
"The predicted outputs with shape [batch_size, class_size], or "
"[batch_size, class_size, D1, D2 , ..., Dk], where K is the number of dimensions.",
@ -3913,14 +4045,26 @@ Return true if all elements are true and false otherwise.
"Weighted loss float Tensor. If reduction is 'none', this has the "
"shape of [batch_size], or [batch_size, D1, D2, ..., Dk] in case of "
"K-dimensional loss. Otherwise, it is a scalar.",
"T")
.Output(1, "log_prob", "Log probability tensor. If the output of softmax is prob, its value is log(prob).", "T")
"TOut")
.Output(1, "log_prob",
"Log probability tensor. If the output of softmax is prob, its value is log(prob).",
"TOut")
.TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Constrain input and output types to float tensors.")
"Constrain input types to float tensors.")
.TypeConstraint("TOut", {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Constrain input types to float tensors.")
.TypeConstraint("Tind", {"tensor(int32)", "tensor(int64)"}, "Constrain target to integer types")
.TypeConstraint("I", {"tensor(int64)"}, "Constrain ignore_index tensor to int64")
.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
auto output_type_attr = ctx.getAttribute("output_type");
if (output_type_attr) {
propagateElemTypeFromAttributeToOutput(ctx, "output_type", 0);
propagateElemTypeFromAttributeToOutput(ctx, "output_type", 1);
} else {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
propagateElemTypeFromInputToOutput(ctx, 0, 1);
}
std::string reduction = getAttribute(ctx, "reduction", "mean");
if (reduction.compare("none") == 0) {
if (hasInputShape(ctx, 1)) {
@ -3929,11 +4073,7 @@ Return true if all elements are true and false otherwise.
} else {
updateOutputShape(ctx, 0, TensorShapeProto());
}
if (ctx.getNumOutputs() == 2) {
propagateElemTypeFromInputToOutput(ctx, 0, 1);
propagateShapeFromInputToOutput(ctx, 0, 1);
}
propagateShapeFromInputToOutput(ctx, 0, 1);
})
.SetContextDependentFunctionBodyBuilder(SCELossInternalFunBuilder)
.SetDoc(R"DOC(SoftmaxCrossEntropyLossInternal)DOC");
@ -3942,6 +4082,11 @@ Return true if all elements are true and false otherwise.
.SetDomain(kMSDomain)
.SinceVersion(1)
.Attr("reduction", reduction_doc, AttributeProto::STRING, std::string("mean"))
.Attr("output_type",
"(Optional) The data type for the output tensor. "
"If not provided, output tensor has the same type as input tensor."
"Strictly must be one of the types from DataType enum in TensorProto",
AttributeProto::INT, OPTIONAL_VALUE)
.Input(0, "dY", "gradient of Y", "T")
.Input(1, "log_prob", "logsoftmax(logits), (N+1)-D input of shape (batch_size).", "T")
.Input(2, "label",
@ -3953,14 +4098,22 @@ Return true if all elements are true and false otherwise.
.Input(4, "ignore_index",
"Scalar tensor to specify a target value that is ignored and does not contribute to the input gradient.",
"I", OpSchema::Optional)
.Input(5, "bias", "data to be non-broadcasting added to the gradient.", "T", OpSchema::Optional)
.Output(0, "d_logits", "gradient of logits", "T")
.Input(5, "bias", "data to be non-broadcasting added to the gradient.", "TOut", OpSchema::Optional)
.Output(0, "d_logits", "gradient of logits", "TOut")
.TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Constrain to float, float16 and double tensors.")
"Constrain input types to float tensors.")
.TypeConstraint("TOut", {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Constrain input types to float tensors.")
.TypeConstraint("Tind", {"tensor(int32)", "tensor(int64)"}, "Constrain indices to integer types")
.TypeConstraint("I", {"tensor(int64)"}, "Constrain ignore_index tensor to int64")
.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 1, 0);
auto output_type_attr = ctx.getAttribute("output_type");
if (output_type_attr) {
propagateElemTypeFromAttributeToOutput(ctx, "output_type", 0);
} else {
propagateElemTypeFromInputToOutput(ctx, 1, 0);
}
propagateShapeFromInputToOutput(ctx, 1, 0);
})
.SetContextDependentFunctionBodyBuilder(

7
orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py
View file

@ -16,6 +16,7 @@ from onnxruntime.training import ortmodule
from . import _logger
from ._fallback import ORTModuleONNXModelException, wrap_exception
from ._custom_op_symbolic_registry import pytorch_type_to_onnx
# Some autograd.Function's shouldn't be exported as PythonOp.
# If CheckpointFunction is exported as PythonOp, the checkpointed computation
@ -50,7 +51,7 @@ def _full_name(klass):
return module + "." + klass.__qualname__
def _pytorch_type_to_onnx(scalar_type: str) -> torch.onnx.TensorProtoDataType:
def pytorch_type_to_onnx(scalar_type: str) -> torch.onnx.TensorProtoDataType:
try:
return torch.onnx.JitScalarType.from_name(scalar_type).onnx_type()
except AttributeError:
@ -131,7 +132,7 @@ def _export_pt_1_10(g, n, *args, **kwargs):
if call_type == "d":
# Got a tensor variable.
tensor_args.append(arg)
scalar_type = _pytorch_type_to_onnx(arg.type().scalarType())
scalar_type = pytorch_type_to_onnx(arg.type().scalarType())
input_tensor_types.append(scalar_type)
input_tensor_ranks.append(arg.type().dim())
elif call_type == "c":
@ -178,7 +179,7 @@ def _export_pt_1_10(g, n, *args, **kwargs):
output_tensor_ranks = []
for arg in n.outputs():
# Type of tensor's elements.
scalar_type = _pytorch_type_to_onnx(arg.type().scalarType())
scalar_type = pytorch_type_to_onnx(arg.type().scalarType())
output_tensor_types.append(scalar_type)
output_tensor_ranks.append(arg.type().dim())

61
orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py
View file

@ -14,6 +14,31 @@ from torch.onnx.symbolic_helper import _get_tensor_dim_size, _get_tensor_sizes,
from onnxruntime.training import ortmodule
# Mapping from pytorch scalar type to onnx scalar type.
_CAST_PYTORCH_TO_ONNX = {
"Byte": torch.onnx.TensorProtoDataType.UINT8,
"Char": torch.onnx.TensorProtoDataType.INT8,
"Double": torch.onnx.TensorProtoDataType.DOUBLE,
"Float": torch.onnx.TensorProtoDataType.FLOAT,
"Half": torch.onnx.TensorProtoDataType.FLOAT16,
"Int": torch.onnx.TensorProtoDataType.INT32,
"Long": torch.onnx.TensorProtoDataType.INT64,
"Short": torch.onnx.TensorProtoDataType.INT16,
"Bool": torch.onnx.TensorProtoDataType.BOOL,
"ComplexFloat": torch.onnx.TensorProtoDataType.COMPLEX64,
"ComplexDouble": torch.onnx.TensorProtoDataType.COMPLEX128,
"BFloat16": torch.onnx.TensorProtoDataType.BFLOAT16,
"Undefined": torch.onnx.TensorProtoDataType.UNDEFINED,
}
def pytorch_type_to_onnx(scalar_type: str) -> torch.onnx.TensorProtoDataType:
try:
return torch.onnx.JitScalarType.from_name(scalar_type).onnx_type()
except AttributeError:
return _CAST_PYTORCH_TO_ONNX[scalar_type]
def wrap_custom_export_function(original_func):
# Starting from PyTorch 1.11, there has been a change to symbolic function signature
# in terms of how additional context is accessed. More info at
@ -84,19 +109,17 @@ def cross_entropy_loss(g, node, logits, target, weight, reduction, ignore_index,
output_type = None
#####################################################################################################
# Workaround: cross_entropy_loss takes fp16 as input and generates fp32 output.
# cross_entropy_loss takes fp16 as input and generates fp32 output.
# sample aten graph:
# %target : Long(16, strides=[1], requires_grad=0, device=cuda:0)
# %input : Half(16, 3, strides=[3, 1], requires_grad=0, device=cuda:0) = aten::linear(%18, %13, %19)
# Float(requires_grad=0, device=cuda:0) = aten::cross_entropy_loss(%input, %target, %21, %22, %23, %24)
# If ORT do the compute with the input data type, the scaled loss gradient will become inf (cannot represented
# with fp16). Here we try to cast the fp16 to fp32 based on the export context (if there is).
# Currently not all type promotion/demotion are considered, only fp16 to fp32 is considered. For others, they
# remain the same behavior as before, but leave a warning message.
#
# So here if we could get node, then explicitly set output type that might be different with input type;
# otherwise, we do the cast (because there is no good way to define a float output type without inheriting from
# existing node)
if not node:
# For lower version torch we cannot get node output types, we do the type promotion for safety.
# Assume a failure will happen if a non-float32 result is expected.
if logits.type().scalarType() == "Half":
logits_casted = g.op("Cast", logits, to_i=torch.onnx.TensorProtoDataType.FLOAT)
@ -105,30 +128,9 @@ def cross_entropy_loss(g, node, logits, target, weight, reduction, ignore_index,
output_type = logits_casted.type()
else:
# For higher version torch we can get node output types, only adding cast for known type promotion cases.
# For higher version torch we can get node output types
loss_output = list(node.outputs())[0]
loss_scalar_type = loss_output.type().scalarType()
logits_scalar_type = logits.type().scalarType()
if loss_scalar_type != logits_scalar_type and logits_scalar_type == "Half" and loss_scalar_type == "Float":
# TODO: remove the cast once SoftmaxCrossEntropyLossInternal supports fp16 input and fp32 output.
logits_casted = g.op("Cast", logits, to_i=torch.onnx.TensorProtoDataType.FLOAT)
if not weight.node().mustBeNone():
if weight.type().scalarType() == "Half":
weight_casted = g.op("Cast", weight, to_i=torch.onnx.TensorProtoDataType.FLOAT)
else:
warnings.warn(
"Unsupported diverged input and output types for weight when export cross_entropy_loss."
f"weight type: {weight.type().scalarType()}, loss type: {loss_scalar_type}"
)
else:
warnings.warn(
"Unsupported diverged input and output types for logits when export cross_entropy_loss."
f"logits type: {logits_scalar_type}, loss type: {loss_scalar_type}"
)
output_type = loss_output.type()
# End of workaround
##################################
# reduction: 0->none, 1->mean, 2->sum
@ -143,6 +145,7 @@ def cross_entropy_loss(g, node, logits, target, weight, reduction, ignore_index,
weight_casted,
ignore_index,
reduction_s=reduction,
output_type_i=pytorch_type_to_onnx(output_type.scalarType()),
outputs=2,
)
output.setType(output_type)

9
orttraining/orttraining/test/training_api/core/training_api_tests.cc
View file

@ -9,6 +9,7 @@
#include "test/framework/test_utils.h"
#include "test/util/include/asserts.h"
#include "test/util/include/test_utils.h"
#include "core/framework/tensorprotoutils.h"
#include "orttraining/training_api/utils.h"
#include "orttraining/training_api/module.h"
@ -45,7 +46,7 @@ void GenerateRandomInput(gsl::span<const int64_t> dims, OrtValue& input) {
TensorShape shape(dims);
std::vector<float> data(shape.Size());
GenerateRandomData(data);
onnxruntime::training::api::utils::CreateInputOrtValue<float>(dims, data, &input);
onnxruntime::test::CreateInputOrtValueOnCPU<float>(dims, data, &input);
}
void TestModuleExport(const std::vector<std::shared_ptr<IExecutionProvider>>& providers) {
@ -145,7 +146,7 @@ void TestLRSchduler(const std::basic_string<ORTCHAR_T>& test_file_name, float in
OrtValue input, target;
GenerateRandomInput(std::array<int64_t, 2>{2, 784}, input);
onnxruntime::training::api::utils::CreateInputOrtValue<int32_t>(
onnxruntime::test::CreateInputOrtValueOnCPU<int32_t>(
std::array<int64_t, 1>{2}, std::vector<int32_t>(2, 1), &target);
/// Load test data for learning rate schedulers.
@ -272,7 +273,7 @@ TEST(TrainingApiTest, ModuleTrainStep) {
ASSERT_EQ(model->GetTrainingModelOutputCount(), 1);
OrtValue input, target;
GenerateRandomInput(std::array<int64_t, 2>{2, 784}, input);
onnxruntime::training::api::utils::CreateInputOrtValue<int32_t>(
onnxruntime::test::CreateInputOrtValueOnCPU<int32_t>(
std::array<int64_t, 1>{2}, std::vector<int32_t>(2, 1), &target);
auto data_loader = std::vector<std::vector<OrtValue>>(4, std::vector<OrtValue>{input, target});
@ -347,7 +348,7 @@ TEST(TrainingApiTest, OptimStep) {
OrtValue input, target;
GenerateRandomInput(std::array<int64_t, 2>{2, 784}, input);
onnxruntime::training::api::utils::CreateInputOrtValue<int32_t>(
onnxruntime::test::CreateInputOrtValueOnCPU<int32_t>(
std::array<int64_t, 1>{2}, std::vector<int32_t>(2, 1), &target);
auto data_loader = std::vector<std::vector<OrtValue>>(4, std::vector<OrtValue>{input, target});

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

File diff suppressed because it is too large Load diff

26
orttraining/orttraining/training_api/utils.h
View file

@ -47,32 +47,6 @@ void WrapInOrtValue(T value,
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
}
// Create OrtValue on CPU out of provided inputs
template <typename T>
static void CreateInputOrtValue(gsl::span<const int64_t> dims,
const std::vector<T>& value,
OrtValue* p_ortvalue,
AllocatorPtr alloc = nullptr) {
static CPUExecutionProviderInfo info;
static CPUExecutionProvider cpu_provider(info);
static AllocatorPtr cpu_allocator = cpu_provider.GetAllocator(OrtMemTypeDefault);
TensorShape shape(dims);
assert(shape.Size() == static_cast<int64_t>(value.size()));
auto element_type = DataTypeImpl::GetType<T>();
auto allocator = alloc ? alloc : cpu_allocator;
auto p_tensor = std::make_unique<Tensor>(element_type, shape, allocator);
// TODO: Handle memcpy for other allocators
if (value.size() > 0 && !alloc) { // using CPU allocator
memcpy(p_tensor->MutableDataRaw(), value.data(), p_tensor->SizeInBytes());
}
p_ortvalue->Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
}
template <typename T>
T GetValue(OrtValue& ort_value) {
const Tensor& tensor = ort_value.Get<Tensor>();

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

@ -403,6 +403,7 @@ Status SoftmaxCrossEntropyLossGrad<T1, T2>::Compute(OpKernelContext* context) co
ONNX_OPERATOR_TWO_TYPED_KERNEL_EX(OpName, kMSDomain, 1, T1, T2, kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("TOut", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<T2>()) \
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()), \
ClassName<T1, T2>);

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

@ -65,12 +65,14 @@ class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDom
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, SoftmaxGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, LogSoftmaxGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, SoftmaxGrad_13);
@ -317,12 +319,14 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float, BatchNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double_double, BatchNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_MLFloat16_MLFloat16, BatchNormalizationGrad)>,

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

@ -13,43 +13,46 @@ namespace cuda {
OrtValue AllocateTensorInMLValue(const MLDataType data_type, const TensorShape& shape, AllocatorPtr& allocator) {
auto new_tensor = Tensor::Create(data_type, shape, allocator);
auto ml_tensor = DataTypeImpl::GetType<Tensor>();
return OrtValue{new_tensor.release(), ml_tensor,
ml_tensor->GetDeleteFunc()};
};
#define REGISTER_KERNEL_VERSIONED_TYPED_TWO_TYPES(Class, T, Tin, domain, startver, endver) \
ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_EX( \
Class, \
domain, \
startver, endver, \
T, Tin, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<Tin>()), \
Class<T, Tin>);
#define REGISTER_KERNEL_VERSIONED_TYPED_TWO_TYPES(Class, T, TLabel, domain, startver, endver) \
ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_EX( \
Class, \
domain, \
startver, endver, \
T, TLabel, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TLabel>()), \
Class<T, TLabel, T>);
#define REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, Tin, domain, version) \
ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \
Class, \
domain, \
version, \
T, Tin, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<Tin>()), \
Class<T, Tin>);
#define REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, TLabel, domain, version) \
ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \
Class, \
domain, \
version, \
T, TLabel, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TLabel>()), \
Class<T, TLabel, T>);
template <typename T, typename TLabel, typename TOut>
Status SoftmaxCrossEntropyLoss<T, TLabel, TOut>::ComputeInternal(OpKernelContext* ctx) const {
typedef typename ToCudaType<T>::MappedType CudaT_IN;
typedef typename ToCudaType<TOut>::MappedType CudaT_OUT;
template <typename T, typename Tin>
Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) const {
typedef typename ToCudaType<T>::MappedType CudaT;
const Tensor& logit = *ctx->Input<Tensor>(0);
const Tensor& label = *ctx->Input<Tensor>(1);
const Tensor* p_weight = ctx->Input<Tensor>(2);
const Tensor* p_ignore_index = ctx->Input<Tensor>(3);
const Tensor* p_weight = ctx->Input<Tensor>(2); // optional input
const Tensor* p_ignore_index = ctx->Input<Tensor>(3); // optional input
// Prefer ignore index from input over attributes.
int64_t ignore_index = ignore_index_;
if (p_ignore_index) {
ORT_ENFORCE(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar.");
@ -61,33 +64,32 @@ Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) co
onnxruntime::contrib::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;
// N_D = N * D1 * D2...Dk
int64_t N_D, C;
onnxruntime::contrib::GetNDCFromLogitAndLabelShape(logit_shape, label_shape, N_D, C);
const TensorShape logit_reshape({N_D, C});
Tensor* total_loss = ctx->Output(0, reduction_ == ReductionType::NONE ? TensorShape(label.Shape()) : TensorShape({}));
T* total_loss_data = total_loss->template MutableData<T>();
T* tmp_loss_sample_buffer = nullptr;
IAllocatorUniquePtr<T> tmp_loss_sample;
TOut* total_loss_data = total_loss->template MutableData<TOut>();
TOut* tmp_loss_sample_buffer = nullptr;
IAllocatorUniquePtr<TOut> tmp_loss_sample;
if (reduction_ == ReductionType::NONE) {
tmp_loss_sample_buffer = total_loss_data;
} else {
tmp_loss_sample = GetScratchBuffer<T>(N_D, ctx->GetComputeStream());
tmp_loss_sample = GetScratchBuffer<TOut>(N_D, ctx->GetComputeStream());
tmp_loss_sample_buffer = tmp_loss_sample.get();
}
const T* logit_data = logit.template Data<T>();
const Tin* label_data = label.template Data<Tin>();
const TLabel* label_data = label.template Data<TLabel>();
T* log_prob_data = nullptr;
TOut* log_prob_data = nullptr;
Tensor* log_prob = nullptr;
IAllocatorUniquePtr<T> log_prob_scratch_buffer;
IAllocatorUniquePtr<TOut> log_prob_scratch_buffer;
if (ctx->OutputCount() > 1) {
log_prob = ctx->Output(1, logit_shape);
log_prob_data = log_prob->template MutableData<T>();
log_prob_data = log_prob->template MutableData<TOut>();
} else {
log_prob_scratch_buffer = GetScratchBuffer<T>(logit_shape.Size(), ctx->GetComputeStream());
log_prob_scratch_buffer = GetScratchBuffer<TOut>(logit_shape.Size(), ctx->GetComputeStream());
log_prob_data = log_prob_scratch_buffer.get();
}
@ -102,16 +104,13 @@ Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) co
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc));
onnxruntime::contrib::GetPermutationAndShape(true, logit_shape, new_shape, permutations);
transpose_output = AllocateTensorInMLValue(logit.DataType(), new_shape, alloc);
ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, logit, *transpose_output.GetMutable<Tensor>()));
ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, logit,
*transpose_output.GetMutable<Tensor>()));
logit_data = (*transpose_output.GetMutable<Tensor>()).template Data<T>();
}
// calculate logsoftmax
auto status = SoftMaxComputeHelper<T, true>(Stream(ctx),
logit_data,
logit_reshape,
log_prob_data,
1);
// Calculate logsoftmax
auto status = SoftMaxComputeHelper<T, TOut, true>(Stream(ctx), logit_data, logit_reshape, log_prob_data, 1);
ORT_RETURN_IF_ERROR(status);
const T* weight_data = nullptr;
@ -120,46 +119,46 @@ Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) co
weight_data = weight.template Data<T>();
}
IAllocatorUniquePtr<T> weight_data_nd = GetScratchBuffer<T>(N_D, ctx->GetComputeStream());
T* weight_data_nd_data = weight_data_nd.get();
IAllocatorUniquePtr<CudaT_OUT> weight_data_nd = GetScratchBuffer<CudaT_OUT>(N_D, ctx->GetComputeStream());
CudaT_OUT* weight_data_nd_data = weight_data_nd.get();
ComputeSoftmaxCrossEntropyWeightsImpl(Stream(ctx),
label_data,
reinterpret_cast<const CudaT*>(weight_data),
reinterpret_cast<const CudaT_IN*>(weight_data),
N_D, C,
ignore_index,
reinterpret_cast<CudaT*>(weight_data_nd_data));
reinterpret_cast<CudaT_OUT*>(weight_data_nd_data));
// Compute buffer size in byte for reduction APIs.
const auto buffer_size =
compute_reduction_buffer_size<CudaT>(static_cast<int>(N_D));
const auto buffer_size = compute_reduction_buffer_size<CudaT_OUT>(static_cast<int>(N_D));
// Allocate reduction buffer whose size is buffer_size bytes, or nullptr if no reduction.
IAllocatorUniquePtr<void> reduction_buffer = GetScratchBuffer<void>(
reduction_ != ReductionType::NONE ? buffer_size : 0, ctx->GetComputeStream());
typedef AccumulationType_t<CudaT> TBuf;
typedef AccumulationType_t<CudaT_OUT> TBuf;
auto normalize_factor_data = GetScratchBuffer<TBuf>(1, ctx->GetComputeStream());
if (reduction_ == ReductionType::MEAN) {
ORT_RETURN_IF_ERROR(reduce_sum(
Stream(ctx),
reinterpret_cast<CudaT*>(weight_data_nd_data),
reinterpret_cast<CudaT_OUT*>(weight_data_nd_data),
normalize_factor_data.get(),
static_cast<int>(N_D),
reduction_buffer.get(),
buffer_size));
} else {
constexpr TBuf normalize_factor = static_cast<TBuf>(1.0f);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(TBuf), cudaMemcpyHostToDevice, Stream(ctx)));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(TBuf),
cudaMemcpyHostToDevice, Stream(ctx)));
}
SoftmaxCrossEntropyLossImpl(Stream(ctx),
reinterpret_cast<CudaT*>(log_prob_data),
reinterpret_cast<CudaT_OUT*>(log_prob_data),
label_data,
reinterpret_cast<CudaT*>(weight_data_nd_data),
reinterpret_cast<CudaT_OUT*>(weight_data_nd_data),
normalize_factor_data.get(),
N_D,
C,
ignore_index,
reinterpret_cast<CudaT*>(tmp_loss_sample_buffer));
reinterpret_cast<CudaT_OUT*>(tmp_loss_sample_buffer));
// Transpose log probability from [N, D1, D2...Dk, C] to [N, C, D1, D2 .. Dk].
if (logit_shape.NumDimensions() > 2 && log_prob != nullptr) {
@ -167,11 +166,13 @@ Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) co
new_shape.clear();
permutations.clear();
onnxruntime::contrib::GetPermutationAndShape(false, log_prob_shape, new_shape, permutations);
auto* transposed_data = (*transpose_output.GetMutable<Tensor>()).template MutableData<T>();
auto* transposed_data = (*transpose_output.GetMutable<Tensor>()).template MutableData<TOut>();
transpose_output.GetMutable<Tensor>()->Reshape(log_prob->Shape());
log_prob->Reshape(log_prob_shape);
ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, *log_prob, *transpose_output.GetMutable<Tensor>()));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(log_prob_data, transposed_data, sizeof(T) * logit_shape.Size(), cudaMemcpyDeviceToDevice, Stream(ctx)));
ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, *log_prob,
*transpose_output.GetMutable<Tensor>()));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(log_prob_data, transposed_data, sizeof(TOut) * logit_shape.Size(),
cudaMemcpyDeviceToDevice, Stream(ctx)));
log_prob->Reshape(new_shape);
}
@ -179,8 +180,8 @@ Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) co
// ReduceSum on loss_per_sample
ORT_RETURN_IF_ERROR(reduce_sum(
Stream(ctx),
reinterpret_cast<CudaT*>(tmp_loss_sample_buffer),
reinterpret_cast<CudaT*>(total_loss_data),
reinterpret_cast<CudaT_OUT*>(tmp_loss_sample_buffer),
reinterpret_cast<CudaT_OUT*>(total_loss_data),
static_cast<int>(N_D),
reduction_buffer.get(),
buffer_size));
@ -189,15 +190,18 @@ Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) co
return Status::OK();
}
template <typename T, typename Tin>
Status SoftmaxCrossEntropyLossGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx) const {
typedef typename ToCudaType<T>::MappedType CudaT;
template <typename T, typename TLabel, typename TOut>
Status SoftmaxCrossEntropyLossGrad<T, TLabel, TOut>::ComputeInternal(OpKernelContext* ctx) const {
typedef typename ToCudaType<T>::MappedType CudaT_IN;
typedef typename ToCudaType<TOut>::MappedType CudaT_OUT;
const Tensor& dY = *ctx->Input<Tensor>(0);
const Tensor& log_prob = *ctx->Input<Tensor>(1);
const Tensor& label = *ctx->Input<Tensor>(2);
const Tensor* p_weight = ctx->Input<Tensor>(3);
const Tensor* p_ignore_index = ctx->Input<Tensor>(4);
const Tensor* p_bias = ctx->Input<Tensor>(5);
int64_t ignore_index = ignore_index_;
if (p_ignore_index) {
ORT_ENFORCE(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar.");
@ -209,15 +213,14 @@ Status SoftmaxCrossEntropyLossGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx
onnxruntime::contrib::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;
// N_D = N * D1 * D2...Dk
int64_t N_D, C;
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>();
const T* log_prob_data = log_prob.Data<T>();
const Tin* label_data = label.Data<Tin>();
T* d_logit_data = d_logit->MutableData<T>();
const TLabel* label_data = label.Data<TLabel>();
TOut* d_logit_data = d_logit->MutableData<TOut>();
const T* weight_data = nullptr;
OrtValue transpose_output;
TensorShapeVector new_shape;
@ -230,7 +233,8 @@ Status SoftmaxCrossEntropyLossGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc));
onnxruntime::contrib::GetPermutationAndShape(true, probability_shape, new_shape, permutations);
transpose_output = AllocateTensorInMLValue(log_prob.DataType(), new_shape, alloc);
ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, log_prob, *transpose_output.GetMutable<Tensor>()));
ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations,
log_prob, *transpose_output.GetMutable<Tensor>()));
log_prob_data = (*transpose_output.GetMutable<Tensor>()).Data<T>();
}
@ -243,44 +247,45 @@ Status SoftmaxCrossEntropyLossGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx
T* weight_data_nd_data = weight_data_nd.get();
ComputeSoftmaxCrossEntropyWeightsImpl(Stream(ctx),
label_data,
reinterpret_cast<const CudaT*>(weight_data),
reinterpret_cast<const CudaT_IN*>(weight_data),
N_D, C,
ignore_index,
reinterpret_cast<CudaT*>(weight_data_nd_data));
typedef AccumulationType_t<CudaT> TBuf;
reinterpret_cast<CudaT_IN*>(weight_data_nd_data));
typedef AccumulationType_t<CudaT_IN> TBuf;
auto normalize_factor_data = GetScratchBuffer<TBuf>(1, ctx->GetComputeStream());
if (reduction_ == ReductionType::MEAN) {
// Compute buffer size in byte for reduction APIs.
const auto buffer_size =
compute_reduction_buffer_size<CudaT>(static_cast<int>(N_D));
compute_reduction_buffer_size<CudaT_IN>(static_cast<int>(N_D));
// Allocate reduction buffer whose size is buffer_size bytes.
IAllocatorUniquePtr<void> reduction_buffer = GetScratchBuffer<void>(
buffer_size, ctx->GetComputeStream());
ORT_RETURN_IF_ERROR(reduce_sum(
Stream(ctx),
reinterpret_cast<const CudaT*>(weight_data_nd_data),
reinterpret_cast<const CudaT_IN*>(weight_data_nd_data),
normalize_factor_data.get(),
static_cast<int>(N_D),
reduction_buffer.get(),
buffer_size));
} else {
constexpr TBuf normalize_factor = static_cast<TBuf>(1.0f);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(TBuf), cudaMemcpyHostToDevice, Stream(ctx)));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(TBuf),
cudaMemcpyHostToDevice, Stream(ctx)));
}
const T* bias_data = p_bias ? p_bias->Data<T>() : nullptr;
const TOut* bias_data = p_bias ? p_bias->Data<TOut>() : nullptr;
SoftmaxCrossEntropyLossGradImpl(Stream(ctx),
reinterpret_cast<const CudaT*>(dY_data),
reinterpret_cast<const CudaT*>(log_prob_data),
reinterpret_cast<const CudaT_IN*>(dY_data),
reinterpret_cast<const CudaT_IN*>(log_prob_data),
label_data,
reinterpret_cast<const CudaT*>(weight_data_nd_data),
reinterpret_cast<const CudaT_IN*>(weight_data_nd_data),
normalize_factor_data.get(),
reinterpret_cast<const CudaT*>(bias_data),
reinterpret_cast<const CudaT_OUT*>(bias_data),
N_D,
C,
ReductionType::NONE == reduction_,
reinterpret_cast<CudaT*>(d_logit_data));
reinterpret_cast<CudaT_OUT*>(d_logit_data));
// Transpose logit from [N, D1, D2...Dk, C] to [N, C, D1, D2 .. Dk]
if (probability_shape.NumDimensions() > 2) {
@ -290,21 +295,23 @@ Status SoftmaxCrossEntropyLossGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx
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>()));
auto* transposed_data = (*transpose_output.GetMutable<Tensor>()).template Data<T>();
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(d_logit_data, transposed_data, sizeof(T) * probability_shape.Size(), cudaMemcpyDeviceToDevice, Stream(ctx)));
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)));
d_logit->Reshape(new_shape);
}
return Status::OK();
}
#define INSTANTIATE_VERSIONED_COMPUTE_SPARSE(Class, T, Tin, domain, startver, endvar) \
REGISTER_KERNEL_VERSIONED_TYPED_TWO_TYPES(Class, T, Tin, domain, startver, endvar)
#define INSTANTIATE_VERSIONED_COMPUTE_SPARSE(Class, T, TLabel, domain, startver, endvar) \
REGISTER_KERNEL_VERSIONED_TYPED_TWO_TYPES(Class, T, TLabel, domain, startver, endvar)
#define INSTANTIATE_COMPUTE_SPARSE(Class, T, Tin, domain, version) \
REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, Tin, domain, version) \
template Status Class<T, Tin>::ComputeInternal(OpKernelContext* ctx) const;
#define INSTANTIATE_COMPUTE_SPARSE(Class, T, TLabel, domain, version) \
REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, TLabel, domain, version) \
template Status Class<T, TLabel, T>::ComputeInternal(OpKernelContext* ctx) const;
INSTANTIATE_VERSIONED_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, float, int64_t, kOnnxDomain, 12, 12)
INSTANTIATE_VERSIONED_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, MLFloat16, int64_t, kOnnxDomain, 12, 12)
@ -315,21 +322,25 @@ INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, float, int64_t, kMSDomai
INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, MLFloat16, int64_t, kMSDomain, 1)
INSTANTIATE_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, BFloat16, int64_t, kMSDomain, 1)
#define REGISTER_KERNEL_INTERNAL_TYPED(OpName, ClassName, T, Tin, CpuInputIndex) \
ONNX_OPERATOR_TWO_TYPED_KERNEL_EX(OpName, kMSDomain, 1, T, Tin, kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, CpuInputIndex) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<Tin>()) \
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()), \
ClassName<T, Tin>);
#define REGISTER_KERNEL_INTERNAL_TYPED(OpName, ClassName, T, TLabel, TOut, CpuInputIndex) \
ONNX_OPERATOR_TYPED_KERNEL_EX(OpName, kMSDomain, 1, T##_##TLabel##_##TOut, kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, CpuInputIndex) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TLabel>()) \
.TypeConstraint("TOut", DataTypeImpl::GetTensorType<TOut>()) \
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()), \
ClassName<T, TLabel, TOut>);
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, float, int64_t, 3)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, MLFloat16, int64_t, 3)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, BFloat16, int64_t, 3)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, float, int64_t, 4)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, MLFloat16, int64_t, 4)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, BFloat16, int64_t, 4)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, MLFloat16, int64_t, float, 3)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, float, int64_t, float, 3)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, MLFloat16, int64_t, MLFloat16, 3)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternal, SoftmaxCrossEntropyLoss, BFloat16, int64_t, BFloat16, 3)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, float, int64_t, float, 4)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, float, int64_t, MLFloat16, 4)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, MLFloat16, int64_t, MLFloat16, 4)
REGISTER_KERNEL_INTERNAL_TYPED(SoftmaxCrossEntropyLossInternalGrad, SoftmaxCrossEntropyLossGrad, BFloat16, int64_t, BFloat16, 4)
} // namespace cuda
} // namespace onnxruntime

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

@ -7,58 +7,60 @@
namespace onnxruntime {
namespace cuda {
template <typename T, typename Tin, bool IsWeighted>
template <typename T, typename TLabel, typename TOut, bool IsWeighted>
struct OpSoftmaxCrossEntropyWeights {
OpSoftmaxCrossEntropyWeights(const Tin* label_data, const T* weight_data, Tin C, Tin ignore_index)
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__ T operator()(CUDA_LONG idx) const {
__device__ __inline__ TOut operator()(CUDA_LONG idx) const {
if (label_data_[idx] != ignore_index_) {
if (IsWeighted) {
CUDA_KERNEL_ASSERT(label_data_[idx] >= 0 && label_data_[idx] < C_);
return weight_data_[label_data_[idx]];
return TOut(weight_data_[label_data_[idx]]);
}
return T(1.f);
return TOut(1.f);
}
return T(0.f);
return TOut(0.f);
}
const Tin* label_data_;
const TLabel* label_data_;
const T* weight_data_;
Tin C_;
Tin ignore_index_;
TLabel C_;
TLabel ignore_index_;
};
template <typename T, typename Tin>
void ComputeSoftmaxCrossEntropyWeightsImpl(cudaStream_t stream, const Tin* label, const T* weight, size_t count,
size_t label_depth, int64_t ignore_index, T* weight_data_nd) {
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, Tin, true> op(label, weight, static_cast<Tin>(label_depth),
static_cast<Tin>(ignore_index));
LaunchElementwiseKernel<T, decltype(op)>(stream, weight_data_nd, op, count);
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);
} else {
OpSoftmaxCrossEntropyWeights<T, Tin, false> op(label, nullptr, static_cast<Tin>(label_depth),
static_cast<Tin>(ignore_index));
LaunchElementwiseKernel<T, decltype(op)>(stream, weight_data_nd, op, count);
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);
}
}
#define INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(T, Tin) \
template void ComputeSoftmaxCrossEntropyWeightsImpl(cudaStream_t stream, const Tin* label, const T* weight, \
size_t count, size_t label_depth, int64_t ignore_index, \
T* weight_data_nd)
#define INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(T, TLabel, TOut) \
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)
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(float, int32_t);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(float, int64_t);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(half, int64_t);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(BFloat16, int64_t);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(float, int32_t, float);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(float, int64_t, float);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(half, int32_t, float);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(half, int64_t, float);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(half, int64_t, half);
INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(BFloat16, int64_t, BFloat16);
#undef INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL
template <typename T, typename TAcc, typename Tin>
template <typename T, typename TAcc, typename TLabel>
struct OpWeightedSoftmaxCrossEntropyLoss {
OpWeightedSoftmaxCrossEntropyLoss(const T* log_prob_data, const Tin* label_data, const T* weight_data,
const TAcc* normalize_factor_data, Tin C, Tin ignore_index)
OpWeightedSoftmaxCrossEntropyLoss(const T* log_prob_data, const TLabel* label_data, const T* weight_data,
const TAcc* normalize_factor_data, TLabel C, TLabel ignore_index)
: log_prob_data_(log_prob_data),
label_data_(label_data),
weight_data_(weight_data),
@ -76,27 +78,27 @@ struct OpWeightedSoftmaxCrossEntropyLoss {
}
const T* log_prob_data_;
const Tin* label_data_;
const TLabel* label_data_;
const T* weight_data_;
const TAcc* normalize_factor_data_;
Tin C_;
Tin ignore_index_;
TLabel C_;
TLabel ignore_index_;
};
template <typename T, typename TAcc, typename Tin>
void SoftmaxCrossEntropyLossImpl(cudaStream_t stream, const T* log_prob, const Tin* label, const T* weight,
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, Tin> op(log_prob, label, weight, normalize_factor,
static_cast<Tin>(label_depth), static_cast<Tin>(ignore_index));
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);
}
template <typename T, typename TAcc, typename Tin, bool IsReductionNone, bool HasBias>
template <typename T, typename TAcc, typename TLabel, typename TOut, bool IsReductionNone, bool HasBias>
struct OpWeightedSoftmaxCrossEntropyLossGrad {
OpWeightedSoftmaxCrossEntropyLossGrad(const T* dY_data, const T* log_prob_data, const Tin* label_data,
const T* weight_data, const TAcc* normalize_factor_data, const T* bias_data,
Tin C)
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,
TLabel C)
: dY_data_(dY_data),
log_prob_data_(log_prob_data),
label_data_(label_data),
@ -107,39 +109,39 @@ struct OpWeightedSoftmaxCrossEntropyLossGrad {
C_fdm_ = fast_divmod(static_cast<int>(C));
}
__device__ __inline__ T operator()(CUDA_LONG idx) const {
__device__ __inline__ TOut operator()(CUDA_LONG 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.
T result = T(0.f);
TAcc result = TAcc(0.f);
if (*normalize_factor_data_ != TAcc(0.f)) {
int 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<T>(static_cast<TAcc>((IsReductionNone ? dY_data_[row] : *dY_data_) * weight_data_[row]) *
(_Exp(static_cast<TAcc>(log_prob_data_[idx])) - (TAcc)(d == label_data_[row])) /
(*normalize_factor_data_));
result = static_cast<TAcc>((IsReductionNone ? dY_data_[row] : *dY_data_) * weight_data_[row]) *
(_Exp(static_cast<TAcc>(log_prob_data_[idx])) - (TAcc)(d == label_data_[row])) /
(*normalize_factor_data_);
}
return HasBias ? result + bias_data_[idx] : result;
return HasBias ? static_cast<TOut>(result + static_cast<TAcc>(bias_data_[idx])) : static_cast<TOut>(result);
}
const T* dY_data_;
const T* log_prob_data_;
const Tin* label_data_;
const TLabel* label_data_;
const T* weight_data_;
const TAcc* normalize_factor_data_;
const T* bias_data_;
Tin C_;
const TOut* bias_data_;
TLabel C_;
fast_divmod C_fdm_;
};
template <typename T, typename TAcc, typename Tin>
void SoftmaxCrossEntropyLossGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const Tin* label,
const T* weight, const TAcc* normalize_factor, const T* bias_data, size_t count,
size_t label_depth, bool reduction_none, T* output_data) {
#define LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL(is_reduction_none, has_bias) \
OpWeightedSoftmaxCrossEntropyLossGrad<T, TAcc, Tin, is_reduction_none, has_bias> op( \
dY, log_prob, label, weight, normalize_factor, bias_data, static_cast<Tin>(label_depth)); \
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count * label_depth)
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)
if (reduction_none) {
if (bias_data) {
LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL(true, true);
@ -156,14 +158,10 @@ void SoftmaxCrossEntropyLossGradImpl(cudaStream_t stream, const T* dY, const T*
#undef LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL
}
#define INSTANTIATE_SCE_LOSS_IMPL(T, TAcc, Tin) \
template void SoftmaxCrossEntropyLossImpl(cudaStream_t stream, const T* log_prob, const Tin* label, const T* weight, \
const TAcc* normalize_factor, size_t count, size_t label_depth, \
int64_t ignore_index, T* output_data); \
template void SoftmaxCrossEntropyLossGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const Tin* label, \
const T* weight, const TAcc* normalize_factor, const T* bias_data, \
size_t count, size_t label_depth, bool reducation_none, \
T* output_data)
#define INSTANTIATE_SCE_LOSS_IMPL(T, TAcc, TLabel) \
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);
INSTANTIATE_SCE_LOSS_IMPL(float, float, int32_t);
INSTANTIATE_SCE_LOSS_IMPL(float, float, int64_t);
@ -172,5 +170,20 @@ INSTANTIATE_SCE_LOSS_IMPL(BFloat16, float, int64_t);
#undef INSTANTIATE_SCE_LOSS_IMPL
#define INSTANTIATE_SCE_LOSS_GRAD_IMPL(T, TAcc, TLabel, TOut) \
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 reducation_none, \
TOut* output_data)
INSTANTIATE_SCE_LOSS_GRAD_IMPL(float, float, int32_t, float);
INSTANTIATE_SCE_LOSS_GRAD_IMPL(float, float, int32_t, half);
INSTANTIATE_SCE_LOSS_GRAD_IMPL(float, float, int64_t, float);
INSTANTIATE_SCE_LOSS_GRAD_IMPL(float, float, int64_t, half);
INSTANTIATE_SCE_LOSS_GRAD_IMPL(half, float, int64_t, half);
INSTANTIATE_SCE_LOSS_GRAD_IMPL(BFloat16, float, int64_t, BFloat16);
#undef INSTANTIATE_SCE_LOSS_GRAD_IMPL
} // namespace cuda
} // namespace onnxruntime

22
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h
View file

@ -10,11 +10,11 @@
namespace onnxruntime {
namespace cuda {
template <typename T, typename TAcc, typename Tin>
template <typename T, typename TAcc, typename TLabel>
void SoftmaxCrossEntropyLossImpl(
cudaStream_t stream,
const T* log_prob,
const Tin* label,
const TLabel* label,
const T* weight,
const TAcc* normalize_factor,
size_t count,
@ -22,31 +22,31 @@ void SoftmaxCrossEntropyLossImpl(
int64_t ignore_index,
T* output_data);
template <typename T, typename TAcc, typename Tin>
template <typename T, typename TAcc, typename TLabel, typename TOut>
void SoftmaxCrossEntropyLossGradImpl(
cudaStream_t stream,
const T* dY,
const T* log_prob,
const Tin* label,
const TLabel* label,
const T* weight,
const TAcc* normalize_factor,
const T* bias_data,
const TOut* bias_data,
size_t count,
size_t label_depth,
bool reduction_none,
T* output_data);
TOut* output_data);
template <typename T, typename Tin>
template <typename T, typename TLabel, typename TOut>
void ComputeSoftmaxCrossEntropyWeightsImpl(
cudaStream_t stream,
const Tin* label,
const TLabel* label,
const T* weight,
size_t count,
size_t label_depth,
int64_t ignore_index,
T* weight_data_nd);
TOut* weight_data_nd);
template <typename T, typename Tin>
template <typename T, typename TLabel, typename TOut>
class SoftmaxCrossEntropyLoss final : public LossBase {
public:
SoftmaxCrossEntropyLoss(const OpKernelInfo& info) : LossBase(info) {
@ -60,7 +60,7 @@ class SoftmaxCrossEntropyLoss final : public LossBase {
int64_t ignore_index_;
};
template <typename T, typename Tin>
template <typename T, typename TLabel, typename TOut>
class SoftmaxCrossEntropyLossGrad final : public LossBase {
public:
SoftmaxCrossEntropyLossGrad(const OpKernelInfo& info) : LossBase(info) {

34
orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cc
View file

@ -26,7 +26,7 @@ namespace cuda {
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<Tin>()), \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<Tin>()), \
Class<T, Tin>);
template <typename T>
@ -49,11 +49,11 @@ Status SoftmaxCrossEntropy<T>::ComputeInternal(OpKernelContext* ctx) const {
T* log_prob_data = log_prob->template MutableData<T>();
// calculate logsoftmax
auto status = SoftMaxComputeHelper<T, true>(Stream(ctx),
logit_data,
logit_reshape,
log_prob_data,
1 /*axis default*/);
auto status = SoftMaxComputeHelper<T, T, true>(Stream(ctx),
logit_data,
logit_reshape,
log_prob_data,
1 /*axis default*/);
ORT_RETURN_IF_ERROR(status);
size_t normalize_factor = N;
@ -151,11 +151,11 @@ Status SparseSoftmaxCrossEntropy<T, Tin>::ComputeInternal(OpKernelContext* ctx)
T* log_prob_data = log_prob->template MutableData<T>();
// calculate logsoftmax
auto status = SoftMaxComputeHelper<T, true>(Stream(ctx),
logit_data,
logit_reshape,
log_prob_data,
1 /*axis default*/);
auto status = SoftMaxComputeHelper<T, T, true>(Stream(ctx),
logit_data,
logit_reshape,
log_prob_data,
1 /*axis default*/);
ORT_RETURN_IF_ERROR(status);
// calculate (label * log(softmax)) for each sample
@ -177,11 +177,13 @@ Status SparseSoftmaxCrossEntropy<T, Tin>::ComputeInternal(OpKernelContext* ctx)
auto normalize_factor_data = GetScratchBuffer<T>(1, ctx->GetComputeStream());
if (reduction_ == ReductionType::SUM) {
constexpr T normalize_factor_one = static_cast<T>(1);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor_one, sizeof(T), cudaMemcpyHostToDevice, Stream(ctx)));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor_one, sizeof(T),
cudaMemcpyHostToDevice, Stream(ctx)));
} else if (reduction_ == ReductionType::MEAN) {
if (weight_data == nullptr) {
const T normalize_factor = static_cast<T>(N);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), cudaMemcpyHostToDevice, Stream(ctx)));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T),
cudaMemcpyHostToDevice, Stream(ctx)));
} else {
ORT_RETURN_IF_ERROR(reduce_sum(
Stream(ctx),
@ -247,11 +249,13 @@ Status SparseSoftmaxCrossEntropyGrad<T, Tin>::ComputeInternal(OpKernelContext* c
auto normalize_factor_data = GetScratchBuffer<T>(1, ctx->GetComputeStream());
if (reduction_ == ReductionType::SUM) {
constexpr T normalize_factor_one = static_cast<T>(1);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor_one, sizeof(T), cudaMemcpyHostToDevice, Stream(ctx)));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor_one, sizeof(T),
cudaMemcpyHostToDevice, Stream(ctx)));
} else if (reduction_ == ReductionType::MEAN) {
if (weight_data == nullptr) {
const T normalize_factor = static_cast<T>(N);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), cudaMemcpyHostToDevice, Stream(ctx)));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T),
cudaMemcpyHostToDevice, Stream(ctx)));
} else {
// Compute buffer size in byte for reduction APIs.
const auto buffer_size =

28
orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc
View file

@ -61,12 +61,14 @@ class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDom
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternal);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternalGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, SoftmaxGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, LogSoftmaxGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, SoftmaxGrad_13);
@ -273,12 +275,14 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternal)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_int64_t_float, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_MLFloat16, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16_int64_t_BFloat16, SoftmaxCrossEntropyLossInternalGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_float_float, BatchNormalizationGrad)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double_double_double, BatchNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_MLFloat16_MLFloat16, BatchNormalizationGrad)>,