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,