From 29c74d3c43923cafa9099cdd523fd5d64c814b7f Mon Sep 17 00:00:00 2001 From: zhijiang <43435212+zhijxu-MS@users.noreply.github.com> Date: Tue, 11 Apr 2023 17:02:40 +0800 Subject: [PATCH] softmax perf improvement pr1 - add more softmax related test (#15176) 1. add fp16 test 2. add test for shape is not power of two. --- .../test/common/tensor_op_test_utils.h | 14 ++ .../test/training_ops/cuda/softmax_test.cc | 188 ++++++++++++++---- .../github/pai/pai-excluded-tests.txt | 8 + 3 files changed, 173 insertions(+), 37 deletions(-) diff --git a/onnxruntime/test/common/tensor_op_test_utils.h b/onnxruntime/test/common/tensor_op_test_utils.h index 0e58dd64f1..0f1de05d0d 100644 --- a/onnxruntime/test/common/tensor_op_test_utils.h +++ b/onnxruntime/test/common/tensor_op_test_utils.h @@ -65,6 +65,20 @@ class RandomValueGenerator { return val; } + // Random values generated are in the range [min, max). + template + typename std::enable_if< + std::is_same_v, + std::vector>::type + Uniform(gsl::span dims, float min, float max) { + std::vector val(detail::SizeFromDims(dims)); + std::uniform_real_distribution distribution(min, max); + for (size_t i = 0; i < val.size(); ++i) { + val[i] = TFloat16(math::floatToHalf(distribution(generator_))); + } + return val; + } + // Random values generated are in the range [min, max). template typename std::enable_if< diff --git a/orttraining/orttraining/test/training_ops/cuda/softmax_test.cc b/orttraining/orttraining/test/training_ops/cuda/softmax_test.cc index 93d1a4be75..1eadc5e155 100644 --- a/orttraining/orttraining/test/training_ops/cuda/softmax_test.cc +++ b/orttraining/orttraining/test/training_ops/cuda/softmax_test.cc @@ -12,24 +12,25 @@ constexpr const char* kGpuExecutionProvider = kCudaExecutionProvider; constexpr const char* kGpuExecutionProvider = kRocmExecutionProvider; #endif +template static void TestSoftmax(const std::vector& X_dims, const std::vector& Y_dims, int axis = 1, bool is_log_softmax=false, double per_sample_tolerance = 1e-4, double relative_per_sample_tolerance = 1e-4) { - + const char* op = is_log_softmax? "LogSoftmax" : "Softmax"; CompareOpTester test(op); test.AddAttribute("axis", axis); // create rand inputs RandomValueGenerator random{}; - std::vector X_data = random.Uniform(X_dims, -10.0f, 10.0f); - test.AddInput("X", X_dims, X_data); + std::vector X_data = random.Uniform(X_dims, -10.0f, 10.0f); + test.AddInput("X", X_dims, X_data); - std::vector Y_data = FillZeros(Y_dims); - test.AddOutput("Y", Y_dims, Y_data); + std::vector Y_data = FillZeros(Y_dims); + test.AddOutput("Y", Y_dims, Y_data); test.CompareWithCPU(kGpuExecutionProvider, per_sample_tolerance, relative_per_sample_tolerance); } @@ -39,63 +40,116 @@ static void TestSoftmax(const std::vector& X_dims, TEST(CudaKernelTest, Softmax_SmallTensor_LastAxis) { std::vector X_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; - TestSoftmax(X_dims, Y_dims, 2, false); + TestSoftmax(X_dims, Y_dims, 2, false); } TEST(CudaKernelTest, Softmax_SmallTensor_AllAxis) { std::vector X_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; - TestSoftmax(X_dims, Y_dims, 0, false); - TestSoftmax(X_dims, Y_dims, 1, false); + TestSoftmax(X_dims, Y_dims, 0, false); + TestSoftmax(X_dims, Y_dims, 1, false); } // large tensor to check cuda DNN softmax forward TEST(CudaKernelTest, Softmax_LargeTensor_LastAxis) { std::vector X_dims{8, 16, 2048}; std::vector Y_dims{8, 16, 2048}; - TestSoftmax(X_dims, Y_dims, 2, false); + TestSoftmax(X_dims, Y_dims, 2, false); +} + +TEST(CudaKernelTest, Softmax_LargeTensor_LastAxis_Float16) { + std::vector X_dims{8, 16, 2048}; + std::vector Y_dims{8, 16, 2048}; + TestSoftmax(X_dims, Y_dims, 2, false, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, Softmax_LargeTensor_LastAxis_Float16_NoPowerOfTwo) { + std::vector X_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + TestSoftmax(X_dims, Y_dims, 2, false, 1e-3, 1e-3); } TEST(CudaKernelTest, Softmax_LargeTensor_AllAxis) { std::vector X_dims{8, 16, 512}; std::vector Y_dims{8, 16, 512}; - TestSoftmax(X_dims, Y_dims, 0, false); - TestSoftmax(X_dims, Y_dims, 1, false); + TestSoftmax(X_dims, Y_dims, 0, false); + TestSoftmax(X_dims, Y_dims, 1, false); +} + +TEST(CudaKernelTest, Softmax_LargeTensor_AllAxis_Float16) { + std::vector X_dims{8, 16, 512}; + std::vector Y_dims{8, 16, 512}; + TestSoftmax(X_dims, Y_dims, 0, false, 1e-3, 1e-3); + TestSoftmax(X_dims, Y_dims, 1, false, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, Softmax_LargeTensor_AllAxis_Float16_NoPowerOfTwo) { + std::vector X_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + TestSoftmax(X_dims, Y_dims, 0, false, 1e-3, 1e-3); + TestSoftmax(X_dims, Y_dims, 1, false, 1e-3, 1e-3); } TEST(CudaKernelTest, LogSoftmax_SmallTensor_LastAxis) { std::vector X_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; - TestSoftmax(X_dims, Y_dims, 2, true); + TestSoftmax(X_dims, Y_dims, 2, true); } TEST(CudaKernelTest, LogSoftmax_SmallTensor_AllAxis) { std::vector X_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; - TestSoftmax(X_dims, Y_dims, 0, true); - TestSoftmax(X_dims, Y_dims, 1, true); + TestSoftmax(X_dims, Y_dims, 0, true); + TestSoftmax(X_dims, Y_dims, 1, true); } TEST(CudaKernelTest, LogSoftmax_LargeTensor_LastAxis) { std::vector X_dims{8, 16, 2048}; std::vector Y_dims{8, 16, 2048}; - TestSoftmax(X_dims, Y_dims, 2, true); + TestSoftmax(X_dims, Y_dims, 2, true); +} + +TEST(CudaKernelTest, LogSoftmax_LargeTensor_LastAxis_Float16) { + std::vector X_dims{8, 16, 2048}; + std::vector Y_dims{8, 16, 2048}; + TestSoftmax(X_dims, Y_dims, 2, true, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, LogSoftmax_LargeTensor_LastAxis_Float16_NoPowerOfTwo) { + std::vector X_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + TestSoftmax(X_dims, Y_dims, 2, true, 1e-3, 1e-3); } TEST(CudaKernelTest, LogSoftmax_LargeTensor_AllAxis) { std::vector X_dims{8, 16, 512}; std::vector Y_dims{8, 16, 512}; - TestSoftmax(X_dims, Y_dims, 0, true); - TestSoftmax(X_dims, Y_dims, 1, true); + TestSoftmax(X_dims, Y_dims, 0, true); + TestSoftmax(X_dims, Y_dims, 1, true); } +TEST(CudaKernelTest, LogSoftmax_LargeTensor_AllAxis_Float16) { + std::vector X_dims{8, 16, 512}; + std::vector Y_dims{8, 16, 512}; + TestSoftmax(X_dims, Y_dims, 0, true, 1e-3, 1e-3); + TestSoftmax(X_dims, Y_dims, 1, true, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, LogSoftmax_LargeTensor_AllAxis_Float16_NoPowerOfTwo) { + std::vector X_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + TestSoftmax(X_dims, Y_dims, 0, true, 1e-3, 1e-3); + TestSoftmax(X_dims, Y_dims, 1, true, 1e-3, 1e-3); +} + +template static void TestSoftmaxGrad(const std::vector& dY_dims, const std::vector& Y_dims, const std::vector& dX_dims, int axis = 1, bool is_log_softmax = false, double per_sample_tolerance = 1e-4, - double relative_per_sample_tolerance = 1e-4) { + double relative_per_sample_tolerance = 1e-4) { const char* op = is_log_softmax? "LogSoftmaxGrad" : "SoftmaxGrad"; CompareOpTester test(op, 1, kMSDomain); @@ -103,15 +157,15 @@ static void TestSoftmaxGrad(const std::vector& dY_dims, // create rand inputs RandomValueGenerator random{}; - std::vector dY_data = random.Uniform(dY_dims, 0.0f, 1.0f); + std::vector dY_data = random.Uniform(dY_dims, -1.0f, 1.0f); // Add 1e-2 for numerical stability to prevent zero probability. - std::vector Y_data = random.Uniform(Y_dims, 0.02f, 1.02f); + std::vector Y_data = random.Uniform(Y_dims, -1.02f, 1.02f); - test.AddInput("dY", dY_dims, dY_data); - test.AddInput("Y", Y_dims, Y_data); + test.AddInput("dY", dY_dims, dY_data); + test.AddInput("Y", Y_dims, Y_data); - std::vector dX_data = FillZeros(dX_dims); - test.AddOutput("dX", dX_dims, dX_data); + std::vector dX_data = FillZeros(dX_dims); + test.AddOutput("dX", dX_dims, dX_data); test.CompareWithCPU(kGpuExecutionProvider, per_sample_tolerance, relative_per_sample_tolerance); } @@ -121,15 +175,15 @@ TEST(CudaKernelTest, SoftmaxGrad_SmallTensor_LastAxis) { std::vector dY_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; std::vector dX_dims{4, 2, 128}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2); } TEST(CudaKernelTest, SoftmaxGrad_SmallTensor_AllAxis) { std::vector dY_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; std::vector dX_dims{4, 2, 128}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0); - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1); } // large tensor to check cuda DNN softmax backward @@ -137,7 +191,21 @@ TEST(CudaKernelTest, SoftmaxGrad_LargeTensor_LastAxis) { std::vector dY_dims{8, 16, 2048}; std::vector Y_dims{8, 16, 2048}; std::vector dX_dims{8, 16, 2048}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2); +} + +TEST(CudaKernelTest, SoftmaxGrad_LargeTensor_LastAxis_Float16) { + std::vector dY_dims{8, 16, 2048}; + std::vector Y_dims{8, 16, 2048}; + std::vector dX_dims{8, 16, 2048}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, false, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, SoftmaxGrad_LargeTensor_LastAxis_Float16_NoPowerOfTwo) { + std::vector dY_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + std::vector dX_dims{8, 16, 1500}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, false, 1e-3, 1e-3); } // large tensor to check cuda DNN softmax backward @@ -145,38 +213,84 @@ TEST(CudaKernelTest, SoftmaxGrad_LargeTensor_AllAxis) { std::vector dY_dims{8, 16, 512}; std::vector Y_dims{8, 16, 512}; std::vector dX_dims{8, 16, 512}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0); - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1); +} + +TEST(CudaKernelTest, SoftmaxGrad_LargeTensor_AllAxis_Float16) { + std::vector dY_dims{8, 16, 512}; + std::vector Y_dims{8, 16, 512}; + std::vector dX_dims{8, 16, 512}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, false, 1e-3, 1e-3); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, false, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, SoftmaxGrad_LargeTensor_AllAxis_Float16_NoPowerOfTwo) { + std::vector dY_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + std::vector dX_dims{8, 16, 1500}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, false, 1e-3, 1e-3); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, false, 1e-3, 1e-3); } TEST(CudaKernelTest, LogSoftmaxGrad_SmallTensor_LastAxis) { std::vector dY_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; std::vector dX_dims{4, 2, 128}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, true); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, true); } TEST(CudaKernelTest, LogSoftmaxGrad_SmallTensor_AllAxis) { std::vector dY_dims{4, 2, 128}; std::vector Y_dims{4, 2, 128}; std::vector dX_dims{4, 2, 128}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, true); - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, true); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, true); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, true); } TEST(CudaKernelTest, LogSoftmaxGrad_LargeTensor_LastAxis) { std::vector dY_dims{8, 16, 2048}; std::vector Y_dims{8, 16, 2048}; std::vector dX_dims{8, 16, 2048}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, true); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, true); +} + +TEST(CudaKernelTest, LogSoftmaxGrad_LargeTensor_LastAxis_Float16) { + std::vector dY_dims{8, 16, 2048}; + std::vector Y_dims{8, 16, 2048}; + std::vector dX_dims{8, 16, 2048}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, true, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, LogSoftmaxGrad_LargeTensor_LastAxis_Float16_NoPowerOfTwo) { + std::vector dY_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + std::vector dX_dims{8, 16, 1500}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 2, true, 1e-3, 1e-3); } TEST(CudaKernelTest, LogSoftmaxGrad_LargeTensor_AllAxis) { std::vector dY_dims{8, 16, 512}; std::vector Y_dims{8, 16, 512}; std::vector dX_dims{8, 16, 512}; - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, true); - TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, true); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, true, 1e-3, 1e-3); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, true, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, LogSoftmaxGrad_LargeTensor_AllAxis_Float16) { + std::vector dY_dims{8, 16, 512}; + std::vector Y_dims{8, 16, 512}; + std::vector dX_dims{8, 16, 512}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, true, 1e-3, 1e-3); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, true, 1e-3, 1e-3); +} + +TEST(CudaKernelTest, LogSoftmaxGrad_LargeTensor_AllAxis_Float16_NoPowerOfTwo) { + std::vector dY_dims{8, 16, 1500}; + std::vector Y_dims{8, 16, 1500}; + std::vector dX_dims{8, 16, 1500}; + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 0, true, 1e-3, 1e-3); + TestSoftmaxGrad(dY_dims, Y_dims, dX_dims, 1, true, 1e-3, 1e-3); } static void TestSoftmaxGrad_13(const std::vector& dY_dims, @@ -185,7 +299,7 @@ static void TestSoftmaxGrad_13(const std::vector& dY_dims, int axis = 1, bool is_log_softmax = false, double per_sample_tolerance = 1e-4, - double relative_per_sample_tolerance = 1e-4) { + double relative_per_sample_tolerance = 1e-4) { const char* op = is_log_softmax? "LogSoftmaxGrad_13" : "SoftmaxGrad_13"; CompareOpTester test(op, 1, kMSDomain); diff --git a/tools/ci_build/github/pai/pai-excluded-tests.txt b/tools/ci_build/github/pai/pai-excluded-tests.txt index 06936399a0..b446dac203 100644 --- a/tools/ci_build/github/pai/pai-excluded-tests.txt +++ b/tools/ci_build/github/pai/pai-excluded-tests.txt @@ -1,6 +1,14 @@ CudaKernelTest.NegativeLogLikelihoodLoss_TinySizeTensor CudaKernelTest.NegativeLogLikelihoodLoss_SmallSizeTensor CudaKernelTest.NegativeLogLikelihoodLoss_MediumSizeTensor +CudaKernelTest.SoftmaxGrad_LargeTensor_LastAxis_Float16 +CudaKernelTest.SoftmaxGrad_LargeTensor_LastAxis_Float16_NoPowerOfTwo +CudaKernelTest.SoftmaxGrad_LargeTensor_AllAxis_Float16 +CudaKernelTest.SoftmaxGrad_LargeTensor_AllAxis_Float16_NoPowerOfTwo +CudaKernelTest.LogSoftmaxGrad_LargeTensor_LastAxis_Float16 +CudaKernelTest.LogSoftmaxGrad_LargeTensor_LastAxis_Float16_NoPowerOfTwo +CudaKernelTest.LogSoftmaxGrad_LargeTensor_AllAxis_Float16 +CudaKernelTest.LogSoftmaxGrad_LargeTensor_AllAxis_Float16_NoPowerOfTwo ReductionOpTest.ReductionVariationTest ReductionOpTest.ReduceLogSumExp_default_axes_keepdims_double ReductionOpTest.ReduceLogSumExp_default_axes_do_not_keep_dims_double