code: def __init__(
      self,
      config: GPT2Config,
      process_group=None,
      device=None,
      dtype=None,
  ):
      super().__init__(config)
      factory_kwargs = {"device": device, "dtype": dtype}
      self.process_group = process_group
      self.sequence_parallel = getattr(
          config, "sequence_parallel", True
      )
      assert config.activation_function in [
          "gelu",
          "gelu_new",
          "gelu_fast",
          "gelu_approx",
          "gelu_pytorch_tanh",
          "relu",
          "sqrelu",
          "glu",
          "swiglu",
          "geglu",
      ]
      pad_vocab_size_multiple = getattr(
          config, "pad_vocab_size_multiple", 1
      )
      vocab_size = (
          math.ceil(config.vocab_size / pad_vocab_size_multiple)
          * pad_vocab_size_multiple
      )
      self.embeddings_multiplier = getattr(
          config, "mup_embeddings_multiplier", 1.0
      )
      # TD [2022-07-30]: Force residual in fp32, seems to make fp16 training more stable
      self.residual_in_fp32 = getattr(config, "residual_in_fp32", False)
      # These 2 options are for OPT-350m
      self.prenorm = getattr(config, "prenorm", True)
      use_rms_norm = getattr(config, "rms_norm", False)
      word_embed_proj_dim = getattr(config, "word_embed_proj_dim", None)
      # For GPT-J, GPT-NeoX
      self.parallel_block = getattr(config, "parallel_block", False)
  
      if process_group is None:
          self.embeddings = GPT2Embeddings(
              config.hidden_size,
              vocab_size,
              config.max_position_embeddings,
              word_embed_proj_dim=word_embed_proj_dim,
              **factory_kwargs,
          )
      else:
          self.embeddings = ParallelGPT2Embeddings(
              config.hidden_size,
              vocab_size,
              config.max_position_embeddings,
              process_group=process_group,
              sequence_parallel=self.sequence_parallel,
              **factory_kwargs,
          )
  
      # We change the order of dropout, residual and layer norm:
      # Instead of LN -> Attn / MLP -> Dropout -> Add, we do:
      # Dropout -> Add -> LN -> Attn / MLP, returning both the residual branch (output of Add) and
      # the main branch (output of MLP). The model definition is unchanged, but the mapping of the
      # nn.Dropout probabilities are changed.
      # This is for performance reason: we can fuse dropout + add + layer_norm.
      self.layers = nn.ModuleList(
          [
              create_block(
                  config,
                  layer_idx=i,
                  process_group=process_group,
                  **factory_kwargs,
              )
              for i in range(config.num_hidden_layers)
          ]
      )
      rotary_emb_fraction = getattr(config, "rotary_emb_fraction", 0.0)
      if (
          rotary_emb_fraction > 0.0
      ):  # Tie all the RotaryEmbedding modules to share the same cos/sin cache
          for layer in self.layers[1:]:
              layer.mixer.rotary_emb = self.layers[0].mixer.rotary_emb
  
      self.fused_dropout_add_ln = getattr(
          config, "fused_dropout_add_ln", False
      )
      if self.fused_dropout_add_ln:
          if layer_norm_fn is None:
              raise ImportError("Triton is not installed")
      if self.prenorm:
          self.drop_f = nn.Dropout(config.resid_pdrop)
          norm_cls = nn.LayerNorm if not use_rms_norm else RMSNorm
          self.ln_f = norm_cls(
              config.hidden_size,
              eps=config.layer_norm_epsilon,
              **factory_kwargs,
          )
      if process_group is not None:
          for p in self.ln_f.parameters():
              # Mark the norm parameters as "shared_params" so that we sync their values at init.
              p._shared_params = True
              # Mark the norm params as "sequence_parallel" so we run all-reduce on their grads.
              if self.sequence_parallel:
                  p._sequence_parallel = True
  
      self.apply(
          partial(
              _init_weights,
              n_layer=config.num_hidden_layers,
              initializer_range=config.initializer_range,
              mup_width_scale=getattr(config, "mup_width_scale", 1.0),
          )
      )
      self.tie_weights()
  
imports: import logging
  import math
  import re
  from collections import OrderedDict, namedtuple
  from collections.abc import Sequence
  from functools import partial
  from typing import Dict, List
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from transformers import GPT2Config
  from flash_attn.models.bigcode import remap_state_dict_hf_bigcode
  from flash_attn.models.falcon import remap_state_dict_hf_falcon
  from flash_attn.models.gpt_neox import remap_state_dict_hf_gpt_neox
  from flash_attn.models.gptj import remap_state_dict_hf_gptj
  from flash_attn.models.llama import remap_state_dict_hf_llama
  from flash_attn.models.opt import remap_state_dict_hf_opt
  from flash_attn.modules.block import Block, ParallelBlock
  from flash_attn.modules.embedding import GPT2Embeddings, ParallelGPT2Embeddings
  from flash_attn.modules.mha import MHA, ParallelMHA
  from flash_attn.modules.mlp import FusedMLP, GatedMlp, Mlp, ParallelFusedMLP, ParallelGatedMlp, ParallelMLP
  from flash_attn.ops.activations import sqrelu_fwd
  from flash_attn.utils.distributed import all_gather, all_gather_raw, get_dim_for_local_rank, sync_shared_params
  from flash_attn.utils.generation import GenerationMixin
  from flash_attn.utils.pretrained import state_dict_from_pretrained
label_desc: Use image node_id 1 to create a new GPTModel object from inherited base classes: GPTPreTrainedModel with example: obj = GPTModel(config, process_group, device, dtype)