easydel.modules.gpt_j.gpt_j_configuration#

class easydel.modules.gpt_j.gpt_j_configuration.GPTJConfig(vocab_size: int = 50400, n_positions: int = 2048, n_embd: int = 4096, n_layer: int = 28, n_head: int = 16, rotary_dim: int = 64, n_inner: int = None, activation_function: str = 'gelu_new', resid_pdrop: float = 0.0, embd_pdrop: float = 0.0, attn_pdrop: float = 0.0, layer_norm_epsilon: float = 1e-05, initializer_range: int = 0.02, use_cache: int = True, bos_token_id: int = 50256, eos_token_id: int = 50256, tie_word_embeddings: bool = False, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, **kwargs)[source]#

Bases: EasyDeLBaseConfig

Configuration objects inherit from [EasyDeLBaseConfig] and can be used to control the model outputs. Read the documentation from [EasyDeLBaseConfig] for more information.

Parameters
  • vocab_size (int, optional, defaults to 50400) โ€“ Vocabulary size of the GPT-J model. Defines the number of different tokens that can be represented by the inputs_ids passed to the forward method.

  • n_positions (int, optional, defaults to 2048) โ€“ The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 2048 or 4096).

  • n_embd (int, optional, defaults to 4096) โ€“ Dimensionality of the encoder layers and the pooler layer.

  • n_layer (int, optional, defaults to 28) โ€“ Number of hidden layers in the Transformer encoder.

  • n_head (int, optional, defaults to 16) โ€“ Number of attention heads for each attention layer in the Transformer encoder.

  • rotary_dim (int, optional, defaults to 64) โ€“ The dimension of the rotary position embedding.

  • n_inner (int, optional) โ€“ Dimensionality of the inner feed-forward layers.

  • activation_function (str, optional, defaults to โ€œgelu_newโ€) โ€“ The non-linear activation function (function or string) to use in the encoder and pooler. If string, โ€œgeluโ€, โ€œreluโ€, โ€œswishโ€ and โ€œgelu_newโ€ are supported.

  • resid_pdrop (float, optional, defaults to 0.0) โ€“ The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.

  • embd_pdrop (float, optional, defaults to 0.0) โ€“ The dropout ratio for the embeddings.

  • attn_pdrop (float, optional, defaults to 0.0) โ€“ The dropout ratio for the attention probabilities.

  • layer_norm_epsilon (float, optional, defaults to 1e-5) โ€“ The epsilon to use in the layer normalization layers.

  • initializer_range (float, optional, defaults to 0.02) โ€“ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.

  • use_cache (bool, optional, defaults to True) โ€“ Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if config.is_decoder=True.

  • bos_token_id (int, optional, defaults to 50256) โ€“ The id of the beginning-of-sequence token.

  • eos_token_id (int, optional, defaults to 50256) โ€“ The id of the end-of-sequence token.

  • tie_word_embeddings (bool, optional, defaults to False) โ€“ Whether to tie the weights of the input embeddings and the output embeddings.

  • gradient_checkpointing (str, optional, defaults to โ€œโ€) โ€“ The gradient checkpointing configuration.

  • bits (int, optional) โ€“ The number of bits to quantize the model to.

add_jax_args(vocab_size: int = 50400, n_positions: int = 2048, n_embd: int = 4096, n_layer: int = 28, n_head: int = 16, rotary_dim: int = 64, n_inner: int = None, activation_function: str = 'gelu_new', resid_pdrop: float = 0.0, embd_pdrop: float = 0.0, attn_pdrop: float = 0.0, layer_norm_epsilon: float = 1e-05, initializer_range: int = 0.02, use_cache: int = True, bos_token_id: int = 50256, eos_token_id: int = 50256, tie_word_embeddings: bool = False, bits: Optional[int] = None, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, **kwargs)[source]#
attribute_map: Dict[str, str] = {'hidden_size': 'n_embd', 'max_position_embeddings': 'n_positions', 'num_attention_heads': 'n_head', 'num_hidden_layers': 'n_layer'}#
static get_mesh_names()[source]#
get_partition_rules(*args, **kwargs)[source]#

Get the partition rules for the model. :returns: The partition rules. :rtype: tp.Tuple[tp.Tuple[str, PartitionSpec]]

model_type: str = 'gptj'#