docstr: Autoregressive Transformer with various bells and whistles. We follow the same
  model-parallel sharding scheme as the one described in https://nv-adlr.github.io/MegatronLM
docstr_tok: ["Autoregressive","Transformer","with","various","bells","and","whistles",".","We","follow","the","same","model-parallel","sharding","scheme","as","the","one","described","in","https",":","\/\/nv-adlr.github.io\/MegatronLM"]
code: def __init__(
      self,
      n_ctx,
      n_vocab,
      d_model=128,
      n_layer=2,
      heads=1,
      attn_dropout=0.0,
      resid_dropout=0.0,
      emb_dropout=0.0,
      zero_out=False,
      init_scale=1.0,
      res_scale=False,
      m_attn=0.25,
      m_mlp=1,
      mp_comm: Comm = None,
      include_pos_embeddings=True,
      include_input_embeddings=True,
      include_output_unembeddings=True,
      # For e.g. reward model training, we want the final layer norm before the extra head
      # but not the output unembeddings, so control this separately.
      include_final_layer_norm=True,
      afn="quick_gelu",
      key_bias=False,
      flatten_multi_index_batch_dims=False,  # allows the first dims to be batch dims and flattens them, outputs will be flattened
      global_idxs_for_resblocks: Optional[List[int]] = None,
      **extra_args,
  ):
      """
      Autoregressive Transformer with various bells and whistles. We follow the same
      model-parallel sharding scheme as the one described in https://nv-adlr.github.io/MegatronLM
      """
      super().__init__()
      self.n_ctx = n_ctx
      self.n_vocab = n_vocab
      self.d_model = d_model
      self.depth = n_layer
      self.include_input_embeddings = include_input_embeddings
      self.include_pos_embeddings = include_pos_embeddings
      self.include_final_layer_norm = include_final_layer_norm
      self.include_output_unembeddings = include_output_unembeddings
      self.mp_comm = (
          mp_comm  # ShardedModel must expose an mp_comm object
      )
      self._global_idxs_for_resblocks = global_idxs_for_resblocks
  
      self.flatten_multi_index_batch_dims = (
          flatten_multi_index_batch_dims
      )
  
      self.is_using_dropout = bool(
          attn_dropout or resid_dropout or emb_dropout
      )
  
      n_shards = 1 if mp_comm is None else mp_comm.size
      assert d_model % n_shards == 0
  
      # Round to 128 for tensorcores
      self.n_vocab_padded = int(np.ceil(n_vocab / 128)) * 128
  
      if include_input_embeddings:
          # No padding needed here, since no matmuls are involved
          # Keeping for backwards-compatibility, but could change this if desired
          input_n_vocab = self.n_vocab_padded
  
          self.embedding = nn.Embedding(
              input_n_vocab, exact_div(d_model, n_shards)
          )
  
          nn.init.normal_(self.embedding.weight, std=0.02 * init_scale)
          self.embed_dropout = nn.Dropout(emb_dropout)
  
          if self.include_pos_embeddings:
              self.position_embedding = PositionEmbedding(
                  n_ctx=n_ctx,
                  d_model=exact_div(d_model, n_shards),
                  init_scale=init_scale,
              )
              self.pos_emb_dropout = nn.Dropout(emb_dropout)
  
      self.torso = TransformerTorso(
          d_model=d_model,
          n_ctx=n_ctx,
          n_head=heads,
          n_depth=n_layer,
          attn_dropout=attn_dropout,
          resid_dropout=resid_dropout,
          afn=afn,
          zero_out=zero_out,
          init_scale=init_scale,
          res_scale=res_scale,
          m_attn=m_attn,
          m_mlp=m_mlp,
          mp_comm=mp_comm,
          key_bias=key_bias,
      )
  
      if include_final_layer_norm:
          self.ln_f = LayerNorm(d_model)
  
      if include_output_unembeddings:
          # Note: We use n_vocab_padded here so that we can use TensorCores for the big matmul and softmax
          x_out_weight = nn.Embedding(
              self.n_vocab_padded, exact_div(d_model, n_shards)
          ).weight
          nn.init.normal_(x_out_weight, std=0.02 * init_scale)
          self.unembedding_weights = x_out_weight
  
imports: import logging
  from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import List, Optional, Dict, TypeVar
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn import LayerNorm
  from summarize_from_feedback.model_layout import ModelLayout
  from summarize_from_feedback.models.attention import Attention
  from summarize_from_feedback.models.ops import ACT_FNS, Conv1D
  from summarize_from_feedback.utils import exact_div
  from summarize_from_feedback.utils.dist_utils import create_model_parallel_comm, Comm
label_desc: Use image node_id 1 to create a new Transformer object from inherited base classes: nn with example: obj = Transformer(n_ctx, n_vocab, d_model, n_layer, heads, attn_dropout, resid_dropout, emb_dropout, zero_out, init_scale, res_scale, m_attn, m_mlp, mp_comm, include_pos_embeddings, include_input_embeddings, include_output_unembeddings, include_final_layer_norm, afn, key_bias, flatten_multi_index_batch_dims, global_idxs_for_resblocks)