From e9ec4c098b07a0f9d7b00888a7c37d59d02b24fe Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Wed, 22 Feb 2023 09:16:22 +0800 Subject: [PATCH] [CUDA] Fix FP16 Precision for Sigmoid Op (#14727) Current Sigmoid's CUDA kernel uses target data type for all computation. For some small negative numbers, if using FP16, it will loss precision. For example, for input [-7.8477, 7.3320, -7.8008, 6.6016], the expected output is [3.9047e-04, 9.9935e-01, 4.0919e-04, 9.9864e-01], but current kernel will generate result [0.0000, 0.9990, 0.0000, 0.9990]. If some sub-graph contains Sigmoid, such as BinaryCrossEntropyWithLogits, it's likely to produce NaN as compute result. The PR fixes this by using FP32 for kernel internal computation. Note that the fix will not have perf regression, as CUDA's _Exp will also do float to half casting, so the fix doesn't introduce extra cast. We move the cast to right begin and end of the whole kernel so that other parts of computation are also in FP32 (instead of only Exp). --- .../cuda/activation/activations_impl.cu | 13 +++++++- .../python/orttraining_test_ortmodule_api.py | 33 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/cuda/activation/activations_impl.cu b/onnxruntime/core/providers/cuda/activation/activations_impl.cu index 05f90ee05b..0f89fe9843 100644 --- a/onnxruntime/core/providers/cuda/activation/activations_impl.cu +++ b/onnxruntime/core/providers/cuda/activation/activations_impl.cu @@ -47,7 +47,18 @@ struct OP_Selu : public CtxSelu { template struct OP_Sigmoid : public CtxSigmoid { __device__ __inline__ T operator()(const T& a) const { - return a > T(0) ? (T)1 / ((T)1. + _Exp(-_Abs(a))) : (T)1 - (T)1 / ((T)1 + _Exp(-_Abs(a))); + return a > T(0) ? (T)1 / ((T)1. + _Exp(-a)) : (T)1 - (T)1 / ((T)1 + _Exp(a)); + } +}; + +template <> +struct OP_Sigmoid : public CtxSigmoid { + __device__ __inline__ half operator()(const half& a) const { + // For some small negative values, it will loss precision if using half for compute. + // The cast between half and float below won't cause perf issue since there is cast inside _Exp if input is half. + float af = static_cast(a); + float res = af > 0.0f ? 1.0f / (1.0f + _Exp(-af)) : 1.0f - 1.0f / (1.0f + _Exp(af)); + return static_cast(res); } }; diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 1a386579f8..f668776d02 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -5681,3 +5681,36 @@ def test_non_contiguous_tensors_as_inputs(training_mode): x_copy = copy.deepcopy(x) assert not x.is_contiguous() _test_helpers.assert_values_are_close(pt_model(x), ort_model(x_copy)) + + +def test_gradient_correctness_bce_with_logits(): + class NeuralNetBCEWithLogitsLoss(torch.nn.Module): + def __init__(self, input_size, hidden_size): + super(NeuralNetBCEWithLogitsLoss, self).__init__() + self.linear = torch.nn.Linear(input_size, hidden_size) + + def forward(self, input, target): + loss_fct = torch.nn.BCEWithLogitsLoss() + return loss_fct(self.linear(input), target) + + N, D, H = 16, 256, 128 + device = "cuda" + pt_model = NeuralNetBCEWithLogitsLoss(D, H).to(device) + ort_model = ORTModule(copy.deepcopy(pt_model)) + + def run_step(model, input, target): + prediction = model(input, target) + loss = prediction.sum() + loss.backward() + return prediction + + for _ in range(10): + pt_input = torch.rand((N, D), device=device, requires_grad=True) + ort_input = copy.deepcopy(pt_input) + pt_target = torch.rand((N, H), device=device) + ort_target = copy.deepcopy(pt_target) + pt_prediction = run_step(pt_model, pt_input, pt_target) + ort_prediction = run_step(ort_model, ort_input, ort_target) + + _test_helpers.assert_values_are_close(ort_prediction, pt_prediction) + _test_helpers.assert_values_are_close(ort_input.grad, pt_input.grad)