class_docstr: A single transformer layer.
  
  Transformer layer takes input with size [s, b, h] and returns an
  output of the same size.
class_docstr_tok: ["A","single","transformer","layer",".","Transformer","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","an","output","of","the","same","size","."]
code: def __init__(
      self,
      config: TransformerConfig,
      submodules: TransformerLayerSubmodules,
      layer_number: int = 1,
      hidden_dropout: float = None,
  ):
      super().__init__(config=config)
  
      self.layer_number = layer_number + self._get_layer_offset()
      self.hidden_dropout = (
          config.hidden_dropout
          if hidden_dropout is None
          else hidden_dropout
      )
  
      ## [Module 1: Input Layernorm] Optional Layernorm on the input data
      # TODO: add pytorch only layernorm
      self.input_layernorm = build_module(
          submodules.input_layernorm,
          config=self.config,
          hidden_size=self.config.hidden_size,
          eps=self.config.layernorm_epsilon,
      )
  
      ## [Module 2: SelfAttention]
      self.self_attention = build_module(
          submodules.self_attention,
          config=self.config,
          layer_number=layer_number,
      )
  
      ## [Module 3: BiasDropoutFusion]
      self.self_attn_bda = build_module(submodules.self_attn_bda)
  
      ## [Module 4: Post SelfAttention] Optional Layernorm after self-attn
      self.pre_cross_attn_layernorm = build_module(
          submodules.pre_cross_attn_layernorm,
          config=self.config,
          hidden_size=self.config.hidden_size,
          eps=self.config.layernorm_epsilon,
      )
  
      ## [Module 5: CrossAttention]
      self.cross_attention = build_module(
          submodules.cross_attention,
          config=self.config,
          layer_number=layer_number,
      )
  
      ## [Module 6: BiasDropoutFusion]
      self.cross_attn_bda = build_module(
          submodules.cross_attn_bda,
          config=self.config,
      )
  
      ## [Module 7: Pre MLP] Optional Layernorm before MLP
      self.pre_mlp_layernorm = build_module(
          submodules.pre_mlp_layernorm,
          config=self.config,
          hidden_size=self.config.hidden_size,
          eps=self.config.layernorm_epsilon,
      )
  
      ## [Module 8: MLP block]
      # TODO how to set the gpt_layer_spec.py when we have moe_frequency > 1,
      #      where MLP and SwitchMLP both appear alternately?
      self.mlp = build_module(submodules.mlp, config=self.config)
  
      ## [Module 9: BiasDropoutFusion]
      self.mlp_bda = build_module(submodules.mlp_bda)
  
      # @jcasper how should we handle nvfuser?
      # Set bias+dropout+add fusion grad_enable execution handler.
      # TORCH_MAJOR = int(torch.__version__.split('.')[0])
      # TORCH_MINOR = int(torch.__version__.split('.')[1])
      # use_nvfuser = TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10)
      # self.bias_dropout_add_exec_handler = nullcontext if use_nvfuser else torch.enable_grad
      self.bias_dropout_add_exec_handler = torch.enable_grad
  
file_path: megatron-lm/megatron/core/transformer/transformer_layer.py
imports: from dataclasses import dataclass
  from typing import Union
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedTensor
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import make_viewless_tensor
label_desc: Use image node_id 1 to create a new TransformerLayer object from inherited base classes: MegatronModule with example: obj = TransformerLayer(config, submodules, layer_number, hidden_dropout)