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)

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.

attach_custom_arguments(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)

Attaches custom arguments to the configuration object.

This method allows adding or overriding configuration attributes dynamically. It iterates through a dictionary of basic configuration parameters and sets them as attributes of the configuration object if they don’t already exist.

Parameters

• vocab_size (int, optional) – Vocabulary size. Defaults to 50400.

• n_positions (int, optional) – Maximum sequence length. Defaults to 2048.

• n_embd (int, optional) – Hidden size. Defaults to 4096.

• n_layer (int, optional) – Number of hidden layers. Defaults to 28.

• n_head (int, optional) – Number of attention heads. Defaults to 16.

• rotary_dim (int, optional) – Dimension for rotary position embeddings. Defaults to 64.

• n_inner (int, optional) – Inner dimension of FFN. Defaults to None.

• activation_function (str, optional) – Activation function. Defaults to “gelu_new”.

• resid_pdrop (float, optional) – Residual dropout probability. Defaults to 0.0.

• embd_pdrop (float, optional) – Embedding dropout probability. Defaults to 0.0.

• attn_pdrop (float, optional) – Attention dropout probability. Defaults to 0.0.

• layer_norm_epsilon (float, optional) – Epsilon for layer normalization. Defaults to 1e-5.

• initializer_range (int, optional) – Initializer range. Defaults to 0.02.

• use_cache (int, optional) – Whether to use KV cache. Defaults to True.

• bos_token_id (int, optional) – Beginning-of-sequence token ID. Defaults to 50256.

• eos_token_id (int, optional) – End-of-sequence token ID. Defaults to 50256.

• tie_word_embeddings (bool, optional) – Whether to tie input/output embeddings. Defaults to False.

• bits (tp.Optional[int], optional) – Quantization bits. Defaults to None.

• gradient_checkpointing (EasyDeLGradientCheckPointers, optional) – Gradient checkpointing strategy. Defaults to EasyDeLGradientCheckPointers.NONE.

• **kwargs – Additional keyword arguments.

Returns

The configuration object itself (self).

Return type

GPTJConfig

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()

Returns the mesh names used for model parallelism. For GPT-J, it returns mesh names for data parallelism (‘dp’), fully sharded data parallelism (‘fsdp’), sequence parallelism (‘sp’), and tensor parallelism (‘tp’).

Returns

A tuple containing the mesh names (“dp”, “fsdp”, “sp”, “sp”).

Return type

tuple

get_partition_rules(*args, **kwargs)

Get the partition rules for the model. This method defines how the model’s parameters are partitioned across devices for distributed training and inference.

Parameters

• *args – Additional positional arguments (unused).

• **kwargs – Additional keyword arguments (unused).

Returns

A tuple of partition rules, where each rule is a tuple

containing a regex pattern for parameter names and the corresponding PartitionSpec.

Return type

tp.Tuple[tp.Tuple[str, PartitionSpec]]

model_type: str = 'gptj'