easydel.modules.mixtral.mixtral_configuration

class easydel.modules.mixtral.mixtral_configuration.MixtralConfig(vocab_size=32000, hidden_size=4096, intermediate_size=14336, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, hidden_act='silu', max_position_embeddings=131072, initializer_range=0.02, rms_norm_eps=1e-05, use_cache=True, pad_token_id=None, bos_token_id=1, eos_token_id=2, tie_word_embeddings=False, rope_theta=1000000.0, sliding_window=4096, attention_dropout=0.0, num_experts_per_tok=2, num_local_experts=8, output_router_logits=False, router_aux_loss_coef=0.001, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, number_rep_kv: int = 1, bits: Optional[int] = None, rope_scaling: Dict[str, Union[str, float]] = None, attention_bias: bool = False, initialization_of_moe: bool = False, router_jitter_noise=0.0, **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 32000) – Vocabulary size of the Mixtral 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 14336) – 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.

• num_key_value_heads (int, optional, defaults to 8) – Number of key and value heads for each attention layer in the Transformer encoder.

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

• max_position_embeddings (int, optional, defaults to 4096 * 32) – 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.

• rms_norm_eps (float, optional, defaults to 1e-5) – The epsilon used by the rms 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.

• pad_token_id (int, optional) – The index of the padding token in the vocabulary.

• bos_token_id (int, optional, defaults to 1) – 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.

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

• sliding_window (int, optional, defaults to 4096) – The sliding window size.

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

• num_experts_per_tok (int, optional, defaults to 2) – The number of experts per token.

• num_local_experts (int, optional, defaults to 8) – The number of local experts.

• output_router_logits (bool, optional, defaults to False) – Whether to output router logits.

• router_aux_loss_coef (float, optional, defaults to 0.001) – The router auxiliary loss coefficient.

• gradient_checkpointing (str, optional, defaults to “nothing_saveable”) – The gradient checkpointing configuration.

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

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

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

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

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

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

• initialization_of_moe (bool, optional, defaults to False) – Whether to initialize the MoE layers.

• router_jitter_noise (float, optional, defaults to 0.0) – The jitter noise for the router.

attach_custom_arguments(gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, number_rep_kv: int = 1, bits: Optional[int] = None, attention_dropout: float = 0.0, rope_scaling: Dict[str, Union[str, float]] = None, attention_bias: bool = False, initialization_of_moe: bool = False, **kwargs)

The attach_custom_arguments function adds the following arguments to the model:

Parameters

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

• use_scan_mlp (bool, optional) – Whether to use scan for MLP layers. Defaults to False.

• scan_mlp_chunk_size (int, optional) – Chunk size for scan MLP. Defaults to 1024.

• number_rep_kv (int, optional) – Number of repetitions for key/value heads. Defaults to 1.

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

• attention_dropout (float, optional) – Dropout probability for attention. Defaults to 0.0.

• rope_scaling (tp.Dict[str, tp.Union[str, float]], optional) – RoPE scaling configuration. Defaults to None.

• attention_bias (bool, optional) – Whether to use bias in attention layers. Defaults to False.

• initialization_of_moe (bool, optional) – Whether MoE layers are being initialized. Defaults to False.

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

Return type

A tuple of the following

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]]

static get_weight_decay_exclusions()

Returns a tuple of parameter names for which weight decay should be excluded.

Returns

An empty tuple, indicating no specific weight decay exclusions for this model.

Return type

tuple

property granted_freq_max_position_embedding: int

Returns the maximum position embedding size specifically for frequency-based position embeddings.

If freq_max_position_embeddings is set, it returns that value. Otherwise, it falls back to max_position_embeddings.

Returns

The granted maximum position embedding size for frequency encoding.

Return type

int

property granted_mask_max_position_embedding: int

Returns the maximum position embedding size specifically for mask-based position embeddings.

If mask_max_position_embeddings is set, it returns that value. Otherwise, it falls back to max_position_embeddings.

Returns

The granted maximum position embedding size for mask encoding.

Return type

int

model_type: str = 'mixtral'

static rng_keys()

Returns the names of the random number generator keys used by the model.

Returns

A tuple containing “params”, “dropout”, and “jitter” as the RNG keys.

Return type

tuple