update

megatron/core/enums.py

+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import enum
+
+class ModelType(enum.Enum):
+    encoder_or_decoder = 1
+    encoder_and_decoder = 2
+    retro_encoder = 3
+    retro_decoder = 4

megatron/core/fusions/fused_bias_dropout.py

+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+from typing import Tuple, Optional
+
+def _bias_dropout_add_func(x, bias, residual, prob, training):
+    # type: (Tensor, Optional[Tensor], Tensor, float, bool) -> Tensor
+    # NOTE: Previously, the argument `bias` used to be passed as
+    # `bias.expand_as(residual)` when the `bias_dropout_func` is called from the
+    # transformer layer but broadcasting should automatically take care of that.
+    # Also, looking at broadcasting semantics, `expand_as` and broadcasting
+    # seem to be identical performance-wise (both just change the view).
+    if bias is not None:
+        x = x + bias
+    out = torch.nn.functional.dropout(x, p=prob, training=training)
+    out = residual + out
+    return out
+
+def get_bias_dropout_add(training, fused):
+
+    def unfused_bias_dropout_add(x_with_bias, residual, prob):
+        x, bias = x_with_bias # unpack
+        return _bias_dropout_add_func(x, bias, residual, prob, training)
+
+    @torch.jit.script
+    def bias_dropout_add_fused_train(
+        x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]],
+        residual: torch.Tensor,
+        prob: float
+    ) -> torch.Tensor:
+        x, bias = x_with_bias # unpack
+        return _bias_dropout_add_func(x, bias, residual, prob, True)
+
+    @torch.jit.script
+    def bias_dropout_add_fused_inference(
+        x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]],
+        residual: torch.Tensor,
+        prob: float
+    ) -> torch.Tensor:
+        x, bias = x_with_bias # unpack
+        return _bias_dropout_add_func(x, bias, residual, prob, False)
+
+    if fused:
+        # jit scripting for a nn.module (with dropout) is not
+        # triggering the fusion kernel. For now, we use two
+        # different nn.functional routines to account for varying
+        # dropout semantics during training and inference phases.
+        if training:
+            return bias_dropout_add_fused_train
+        else:
+            return bias_dropout_add_fused_inference
+    else:
+        return unfused_bias_dropout_add

megatron/core/fusions/fused_bias_gelu.py

+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+
+
+###### BIAS GELU FUSION/ NO AUTOGRAD ################
+# 1/sqrt(2*pi)-> 0.3989423
+# 1/sqrt(2)   -> 0.70710678
+# sqrt(2/pi)  -> 0.79788456
+# this function is tanh approximation of gelu
+# actual gelu is:
+# x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
+
+@torch.jit.script
+def bias_gelu(bias, y):
+    x = bias + y
+    return  x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+# gradient of tanh approximation of gelu
+# gradient of actual gelu is:
+# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+@torch.jit.script
+def bias_gelu_back(g, bias, y):
+    x = bias + y
+    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
+    return ff*g
+
+class GeLUFunction(torch.autograd.Function):
+    @staticmethod
+    # bias is an optional argument
+    def forward(ctx, input, bias):
+        ctx.save_for_backward(input, bias)
+        return bias_gelu(bias, input)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, bias = ctx.saved_tensors
+        tmp = bias_gelu_back(grad_output, bias, input)
+        return tmp, tmp
+
+bias_gelu_impl = GeLUFunction.apply