easydel.modules.gpt_neox.gpt_neox_configuration

easydel.modules.gpt_neox.gpt_neox_configuration#

class easydel.modules.gpt_neox.gpt_neox_configuration.GPTNeoXConfig(vocab_size=50432, hidden_size=6144, num_hidden_layers=44, num_attention_heads=64, intermediate_size=24576, hidden_act='gelu', rotary_pct=0.25, rotary_emb_base=10000, attention_dropout=0.0, hidden_dropout=0.0, classifier_dropout=0.1, max_position_embeddings=2048, initializer_range=0.02, layer_norm_eps=1e-05, use_cache=True, bos_token_id=0, eos_token_id=2, tie_word_embeddings=False, use_parallel_residual=True, rope_scaling=None, attention_bias=True, gradient_checkpointing=EasyDeLGradientCheckPointers.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 50432) – Vocabulary size of the GPT NeoX model. Defines the number of different tokens that can be represented by the inputs_ids passed to the forward method.

  • hidden_size (int, optional, defaults to 6144) – Dimensionality of the encoder layers and the pooler layer.

  • num_hidden_layers (int, optional, defaults to 44) – Number of hidden layers in the Transformer encoder.

  • num_attention_heads (int, optional, defaults to 64) – Number of attention heads for each attention layer in the Transformer encoder.

  • intermediate_size (int, optional, defaults to 24576) – Dimensionality of the “intermediate” (i.e., feed-forward) layer in the Transformer encoder.

  • hidden_act (str or function, optional, defaults to “gelu”) – 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.

  • rotary_pct (float, optional, defaults to 0.25) – The percentage of hidden dimensions to allocate to rotary embeddings.

  • rotary_emb_base (int, optional, defaults to 10000) – The base for the rotary position embedding.

  • classifier_dropout (float, optional, defaults to 0.1) – The dropout ratio for the classifier layer.

  • max_position_embeddings (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).

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

  • layer_norm_eps (float, optional, defaults to 1e-5) – The epsilon used by the layer normalization layers.

  • 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 0) – The id of the beginning-of-sequence token.

  • eos_token_id (int, optional, defaults to 2) – 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 “everything_saveable”) – The gradient checkpointing configuration.

  • use_parallel_residual (bool, optional, defaults to True) – Whether to use a parallel residual connection in the attention layer.

add_jax_args(**kwargs)[source]#
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 = 'gpt_neox'#
Contents