text: # coding=utf-8
  # Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Transformer."""
  import math
  
  import torch
  import torch.nn.functional as F
  
  from nemo.collections.nlp.modules.common.megatron.fused_bias_dropout_add import (
      bias_dropout_add,
      bias_dropout_add_fused_inference,
      bias_dropout_add_fused_train,
  )
  from nemo.collections.nlp.modules.common.megatron.fused_bias_gelu import fused_bias_gelu
  from nemo.collections.nlp.modules.common.megatron.fused_layer_norm import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, attention_mask_func, erf_gelu
  
  try:
      from apex.transformer import parallel_state, tensor_parallel
      from apex.transformer.enums import AttnMaskType, AttnType, LayerType, ModelType
      from apex.transformer.functional.fused_softmax import FusedScaleMaskSoftmax
      from apex.transformer.utils import divide as safe_divide
  
      HAVE_APEX = True
  except (ImportError, ModuleNotFoundError):
      HAVE_APEX = False
  
      # fake missing classes with None attributes
      ModelType = AttnMaskType = AttnType = LayerType = ApexGuardDefaults()
  
  """ We use the following notation throughout this file:
       h: hidden size
       n: number of attention heads
       p: number of model parallel partitions
       np: n/p
       hp: h/p
       hn: h/n
       b: batch size
       s: sequence length
       l: number of layers
      Transformer takes input of size [s, b, h] and returns a
      tensor of the same size. We use the following arguments:
          hyperparameters: transformer hyperparameters
  """
  
  
  class ParallelMLP(MegatronModule):
      """MLP.
  
      MLP will take the input with h hidden state, project it to 4*h
      hidden dimension, perform nonlinear transformation, and project the
      state back into h hidden dimension.
      """
  
      def __init__(
          self,
          init_method,
          output_layer_init_method,
          hidden_size,
          ffn_hidden_size,
          use_cpu_initialization=False,
          bias_gelu_fusion=True,
          openai_gelu=False,
          onnx_safe=False,
          activation='gelu',
      ):
          super(ParallelMLP, self).__init__()
          self.activation = activation
  
          if activation not in ['gelu', 'geglu', 'reglu', 'swiglu']:
              raise ValueError(f"Activation {activation} not supported. Only gelu, geglu, reglu, swiglu are supported.")
  
          # Project to 4h.
          self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(
              hidden_size,
              ffn_hidden_size,  # NOTE: When using geglu, divide ffn dim by 2/3 to keep overall params the same.
              gather_output=False,
              init_method=init_method,
              skip_bias_add=True,
              use_cpu_initialization=use_cpu_initialization,
          )
  
          if activation in ['geglu', 'reglu', 'swiglu']:
              # Separate linear layer for *GLU activations.
              # Source: https://github.com/huggingface/transformers/blob/bee361c6f1f7704f8c688895f2f86f6e5ff84727/src/transformers/models/t5/modeling_t5.py#L292
              self.dense_h_to_4h_2 = tensor_parallel.ColumnParallelLinear(
                  hidden_size,
                  ffn_hidden_size,  # NOTE: When using *glu, divide ffn dim by 2/3 to keep overall params the same.
                  gather_output=False,
                  init_method=init_method,
                  skip_bias_add=True,
                  use_cpu_initialization=use_cpu_initialization,
              )
              glu_activation_family = True
          else:
              glu_activation_family = False
  
          if glu_activation_family and bias_gelu_fusion:
              raise ValueError(
                  f"Cannot use bias_gelu_fusion with {activation} activation. Please turn bias gelu fusion off."
              )
  
          if glu_activation_family and openai_gelu:
              raise ValueError(
                  f"Cannot use openai_gelu with specificed activation function : {activation} Please turn openai gelu off."
              )
  
          if glu_activation_family and onnx_safe:
              raise ValueError(
                  f"Cannot use onnx_safe with specificed activation function : {activation} Please turn onnx safe off."
              )
  
          self.bias_gelu_fusion = bias_gelu_fusion
  
          if activation == "gelu":
              self.activation_func = F.gelu
          elif openai_gelu:
              self.activation_func = openai_gelu
          elif onnx_safe:
              self.activation_func = erf_gelu
          else:
              # Remaining acitvations are implemeted in the forward function.
              self.activation_func = None
  
          # Project back to h.
          self.dense_4h_to_h = tensor_parallel.RowParallelLinear(
              ffn_hidden_size,
              hidden_size,
              input_is_parallel=True,
              init_method=output_layer_init_method,
              skip_bias_add=True,
              use_cpu_initialization=use_cpu_initialization,
          )
  
      def forward(self, hidden_states):
  
          # [s, b, 4hp]
          intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states)
  
          if self.activation in ['geglu', 'reglu', 'swiglu']:
              intermediate_parallel_2, bias_parallel_2 = self.dense_h_to_4h_2(hidden_states)
  
          if self.activation == 'geglu':
              intermediate_parallel = F.gelu(intermediate_parallel + bias_parallel) * (
                  intermediate_parallel_2 + bias_parallel_2
              )
          elif self.activation == 'swiglu':
              # SiLU or sigmoid linear unit is the same as swish with beta = 1 (which is what https://arxiv.org/pdf/2002.05202.pdf uses.)
              intermediate_parallel = F.silu(intermediate_parallel + bias_parallel) * (
                  intermediate_parallel_2 + bias_parallel_2
              )
          elif self.activation == 'reglu':
              intermediate_parallel = F.relu(intermediate_parallel + bias_parallel) * (
                  intermediate_parallel_2 + bias_parallel_2
              )
          elif self.bias_gelu_fusion and self.activation == 'gelu':
              intermediate_parallel = fused_bias_gelu(intermediate_parallel, bias_parallel)
          else:
              intermediate_parallel = self.activation_func(intermediate_parallel + bias_parallel)
  
          # [s, b, h]
          output, output_bias = self.dense_4h_to_h(intermediate_parallel)
          return output, output_bias
  
  
  class ParallelAttention(MegatronModule):
      """Parallel self-attention layer abstract class.
  
      Self-attention layer takes input with size [b, s, h]
      and returns output of the same size.
      """
  
      def __init__(
          self,
          init_method,
          output_layer_init_method,
          layer_number,
          num_attention_heads,
          hidden_size,
          attention_type=AttnType.self_attn,
          attn_mask_type=AttnMaskType.padding,
          precision=16,
          apply_query_key_layer_scaling=True,
          kv_channels=None,
          use_cpu_initialization=False,
          masked_softmax_fusion=True,
          attention_dropout=0.1,
          megatron_legacy=False,
      ):
          super(ParallelAttention, self).__init__()
  
          self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
          self.attention_softmax_in_fp32 = False
          if self.apply_query_key_layer_scaling:
              self.attention_softmax_in_fp32 = True
          self.layer_number = max(1, layer_number)
          self.attention_type = attention_type
          self.attn_mask_type = attn_mask_type
          self.megatron_legacy = megatron_legacy
  
          if kv_channels is None:
              assert (
                  hidden_size % num_attention_heads == 0
              ), 'hidden_size must be divisible by num_attention_heads if kv_channels is None'
              kv_channels = hidden_size // num_attention_heads
          projection_size = kv_channels * num_attention_heads
  
          # Per attention head and per partition values.
          world_size = parallel_state.get_tensor_model_parallel_world_size()
          self.hidden_size_per_partition = safe_divide(projection_size, world_size)
          self.hidden_size_per_attention_head = safe_divide(projection_size, num_attention_heads)
          self.num_attention_heads_per_partition = safe_divide(num_attention_heads, world_size)
  
          # Strided linear layer.
          if attention_type == AttnType.self_attn:
              self.query_key_value = tensor_parallel.ColumnParallelLinear(
                  hidden_size,
                  3 * projection_size,
                  gather_output=False,
                  init_method=init_method,
                  use_cpu_initialization=use_cpu_initialization,
              )
          else:
              assert attention_type == AttnType.cross_attn
              self.query = tensor_parallel.ColumnParallelLinear(
                  hidden_size, projection_size, gather_output=False, init_method=init_method
              )
  
              self.key_value = tensor_parallel.ColumnParallelLinear(
                  hidden_size, 2 * projection_size, gather_output=False, init_method=init_method
              )
  
          coeff = None
          self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
          if self.apply_query_key_layer_scaling:
              coeff = self.layer_number
              self.norm_factor *= coeff
  
          fused_fp16 = precision == 16
          fused_bf16 = precision == 'bf16'
          self.scale_mask_softmax = FusedScaleMaskSoftmax(
              fused_fp16,
              fused_bf16,
              self.attn_mask_type,
              masked_softmax_fusion,
              attention_mask_func,
              self.attention_softmax_in_fp32,
              coeff,
          )
  
          # Dropout. Note that for a single iteration, this layer will generate
          # different outputs on different number of parallel partitions but
          # on average it should not be partition dependent.
          self.attention_dropout = torch.nn.Dropout(attention_dropout)
  
          # Output.
          self.dense = tensor_parallel.RowParallelLinear(
              projection_size,
              hidden_size,
              input_is_parallel=True,
              init_method=output_layer_init_method,
              skip_bias_add=True,
              use_cpu_initialization=use_cpu_initialization,
          )
  
          # Inference key-value memory
          self.inference_key_memory = None
          self.inference_value_memory = None
          self.inference_current_sequence_len = 0
  
      def _allocate_memory(self, inference_max_sequence_len, batch_size, dtype):
          return torch.empty(
              inference_max_sequence_len,
              batch_size,
              self.num_attention_heads_per_partition,
              self.hidden_size_per_attention_head,
              dtype=dtype,
              device=torch.cuda.current_device(),
          )
  
      def _transpose_last_dim(self, mixed_layer, num_splits, num_splits_first):
          input_shape = mixed_layer.size()
          if num_splits_first:
              """[s, b, num_splits * np * hn]
              -->(view) [s, b, num_splits, np, hn]
              -->(tranpose) [s, b, np, num_splits, hn]
              -->(view) [s, b, np * num_splits * hn] """
  
              intermediate_shape = input_shape[:-1] + (
                  num_splits,
                  self.num_attention_heads_per_partition,
                  self.hidden_size_per_attention_head,
              )
  
              mixed_layer = mixed_layer.view(*intermediate_shape)
              mixed_layer = mixed_layer.transpose(-2, -3).contiguous()
          else:
              """[s, b, np * hn * num_splits]
              -->(view) [s, b, np, hn, num_splits]
              -->(tranpose) [s, b, np, num_splits, hn]
              -->(view) [s, b, np * num_splits * hn] """
  
              intermediate_shape = input_shape[:-1] + (
                  self.num_attention_heads_per_partition,
                  self.hidden_size_per_attention_head,
                  num_splits,
              )
  
              mixed_layer = mixed_layer.view(*intermediate_shape)
              mixed_layer = mixed_layer.transpose(-1, -2).contiguous()
          mixed_layer = mixed_layer.view(*input_shape)
  
          return mixed_layer
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          layer_past=None,
          get_key_value=False,
          encoder_output=None,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=None,
      ):
          # hidden_states: [sq, b, h]
  
          # =================================================
          # Pre-allocate memory for key-values for inference.
          # =================================================
          if set_inference_key_value_memory:
              assert inference_max_sequence_len and inference_max_sequence_len > 0
              self.inference_key_memory = self._allocate_memory(
                  inference_max_sequence_len, hidden_states.size(1), hidden_states.dtype
              )
              self.inference_value_memory = self._allocate_memory(
                  inference_max_sequence_len, hidden_states.size(1), hidden_states.dtype
              )
              self.inference_current_sequence_len = 0
          # Some consistency check.
          if inference_max_sequence_len:
              assert self.inference_current_sequence_len < self.inference_key_memory.size(0)
              assert inference_max_sequence_len == self.inference_key_memory.size(0)
          # This is added for safety. In case inference_max_sequence_len
          # is not provided, make sure there is no potential memory left
          # from previous inference.
          if not inference_max_sequence_len:
              self.inference_key_memory = None
              self.inference_value_memory = None
          # =====================
          # Query, Key, and Value
          # =====================
  
          if self.attention_type == AttnType.self_attn:
              # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
              mixed_x_layer, _ = self.query_key_value(hidden_states)
  
              # [sq, b, (np * 3 * hn)] --> [sq, b, np, 3 * hn]
              new_tensor_shape = mixed_x_layer.size()[:-1] + (
                  self.num_attention_heads_per_partition,
                  3 * self.hidden_size_per_attention_head,
              )
              if self.megatron_legacy:
                  mixed_x_layer = self._transpose_last_dim(mixed_x_layer, 3, True)
              mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
  
              # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
              (query_layer, key_layer, value_layer) = tensor_parallel.split_tensor_along_last_dim(mixed_x_layer, 3)
          else:
              # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)]
              mixed_kv_layer, _ = self.key_value(encoder_output)
  
              # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn]
              new_tensor_shape = mixed_kv_layer.size()[:-1] + (
                  self.num_attention_heads_per_partition,
                  2 * self.hidden_size_per_attention_head,
              )
              mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape)
  
              # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn]
              (key_layer, value_layer) = tensor_parallel.split_tensor_along_last_dim(mixed_kv_layer, 2)
  
              # Attention head [sq, b, h] --> [sq, b, hp]
              query_layer, _ = self.query(hidden_states)
              # [sq, b, hp] --> [sq, b, np, hn]
              new_tensor_shape = query_layer.size()[:-1] + (
                  self.num_attention_heads_per_partition,
                  self.hidden_size_per_attention_head,
              )
              query_layer = query_layer.view(*new_tensor_shape)
  
          # ===================================================
          # Adjust key, value, and attention mask for inference
          # ===================================================
  
          if inference_max_sequence_len:
              # Adjust the range variables.
              start = self.inference_current_sequence_len
              self.inference_current_sequence_len += key_layer.size(0)
              end = self.inference_current_sequence_len
              # Copy key and values.
              self.inference_key_memory[start:end, ...] = key_layer
              self.inference_value_memory[start:end, ...] = value_layer
              key_layer = self.inference_key_memory[:end, ...]
              value_layer = self.inference_value_memory[:end, ...]
              # Adjust attention mask
              attention_mask = attention_mask[..., start:end, :end]
  
          if layer_past is not None:
              past_key, past_value = layer_past
              key_layer = torch.cat((past_key.type_as(key_layer), key_layer), dim=0)
              value_layer = torch.cat((past_value.type_as(value_layer), value_layer), dim=0)
          if get_key_value:
              present = (key_layer, value_layer)
  
          # ===================================
          # Raw attention scores. [b, np, s, s]
          # ===================================
  
          # [b, np, sq, sk]
          output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
  
          # [sq, b, np, hn] -> [sq, b * np, hn]
          query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
          # [sk, b, np, hn] -> [sk, b * np, hn]
          key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
  
          # preallocting result tensor: [b * np, sq, sk]
          matmul_result = torch.empty(
              output_size[0] * output_size[1],
              output_size[2],
              output_size[3],
              dtype=query_layer.dtype,
              device=torch.cuda.current_device(),
          )
  
          # Raw attention scores. [b * np, sq, sk]
          matmul_result = torch.baddbmm(
              matmul_result,
              query_layer.transpose(0, 1),  # [b * np, sq, hn]
              key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
              beta=0.0,
              alpha=(1.0 / self.norm_factor),
          )
  
          # change view to [b, np, sq, sk]
          attention_scores = matmul_result.view(*output_size)
  
          # ==================================================
          # Update attention mask for inference. [b, np, sq, sk]
          # ==================================================
  
          if get_key_value:
              with torch.no_grad():
                  if layer_past is not None:
                      attention_mask = attention_mask[
                          ..., attention_scores.size(3) - 1, : attention_scores.size(3)
                      ].unsqueeze(2)
                  else:
                      attention_mask = attention_mask[..., : attention_scores.size(3), : attention_scores.size(3)]
  
          # ===========================
          # Attention probs and dropout
          # ===========================
  
          # attention scores and attention mask [b, np, sq, sk]
          attention_probs = self.scale_mask_softmax(attention_scores, attention_mask)
  
          # This is actually dropping out entire tokens to attend to, which might
          # seem a bit unusual, but is taken from the original Transformer paper.
          with tensor_parallel.random.get_cuda_rng_tracker().fork():
              attention_probs = self.attention_dropout(attention_probs)
  
          # =========================
          # Context layer. [sq, b, hp]
          # =========================
  
          # value_layer -> context layer.
          # [sk, b, np, hn] --> [b, np, sq, hn]
  
          # context layer shape: [b, np, sq, hn]
          output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
  
          # change view [sk, b * np, hn]
          value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
  
          # change view [b * np, sq, sk]
          attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
  
          # matmul: [b * np, sq, hn]
          context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
  
          # change view [b, np, sq, hn]
          context_layer = context_layer.view(*output_size)
  
          # [b, np, sq, hn] --> [sq, b, np, hn]
          context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
  
          # [sq, b, np, hn] --> [sq, b, hp]
          new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
          context_layer = context_layer.view(*new_context_layer_shape)
  
          # =================
          # Output. [sq, b, h]
          # =================
  
          output, bias = self.dense(context_layer)
  
          if get_key_value:
              output = [output, present]
  
          return output, bias
  
  
  def get_bias_dropout_add(training):
      def _bias_dropout_add(x, bias, residual, prob):
          return bias_dropout_add(x, bias, residual, prob, training)
  
      return _bias_dropout_add
  
  
  class ParallelTransformerLayer_(MegatronModule):
      """A single transformer layer.
  
      Transformer layer takes input with size [b, s, h] and returns an
      output of the same size.
      """
  
      def __init__(
          self,
          init_method,
          output_layer_init_method,
          layer_number,
          hidden_size,
          ffn_hidden_size,
          num_attention_heads,
          layer_type=LayerType.encoder,
          self_attn_mask_type=AttnMaskType.padding,
          fp32_residual_connection=False,
          apply_residual_connection_post_layernorm=False,
          precision=16,
          apply_query_key_layer_scaling=True,
          kv_channels=None,
          layernorm_epsilon=1e-5,
          hidden_dropout=0.1,
          bias_dropout_fusion=True,
          persist_layer_norm=False,
          use_cpu_initialization=False,
          bias_gelu_fusion=True,
          openai_gelu=False,
          onnx_safe=False,
          masked_softmax_fusion=True,
          attention_dropout=0.1,
          activation='gelu',
          megatron_legacy=False,
      ):
          super(ParallelTransformerLayer_, self).__init__()
  
          if kv_channels is None:
              assert (
                  hidden_size % num_attention_heads == 0
              ), 'hidden_size must be divisible by num_attention_heads if kv_channels is None'
              kv_channels = hidden_size // num_attention_heads
  
          self.layer_number = layer_number
          self.layer_type = layer_type
  
          # if true apply residual connection post layer norm (like original bert)
          self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
  
          self.fp32_residual_connection = fp32_residual_connection  # if true move residual connections to fp32
  
          # Layernorm on the input data.
          self.input_layernorm = get_layer_norm(hidden_size, layernorm_epsilon, persist_layer_norm)
  
          # Self attention.
          self.self_attention = ParallelAttention(
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              layer_number=layer_number,
              num_attention_heads=num_attention_heads,
              hidden_size=hidden_size,
              attention_type=AttnType.self_attn,
              attn_mask_type=self_attn_mask_type,
              precision=precision,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              use_cpu_initialization=use_cpu_initialization,
              masked_softmax_fusion=masked_softmax_fusion,
              attention_dropout=attention_dropout,
              megatron_legacy=megatron_legacy,
          )
          self.hidden_dropout = hidden_dropout
          self.attention_dropout = attention_dropout
          self.bias_dropout_fusion = bias_dropout_fusion  # if true, enable bias dropout fusion
  
          # Layernorm on the attention output
          self.post_attention_layernorm = get_layer_norm(hidden_size, layernorm_epsilon, persist_layer_norm)
  
          if self.layer_type == LayerType.decoder:
              self.inter_attention = ParallelAttention(
                  init_method=init_method,
                  output_layer_init_method=output_layer_init_method,
                  layer_number=layer_number,
                  num_attention_heads=num_attention_heads,
                  hidden_size=hidden_size,
                  attention_type=AttnType.cross_attn,
                  attn_mask_type=AttnMaskType.padding,
                  precision=precision,
                  apply_query_key_layer_scaling=apply_query_key_layer_scaling,
                  kv_channels=kv_channels,
                  use_cpu_initialization=use_cpu_initialization,
                  masked_softmax_fusion=masked_softmax_fusion,
                  attention_dropout=attention_dropout,
                  megatron_legacy=megatron_legacy,
              )
              # Layernorm on the attention output.
              self.post_inter_attention_layernorm = get_layer_norm(hidden_size, layernorm_epsilon, persist_layer_norm)
  
          # MLP
          self.mlp = ParallelMLP(
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              hidden_size=hidden_size,
              ffn_hidden_size=ffn_hidden_size,
              use_cpu_initialization=use_cpu_initialization,
              bias_gelu_fusion=bias_gelu_fusion,
              openai_gelu=openai_gelu,
              onnx_safe=onnx_safe,
              activation=activation,
          )
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          encoder_output=None,
          enc_dec_attn_mask=None,
          layer_past=None,
          get_key_value=False,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=None,
      ):
          # hidden_states: [b, s, h]
  
          # Layer norm at the beginning of the transformer layer.
          layernorm_output = self.input_layernorm(hidden_states)
          # Self attention.
          attention_output, attention_bias = self.self_attention(
              layernorm_output,
              attention_mask,
              layer_past=layer_past,
              get_key_value=get_key_value,
              set_inference_key_value_memory=set_inference_key_value_memory,
              inference_max_sequence_len=inference_max_sequence_len,
          )
  
          if get_key_value:
              attention_output, presents = attention_output
  
          # Residual connection.
          if self.apply_residual_connection_post_layernorm:
              residual = layernorm_output
          else:
              residual = hidden_states
  
          # jit scripting for a nn.module (with dropout) is not
          # trigerring the fusion kernel. For now, we use two
          # different nn.functional routines to account for varying
          # dropout semantics during training and inference phases.
          if self.bias_dropout_fusion:
              if self.training:
                  bias_dropout_add_func = bias_dropout_add_fused_train
              else:
                  bias_dropout_add_func = bias_dropout_add_fused_inference
          else:
              bias_dropout_add_func = get_bias_dropout_add(self.training)
  
          layernorm_input = bias_dropout_add_func(
              attention_output, attention_bias.expand_as(residual), residual, self.hidden_dropout
          )
  
          # Layer norm post the self attention.
          layernorm_output = self.post_attention_layernorm(layernorm_input)
  
          if self.layer_type == LayerType.decoder:
              attention_output, attention_bias = self.inter_attention(
                  layernorm_output, enc_dec_attn_mask, encoder_output=encoder_output
              )
              # residual connection
              if self.apply_residual_connection_post_layernorm:
                  residual = layernorm_output
              else:
                  residual = layernorm_input
  
              layernorm_input = bias_dropout_add_func(
                  attention_output, attention_bias.expand_as(residual), residual, self.hidden_dropout
              )
  
              # Layer norm post the decoder attention
              layernorm_output = self.post_inter_attention_layernorm(layernorm_input)
  
          # MLP.
          mlp_output, mlp_bias = self.mlp(layernorm_output)
  
          # Second residual connection.
          if self.apply_residual_connection_post_layernorm:
              residual = layernorm_output
          else:
              residual = layernorm_input
  
          output = bias_dropout_add_func(mlp_output, mlp_bias.expand_as(residual), residual, self.hidden_dropout)
  
          if get_key_value:
              output = [output, presents]
  
          return output
  
  
  class ParallelTransformerLayer(ParallelTransformerLayer_):
      def __init__(self, **kwargs):
          super(ParallelTransformerLayer, self).__init__(**kwargs)
  
          if kwargs['precision'] == 32:
              self.dtype = torch.float32
          elif kwargs['precision'] == 16:
              self.dtype = torch.float16
          elif kwargs['precision'] == 'bf16':
              self.dtype = torch.bfloat16
          else:
              raise ValueError
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          encoder_output=None,
          enc_dec_attn_mask=None,
          layer_past=None,
          get_key_value=False,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=None,
      ):
          if self.dtype == torch.float32:
              return super().forward(
                  hidden_states,
                  attention_mask,
                  encoder_output,
                  enc_dec_attn_mask,
                  layer_past,
                  get_key_value,
                  set_inference_key_value_memory,
                  inference_max_sequence_len,
              )
          with torch.autocast(device_type="cuda", dtype=self.dtype):
              return super().forward(
                  hidden_states,
                  attention_mask,
                  encoder_output,
                  enc_dec_attn_mask,
                  layer_past,
                  get_key_value,
                  set_inference_key_value_memory,
                  inference_max_sequence_len,
              )
  
  
  class ParallelTransformer(MegatronModule):
      """Transformer class."""
  
      def __init__(
          self,
          init_method,
          output_layer_init_method,
          num_layers,
          hidden_size,
          ffn_hidden_size,
          num_attention_heads,
          apply_query_key_layer_scaling=True,
          kv_channels=None,
          layer_type=LayerType.encoder,
          self_attn_mask_type=AttnMaskType.padding,
          pre_process=True,
          post_process=True,
          precision=16,
          fp32_residual_connection=False,
          activations_checkpoint_method=None,
          activations_checkpoint_num_layers=1,
          layernorm_epsilon=1e-5,
          hidden_dropout=0.1,
          attention_dropout=0.1,
          use_cpu_initialization=False,
          bias_gelu_fusion=True,
          masked_softmax_fusion=True,
          persist_layer_norm=False,
          openai_gelu=False,
          onnx_safe=False,
          activation='gelu',
          model_type=ModelType.encoder_or_decoder,
          megatron_legacy=False,
      ):
          super(ParallelTransformer, self).__init__()
  
          if kv_channels is None:
              assert (
                  hidden_size % num_attention_heads == 0
              ), 'hidden_size must be divisible by num_attention_heads if kv_channels is None'
              kv_channels = hidden_size // num_attention_heads
  
          self.fp32_residual_connection = fp32_residual_connection
          self.pre_process = pre_process
          self.post_process = post_process
          self.input_tensor = None
          self.self_attn_mask_type = self_attn_mask_type
          self.model_type = model_type
  
          # Store activation checkpointing flag.
          self.activations_checkpoint_method = activations_checkpoint_method
          self.activations_checkpoint_num_layers = activations_checkpoint_num_layers
  
          if self.model_type == ModelType.encoder_or_decoder:
              assert (
                  num_layers % parallel_state.get_pipeline_model_parallel_world_size() == 0
              ), 'num_layers must be divisible by pipeline_model_parallel_size'
  
          # TODO: Add similar assert for encoder-decoder.
  
          self.num_layers = self.get_num_layers(num_layers)
          # Transformer layers.
          def build_layer(layer_number):
              return ParallelTransformerLayer(
                  init_method=init_method,
                  output_layer_init_method=output_layer_init_method,
                  layer_number=layer_number,
                  hidden_size=hidden_size,
                  ffn_hidden_size=ffn_hidden_size,
                  num_attention_heads=num_attention_heads,
                  apply_query_key_layer_scaling=apply_query_key_layer_scaling,
                  kv_channels=kv_channels,
                  layer_type=layer_type,
                  self_attn_mask_type=self_attn_mask_type,
                  precision=precision,
                  fp32_residual_connection=fp32_residual_connection,
                  layernorm_epsilon=layernorm_epsilon,
                  hidden_dropout=hidden_dropout,
                  attention_dropout=attention_dropout,
                  use_cpu_initialization=use_cpu_initialization,
                  bias_gelu_fusion=bias_gelu_fusion,
                  masked_softmax_fusion=masked_softmax_fusion,
                  persist_layer_norm=persist_layer_norm,
                  openai_gelu=openai_gelu,
                  onnx_safe=onnx_safe,
                  activation=activation,
                  megatron_legacy=megatron_legacy,
              )
  
          if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:
              assert num_layers % parallel_state.get_virtual_pipeline_model_parallel_world_size() == 0, (
                  'num_layers_per_stage must be divisible by ' 'virtual_pipeline_model_parallel_size'
              )
              assert self.model_type != ModelType.encoder_or_decoder
              # Number of layers in each model chunk is the number of layers in the stage,
              # divided by the number of model chunks in a stage.
              self.num_layers = self.num_layers // parallel_state.get_virtual_pipeline_model_parallel_world_size()
              # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of
              # layers to stages like (each list is a model chunk):
              # Stage 0: [0]  [2]  [4]  [6]
              # Stage 1: [1]  [3]  [5]  [7]
              # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
              # layers to stages like (each list is a model chunk):
              # Stage 0: [0, 1]  [4, 5]
              # Stage 1: [2, 3]  [6, 7]
              offset = parallel_state.get_virtual_pipeline_model_parallel_rank() * (
                  num_layers // parallel_state.get_virtual_pipeline_model_parallel_world_size()
              ) + (parallel_state.get_pipeline_model_parallel_rank() * self.num_layers)
          else:
              # Each stage gets a contiguous set of layers.
              if (
                  self.model_type == ModelType.encoder_and_decoder
                  and parallel_state.get_pipeline_model_parallel_world_size() > 1
              ):
                  pipeline_rank = parallel_state.get_pipeline_model_parallel_rank()
                  if layer_type == LayerType.encoder:
                      offset = pipeline_rank * self.num_layers
                  else:
                      num_ranks_in_enc = parallel_state.get_pipeline_model_parallel_split_rank()
                      offset = (pipeline_rank - num_ranks_in_enc) * self.num_layers
              else:
                  offset = parallel_state.get_pipeline_model_parallel_rank() * self.num_layers
  
          self.layers = torch.nn.ModuleList([build_layer(i + 1 + offset) for i in range(self.num_layers)])
  
          if self.post_process:
              # Final layer norm before output.
              self.final_layernorm = get_layer_norm(hidden_size, layernorm_epsilon, persist_layer_norm)
  
      def _get_layer(self, layer_number):
          return self.layers[layer_number]
  
      def get_num_layers(self, num_layers):
          """Compute the number of transformer layers resident on the current rank."""
          if parallel_state.get_pipeline_model_parallel_world_size() > 1:
              if self.model_type == ModelType.encoder_and_decoder:
                  assert parallel_state.get_pipeline_model_parallel_split_rank() is not None
                  num_ranks_in_encoder = parallel_state.get_pipeline_model_parallel_split_rank()
                  num_ranks_in_decoder = parallel_state.get_pipeline_model_parallel_world_size() - num_ranks_in_encoder
                  assert (
                      num_layers % num_ranks_in_encoder == 0
                  ), 'num_layers must be divisible by number of ranks given to encoder'
                  assert (
                      num_layers % num_ranks_in_decoder == 0
                  ), 'num_layers must be divisible by number of ranks given to decoder'
                  if parallel_state.is_pipeline_stage_before_split():
                      num_layers = num_layers // num_ranks_in_encoder
                  else:
                      num_layers = num_layers // num_ranks_in_decoder
              else:
                  assert (
                      num_layers % parallel_state.get_pipeline_model_parallel_world_size() == 0
                  ), 'num_layers must be divisible by pipeline_model_parallel_size'
                  num_layers = num_layers // parallel_state.get_pipeline_model_parallel_world_size()
  
          return num_layers
  
      def _checkpointed_forward(self, hidden_states, attention_mask, encoder_output, enc_dec_attn_mask):
          """Forward method with activation checkpointing."""
  
          def custom(start, end):
              def custom_forward(*inputs):
                  x_ = inputs[0]
                  attention_mask = inputs[1]
                  encoder_output = inputs[2]
                  enc_dec_attn_mask = inputs[3]
                  for index in range(start, end):
                      layer = self._get_layer(index)
                      x_ = layer(x_, attention_mask, encoder_output, enc_dec_attn_mask)
                  return x_
  
              return custom_forward
  
          # Make sure memory is freed.
          tensor_parallel.reset_checkpointed_activations_memory_buffer()
  
          if self.activations_checkpoint_method == 'uniform':
              # Uniformly divide the total number of Transformer layers and checkpoint
              # the input activation of each divided chunk.
              # A method to further reduce memory usage reducing checkpoints.
              l = 0
              while l < self.num_layers:
                  hidden_states = tensor_parallel.checkpoint(
                      custom(l, l + self.activations_checkpoint_num_layers),
                      hidden_states,
                      attention_mask,
                      encoder_output,
                      enc_dec_attn_mask,
                  )
                  l += self.activations_checkpoint_num_layers
          elif self.activations_checkpoint_method == 'block':
              # Checkpoint the input activation of only a set number of individual
              # Transformer layers and skip the rest.
              # A method fully use the device memory removing redundant re-computation.
              for l in range(self.num_layers):
                  if l < self.activations_checkpoint_num_layers:
                      hidden_states = tensor_parallel.checkpoint(
                          custom(l, l + 1), hidden_states, attention_mask, encoder_output, enc_dec_attn_mask
                      )
                  else:
                      hidden_states = custom(l, l + 1)(hidden_states, attention_mask, encoder_output, enc_dec_attn_mask)
          else:
              raise ValueError("Invalid activation checkpoint method.")
  
          return hidden_states
  
      def set_input_tensor(self, input_tensor):
          """Set input tensor to be used instead of forward()'s input.
  
          When doing pipeline parallelism the input from the previous
          stage comes from communication, not from the input, so the
          model's forward_step_func won't have it. This function is thus
          used by internal code to bypass the input provided by the
          forward_step_func"""
          self.input_tensor = input_tensor
  
      def forward(
          self,
          hidden_states,
          attention_mask,
          layer_past=None,
          get_key_value=False,
          encoder_output=None,
          enc_dec_attn_mask=None,
          set_inference_key_value_memory=False,
          inference_max_sequence_len=None,
      ):
          # Checks.
          if inference_max_sequence_len:
              assert self.activations_checkpoint_method is None, 'inference does not work with activation checkpointing'
  
          if layer_past is not None:
              assert get_key_value, 'for not None values in layer_past, ' 'expected get_key_value to be set'
          if get_key_value:
              assert self.activations_checkpoint_method is None, (
                  'get_key_value does not work with ' 'activation checkpointing'
              )
  
          if self.pre_process:
              # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
              # If the input flag for fp32 residual connection is set, convert for float.
              if self.fp32_residual_connection:
                  hidden_states = hidden_states.transpose(0, 1).contiguous().float()
              # Otherwise, leave it as is.
              else:
                  hidden_states = hidden_states.transpose(0, 1).contiguous()
          else:
              # See set_input_tensor()
              hidden_states = self.input_tensor
  
          if encoder_output is not None:
              encoder_output = encoder_output.transpose(0, 1).contiguous()
  
          if self.activations_checkpoint_method is not None:
              hidden_states = self._checkpointed_forward(
                  hidden_states, attention_mask, encoder_output, enc_dec_attn_mask
              )
          else:
              if get_key_value:
                  presents = []
              for index in range(self.num_layers):
                  layer = self._get_layer(index)
                  past = None
                  if layer_past is not None:
                      past = layer_past[index]
                  hidden_states = layer(
                      hidden_states,
                      attention_mask,
                      encoder_output=encoder_output,
                      enc_dec_attn_mask=enc_dec_attn_mask,
                      layer_past=past,
                      get_key_value=get_key_value,
                      set_inference_key_value_memory=set_inference_key_value_memory,
                      inference_max_sequence_len=inference_max_sequence_len,
                  )
                  if get_key_value:
                      hidden_states, present = hidden_states
                      presents.append(present)
  
          # Final layer norm.
          if self.post_process:
              # Reverting data format change [s b h] --> [b s h].
              hidden_states = hidden_states.transpose(0, 1).contiguous()
              output = self.final_layernorm(hidden_states)
          else:
              output = hidden_states
          if get_key_value:
              output = [output, presents]
  
          return output
  