From 32295b15a131c07884b409e485c3b823bcb5a3dc Mon Sep 17 00:00:00 2001 From: Funtowicz Morgan Date: Tue, 22 Feb 2022 18:21:16 +0100 Subject: [PATCH] Gelu10 (#15676) * Add GeLU10 (clipped version of GeLU) to transformers to improve quantization performances. * Add unittests. * Import tensorflow after `is_tf_available` check. * Fix tensorflow wrong function `tf.tensor` to `tf.constant` * style. * use `tf.math.max` * Fix tf tests. * style. * style style style style style style * style style style style style style * Address @sgugger comments. * Fix wrong operator for raising ValueError for ClippedGELUActivation. --- src/transformers/activations.py | 26 ++++++++++++++++++++++++++ src/transformers/activations_tf.py | 15 +++++++++++++++ tests/test_activations.py | 14 ++++++++++++++ tests/test_activations_tf.py | 18 ++++++++++++++++++ 4 files changed, 73 insertions(+) diff --git a/src/transformers/activations.py b/src/transformers/activations.py index 998534709f..07421bfa55 100644 --- a/src/transformers/activations.py +++ b/src/transformers/activations.py @@ -74,6 +74,31 @@ class QuickGELUActivation(nn.Module): return input * torch.sigmoid(1.702 * input) +class ClippedGELUActivation(nn.Module): + """ + Clip the range of possible GeLU outputs between [min, max]. This is especially useful for quantization purpose, as + it allows mapping negatives values in the GeLU spectrum. For more information on this trick, please refer to + https://arxiv.org/abs/2004.09602. + + Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when + initially created. + + For information: OpenAI GPT's gelu is slightly different (and gives slightly different results): 0.5 * x * (1 + + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))). See https://arxiv.org/abs/1606.08415 + """ + + def __init__(self, min: float, max: float): + if min > max: + raise ValueError(f"min should be < max (got min: {min}, max: {max})") + + super().__init__() + self.min = min + self.max = max + + def forward(self, x: Tensor) -> Tensor: + return torch.clip(gelu(x), self.min, self.max) + + class SiLUActivation(nn.Module): """ See Gaussian Error Linear Units (Hendrycks et al., https://arxiv.org/abs/1606.08415) where the SiLU (Sigmoid Linear @@ -136,6 +161,7 @@ ACT2FN = { "gelu_new": NewGELUActivation(), "gelu_fast": FastGELUActivation(), "quick_gelu": QuickGELUActivation(), + "gelu_10": ClippedGELUActivation(-10, 10), "mish": MishActivation(), "linear": LinearActivation(), "sigmoid": nn.Sigmoid(), diff --git a/src/transformers/activations_tf.py b/src/transformers/activations_tf.py index 32c333e46a..ba74e9850e 100644 --- a/src/transformers/activations_tf.py +++ b/src/transformers/activations_tf.py @@ -69,6 +69,20 @@ def quick_gelu(x): return x * tf.math.sigmoid(coeff * x) +def gelu_10(x): + """ + Clip the range of possible GeLU outputs between [-10, 10]. This is especially useful for quantization purpose, as + it allows mapping 2 negatives values in the GeLU spectrum. For more information on this trick, please refer to + https://arxiv.org/abs/2004.09602 + + Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when + initially created. For information: OpenAI GPT's gelu is slightly different (and gives slightly different results): + 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) Also see + https://arxiv.org/abs/1606.08415 :param x: :return: + """ + return tf.clip_by_value(_gelu(x), -10, 10) + + def glu(x, axis=-1): """ Gated Linear Unit. Implementation as defined in the original paper (see https://arxiv.org/abs/1612.08083), where @@ -107,6 +121,7 @@ ACT2FN = { "tanh": tf.keras.activations.tanh, "gelu_fast": gelu_fast, "quick_gelu": quick_gelu, + "gelu_10": gelu_10, "glu": glu, } diff --git a/tests/test_activations.py b/tests/test_activations.py index 71b2991310..339d2eda16 100644 --- a/tests/test_activations.py +++ b/tests/test_activations.py @@ -32,6 +32,19 @@ class TestActivations(unittest.TestCase): self.assertTrue(torch.allclose(gelu_python(x), torch_builtin(x))) self.assertFalse(torch.allclose(gelu_python(x), gelu_new(x))) + def test_gelu_10(self): + x = torch.tensor([-100, -1, -0.1, 0, 0.1, 1.0, 100]) + torch_builtin = get_activation("gelu") + gelu10 = get_activation("gelu_10") + + y_gelu = torch_builtin(x) + y_gelu_10 = gelu10(x) + + clipped_mask = torch.where(y_gelu_10 < 10.0, 1, 0) + + self.assertTrue(torch.max(y_gelu_10).item() == 10.0) + self.assertTrue(torch.allclose(y_gelu * clipped_mask, y_gelu_10 * clipped_mask)) + def test_get_activation(self): get_activation("swish") get_activation("silu") @@ -40,6 +53,7 @@ class TestActivations(unittest.TestCase): get_activation("gelu_new") get_activation("gelu_fast") get_activation("gelu_python") + get_activation("gelu_10") get_activation("quick_gelu") get_activation("mish") get_activation("linear") diff --git a/tests/test_activations_tf.py b/tests/test_activations_tf.py index 236c37c333..9c02639835 100644 --- a/tests/test_activations_tf.py +++ b/tests/test_activations_tf.py @@ -14,16 +14,33 @@ import unittest +import numpy as np + from transformers import is_tf_available from transformers.testing_utils import require_tf if is_tf_available(): + import tensorflow as tf + from transformers.activations_tf import get_tf_activation @require_tf class TestTFActivations(unittest.TestCase): + def test_gelu_10(self): + x = tf.constant([-100, -1.0, -0.1, 0, 0.1, 1.0, 100.0]) + gelu = get_tf_activation("gelu") + gelu10 = get_tf_activation("gelu_10") + + y_gelu = gelu(x) + y_gelu_10 = gelu10(x) + + clipped_mask = tf.where(y_gelu_10 < 10.0, 1.0, 0.0) + + self.assertEqual(tf.math.reduce_max(y_gelu_10).numpy().item(), 10.0) + self.assertTrue(np.allclose(y_gelu * clipped_mask, y_gelu_10 * clipped_mask)) + def test_get_activation(self): get_tf_activation("swish") get_tf_activation("silu") @@ -32,6 +49,7 @@ class TestTFActivations(unittest.TestCase): get_tf_activation("tanh") get_tf_activation("gelu_new") get_tf_activation("gelu_fast") + get_tf_activation("gelu_10") get_tf_activation("mish") get_tf_activation("quick_gelu") get_tf_activation("glu")