diff --git a/onnxruntime/test/providers/cpu/activation/activation_op_test.cc b/onnxruntime/test/providers/cpu/activation/activation_op_test.cc index 7333800b8c..fb6ae200ab 100644 --- a/onnxruntime/test/providers/cpu/activation/activation_op_test.cc +++ b/onnxruntime/test/providers/cpu/activation/activation_op_test.cc @@ -57,7 +57,11 @@ float ReluGrad(float dy, float x) { float SigmoidGrad(float dy, float y) { return dy * y * (1 - y); } + +float TanhGrad(float dy, float y) { + return dy * (1 - y * y); } +} // namespace #endif TEST_F(ActivationOpTest, Sigmoid) { @@ -303,6 +307,22 @@ TEST(SigmoidGradInferenceTest, Basic) { }, {}, 1, kMSDomain); } + +TEST(TanhGradInferenceTest, Basic) { + const std::vector y_vals = {-1.0f, 0, 1.0f, 100.0f, -100.0f, 1000.0f, -1000.0f}; + const std::vector dY(7, 1.0f); + + TestElementwiseGradientOp( + "TanhGrad", + {{"dY", dY}, {"Y", y_vals}}, + [](const std::vector& params) { + ORT_ENFORCE(params.size() == 2); + const auto dy = params[0], y = params[1]; + + return TanhGrad(dy, y); + }, + {}, 1, kMSDomain); +} #endif } // namespace test diff --git a/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json b/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json index 17f16e260e..764b882112 100644 --- a/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json +++ b/onnxruntime/test/testdata/kernel_def_hashes/training_ops.cpu.json @@ -251,8 +251,12 @@ "SplitTraining com.microsoft CPUExecutionProvider", 12689204749897364688 ], + [ + "TanhGrad com.microsoft CPUExecutionProvider", + 7147744030478490408 + ], [ "ZeroGradient com.microsoft CPUExecutionProvider", 3284255990062374928 ] -] +] \ No newline at end of file diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index cbb8fe716c..128312dccb 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -38,7 +38,7 @@ static bool SimplifyReshape(const std::vector& target_shape, // the return false; } } - //trim empty strings in the tail of list + // trim empty strings in the tail of list while (!dim_params.empty() && dim_params.back().empty()) { dim_params.pop_back(); } @@ -90,15 +90,10 @@ IMPLEMENT_GRADIENT_BUILDER(GetLogGradient) { } IMPLEMENT_GRADIENT_BUILDER(GetTanhGradient) { - ArgDef Y = O(0); - std::vector result; - NodeDef one_constant_node = OneConstantNode(OElemType(0)); - ArgDef one_arg = one_constant_node.output_args[0]; - result.push_back(one_constant_node); - result.push_back(NodeDef("Mul", {Y, Y}, {IA("Squared_Y")})); - result.push_back(NodeDef("Sub", {one_arg, IA("Squared_Y")}, {IA("Sub_Squared_Y")})); - result.push_back(NodeDef("Mul", {GO(0), IA("Sub_Squared_Y")}, {GI(0)})); - return result; + return std::vector{ + NodeDef(OpDef{"TanhGrad", kMSDomain, 1}, + {GO(0), O(0)}, + {GI(0)})}; } IMPLEMENT_GRADIENT_BUILDER(GetSqrtGradient) { @@ -241,7 +236,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetMatMulGradient) { NodeDef(OpDef{"FusedMatMul", kMSDomain, 1}, {GO(0), B}, {matmul_out}, - {{"transB", MakeAttribute("transB", int64_t(1))}})); + {{"transB", MakeAttribute("transB", int64_t(1))}})); if (A_axes.size() > 0) { AddReduceSumNode(IA("PreReduceGrad0"), IA("ReduceGrad0"), A_axes, true, result); result.push_back(NodeDef("Shape", {A}, {IA("A_shape")})); @@ -281,7 +276,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetMatMulGradient) { } } } else { - //GetShape failed, build shape-independent gradient graph + // GetShape failed, build shape-independent gradient graph ArgDef a_axes, b_axes, a_shape, b_shape, ia_shape; a_shape = IA("Shape_" + A.name); b_shape = IA("Shape_" + B.name); @@ -451,7 +446,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetGemmGradient) { } } } else { - //GetShape failed, build shape-independent gradient graph + // GetShape failed, build shape-independent gradient graph ArgDef c_axes = IA("ReduceAxes_" + C.name); ArgDef c_shape = IA("Shape_" + C.name); ArgDef dy_shape = IA("Shape_" + dY.name); @@ -617,7 +612,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetTransposeGradient) { std::vector new_attributes; if (attributes.empty()) { const TensorShapeProto& input_shape = I(0).type_proto->tensor_type().shape(); - if (input_shape.dim_size() > 0) { //input_shape is available + if (input_shape.dim_size() > 0) { // input_shape is available int n = input_shape.dim_size() - 1; bw_perm.resize(n + 1); std::generate(bw_perm.begin(), bw_perm.end(), [&n] { return n--; }); @@ -694,7 +689,6 @@ IMPLEMENT_GRADIENT_BUILDER(GetConvGradient) { } IMPLEMENT_GRADIENT_BUILDER(GetSigmoidGradient) { - auto const_one = OneConstantNode(OElemType(0)); return std::vector{ NodeDef(OpDef{"SigmoidGrad", kMSDomain, 1}, {GO(0), O(0)}, @@ -860,7 +854,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetAddSubGradient) { } } } else { - //GetShape failed, build shape-independent gradient graph + // GetShape failed, build shape-independent gradient graph ArgDef a_axes = IA("ReduceAxes_" + a.name); ArgDef b_axes = IA("ReduceAxes_" + b.name); ArgDef A_shape = IA("Shape_" + a.name); @@ -944,7 +938,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetMulGradient) { } } } else { - //GetShape failed, build shape-independent gradient graph + // GetShape failed, build shape-independent gradient graph ArgDef a_axes = IA("ReduceAxes_" + a.name); ArgDef b_axes = IA("ReduceAxes_" + b.name); ArgDef A_shape = IA("Shape_" + a.name); @@ -1001,7 +995,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetDivGradient) { output.push_back(NodeDef("Identity", {tmp_grad}, {GI(0)})); } } else { - //GetShape failed, build shape-independent gradient graph + // GetShape failed, build shape-independent gradient graph ArgDef a_axes = IA("ReduceAxes_" + a.name); ArgDef A_shape = IA("Shape_" + a.name); ArgDef B_shape = IA("Shape_" + b.name); @@ -1133,17 +1127,17 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceSumGradient) { ArgDef grad = GO(0); if (!keepdims) { size_t numInputs = GetSrcNodeInputSize(); - if (SrcNodeOpsetVersion() < 13) { //axes is attribute + if (SrcNodeOpsetVersion() < 13) { // axes is attribute if (attributes.find("axes") != attributes.end()) { std::vector axes_values = RetrieveValues(attributes.at("axes")); grad = IA("Unqueezed_Grad"); result.push_back(NodeDef("Unsqueeze", {GO(0)}, {grad}, {MakeAttribute("axes", axes_values)})); } - } else if (numInputs == 2) { //optional input 'axes' is available as input I(1) + } else if (numInputs == 2) { // optional input 'axes' is available as input I(1) grad = IA("Unqueezed_Grad"); result.push_back(NodeDef(OpDef{"Unsqueeze", kOnnxDomain, 13}, {GO(0), I(1)}, {grad})); - } //axes is not available, the GO(0) is a scalar which can be expanded to required shape + } // axes is not available, the GO(0) is a scalar which can be expanded to required shape } result.push_back(NodeDef("Shape", {I(0)}, {IA("Shaped_X")})); @@ -1443,7 +1437,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetExpandGradient) { {GI(0)})); } } else { - //GetShape failed, build shape-independent gradient graph + // GetShape failed, build shape-independent gradient graph ArgDef a_axes = IA("ReduceAxes_" + a.name); ArgDef A_shape = IA("Shape_" + a.name); ArgDef Y_shape = IA("Shape_" + y.name); @@ -1549,10 +1543,10 @@ IMPLEMENT_GRADIENT_BUILDER(GetTileGradient) { NodeDef unsqueeze_axes = ConstantVectorNode(axes_values, Name("unsqueeze_axes")); result.push_back(unsqueeze_axes); result.push_back(NodeDef("Unsqueeze", {IA("orig_shape"), unsqueeze_axes.output_args[0]}, {IA("2d_orig_shape")})); // M, N, K - result.push_back(NodeDef("Unsqueeze", {I(1), unsqueeze_axes.output_args[0]}, {IA("2d_repeats")})); //a, b, c + result.push_back(NodeDef("Unsqueeze", {I(1), unsqueeze_axes.output_args[0]}, {IA("2d_repeats")})); // a, b, c } else { result.push_back(NodeDef("Unsqueeze", {IA("orig_shape")}, {IA("2d_orig_shape")}, {MakeAttribute("axes", axes_values)})); // M, N, K - result.push_back(NodeDef("Unsqueeze", {I(1)}, {IA("2d_repeats")}, {MakeAttribute("axes", axes_values)})); //a, b, c + result.push_back(NodeDef("Unsqueeze", {I(1)}, {IA("2d_repeats")}, {MakeAttribute("axes", axes_values)})); // a, b, c } result.push_back(NodeDef("Concat", {IA("2d_repeats"), IA("2d_orig_shape")}, {IA("concated_dims_T")}, {MakeAttribute("axis", int64_t(1))})); // [[a, M], [b, N], [c, K]] diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 40d9da4001..e14f0f2d9e 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -623,7 +623,7 @@ void RegisterTrainingOpSchemas() { .AddOpset("", 13) .Const("one", int64_t(1)) .Const("k", axis) - .Const("axis_zero", std::vector({0})) // a 1D tensor constant + .Const("axis_zero", std::vector({0})) // a 1D tensor constant .Add(R"( shape = Shape (dY) n_as_vector = Shape (shape) @@ -835,8 +835,8 @@ void RegisterTrainingOpSchemas() { } }); - //TODO: Move this to the right location. Its only here for quick experimentation. - //TODO: Use the mutli weight / grad version. + // TODO: Move this to the right location. Its only here for quick experimentation. + // TODO: Use the mutli weight / grad version. ONNX_CONTRIB_OPERATOR_SCHEMA(SGDOptimizer) .SetDomain(kMSDomain) .SinceVersion(1) @@ -2081,7 +2081,6 @@ Example 4: return true; }); - ONNX_CONTRIB_OPERATOR_SCHEMA(SigmoidGrad) .SetDomain(kMSDomain) .SinceVersion(1) @@ -2112,7 +2111,35 @@ Example 4: onnx_opset_13.set_version(13); return ONNX_NAMESPACE::FunctionBodyHelper::BuildFunctionProto(functionProto, schema, body, {onnx_opset_13}); + }); + ONNX_CONTRIB_OPERATOR_SCHEMA(TanhGrad) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) + .SetDoc("TanhGrad") + .AllowUncheckedAttributes() + .Input(0, "dY", "The gradient tensor from output.", "T") + .Input(1, "Y", "The input tensor. ", "T") + .Output(0, "dX", "Gradient of the input.", "T") + .TypeConstraint( + "T", + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, + "Constrain input and output types to float tensors.") + .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput) + .SetContextDependentFunctionBodyBuilder( + [](const FunctionBodyBuildContext& ctx, const OpSchema& schema, FunctionProto& functionProto) { + auto* tp = ctx.getInputType(0); + if ((tp == nullptr) || (!tp->has_tensor_type())) + return false; + auto elem_type = (ONNX_NAMESPACE::TensorProto_DataType)tp->tensor_type().elem_type(); + std::vector body{ + ONNX_NAMESPACE::Const("C_One", 1.0f, elem_type), + {{"YSquare"}, "Mul", {"Y", "Y"}}, + {{"dTanhX"}, "Sub", {"C_One", "YSquare"}}, + {{"dX"}, "Mul", {"dY", "dTanhX"}}}; + + return ONNX_NAMESPACE::FunctionBodyHelper::BuildFunctionProto(functionProto, schema, body, {}); }); ONNX_CONTRIB_OPERATOR_SCHEMA(LayerNormalizationGrad) diff --git a/orttraining/orttraining/test/gradient/function_ops_test.cc b/orttraining/orttraining/test/gradient/function_ops_test.cc index e396bb350c..80174d6863 100644 --- a/orttraining/orttraining/test/gradient/function_ops_test.cc +++ b/orttraining/orttraining/test/gradient/function_ops_test.cc @@ -235,5 +235,9 @@ TEST_F(FunExpansionTest, SigmoidGrad_float) { TestUnaryOpGrad("SigmoidGrad"); } +TEST_F(FunExpansionTest, TanhGrad_float) { + TestUnaryOpGrad("TanhGrad"); +} + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index f069b8c01c..5739cc0bd9 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -4213,3 +4213,37 @@ def test_sigmoid_grad(): _test_helpers.assert_values_are_close(ort_prediction, pt_prediction) _test_helpers.assert_values_are_close(ort_x.grad, pt_x.grad) _test_helpers.assert_values_are_close(ort_loss, pt_loss) + + +def test_tanh_grad(): + class NeuralNetTanh(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(NeuralNetTanh, self).__init__() + + self.fc1 = torch.nn.Linear(input_size, hidden_size) + self.tanh = torch.nn.Tanh() + + def forward(self, input1): + out = self.fc1(input1) + out = self.tanh(out) + return out + + def run_step(model, x): + prediction = model(x) + loss = prediction.sum() + loss.backward() + return prediction, loss + device = 'cuda' + + N, D_in, H, D_out = 120, 1536, 500, 1536 + pt_model = NeuralNetTanh(D_in, H, D_out).to(device) + ort_model = ORTModule(copy.deepcopy(pt_model)) + + for step in range(10): + pt_x = torch.randn(N, D_in, device=device, requires_grad=True) + ort_x = copy.deepcopy(pt_x) + ort_prediction, ort_loss = run_step(ort_model, ort_x) + pt_prediction, pt_loss = run_step(pt_model, pt_x) + _test_helpers.assert_values_are_close(ort_prediction, pt_prediction) + _test_helpers.assert_values_are_close(ort_x.grad, pt_x.grad) + _test_helpers.assert_values_are_close(ort_loss, pt_loss) \ No newline at end of file diff --git a/orttraining/orttraining/test/training_ops/cpu/activation/activation_op_test.cc b/orttraining/orttraining/test/training_ops/cpu/activation/activation_op_test.cc index f7cbec1c2f..a0f3685c01 100644 --- a/orttraining/orttraining/test/training_ops/cpu/activation/activation_op_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/activation/activation_op_test.cc @@ -80,6 +80,9 @@ float SigmoidGrad(float dy, float y) { return dy * y * (1 - y); } +float TanhGrad(float dy, float y) { + return dy * (1 - y * y); +} } // namespace TEST(GeluGradTest, Basic) { @@ -180,6 +183,22 @@ TEST(SigmoidGradTest, Basic) { {}, 1, kMSDomain); } +TEST(TanhGradTest, Basic) { + const std::vector y_vals = {-1.0f, 0, 1.0f, 100.0f, -100.0f, 1000.0f, -1000.0f}; + const std::vector dY(7, 1.0f); + + TestElementwiseGradientOp( + "TanhGrad", + {{"dY", dY}, {"Y", y_vals}}, + [](const std::vector& params) { + ORT_ENFORCE(params.size() == 2); + const auto dy = params[0], y = params[1]; + + return TanhGrad(dy, y); + }, + {}, 1, kMSDomain); +} + namespace { template void TestBiasGeluGradBroadcastBias(const std::string& op, int opset_version, const std::string& domain, diff --git a/orttraining/orttraining/test/training_ops/cuda/activations_test.cc b/orttraining/orttraining/test/training_ops/cuda/activations_test.cc index 512902c35e..9035ca5019 100644 --- a/orttraining/orttraining/test/training_ops/cuda/activations_test.cc +++ b/orttraining/orttraining/test/training_ops/cuda/activations_test.cc @@ -83,6 +83,13 @@ TEST(CudaKernelTest, SigmoidGrad_basic) { } } +TEST(CudaKernelTest, TanhGrad_basic) { + std::vector> test_dims{{4}, {16, 2}, {8, 2, 128, 128}}; + for (const auto& test_dim : test_dims) { + TestActivations(test_dim, "TanhGrad", true /* grad_op */); + } +} + static void TestActivationsWithBroadcastBias( const std::vector& tensor_dim, const std::string& operator_name, diff --git a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc index 1594a8f5e8..76d4f9dde5 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc @@ -46,12 +46,13 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Gathe class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherElementsGrad); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GeluGrad); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SigmoidGrad); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, TanhGrad); // REVIEW(mzs): ConstEigenVectorArrayMap.cast, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, // REVIEW(mzs): ConstEigenVectorArrayMap.cast, - //BuildKernelCreateInfo, - //BuildKernelCreateInfo, + // BuildKernelCreateInfo, + // BuildKernelCreateInfo, + // BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cpu/op_gradients.cc b/orttraining/orttraining/training_ops/cpu/op_gradients.cc index 005e8d86fd..bbcfdc6ec8 100644 --- a/orttraining/orttraining/training_ops/cpu/op_gradients.cc +++ b/orttraining/orttraining/training_ops/cpu/op_gradients.cc @@ -167,5 +167,25 @@ Status SigmoidGrad::Compute(OpKernelContext* context) const { dx = dy * y * (1 - y); return Status::OK(); } + +ONNX_OPERATOR_KERNEL_EX( + TanhGrad, + kMSDomain, + 1, + kCpuExecutionProvider, + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), + TanhGrad); + +template +Status TanhGrad::Compute(OpKernelContext* context) const { + auto& dY = *context->Input(0); + auto& Y = *context->Input(1); + auto& dX = *context->Output(0, dY.Shape()); + EigenVectorArrayMap dx = EigenVectorArrayMap(dX.template MutableData(), dX.Shape().Size()); + ConstEigenVectorArrayMap y = ConstEigenVectorArrayMap(Y.template Data(), Y.Shape().Size()); + ConstEigenVectorArrayMap dy = ConstEigenVectorArrayMap(dY.template Data(), dY.Shape().Size()); + dx = dy * (1 - y * y); + return Status::OK(); +} } // namespace contrib } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/op_gradients.h b/orttraining/orttraining/training_ops/cpu/op_gradients.h index 2853a1dfec..80081c15c4 100644 --- a/orttraining/orttraining/training_ops/cpu/op_gradients.h +++ b/orttraining/orttraining/training_ops/cpu/op_gradients.h @@ -45,6 +45,18 @@ class SigmoidGrad final : public OpKernel { ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(SigmoidGrad); }; +template +class TanhGrad final : public OpKernel { + public: + explicit TanhGrad(const OpKernelInfo& info) : OpKernel(info) { + } + + Status Compute(OpKernelContext* context) const override; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TanhGrad); +}; + template class SoftmaxGrad final : public OpKernel { public: diff --git a/orttraining/orttraining/training_ops/cuda/activation/activations_grad.cc b/orttraining/orttraining/training_ops/cuda/activation/activations_grad.cc index 3cfcda2a9f..f42a523848 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/activations_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/activation/activations_grad.cc @@ -47,6 +47,7 @@ ACTIVATION_GRAD_OP_HFD(GeluGrad, 1, kMSDomain); ACTIVATION_GRAD_OP_HFD(FastGeluGrad, 1, kMSDomain); ACTIVATION_GRAD_OP_HFD(ReluGrad, 1, kMSDomain); ACTIVATION_GRAD_OP_HFD(SigmoidGrad, 1, kMSDomain); +ACTIVATION_GRAD_OP_HFD(TanhGrad, 1, kMSDomain); -} //namespace cuda +} // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/activation/activations_grad.h b/orttraining/orttraining/training_ops/cuda/activation/activations_grad.h index bb4c521a8e..444565a371 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/activations_grad.h +++ b/orttraining/orttraining/training_ops/cuda/activation/activations_grad.h @@ -55,5 +55,15 @@ class SigmoidGrad final : public BinaryElementwise { MAKE_FUNC_CTX_NULL() }; +template +class TanhGrad final : public BinaryElementwise { + public: + TanhGrad(const OpKernelInfo& info) : BinaryElementwise(info) {} + + Status ComputeInternal(OpKernelContext* context) const override; + + private: + MAKE_FUNC_CTX_NULL() +}; } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.cu index 6f8d5e7de1..befcda3dec 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.cu @@ -27,7 +27,7 @@ struct OP_FastGeluGrad : public CtxGeluGrad { template struct OP_ReluGrad : public CtxReluGrad { __device__ __inline__ T operator()(const T& dy, const T& x) const { - return x > T {0} ? dy : T {0}; + return x > T{0} ? dy : T{0}; } }; @@ -38,6 +38,13 @@ struct OP_SigmoidGrad : public CtxSigmoidGrad { } }; +template +struct OP_TanhGrad : public CtxTanhGrad { + __device__ __inline__ T operator()(const T& dy, const T& y) const { + return dy * ((T)1 - y * y); + } +}; + #define BINARY_ELEMENTWISE_IMPL(name) \ BINARY_ELEMENTWISE_IMPL_DECLARATION(name) { \ BinaryElementWiseNoBroadcastImpl(stream, \ diff --git a/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.h b/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.h index 03ad73a555..af18144377 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.h +++ b/orttraining/orttraining/training_ops/cuda/activation/activations_grad_impl.h @@ -11,16 +11,18 @@ typedef onnxruntime::cuda::CtxNull CtxGeluGrad; typedef onnxruntime::cuda::CtxNull CtxFastGeluGrad; typedef onnxruntime::cuda::CtxNull CtxReluGrad; typedef onnxruntime::cuda::CtxNull CtxSigmoidGrad; +typedef onnxruntime::cuda::CtxNull CtxTanhGrad; -#define ACTIVATION_GRAD_OPS() \ - ACTIVATION_GRAD_OP_NAME(GeluGrad) \ +#define ACTIVATION_GRAD_OPS() \ + ACTIVATION_GRAD_OP_NAME(GeluGrad) \ ACTIVATION_GRAD_OP_NAME(FastGeluGrad) \ - ACTIVATION_GRAD_OP_NAME(ReluGrad) \ - ACTIVATION_GRAD_OP_NAME(SigmoidGrad) + ACTIVATION_GRAD_OP_NAME(ReluGrad) \ + ACTIVATION_GRAD_OP_NAME(SigmoidGrad) \ + ACTIVATION_GRAD_OP_NAME(TanhGrad) #define BINARY_ELEMENTWISE_IMPL_DECLARATION(name) \ template \ - void Impl_##name(cudaStream_t stream, \ + void Impl_##name(cudaStream_t stream, \ const T* lhs_data, \ const T* rhs_data, \ T* output_data, \ diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 2d5e871ceb..394b8ad7c2 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -106,6 +106,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, SigmoidGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, SigmoidGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, TanhGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, TanhGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, TanhGrad); + class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, IsFinite); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, IsFinite); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, IsFinite); @@ -313,6 +317,9 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc index 2fce0f5ef0..0fcb36f161 100644 --- a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc +++ b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc @@ -106,6 +106,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, SigmoidGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, SigmoidGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, TanhGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, TanhGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, TanhGrad); + class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, IsFinite); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, IsFinite); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, IsFinite); @@ -279,6 +283,9 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -337,7 +344,7 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) { #endif #ifdef USE_MPI - // BuildKernelCreateInfo, + // BuildKernelCreateInfo, #endif BuildKernelCreateInfo,