From 79aa0acdd08e7ab84a95ce9f6d985565e68960ac Mon Sep 17 00:00:00 2001 From: pengwa Date: Tue, 28 Feb 2023 18:02:08 +0800 Subject: [PATCH] SCELoss(SCELossGrad) support half(float) input float(half) output (#13972) ### 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 --- .../core/providers/cpu/cpu_provider_shared.h | 4 +- .../core/providers/cuda/math/softmax.cc | 47 +- .../core/providers/cuda/math/softmax.h | 4 +- .../core/providers/cuda/math/softmax_impl.cu | 58 +- .../core/providers/rocm/math/softmax.cc | 51 +- .../core/providers/rocm/math/softmax.h | 4 +- .../core/providers/rocm/math/softmax_impl.cu | 7 +- .../python/tools/symbolic_shape_infer.py | 6 + onnxruntime/test/util/include/test_utils.h | 25 + .../core/graph/gradient_builder.cc | 13 +- .../core/graph/training_op_defs.cc | 229 ++++- .../_custom_autograd_function_exporter.py | 7 +- .../ortmodule/_custom_op_symbolic_registry.py | 61 +- .../training_api/core/training_api_tests.cc | 9 +- .../training_ops/cuda/cross_entropy_test.cc | 800 +++++++++++++----- orttraining/orttraining/training_api/utils.h | 26 - .../cpu/loss/softmax_cross_entropy_loss.cc | 1 + .../cuda/cuda_training_kernels.cc | 28 +- .../loss/softmax_cross_entropy_loss_impl.cc | 225 ++--- .../loss/softmax_cross_entropy_loss_impl.cu | 143 ++-- .../loss/softmax_cross_entropy_loss_impl.h | 22 +- .../cuda/loss/softmaxcrossentropy_impl.cc | 34 +- .../rocm/rocm_training_kernels.cc | 28 +- 23 files changed, 1193 insertions(+), 639 deletions(-) diff --git a/onnxruntime/core/providers/cpu/cpu_provider_shared.h b/onnxruntime/core/providers/cpu/cpu_provider_shared.h index f12e080adf..1ecbca7a55 100644 --- a/onnxruntime/core/providers/cpu/cpu_provider_shared.h +++ b/onnxruntime/core/providers/cpu/cpu_provider_shared.h @@ -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& axes, bool keepdims) { return g_host_cpu.contrib__ExecuteReduceSumATen(p_ctx, axes, keepdims); } +inline Status ExecuteReduceSumATen(OpKernelContext* p_ctx, const gsl::span& axes, bool keepdims) { + return g_host_cpu.contrib__ExecuteReduceSumATen(p_ctx, axes, keepdims); +} } // namespace contrib #endif // ENABLE_TRAINING #endif // USE_CUDA || USE_ROCM diff --git a/onnxruntime/core/providers/cuda/math/softmax.cc b/onnxruntime/core/providers/cuda/math/softmax.cc index 5047a70242..56e4e052bb 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.cc +++ b/onnxruntime/core/providers/cuda/math/softmax.cc @@ -11,37 +11,44 @@ namespace onnxruntime { namespace cuda { -template +template Status SoftMaxComputeHelper( cudaStream_t stream, const T* X, const TensorShape& input_shape, - T* Y, + TOut* Y, int64_t axis) { - typedef typename ToCudaType::MappedType CudaT; + typedef typename ToCudaType::MappedType CudaT_IN; + typedef typename ToCudaType::MappedType CudaT_OUT; + typedef typename ToCudaType::MappedType CudaT_ACCUM; int64_t N = input_shape.SizeToDimension(axis); int64_t D = input_shape.SizeFromDimension(axis); - auto Y_data = reinterpret_cast(Y); - auto X_data = reinterpret_cast(X); + auto Y_data = reinterpret_cast(Y); + auto X_data = reinterpret_cast(X); if (D <= 1024 && D * sizeof(T) <= 4096) { - return dispatch_warpwise_softmax_forward, is_log_softmax>( + return dispatch_warpwise_softmax_forward< + CudaT_IN, CudaT_OUT, AccumulationType_t, is_log_softmax>( stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); } - return dispatch_blockwise_softmax_forward, is_log_softmax>( + + return dispatch_blockwise_softmax_forward, is_log_softmax>( stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); } -#define SPECIALIZED_SOFTMAX_HELPER_IMPL(T) \ - template Status SoftMaxComputeHelper(cudaStream_t stream, const T* input, const TensorShape& shape, T* Y, int64_t axis); \ - template Status SoftMaxComputeHelper(cudaStream_t stream, const T* input, const TensorShape& shape, T* Y, int64_t axis); +#define SPECIALIZED_SOFTMAX_HELPER_IMPL(T, TOut) \ + template Status SoftMaxComputeHelper(cudaStream_t stream, const T* input, \ + const TensorShape& shape, TOut* Y, int64_t axis); \ + template Status SoftMaxComputeHelper(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::ComputeInternal(OpKernelContext* ctx) const { Status status; if (log_softmax_) { - status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, - is_transpose_required ? static_cast(rank) - 1 - : static_cast(axis)); + status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, + is_transpose_required ? static_cast(rank) - 1 + : static_cast(axis)); } else { - status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, - is_transpose_required ? static_cast(rank) - 1 - : static_cast(axis)); + status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, + is_transpose_required ? static_cast(rank) - 1 + : static_cast(axis)); } if (!status.IsOK()) diff --git a/onnxruntime/core/providers/cuda/math/softmax.h b/onnxruntime/core/providers/cuda/math/softmax.h index b66ad32517..8e96857bc7 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.h +++ b/onnxruntime/core/providers/cuda/math/softmax.h @@ -9,12 +9,12 @@ namespace onnxruntime { namespace cuda { -template +template Status SoftMaxComputeHelper( cudaStream_t stream, const T* input, const TensorShape& shape, - T* Y, + TOut* Y, int64_t axis); template diff --git a/onnxruntime/core/providers/cuda/math/softmax_impl.cu b/onnxruntime/core/providers/cuda/math/softmax_impl.cu index 4c097f714b..ddf07803fc 100644 --- a/onnxruntime/core/providers/cuda/math/softmax_impl.cu +++ b/onnxruntime/core/providers/cuda/math/softmax_impl.cu @@ -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 -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(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(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(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(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 diff --git a/onnxruntime/core/providers/rocm/math/softmax.cc b/onnxruntime/core/providers/rocm/math/softmax.cc index 013c810f84..1dae8f3ccf 100644 --- a/onnxruntime/core/providers/rocm/math/softmax.cc +++ b/onnxruntime/core/providers/rocm/math/softmax.cc @@ -11,42 +11,45 @@ namespace onnxruntime { namespace rocm { -template +template 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::MappedType HipT; + typedef typename ToHipType::MappedType HipT_IN; + typedef typename ToHipType::MappedType HipT_OUT; + typedef typename ToHipType::MappedType HipT_ACCUM; int64_t N = input_shape.SizeToDimension(axis); int64_t D = input_shape.SizeFromDimension(axis); - auto Y_data = reinterpret_cast(Y); - auto X_data = reinterpret_cast(X); + auto Y_data = reinterpret_cast(Y); + auto X_data = reinterpret_cast(X); if (D <= 1024 && D * sizeof(T) <= 4096) { - return dispatch_warpwise_softmax_forward, IsLogSoftmax>( + return dispatch_warpwise_softmax_forward, IsLogSoftmax>( stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N), tuning_ctx); } - return dispatch_blockwise_softmax_forward, IsLogSoftmax>( + return dispatch_blockwise_softmax_forward, IsLogSoftmax>( stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N), tuning_ctx); } -#define SPECIALIZED_SOFTMAX_HELPER_IMPL(T) \ - template Status SoftMaxComputeHelper(hipStream_t stream, const T* input, const TensorShape& shape, T* Y, \ - int64_t axis, RocmTuningContext* tuning_ctx); \ - template Status SoftMaxComputeHelper(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(hipStream_t stream, const T* input, const TensorShape& shape, TOut* Y, \ + int64_t axis, RocmTuningContext* tuning_ctx); \ + template Status SoftMaxComputeHelper(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::ComputeInternal(OpKernelContext* ctx) const { Status status; if (log_softmax_) { - status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, - is_transpose_required ? static_cast(rank) - 1 - : static_cast(axis), - GetTuningContext()); + status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, + is_transpose_required ? static_cast(rank) - 1 + : static_cast(axis), + GetTuningContext()); } else { - status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, - is_transpose_required ? static_cast(rank) - 1 - : static_cast(axis), - GetTuningContext()); + status = SoftMaxComputeHelper(Stream(ctx), X_data, *compute_input_shape, Y_data, + is_transpose_required ? static_cast(rank) - 1 + : static_cast(axis), + GetTuningContext()); } if (!status.IsOK()) diff --git a/onnxruntime/core/providers/rocm/math/softmax.h b/onnxruntime/core/providers/rocm/math/softmax.h index 89f75c5eb1..8443cfe593 100644 --- a/onnxruntime/core/providers/rocm/math/softmax.h +++ b/onnxruntime/core/providers/rocm/math/softmax.h @@ -11,12 +11,12 @@ namespace rocm { using tunable::RocmTuningContext; -template +template Status SoftMaxComputeHelper( hipStream_t stream, const T* input, const TensorShape& shape, - T* Y, + TOut* Y, int64_t axis, RocmTuningContext* tuning_ctx = nullptr); diff --git a/onnxruntime/core/providers/rocm/math/softmax_impl.cu b/onnxruntime/core/providers/rocm/math/softmax_impl.cu index 26fe515c77..f7a4a60f31 100644 --- a/onnxruntime/core/providers/rocm/math/softmax_impl.cu +++ b/onnxruntime/core/providers/rocm/math/softmax_impl.cu @@ -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 diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 9daada5d53..a20757d800 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -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()) diff --git a/onnxruntime/test/util/include/test_utils.h b/onnxruntime/test/util/include/test_utils.h index 05555d1bec..2fd51701a9 100644 --- a/onnxruntime/test/util/include/test_utils.h +++ b/onnxruntime/test/util/include/test_utils.h @@ -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 #include @@ -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 +void CreateInputOrtValueOnCPU(gsl::span dims, const std::vector& 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(value.size())); + auto element_type = DataTypeImpl::GetType(); + auto allocator = alloc ? alloc : cpu_allocator; + auto p_tensor = std::make_unique(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(), + DataTypeImpl::GetType()->GetDeleteFunc()); +} + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 668a44d985..fe5d3d1687 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -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 attrs; + for (auto& attr : src_attrs) { + if (attr.first == "output_type") { + attrs.push_back(MakeAttribute("output_type", static_cast(IElemType(0)))); + continue; + } + attrs.push_back(attr.second); + } + return std::vector{ - 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) { diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index fb02924d72..73f6952c2c 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -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({0, 0, -1})) .Add(R"( - X_NCD = Reshape (scores, Shape3D) + X_NCD = Reshape (scores_casted, Shape3D) X_NDC = Transpose (X_NCD) X_LogSM = LogSoftmax (X_NDC) X_LogSM_NCD = Transpose (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 (X_Log, labels, weights, ignore_index)"); + builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, weights_casted, ignore_index)"); else - builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, weights)"); + builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, weights_casted)"); else if (hasIgnoreIndex) - builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, , ignore_index)"); + builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, , ignore_index)"); else - builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels)"); + builder.Add("intermediate_output = com.microsoft.NegativeLogLikelihoodLossInternal2 (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 (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 (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 (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( diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py index 887066d2a3..96f9ccd4ae 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py @@ -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()) diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py index 7cd889a156..24b13b4ed1 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py @@ -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) diff --git a/orttraining/orttraining/test/training_api/core/training_api_tests.cc b/orttraining/orttraining/test/training_api/core/training_api_tests.cc index 534df6d6c3..2efa8b4bc7 100644 --- a/orttraining/orttraining/test/training_api/core/training_api_tests.cc +++ b/orttraining/orttraining/test/training_api/core/training_api_tests.cc @@ -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 dims, OrtValue& input) { TensorShape shape(dims); std::vector data(shape.Size()); GenerateRandomData(data); - onnxruntime::training::api::utils::CreateInputOrtValue(dims, data, &input); + onnxruntime::test::CreateInputOrtValueOnCPU(dims, data, &input); } void TestModuleExport(const std::vector>& providers) { @@ -145,7 +146,7 @@ void TestLRSchduler(const std::basic_string& test_file_name, float in OrtValue input, target; GenerateRandomInput(std::array{2, 784}, input); - onnxruntime::training::api::utils::CreateInputOrtValue( + onnxruntime::test::CreateInputOrtValueOnCPU( std::array{2}, std::vector(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{2, 784}, input); - onnxruntime::training::api::utils::CreateInputOrtValue( + onnxruntime::test::CreateInputOrtValueOnCPU( std::array{2}, std::vector(2, 1), &target); auto data_loader = std::vector>(4, std::vector{input, target}); @@ -347,7 +348,7 @@ TEST(TrainingApiTest, OptimStep) { OrtValue input, target; GenerateRandomInput(std::array{2, 784}, input); - onnxruntime::training::api::utils::CreateInputOrtValue( + onnxruntime::test::CreateInputOrtValueOnCPU( std::array{2}, std::vector(2, 1), &target); auto data_loader = std::vector>(4, std::vector{input, target}); diff --git a/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc b/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc index d79fcb7db3..ae752dfe29 100644 --- a/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc +++ b/orttraining/orttraining/test/training_ops/cuda/cross_entropy_test.cc @@ -1,6 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include "test/compare_ortvalue.h" +#include "test/util/include/default_providers.h" +#include "test/util/include/test_utils.h" #include "test/providers/compare_provider_test_utils.h" using namespace std; @@ -64,7 +68,7 @@ static void TestSoftmaxCrossEntropyGrad(const std::vector& dY_dims, test.CompareWithCPU(kGpuExecutionProvider); } -TEST(CudaKernelTest, SoftmaxCrossEntropy_TinySizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropy_TinySizeTensor) { std::vector X_dims{8, 2}; std::vector label_dims{8, 2}; std::vector Y_dims{}; @@ -73,7 +77,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropy_TinySizeTensor) { TestSoftmaxCrossEntropy(X_dims, label_dims, Y_dims, log_prob_dims, "sum"); } -TEST(CudaKernelTest, SoftmaxCrossEntropy_SmallSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropy_SmallSizeTensor) { std::vector X_dims{8, 20, 10}; std::vector label_dims{8, 20, 10}; std::vector Y_dims{}; @@ -82,7 +86,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropy_SmallSizeTensor) { TestSoftmaxCrossEntropy(X_dims, label_dims, Y_dims, log_prob_dims, "sum"); } -TEST(CudaKernelTest, SoftmaxCrossEntropy_MediumSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropy_MediumSizeTensor) { std::vector X_dims{7, 1024}; std::vector label_dims{7, 1024}; std::vector Y_dims{}; @@ -91,7 +95,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropy_MediumSizeTensor) { TestSoftmaxCrossEntropy(X_dims, label_dims, Y_dims, log_prob_dims, "sum"); } -TEST(CudaKernelTest, SoftmaxCrossEntropy_LargeSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropy_LargeSizeTensor) { std::vector X_dims{2, 512, 30528}; std::vector label_dims{2, 512, 30528}; std::vector Y_dims{}; @@ -100,7 +104,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropy_LargeSizeTensor) { TestSoftmaxCrossEntropy(X_dims, label_dims, Y_dims, log_prob_dims, "sum"); } -TEST(CudaKernelTest, SoftmaxCrossEntropyGrad_TinySizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyGrad_TinySizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{8, 2}; std::vector label_dims{8, 2}; @@ -109,7 +113,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyGrad_TinySizeTensor) { TestSoftmaxCrossEntropyGrad(dY_dims, log_prob_dims, label_dims, dX_dims, "sum"); } -TEST(CudaKernelTest, SoftmaxCrossEntropyGrad_SmallSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyGrad_SmallSizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{8, 20, 10}; std::vector label_dims{8, 20, 10}; @@ -118,7 +122,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyGrad_SmallSizeTensor) { TestSoftmaxCrossEntropyGrad(dY_dims, log_prob_dims, label_dims, dX_dims, "sum"); } -TEST(CudaKernelTest, SoftmaxCrossEntropyGrad_LargeSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyGrad_LargeSizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{2, 512, 30528}; std::vector label_dims{2, 512, 30528}; @@ -127,7 +131,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyGrad_LargeSizeTensor) { TestSoftmaxCrossEntropyGrad(dY_dims, log_prob_dims, label_dims, dX_dims, "sum"); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_Basic) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropy_Basic) { OpTester test("SparseSoftmaxCrossEntropy", 9); test.AddAttribute("reduction", "mean"); @@ -179,7 +183,7 @@ static void TestSparseSoftmaxCrossEntropy(const std::vector* X_dims, test.CompareWithCPU(kGpuExecutionProvider); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_TinySizeTensor) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropy_TinySizeTensor) { std::vector X_dims{8, 2}; std::vector index_dims{8}; std::vector weight_dims{8}; @@ -191,7 +195,7 @@ TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_TinySizeTensor) { TestSparseSoftmaxCrossEntropy(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_SmallSizeTensor) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropy_SmallSizeTensor) { std::vector X_dims{8, 20, 10}; std::vector index_dims{8, 20}; std::vector weight_dims{8, 20}; @@ -203,7 +207,7 @@ TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_SmallSizeTensor) { TestSparseSoftmaxCrossEntropy(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_MediumSizeTensor) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropy_MediumSizeTensor) { std::vector X_dims{8, 1024}; std::vector index_dims{8}; std::vector weight_dims{8}; @@ -215,7 +219,7 @@ TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_MediumSizeTensor) { TestSparseSoftmaxCrossEntropy(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropy_LargeSizeTensor) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropy_LargeSizeTensor) { std::vector X_dims{4, 512, 30528}; std::vector index_dims{4, 512}; std::vector weight_dims{4, 512}; @@ -252,7 +256,7 @@ static void TestSparseSoftmaxCrossEntropyGrad(const std::vector& dY_dim test.CompareWithCPU(kGpuExecutionProvider); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropyGrad_TinySizeTensor) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropyGrad_TinySizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{8, 2}; std::vector index_dims{8}; @@ -261,7 +265,7 @@ TEST(CudaKernelTest, SparseSoftmaxCrossEntropyGrad_TinySizeTensor) { TestSparseSoftmaxCrossEntropyGrad(dY_dims, log_prob_dims, index_dims, dX_dims, "sum"); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropyGrad_SmallSizeTensor) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropyGrad_SmallSizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{8, 20, 10}; std::vector index_dims{8, 20}; @@ -270,7 +274,7 @@ TEST(CudaKernelTest, SparseSoftmaxCrossEntropyGrad_SmallSizeTensor) { TestSparseSoftmaxCrossEntropyGrad(dY_dims, log_prob_dims, index_dims, dX_dims, "sum"); } -TEST(CudaKernelTest, SparseSoftmaxCrossEntropyGrad_LargeSizeTensor) { +TEST(CrossEntropyTest, SparseSoftmaxCrossEntropyGrad_LargeSizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{2, 512, 30528}; std::vector index_dims{2, 512}; @@ -279,25 +283,73 @@ TEST(CudaKernelTest, SparseSoftmaxCrossEntropyGrad_LargeSizeTensor) { TestSparseSoftmaxCrossEntropyGrad(dY_dims, log_prob_dims, index_dims, dX_dims, "sum"); } -static void TestSoftmaxCrossEntropyLoss(CompareOpTester& test, const std::vector* X_dims, - const std::vector* index_dims, const std::vector* weight_dims, - const std::vector* Y_dims, const std::vector* log_prob_dims, - const std::string& reduction, const std::int64_t ignore_index, - const bool test_fp16, const bool is_internal_op, const double error_tolerance) { +static void PrepareSCELossTestData(const std::vector* X_dims, + const std::vector* index_dims, const std::vector* weight_dims, + const std::int64_t ignore_index, + std::vector& X_data, std::vector& index_data, + std::vector& weight_data) { + // create rand inputs + RandomValueGenerator random{2333}; + X_data = random.Uniform(*X_dims, -200.0f, 200.0f); + index_data = random.Uniform(*index_dims, 0, (*X_dims)[1]); + // Add one data point that has ignore_index. + if (index_data.size() > 0) { + index_data[0] = ignore_index; + } + + if (weight_dims) { + weight_data = random.Uniform(*weight_dims, 0.0f, 1.0f); + } +} + +template +static std::vector RunSCELossWithEP(const char* op, + int opset_version, + const char* domain, + std::function()> + ep_creator, + const std::string& reduction, + const std::int64_t ignore_index, + const double error_tolerance, + const std::vector* X_dims, + const std::vector* index_dims, + const std::vector* weight_dims, + const std::vector* Y_dims, + const std::vector* log_prob_dims, + std::vector& X_data, + std::vector& index_data, + std::vector& weight_data) { + /** + * OpTester's atol/rtol check is too strict for our testing cases. Imagine expected value is 4.7683704451628728e-07, + * real value is 0, even we set rtol=1e-1, atol = 1e-4. The check still fail. + * The difference between expected[i] and output[i] is 4.7683704451628728e-07, + * which exceeds *(params.relative_error_) * std::abs(expected[i]) + * (params.relative_error_) * std::abs(expected[i]) evaluates to 4.7683705872714199e-08. + * + * CompareOrtValue called by test.CompareWithCPU looses the check by compare diff with atol + rtol * expected[i]. + * So here we disable OpTester's check by default, and do the check externally. + */ + bool need_verify_outputs = false; + std::vector> expected_values; + // Still need feed the output data even we don't want to verify it. + expected_values.push_back(FillZeros(*Y_dims)); + if (log_prob_dims) { + expected_values.push_back(FillZeros(*log_prob_dims)); + } + + OpTester test(op, opset_version, domain, need_verify_outputs /*verify_output*/); + const bool is_internal_op = (std::string(op).compare("SoftmaxCrossEntropyLossInternal") == 0); + + if (is_internal_op && !std::is_same::value) { + test.AddAttribute("output_type", static_cast(utils::ToTensorProtoElementType())); + } + test.AddAttribute("reduction", reduction); if (!is_internal_op) { test.AddAttribute("ignore_index", ignore_index); } - // create rand inputs - RandomValueGenerator random{2333}; - std::vector X_data = random.Uniform(*X_dims, -200.0f, 200.0f); - std::vector index_data = random.Uniform(*index_dims, 0, (*X_dims)[1]); - // Add one data point that has ignore_index. - if (index_data.size() > 0) { - index_data[0] = ignore_index; - } - if (test_fp16) { + if (std::is_same::value) { std::vector X_data_half(X_data.size()); ConvertFloatToMLFloat16(X_data.data(), X_data_half.data(), static_cast(X_data.size())); test.AddInput("X", *X_dims, X_data_half); @@ -308,8 +360,7 @@ static void TestSoftmaxCrossEntropyLoss(CompareOpTester& test, const std::vector test.AddInput("index", *index_dims, index_data); if (weight_dims) { - std::vector weight_data = random.Uniform(*weight_dims, 0.0f, 1.0f); - if (test_fp16) { + if (std::is_same::value) { std::vector weight_data_half(weight_data.size()); ConvertFloatToMLFloat16(weight_data.data(), weight_data_half.data(), static_cast(weight_data.size())); test.AddInput("weight", *weight_dims, weight_data_half); @@ -322,172 +373,286 @@ static void TestSoftmaxCrossEntropyLoss(CompareOpTester& test, const std::vector test.AddInput("ignore_index", {}, &ignore_index, 1); } - if (test_fp16) { - std::vector Y_data = FillZeros(*Y_dims); - test.AddOutput("output", *Y_dims, Y_data); + if (std::is_same::value) { + std::vector output_half(expected_values[0].size()); + ConvertFloatToMLFloat16(expected_values[0].data(), output_half.data(), static_cast(expected_values[0].size())); + test.AddOutput("output", *Y_dims, output_half); if (log_prob_dims) { - std::vector log_prob_data = FillZeros(*log_prob_dims); - test.AddOutput("log_prob", *log_prob_dims, log_prob_data); + std::vector log_prob_half(expected_values[1].size()); + ConvertFloatToMLFloat16(expected_values[1].data(), log_prob_half.data(), static_cast(expected_values[1].size())); + test.AddOutput("log_prob", *log_prob_dims, log_prob_half); } - // Need to add Cast node if it's FP16 test. For internal Op test, we need to add ONNX domain. - std::unordered_map extra_domain_to_version; - if (is_internal_op) { - extra_domain_to_version[onnxruntime::kOnnxDomain] = 12; - } - - test.CompareWithCPU(kGpuExecutionProvider, error_tolerance, error_tolerance, true, extra_domain_to_version); } else { - std::vector Y_data = FillZeros(*Y_dims); - test.AddOutput("output", *Y_dims, Y_data); - + test.AddOutput("output", *Y_dims, expected_values[0]); if (log_prob_dims) { - std::vector log_prob_data = FillZeros(*log_prob_dims); - test.AddOutput("log_prob", *log_prob_dims, log_prob_data); + test.AddOutput("log_prob", *log_prob_dims, expected_values[1]); + } + } + + std::vector> eps; + eps.emplace_back(ep_creator()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + return test.GetFetches(); +} + +template +static void TestSCELoss(const char* op, int opset_version, + const char* domain, const std::vector* X_dims, + const std::vector* index_dims, const std::vector* weight_dims, + const std::vector* Y_dims, const std::vector* log_prob_dims, + const std::string& reduction, const std::int64_t ignore_index, + const double error_tolerance) { + ASSERT_TRUE((std::is_same::value || std::is_same::value)); + ASSERT_TRUE((std::is_same::value || std::is_same::value)); + + std::vector X_data, weight_data; + std::vector index_data; + PrepareSCELossTestData(X_dims, index_dims, weight_dims, ignore_index, X_data, index_data, weight_data); + + // Run on CPU using float input and output + // (because the CPU implementation doesn't support variant input output types.) + // Be noted, no result comparison is done here. + std::vector cpu_fetches = RunSCELossWithEP( + op, opset_version, domain, + []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, + reduction, ignore_index, error_tolerance, + X_dims, index_dims, weight_dims, + Y_dims, log_prob_dims, + X_data, index_data, weight_data); + + // Run on CUDA. + // Be noted, no result comparison is done here because OpTester's check is too strick for our test cases. + // Check more details in comment of RunSCELossWithEP. + std::vector target_fetches = RunSCELossWithEP( + op, opset_version, domain, + []() -> std::unique_ptr { +#ifdef USE_CUDA + return DefaultCudaExecutionProvider(); +#elif USE_ROCM + return DefaultRocmExecutionProvider(); +#endif + }, + reduction, ignore_index, error_tolerance, + X_dims, index_dims, weight_dims, + Y_dims, log_prob_dims, + X_data, index_data, weight_data); + + // Compare + ASSERT_EQ(cpu_fetches.size(), target_fetches.size()); + for (size_t i = 0; i < cpu_fetches.size(); i++) { + if (std::is_same::value) { + auto y_data_size = cpu_fetches[i].Get().Shape().Size(); + std::vector cpu_temp_buffer; + cpu_temp_buffer.resize(y_data_size); + const float* y_buffer = cpu_fetches[i].Get().Data(); + std::copy(y_buffer, y_buffer + y_data_size, cpu_temp_buffer.begin()); + + std::vector ret_half(cpu_temp_buffer.size()); + ConvertFloatToMLFloat16(cpu_temp_buffer.data(), ret_half.data(), static_cast(cpu_temp_buffer.size())); + + OrtValue target; + test::CreateInputOrtValueOnCPU(cpu_fetches[i].Get().Shape().GetDims(), ret_half, &target); + auto ret = CompareOrtValue(target_fetches[i], target, error_tolerance /*per_sample_tolerance*/, + error_tolerance /*relative_per_sample_tolerance*/, false); + EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second; + + } else { + auto ret = CompareOrtValue(target_fetches[i], cpu_fetches[i], error_tolerance /*per_sample_tolerance*/, + error_tolerance /*relative_per_sample_tolerance*/, false); + EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second; } - test.CompareWithCPU(kGpuExecutionProvider); } } +template static void TestSoftmaxCrossEntropyLoss(const std::vector* X_dims, const std::vector* index_dims, const std::vector* weight_dims, const std::vector* Y_dims, const std::vector* log_prob_dims, const std::string& reduction, - const std::int64_t ignore_index = -1, const bool test_fp16 = false, + const std::int64_t ignore_index = -1, const double error_tolerance = 1e-4) { - CompareOpTester test("SoftmaxCrossEntropyLoss", 12, onnxruntime::kOnnxDomain); - TestSoftmaxCrossEntropyLoss(test, X_dims, index_dims, weight_dims, Y_dims, log_prob_dims, reduction, ignore_index, - test_fp16, false, error_tolerance); + // Only test SoftmaxCrossEntropyLoss if input and output type matches. + if (std::is_same::value) { + TestSCELoss("SoftmaxCrossEntropyLoss", 12, onnxruntime::kOnnxDomain, + X_dims, index_dims, weight_dims, Y_dims, log_prob_dims, + reduction, ignore_index, error_tolerance); + } // Can we add a empty optional input before a non-empty input? if (weight_dims || ignore_index == -1) { - CompareOpTester test_internal("SoftmaxCrossEntropyLossInternal", 1, onnxruntime::kMSDomain); - TestSoftmaxCrossEntropyLoss(test_internal, X_dims, index_dims, weight_dims, Y_dims, log_prob_dims, reduction, - ignore_index, test_fp16, true, error_tolerance); + TestSCELoss("SoftmaxCrossEntropyLossInternal", 1, onnxruntime::kMSDomain, + X_dims, index_dims, weight_dims, Y_dims, log_prob_dims, reduction, + ignore_index, error_tolerance); } } -TEST(CudaKernelTest, SoftmaxCrossEntropyLoss_TinySizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_TinySizeTensor) { std::vector X_dims{8, 2}; std::vector index_dims{8}; std::vector weight_dims{2}; std::vector Y_dims{}; std::vector Y_dims_none{8}; std::vector log_prob_dims{8, 2}; - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); // Just test ignore_index for small tensor because it will increase test time a lot with little verification gain. - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean", 0); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean", 0); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum", 0); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum", 0); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none", 0); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none", 0); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean", 0); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean", 0); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum", 0); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum", 0); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none", 0); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none", 0); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLoss_TinySizeTensor_half) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_TinySizeTensor_half) { std::vector X_dims{8, 2}; std::vector index_dims{8}; std::vector weight_dims{2}; std::vector Y_dims{}; std::vector Y_dims_none{8}; std::vector log_prob_dims{8, 2}; - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none", -1, true, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); // Just test ignore_index for small tensor because it will increase test time a lot with little verification gain. - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean", 0, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean", 0, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum", 0, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum", 0, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none", 0, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none", 0, true, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "mean", 0, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "mean", 0, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "sum", 0, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "sum", 0, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, + "none", 0, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, + "none", 0, 5e-2); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLoss_SmallSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_SmallSizeTensor) { std::vector X_dims{8, 20, 10}; std::vector index_dims{8, 10}; std::vector weight_dims{20}; std::vector Y_dims{}; std::vector Y_dims_none{8, 10}; std::vector log_prob_dims{8, 20, 10}; - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLoss_SmallSizeTensor_half) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_SmallSizeTensor_half) { std::vector X_dims{8, 20, 10}; std::vector index_dims{8, 10}; std::vector weight_dims{20}; std::vector Y_dims{}; std::vector Y_dims_none{8, 10}; std::vector log_prob_dims{8, 20, 10}; - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none", -1, true, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLoss_MediumSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_MediumSizeTensor) { std::vector X_dims{8, 1024}; std::vector index_dims{8}; std::vector weight_dims{1024}; std::vector Y_dims{}; std::vector Y_dims_none{8}; std::vector log_prob_dims{8, 1024}; - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLoss_MediumSizeTensor_half) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_MediumSizeTensor_half) { std::vector X_dims{8, 1024}; std::vector index_dims{8}; std::vector weight_dims{1024}; std::vector Y_dims{}; std::vector Y_dims_none{8}; std::vector log_prob_dims{8, 1024}; - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none", -1, true, 5e-2); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none", -1, true, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_MediumSizeTensor_half_input_float_output) { + std::vector X_dims{8, 1024}; + std::vector index_dims{8}; + std::vector weight_dims{1024}; + std::vector Y_dims{}; + std::vector Y_dims_none{8}; + std::vector log_prob_dims{8, 1024}; + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, + "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, + "none", -1, 5e-2); } // TODO fix flaky test // failing random seed: 2873512643 -TEST(CudaKernelTest, DISABLED_SoftmaxCrossEntropyLoss_LargeSizeTensor) { +TEST(CrossEntropyTest, DISABLED_SoftmaxCrossEntropyLoss_LargeSizeTensor) { std::vector X_dims{4, 512, 30528}; std::vector index_dims{4, 30528}; std::vector weight_dims{512}; std::vector Y_dims{}; std::vector Y_dims_none{4, 30528}; std::vector log_prob_dims{4, 512, 30528}; - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); - TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "mean"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims, &log_prob_dims, "sum"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, &weight_dims, &Y_dims_none, &log_prob_dims, "none"); + TestSoftmaxCrossEntropyLoss(&X_dims, &index_dims, nullptr, &Y_dims_none, &log_prob_dims, "none"); } static void TestSoftmaxCrossEntropyLossGrad(const std::vector& dY_dims, @@ -507,7 +672,7 @@ static void TestSoftmaxCrossEntropyLossGrad(const std::vector& dY_dims, std::vector dY_data = random.Uniform(dY_dims, -10.0f, 10.0f); std::vector log_prob_data = random.Uniform(log_prob_dims, -10.0f, 10.0f); std::vector index_data = random.Uniform(index_dims, 0, dX_dims[1]); - //Add one data point that has ignore_index. + // Add one data point that has ignore_index. if (index_data.size() > 0) { index_data[0] = ignore_index; } @@ -538,7 +703,7 @@ static void TestSoftmaxCrossEntropyLossGrad(const std::vector& dY_dims, } } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_TinySizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_TinySizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{8, 2}; std::vector index_dims{8}; @@ -553,7 +718,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_TinySizeTensor) { TestSoftmaxCrossEntropyLossGrad({8}, log_prob_dims, index_dims, dX_dims, "none", 0); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_SmallSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_SmallSizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{8, 20, 10}; std::vector index_dims{8, 10}; @@ -563,7 +728,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_SmallSizeTensor) { TestSoftmaxCrossEntropyLossGrad({8, 10}, log_prob_dims, index_dims, dX_dims, "none"); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_LargeSizeTensor) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LargeSizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{2, 512, 30528}; std::vector index_dims{2, 30528}; @@ -573,7 +738,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_LargeSizeTensor) { TestSoftmaxCrossEntropyLossGrad({2, 30528}, log_prob_dims, index_dims, dX_dims, "none"); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_TinySizeTensor_half) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_TinySizeTensor_half) { std::vector dY_dims{}; std::vector log_prob_dims{8, 2}; std::vector index_dims{8}; @@ -588,7 +753,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_TinySizeTensor_half) { TestSoftmaxCrossEntropyLossGrad({8}, log_prob_dims, index_dims, dX_dims, "none", 0, true, 5e-2); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_SmallSizeTensor_half) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_SmallSizeTensor_half) { std::vector dY_dims{}; std::vector log_prob_dims{8, 20, 10}; std::vector index_dims{8, 10}; @@ -598,7 +763,7 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_SmallSizeTensor_half) { TestSoftmaxCrossEntropyLossGrad({8, 10}, log_prob_dims, index_dims, dX_dims, "none", -1, true, 5e-2); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_LargeSizeTensor_half) { +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LargeSizeTensor_half) { std::vector dY_dims{}; std::vector log_prob_dims{2, 512, 30528}; std::vector index_dims{2, 30528}; @@ -608,139 +773,304 @@ TEST(CudaKernelTest, SoftmaxCrossEntropyLossGrad_LargeSizeTensor_half) { TestSoftmaxCrossEntropyLossGrad({2, 30528}, log_prob_dims, index_dims, dX_dims, "none", -1, true, 5e-2); } -static void TestSoftmaxCrossEntropyLossInternalGrad(const std::vector& dY_dims, - const std::vector& log_prob_dims, - const std::vector& index_dims, - const std::vector& weight_dims, - const std::vector& dX_dims, const std::string& reduction, - const std::int64_t ignore_index = -1, const bool test_fp16 = false, - const double error_tolerance = 1e-4, const bool has_bias = false) { - CompareOpTester test("SoftmaxCrossEntropyLossInternalGrad", 1, onnxruntime::kMSDomain); - test.AddAttribute("reduction", reduction); - - // create rand inputs +static void PrepareSCELossInternalGradTestData( + const std::vector& dY_dims, const std::vector& log_prob_dims, + const std::vector& index_dims, const std::vector& weight_dims, + const std::vector& dX_dims, const std::int64_t ignore_index, bool has_bias, + std::vector& dY_data, std::vector& log_prob_data, + std::vector& index_data, std::vector& weight_data, + std::vector& bias_data) { + // Create rand inputs RandomValueGenerator random{2333}; - std::vector dY_data = random.Uniform(dY_dims, -10.0f, 10.0f); - std::vector log_prob_data = random.Uniform(log_prob_dims, -10.0f, 10.0f); - std::vector index_data = random.Uniform(index_dims, 0, dX_dims[1]); + dY_data = random.Uniform(dY_dims, -10.0f, 10.0f); + log_prob_data = random.Uniform(log_prob_dims, -10.0f, 10.0f); + index_data = random.Uniform(index_dims, 0, dX_dims[1]); // Add one data point that has ignore_index. if ((ignore_index != -1) && (index_data.size() > 0)) { index_data[0] = ignore_index; } - std::vector weight_data = random.Uniform(weight_dims, 0.0f, 1.0f); - if (test_fp16) { - std::vector dY_data_half(dY_data.size()); - ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), static_cast(dY_data.size())); - test.AddInput("dY", dY_dims, dY_data_half); + weight_data = random.Uniform(weight_dims, 0.0f, 1.0f); - std::vector log_prob_data_half(log_prob_data.size()); - ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), static_cast(log_prob_data.size())); - test.AddInput("log_prob", log_prob_dims, log_prob_data_half); - - test.AddInput("index", index_dims, index_data); - - std::vector weight_data_half(weight_data.size()); - ConvertFloatToMLFloat16(weight_data.data(), weight_data_half.data(), static_cast(weight_data.size())); - test.AddInput("weight", weight_dims, weight_data_half); - - if (ignore_index != -1 || has_bias) { - test.AddInput("ignore_index", {}, &ignore_index, 1); - } - - if (has_bias) { - std::vector bias_data = random.Uniform(dX_dims, 0.0f, 1.0f); - std::vector bias_data_half(bias_data.size()); - ConvertFloatToMLFloat16(bias_data.data(), bias_data_half.data(), static_cast(bias_data.size())); - test.AddInput("bias", dX_dims, bias_data_half); - } - - std::vector dX_data = FillZeros(dX_dims); - - test.AddOutput("dX", dX_dims, dX_data); - test.CompareWithCPU(kGpuExecutionProvider, error_tolerance, error_tolerance); - } else { - test.AddInput("dY", dY_dims, dY_data); - test.AddInput("log_prob", log_prob_dims, log_prob_data); - test.AddInput("index", index_dims, index_data); - test.AddInput("weight", weight_dims, weight_data); - if (ignore_index != -1 || has_bias) { - test.AddInput("ignore_index", {}, &ignore_index, 1); - } - - if (has_bias) { - std::vector bias_data = random.Uniform(dX_dims, 0.0f, 1.0f); - test.AddInput("bias", dX_dims, bias_data); - } - - std::vector dX_data = FillZeros(dX_dims); - - test.AddOutput("dX", dX_dims, dX_data); - test.CompareWithCPU(kGpuExecutionProvider); + if (has_bias) { + bias_data = random.Uniform(dX_dims, 0.0f, 1.0f); } } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossInternalGrad_TinySizeTensor) { - std::vector dY_dims{}; - std::vector log_prob_dims{8, 2}; - std::vector index_dims{8}; - std::vector weight_dims{2}; - std::vector dX_dims{8, 2}; - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean"); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum"); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none"); +template +static std::vector RunSCELossInternalGradWithEP( + std::function()> + ep_creator, + const std::string& reduction, + const std::int64_t ignore_index, + const double error_tolerance, + const bool has_bias, + const std::vector& dY_dims, + const std::vector& log_prob_dims, + const std::vector& index_dims, + const std::vector& weight_dims, + const std::vector& dX_dims, + std::vector& dY_data, + std::vector& log_prob_data, + std::vector& index_data, + std::vector& weight_data, + std::vector& bias_data) { + /** + * OpTester's atol/rtol check is too strict for our testing cases. Imagine expected value is 4.7683704451628728e-07, + * real value is 0, even we set rtol=1e-1, atol = 1e-4. The check still fail with the following check: + * The difference between expected[i] and output[i] is 4.7683704451628728e-07, + * which exceeds *(params.relative_error_) * std::abs(expected[i]) + * (params.relative_error_) * std::abs(expected[i]) evaluates to 4.7683705872714199e-08. + * + * CompareOrtValue called by test.CompareWithCPU looses the check by compare diff with atol + rtol * expected[i]. + * So here we disable OpTester's check by default, and do the check externally. + */ + bool need_verify_outputs = false; + std::vector expected_value = FillZeros(dX_dims); - // Just test ignore_index for small tensor because it will increase test time a lot with little verification gain. - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean", 0); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum", 0); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none", 0); + ORT_ENFORCE((std::is_same::value || std::is_same::value)); + ORT_ENFORCE((std::is_same::value || std::is_same::value)); - // Bias. - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean", -1, false, - 1e-4, true); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum", -1, false, - 1e-4, true); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none", -1, false, 1e-4, - true); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean", 0, false, - 1e-4, true); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum", 0, false, - 1e-4, true); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none", 0, false, 1e-4, - true); + std::unique_ptr test_unique_ptr = std::make_unique( + "SoftmaxCrossEntropyLossInternalGrad", 1, onnxruntime::kMSDomain, need_verify_outputs /*verify_output*/); + + OpTester* test = test_unique_ptr.get(); + + test->AddAttribute("reduction", reduction); + if (!std::is_same::value) { + test->AddAttribute("output_type", static_cast(utils::ToTensorProtoElementType())); + } + + if (std::is_same::value) { + std::vector dY_data_half(dY_data.size()); + ConvertFloatToMLFloat16(dY_data.data(), dY_data_half.data(), static_cast(dY_data.size())); + test->AddInput("dY", dY_dims, dY_data_half); + + std::vector log_prob_data_half(log_prob_data.size()); + ConvertFloatToMLFloat16(log_prob_data.data(), log_prob_data_half.data(), static_cast(log_prob_data.size())); + test->AddInput("log_prob", log_prob_dims, log_prob_data_half); + + test->AddInput("index", index_dims, index_data); + + std::vector weight_data_half(weight_data.size()); + ConvertFloatToMLFloat16(weight_data.data(), weight_data_half.data(), static_cast(weight_data.size())); + test->AddInput("weight", weight_dims, weight_data_half); + + if (ignore_index != -1 || has_bias) { + test->AddInput("ignore_index", {}, &ignore_index, 1); + } + } else { + test->AddInput("dY", dY_dims, dY_data); + test->AddInput("log_prob", log_prob_dims, log_prob_data); + test->AddInput("index", index_dims, index_data); + test->AddInput("weight", weight_dims, weight_data); + if (ignore_index != -1 || has_bias) { + test->AddInput("ignore_index", {}, &ignore_index, 1); + } + } + + if (std::is_same::value) { + // Be noted, bias should be aligned with output's data type. + if (has_bias) { + std::vector bias_data_half(bias_data.size()); + ConvertFloatToMLFloat16(bias_data.data(), bias_data_half.data(), static_cast(bias_data.size())); + test->AddInput("bias", dX_dims, bias_data_half); + } + + std::vector expected_data_half(expected_value.size()); + ConvertFloatToMLFloat16(expected_value.data(), expected_data_half.data(), static_cast(expected_value.size())); + test->AddOutput("dX", dX_dims, expected_data_half, false /*sort_output*/, + error_tolerance /*rel_error*/, error_tolerance /*abs_error*/); + + } else { + if (has_bias) { + test->AddInput("bias", dX_dims, bias_data); + } + + test->AddOutput("dX", dX_dims, expected_value, false /*sort_output*/, + error_tolerance /*rel_error*/, error_tolerance /*abs_error*/); + } + + std::vector> eps; + eps.emplace_back(ep_creator()); + test->Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + return test->GetFetches(); } -TEST(CudaKernelTest, SoftmaxCrossEntropyLossInternalGrad_TinySizeTensor_half) { +template +static void TestSoftmaxCrossEntropyLossInternalGrad(const std::vector& dY_dims, + const std::vector& log_prob_dims, + const std::vector& index_dims, + const std::vector& weight_dims, + const std::vector& dX_dims, + const std::string& reduction, + const std::int64_t ignore_index = -1, + const double error_tolerance = 1e-4, + const bool has_bias = false) { + std::vector dY_data, log_prob_data, weight_data, bias_data; + std::vector index_data; + PrepareSCELossInternalGradTestData(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, ignore_index, has_bias, + dY_data, log_prob_data, index_data, weight_data, bias_data); + + // Run on CPU using float input and output + // (because the CPU implementation doesn't support variant input output types.) + // Be noted, no result comparison is done here. + std::vector cpu_fetches = + RunSCELossInternalGradWithEP( + []() -> std::unique_ptr { return DefaultCpuExecutionProvider(); }, + reduction, ignore_index, error_tolerance, has_bias, + dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + dY_data, log_prob_data, index_data, weight_data, bias_data); + + // Run on CUDA and compare results with cpu results. + std::vector target_fetches = + RunSCELossInternalGradWithEP( + []() -> std::unique_ptr { +#ifdef USE_CUDA + return DefaultCudaExecutionProvider(); +#elif USE_ROCM + return DefaultRocmExecutionProvider(); +#endif + }, + reduction, ignore_index, error_tolerance, has_bias, + dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + dY_data, log_prob_data, index_data, weight_data, bias_data); + + // Compare + ASSERT_EQ(cpu_fetches.size(), target_fetches.size()); + for (size_t i = 0; i < cpu_fetches.size(); i++) { + if (std::is_same::value) { + auto y_data_size = cpu_fetches[i].Get().Shape().Size(); + std::vector cpu_temp_buffer; + cpu_temp_buffer.resize(y_data_size); + const float* y_buffer = cpu_fetches[i].Get().Data(); + std::copy(y_buffer, y_buffer + y_data_size, cpu_temp_buffer.begin()); + + std::vector ret_half(cpu_temp_buffer.size()); + ConvertFloatToMLFloat16(cpu_temp_buffer.data(), ret_half.data(), static_cast(cpu_temp_buffer.size())); + + OrtValue target; + test::CreateInputOrtValueOnCPU(cpu_fetches[i].Get().Shape().GetDims(), ret_half, &target); + auto ret = CompareOrtValue(target_fetches[i], target, error_tolerance /*per_sample_tolerance*/, + error_tolerance /*relative_per_sample_tolerance*/, false); + EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second; + + } else { + auto ret = CompareOrtValue(target_fetches[i], cpu_fetches[i], error_tolerance /*per_sample_tolerance*/, + error_tolerance /*relative_per_sample_tolerance*/, false); + EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second; + } + } +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternalGrad_TinySizeTensor) { std::vector dY_dims{}; std::vector log_prob_dims{8, 2}; std::vector index_dims{8}; std::vector weight_dims{2}; std::vector dX_dims{8, 2}; - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean", -1, true, - 5e-2); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum", -1, true, - 5e-2); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none", -1, true, 5e-2); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean"); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum"); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none"); // Just test ignore_index for small tensor because it will increase test time a lot with little verification gain. - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean", 0, true, - 5e-2); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum", 0, true, - 5e-2); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none", 0, true, 5e-2); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + "mean", 0); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + "sum", 0); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", 0); // Bias. - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean", -1, true, - 5e-2, true); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum", -1, true, - 5e-2, true); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none", -1, true, 5e-2, - true); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "mean", 0, true, - 5e-2, true); - TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, "sum", 0, true, - 5e-2, true); - TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, "none", 0, true, 5e-2, - true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + "mean", -1, 1e-4, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + "sum", -1, 1e-4, true); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", -1, 1e-4, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + "mean", 0, 1e-4, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, dX_dims, + "sum", 0, 1e-4, true); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", 0, 1e-4, true); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternalGrad_TinySizeTensor_half) { + std::vector dY_dims{}; + std::vector log_prob_dims{8, 2}; + std::vector index_dims{8}; + std::vector weight_dims{2}; + std::vector dX_dims{8, 2}; + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", -1, 5e-2); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", -1, 5e-2); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", -1, 5e-2); + + // Just test ignore_index for small tensor because it will increase test time a lot with little verification gain. + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", 0, 5e-2); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", 0, 5e-2); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", 0, 5e-2); + + // Bias. + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", -1, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", -1, + 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", -1, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", 0, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", 0, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", 0, 5e-2, true); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternalGrad_TinySizeTensorFloatInputHalfOutput) { + std::vector dY_dims{}; + std::vector log_prob_dims{8, 2}; + std::vector index_dims{8}; + std::vector weight_dims{2}; + std::vector dX_dims{8, 2}; + // Set run_cpu_baseline_seperately = True because CPU kernel did not support multiple type support + // for input and output. + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", -1, 5e-2, false /*has_bias*/); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", -1, 5e-2, false /*has_bias*/); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", -1, 5e-2, false /*has_bias*/); + + // Just test ignore_index for small tensor because it will increase test time a lot with little verification gain. + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", 0, 5e-2, false /*has_bias*/); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", 0, 5e-2, false /*has_bias*/); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", 0, 5e-2, false /*has_bias*/); + + // Bias. + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", -1, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", -1, 1e0, true); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", -1, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "mean", 0, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad(dY_dims, log_prob_dims, index_dims, weight_dims, + dX_dims, "sum", 0, 5e-2, true); + TestSoftmaxCrossEntropyLossInternalGrad({8}, log_prob_dims, index_dims, weight_dims, dX_dims, + "none", 0, 5e-2, true); } } // namespace test diff --git a/orttraining/orttraining/training_api/utils.h b/orttraining/orttraining/training_api/utils.h index 5f3fe2979c..27391e596c 100644 --- a/orttraining/orttraining/training_api/utils.h +++ b/orttraining/orttraining/training_api/utils.h @@ -47,32 +47,6 @@ void WrapInOrtValue(T value, DataTypeImpl::GetType()->GetDeleteFunc()); } -// Create OrtValue on CPU out of provided inputs -template -static void CreateInputOrtValue(gsl::span dims, - const std::vector& 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(value.size())); - auto element_type = DataTypeImpl::GetType(); - auto allocator = alloc ? alloc : cpu_allocator; - auto p_tensor = std::make_unique(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(), - DataTypeImpl::GetType()->GetDeleteFunc()); -} - template T GetValue(OrtValue& ort_value) { const Tensor& tensor = ort_value.Get(); diff --git a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc index 165910d39c..ef6d7a2b0a 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc +++ b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc @@ -403,6 +403,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co ONNX_OPERATOR_TWO_TYPED_KERNEL_EX(OpName, kMSDomain, 1, T1, T2, kCpuExecutionProvider, \ KernelDefBuilder() \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TOut", DataTypeImpl::GetTensorType()) \ .TypeConstraint("Tind", DataTypeImpl::GetTensorType()) \ .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ ClassName); diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 342c747564..ef8f9da6b2 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -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, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc index 2bceb0be35..d0dec4fae6 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc @@ -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(); 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()) \ - .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ - Class); +#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()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ + Class); -#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()) \ - .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ - Class); +#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()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ + Class); + +template +Status SoftmaxCrossEntropyLoss::ComputeInternal(OpKernelContext* ctx) const { + typedef typename ToCudaType::MappedType CudaT_IN; + typedef typename ToCudaType::MappedType CudaT_OUT; -template -Status SoftmaxCrossEntropyLoss::ComputeInternal(OpKernelContext* ctx) const { - typedef typename ToCudaType::MappedType CudaT; const Tensor& logit = *ctx->Input(0); const Tensor& label = *ctx->Input(1); - const Tensor* p_weight = ctx->Input(2); - const Tensor* p_ignore_index = ctx->Input(3); + const Tensor* p_weight = ctx->Input(2); // optional input + const Tensor* p_ignore_index = ctx->Input(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::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* tmp_loss_sample_buffer = nullptr; - IAllocatorUniquePtr tmp_loss_sample; + TOut* total_loss_data = total_loss->template MutableData(); + TOut* tmp_loss_sample_buffer = nullptr; + IAllocatorUniquePtr tmp_loss_sample; if (reduction_ == ReductionType::NONE) { tmp_loss_sample_buffer = total_loss_data; } else { - tmp_loss_sample = GetScratchBuffer(N_D, ctx->GetComputeStream()); + tmp_loss_sample = GetScratchBuffer(N_D, ctx->GetComputeStream()); tmp_loss_sample_buffer = tmp_loss_sample.get(); } const T* logit_data = logit.template Data(); - const Tin* label_data = label.template Data(); + const TLabel* label_data = label.template Data(); - T* log_prob_data = nullptr; + TOut* log_prob_data = nullptr; Tensor* log_prob = nullptr; - IAllocatorUniquePtr log_prob_scratch_buffer; + IAllocatorUniquePtr log_prob_scratch_buffer; if (ctx->OutputCount() > 1) { log_prob = ctx->Output(1, logit_shape); - log_prob_data = log_prob->template MutableData(); + log_prob_data = log_prob->template MutableData(); } else { - log_prob_scratch_buffer = GetScratchBuffer(logit_shape.Size(), ctx->GetComputeStream()); + log_prob_scratch_buffer = GetScratchBuffer(logit_shape.Size(), ctx->GetComputeStream()); log_prob_data = log_prob_scratch_buffer.get(); } @@ -102,16 +104,13 @@ Status SoftmaxCrossEntropyLoss::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())); + ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, logit, + *transpose_output.GetMutable())); logit_data = (*transpose_output.GetMutable()).template Data(); } - // calculate logsoftmax - auto status = SoftMaxComputeHelper(Stream(ctx), - logit_data, - logit_reshape, - log_prob_data, - 1); + // Calculate logsoftmax + auto status = SoftMaxComputeHelper(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::ComputeInternal(OpKernelContext* ctx) co weight_data = weight.template Data(); } - IAllocatorUniquePtr weight_data_nd = GetScratchBuffer(N_D, ctx->GetComputeStream()); - T* weight_data_nd_data = weight_data_nd.get(); + IAllocatorUniquePtr weight_data_nd = GetScratchBuffer(N_D, ctx->GetComputeStream()); + CudaT_OUT* weight_data_nd_data = weight_data_nd.get(); ComputeSoftmaxCrossEntropyWeightsImpl(Stream(ctx), label_data, - reinterpret_cast(weight_data), + reinterpret_cast(weight_data), N_D, C, ignore_index, - reinterpret_cast(weight_data_nd_data)); + reinterpret_cast(weight_data_nd_data)); // Compute buffer size in byte for reduction APIs. - const auto buffer_size = - compute_reduction_buffer_size(static_cast(N_D)); + const auto buffer_size = compute_reduction_buffer_size(static_cast(N_D)); // Allocate reduction buffer whose size is buffer_size bytes, or nullptr if no reduction. IAllocatorUniquePtr reduction_buffer = GetScratchBuffer( reduction_ != ReductionType::NONE ? buffer_size : 0, ctx->GetComputeStream()); - typedef AccumulationType_t TBuf; + typedef AccumulationType_t TBuf; auto normalize_factor_data = GetScratchBuffer(1, ctx->GetComputeStream()); if (reduction_ == ReductionType::MEAN) { ORT_RETURN_IF_ERROR(reduce_sum( Stream(ctx), - reinterpret_cast(weight_data_nd_data), + reinterpret_cast(weight_data_nd_data), normalize_factor_data.get(), static_cast(N_D), reduction_buffer.get(), buffer_size)); } else { constexpr TBuf normalize_factor = static_cast(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(log_prob_data), + reinterpret_cast(log_prob_data), label_data, - reinterpret_cast(weight_data_nd_data), + reinterpret_cast(weight_data_nd_data), normalize_factor_data.get(), N_D, C, ignore_index, - reinterpret_cast(tmp_loss_sample_buffer)); + reinterpret_cast(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::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()).template MutableData(); + auto* transposed_data = (*transpose_output.GetMutable()).template MutableData(); transpose_output.GetMutable()->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())); - 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())); + 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::ComputeInternal(OpKernelContext* ctx) co // ReduceSum on loss_per_sample ORT_RETURN_IF_ERROR(reduce_sum( Stream(ctx), - reinterpret_cast(tmp_loss_sample_buffer), - reinterpret_cast(total_loss_data), + reinterpret_cast(tmp_loss_sample_buffer), + reinterpret_cast(total_loss_data), static_cast(N_D), reduction_buffer.get(), buffer_size)); @@ -189,15 +190,18 @@ Status SoftmaxCrossEntropyLoss::ComputeInternal(OpKernelContext* ctx) co return Status::OK(); } -template -Status SoftmaxCrossEntropyLossGrad::ComputeInternal(OpKernelContext* ctx) const { - typedef typename ToCudaType::MappedType CudaT; +template +Status SoftmaxCrossEntropyLossGrad::ComputeInternal(OpKernelContext* ctx) const { + typedef typename ToCudaType::MappedType CudaT_IN; + typedef typename ToCudaType::MappedType CudaT_OUT; + const Tensor& dY = *ctx->Input(0); const Tensor& log_prob = *ctx->Input(1); const Tensor& label = *ctx->Input(2); const Tensor* p_weight = ctx->Input(3); const Tensor* p_ignore_index = ctx->Input(4); const Tensor* p_bias = ctx->Input(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::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(); const T* log_prob_data = log_prob.Data(); - const Tin* label_data = label.Data(); - T* d_logit_data = d_logit->MutableData(); + const TLabel* label_data = label.Data(); + TOut* d_logit_data = d_logit->MutableData(); const T* weight_data = nullptr; OrtValue transpose_output; TensorShapeVector new_shape; @@ -230,7 +233,8 @@ Status SoftmaxCrossEntropyLossGrad::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())); + ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, + log_prob, *transpose_output.GetMutable())); log_prob_data = (*transpose_output.GetMutable()).Data(); } @@ -243,44 +247,45 @@ Status SoftmaxCrossEntropyLossGrad::ComputeInternal(OpKernelContext* ctx T* weight_data_nd_data = weight_data_nd.get(); ComputeSoftmaxCrossEntropyWeightsImpl(Stream(ctx), label_data, - reinterpret_cast(weight_data), + reinterpret_cast(weight_data), N_D, C, ignore_index, - reinterpret_cast(weight_data_nd_data)); - typedef AccumulationType_t TBuf; + reinterpret_cast(weight_data_nd_data)); + typedef AccumulationType_t TBuf; auto normalize_factor_data = GetScratchBuffer(1, ctx->GetComputeStream()); if (reduction_ == ReductionType::MEAN) { // Compute buffer size in byte for reduction APIs. const auto buffer_size = - compute_reduction_buffer_size(static_cast(N_D)); + compute_reduction_buffer_size(static_cast(N_D)); // Allocate reduction buffer whose size is buffer_size bytes. IAllocatorUniquePtr reduction_buffer = GetScratchBuffer( buffer_size, ctx->GetComputeStream()); ORT_RETURN_IF_ERROR(reduce_sum( Stream(ctx), - reinterpret_cast(weight_data_nd_data), + reinterpret_cast(weight_data_nd_data), normalize_factor_data.get(), static_cast(N_D), reduction_buffer.get(), buffer_size)); } else { constexpr TBuf normalize_factor = static_cast(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() : nullptr; + const TOut* bias_data = p_bias ? p_bias->Data() : nullptr; SoftmaxCrossEntropyLossGradImpl(Stream(ctx), - reinterpret_cast(dY_data), - reinterpret_cast(log_prob_data), + reinterpret_cast(dY_data), + reinterpret_cast(log_prob_data), label_data, - reinterpret_cast(weight_data_nd_data), + reinterpret_cast(weight_data_nd_data), normalize_factor_data.get(), - reinterpret_cast(bias_data), + reinterpret_cast(bias_data), N_D, C, ReductionType::NONE == reduction_, - reinterpret_cast(d_logit_data)); + reinterpret_cast(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::ComputeInternal(OpKernelContext* ctx onnxruntime::contrib::GetPermutationAndShape(false, logit_shape, new_shape, permutations); transpose_output.GetMutable()->Reshape(d_logit->Shape()); d_logit->Reshape(logit_shape); - ORT_RETURN_IF_ERROR(cuda::Transpose::DoTranspose(cuda::Transpose(info), ctx->GetComputeStream(), permutations, *d_logit, *transpose_output.GetMutable())); - auto* transposed_data = (*transpose_output.GetMutable()).template Data(); - 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())); + auto* transposed_data = (*transpose_output.GetMutable()).template Data(); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(d_logit_data, transposed_data, sizeof(TOut) * probability_shape.Size(), + cudaMemcpyDeviceToDevice, Stream(ctx))); 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::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::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()) \ - .TypeConstraint("Tind", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ - ClassName); +#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()) \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TOut", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ + ClassName); -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 diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu index 4e7171a072..e2afa69ec1 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu @@ -7,58 +7,60 @@ namespace onnxruntime { namespace cuda { -template +template 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 -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 +void ComputeSoftmaxCrossEntropyWeightsImpl(cudaStream_t stream, const TLabel* label, const T* weight, size_t count, + size_t label_depth, int64_t ignore_index, TOut* weight_data_nd) { if (weight) { - OpSoftmaxCrossEntropyWeights op(label, weight, static_cast(label_depth), - static_cast(ignore_index)); - LaunchElementwiseKernel(stream, weight_data_nd, op, count); + OpSoftmaxCrossEntropyWeights op(label, weight, static_cast(label_depth), + static_cast(ignore_index)); + LaunchElementwiseKernel(stream, weight_data_nd, op, count); } else { - OpSoftmaxCrossEntropyWeights op(label, nullptr, static_cast(label_depth), - static_cast(ignore_index)); - LaunchElementwiseKernel(stream, weight_data_nd, op, count); + OpSoftmaxCrossEntropyWeights op(label, nullptr, static_cast(label_depth), + static_cast(ignore_index)); + LaunchElementwiseKernel(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 +template 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 -void SoftmaxCrossEntropyLossImpl(cudaStream_t stream, const T* log_prob, const Tin* label, const T* weight, +template +void SoftmaxCrossEntropyLossImpl(cudaStream_t stream, const T* log_prob, const TLabel* label, const T* weight, const TAcc* normalize_factor, size_t count, size_t label_depth, int64_t ignore_index, T* output_data) { - OpWeightedSoftmaxCrossEntropyLoss op(log_prob, label, weight, normalize_factor, - static_cast(label_depth), static_cast(ignore_index)); + OpWeightedSoftmaxCrossEntropyLoss op(log_prob, label, weight, normalize_factor, + static_cast(label_depth), static_cast(ignore_index)); LaunchElementwiseKernel(stream, output_data, op, count); } -template +template struct OpWeightedSoftmaxCrossEntropyLossGrad { - OpWeightedSoftmaxCrossEntropyLossGrad(const T* dY_data, const T* log_prob_data, const 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(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(static_cast((IsReductionNone ? dY_data_[row] : *dY_data_) * weight_data_[row]) * - (_Exp(static_cast(log_prob_data_[idx])) - (TAcc)(d == label_data_[row])) / - (*normalize_factor_data_)); + result = static_cast((IsReductionNone ? dY_data_[row] : *dY_data_) * weight_data_[row]) * + (_Exp(static_cast(log_prob_data_[idx])) - (TAcc)(d == label_data_[row])) / + (*normalize_factor_data_); } - return HasBias ? result + bias_data_[idx] : result; + return HasBias ? static_cast(result + static_cast(bias_data_[idx])) : static_cast(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 -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 op( \ - dY, log_prob, label, weight, normalize_factor, bias_data, static_cast(label_depth)); \ - LaunchElementwiseKernel(stream, output_data, op, count * label_depth) +template +void SoftmaxCrossEntropyLossGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const TLabel* label, + const T* weight, const TAcc* normalize_factor, const TOut* bias_data, size_t count, + size_t label_depth, bool reduction_none, TOut* output_data) { +#define LAUNCH_WEIGHTED_SOFTMAX_CROSS_ENTROPY_LOSS_GRAD_KERNEL(is_reduction_none, has_bias) \ + OpWeightedSoftmaxCrossEntropyLossGrad op( \ + dY, log_prob, label, weight, normalize_factor, bias_data, static_cast(label_depth)); \ + LaunchElementwiseKernel(stream, output_data, op, count * label_depth) 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 diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h index 85b353cd56..e06a6b7633 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h @@ -10,11 +10,11 @@ namespace onnxruntime { namespace cuda { -template +template 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 +template 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 +template 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 +template 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 +template class SoftmaxCrossEntropyLossGrad final : public LossBase { public: SoftmaxCrossEntropyLossGrad(const OpKernelInfo& info) : LossBase(info) { diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cc b/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cc index bfe0acc0c5..f6254e0c11 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cc +++ b/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cc @@ -26,7 +26,7 @@ namespace cuda { kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ + .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ Class); template @@ -49,11 +49,11 @@ Status SoftmaxCrossEntropy::ComputeInternal(OpKernelContext* ctx) const { T* log_prob_data = log_prob->template MutableData(); // calculate logsoftmax - auto status = SoftMaxComputeHelper(Stream(ctx), - logit_data, - logit_reshape, - log_prob_data, - 1 /*axis default*/); + auto status = SoftMaxComputeHelper(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::ComputeInternal(OpKernelContext* ctx) T* log_prob_data = log_prob->template MutableData(); // calculate logsoftmax - auto status = SoftMaxComputeHelper(Stream(ctx), - logit_data, - logit_reshape, - log_prob_data, - 1 /*axis default*/); + auto status = SoftMaxComputeHelper(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::ComputeInternal(OpKernelContext* ctx) auto normalize_factor_data = GetScratchBuffer(1, ctx->GetComputeStream()); if (reduction_ == ReductionType::SUM) { constexpr T normalize_factor_one = static_cast(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(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::ComputeInternal(OpKernelContext* c auto normalize_factor_data = GetScratchBuffer(1, ctx->GetComputeStream()); if (reduction_ == ReductionType::SUM) { constexpr T normalize_factor_one = static_cast(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(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 = diff --git a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc index cb10290449..daa43fdde9 100644 --- a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc +++ b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc @@ -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, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, // BuildKernelCreateInfo, BuildKernelCreateInfo,