easydel.modules.llama.llama_configuration

easydel.modules.llama.llama_configuration#

class easydel.modules.llama.llama_configuration.LlamaConfig(vocab_size: int = 32000, hidden_size: int = 4096, intermediate_size: int = 11008, num_hidden_layers: int = 32, num_attention_heads: int = 32, number_rep_kv: int = 1, head_dim: Optional[int] = None, num_key_value_heads: Optional[int] = None, max_position_embeddings: int = 2048, rms_norm_eps: float = 1e-06, initializer_range: float = 0.02, use_cache: bool = True, bos_token_id: int = 0, eos_token_id: int = 1, resid_pdrop: float = 0.0, embd_pdrop: float = 0.0, attention_dropout: float = 0.0, rope_theta: float = 10000.0, attention_bias: bool = False, tie_word_embeddings: bool = False, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, fcm_min_ratio: float = -1, fcm_max_ratio: float = -1, rope_scaling: Dict[str, Union[str, float]] = None, scan_mlp_chunk_size: int = 1024, bits: Optional[int] = None, hidden_act: str = 'silu', pretraining_tp: int = 1, mlp_bias: bool = False, scan_layers: bool = False, **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 32000) – Vocabulary size of the Llama 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 4096) – Dimensionality of the encoder layers and the pooler layer.

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

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

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

  • number_rep_kv (int, optional, defaults to 1) – Number of repetitions for the key and value vectors.

  • num_key_value_heads (int, optional) – Number of key and value heads for each attention layer in the Transformer encoder. Will default to number_rep_kv * num_attention_heads if not set.

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

    head_dim (int, optional):

    head_dim for attention qkv.

  • rms_norm_eps (float, optional, defaults to 1e-6) – The epsilon used by the rms 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 0) – The id of the beginning-of-sequence token.

  • eos_token_id (int, optional, defaults to 1) – The id of the end-of-sequence token.

  • 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.

  • attention_dropout (float, optional, defaults to 0.0) – The dropout ratio for the attention probabilities.

  • rope_theta (float, optional, defaults to 10000.0) – The theta value to use for rotary position embeddings.

  • attention_bias (bool, optional, defaults to False) – Whether to use attention bias.

  • 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 “nothing_saveable”) – The gradient checkpointing configuration.

  • fcm_min_ratio (float, optional, defaults to -1) – The minimum ratio for Flash Attention.

  • fcm_max_ratio (float, optional, defaults to -1) – The maximum ratio for Flash Attention.

  • rope_scaling (tp.Dict[str, tp.Union[str, float]], optional) – The configuration for rope scaling.

  • scan_mlp_chunk_size (int, optional, defaults to 1024) – The chunk size to use when scanning the MLP.

  • bits (int, optional) – The number of bits to quantize the model to.

  • hidden_act (str, optional, defaults to “silu”) – The hidden activation function to use.

  • pretraining_tp (int, optional, defaults to 1) – The tensor parallelism degree used during pretraining.

  • mlp_bias (bool, optional, defaults to False) – Whether to use bias in the MLP.

  • scan_layers (bool, optional, defaults to False) – Whether to use the scan implementation for the layers.

add_jax_args(resid_pdrop: float = 0.0, embd_pdrop: float = 0.0, attention_dropout: float = 0.0, tie_word_embeddings: bool = False, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, fcm_min_ratio: float = 0.0, fcm_max_ratio: float = 0.0, number_rep_kv: int = 1, bits: Optional[int] = None, rope_theta: float = 10000.0, attention_bias: bool = False, hidden_act: str = 'silu', scan_layers: bool = True, **kwargs)[source]#

The add_jax_args function adds the following arguments to the Transformer class:

Parameters

  • self – Refer to the current object

  • resid_pdrop – float: Set the dropout rate for residual connections

  • embd_pdrop – float: Set the probability of dropping an embedding

  • attention_dropout – float: Set the probability of dropping out the attention layer

  • tie_word_embeddings – bool: Tie the word embeddings to the decoder

  • gradient_checkpointing – str: Control the amount of memory used by jax

  • fcm_min_ratio – float: Control the minimum ratio of the number of chunks to be used in flash-based computation

  • fcm_max_ratio – float: Set the maximum ratio of the number of input tokens to output tokens

  • number_rep_kv – int: Determine how many times the key and value vectors are repeated

  • bits – tp.Optional[int]: Determine the number of bits used in the quantization

  • rope_theta – float : rope_theta for compute rope

  • attention_bias – bool : whenever to use attention bias or no

  • hidden_act – str : hidden_act for mlp

  • scan_layers – bool: Determine whether to use scan layers or not

get_partition_rules(*args, **kwargs)[source]#

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

static get_weight_decay_exclusions()[source]#
property granted_freq_max_position_embedding: int#
property granted_mask_max_position_embedding: int#
model_type: str = 'llama'#
static rng_keys()[source]#
class easydel.modules.llama.llama_configuration.VisionLlamaConfig(vision_vocab_size=8448, tie_vision_embeddings=False, sample_mode='all', **kwargs)[source]#

Bases: LlamaConfig

get_partition_rules(*args, **kwargs)[source]#

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

property granted_freq_max_position_embedding: int#
property granted_mask_max_position_embedding: int#
Contents