easydel.modules.phimoe.__init__

class easydel.modules.phimoe.__init__.PhiMoeConfig(vocab_size=32064, hidden_size=4096, intermediate_size=6400, 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, rope_scaling=None, sliding_window=None, attention_dropout=0.0, num_experts_per_tok=2, num_local_experts=16, output_router_logits=False, router_aux_loss_coef=0.001, router_jitter_noise=0.01, input_jitter_noise=0.0, attention_bias=False, embd_pdrop: float = 0.0, lm_head_bias=False, bits: Optional[int] = None, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.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 32064) – Vocabulary size of the PhiMoE model. Defines the number of different tokens that can be represented by the inputs_ids passed when calling [PhiMoEModel]

• hidden_size (int, optional, defaults to 4096) – Dimension of the hidden representations.

• intermediate_size (int, optional, defaults to 6400) – Dimension of the MLP representations.

• 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) – This is the number of key_value heads that should be used to implement Grouped Query Attention. If num_key_value_heads=num_attention_heads, the model will use Multi Head Attention (MHA), if num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details checkout [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8.

• hidden_act (str or function, optional, defaults to “silu”) – The non-linear activation function (function or string) in the decoder.

• max_position_embeddings (int, optional, defaults to 4096*32) – The maximum sequence length that this model might ever be used with. Mixtral’s sliding window attention allows sequence of up to 4096*32 tokens.

• 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-05) – 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 id of the padding token.

• 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 the model’s input and output word embeddings should be tied.

• rope_theta (float, optional, defaults to 10000.0) – The base period of the RoPE embeddings.

• rope_scaling (dict, optional) – The scaling strategy for the RoPE embeddings. If None, no scaling is applied. If a dictionary, it must contain the following keys: type, short_factor, long_factor, short_mscale, long_mscale and original_max_position_embeddings. The type must be longrope, the short_mscale and long_scale must be numbers, the short_factor and long_factor must be lists of numbers with the same length as half of the attention head size and the original_max_position_embeddings must be an integer.

• sliding_window (int, optional) – Sliding window attention window size. If not specified, will default to 262144.

• 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 to root per-token, can be also interpreted as the top-p routing parameter

• num_local_experts (int, optional, defaults to 16) – Number of experts per Sparse MLP layer.

• output_router_logits (bool, optional, defaults to False) – Whether or not the router logits should be returned by the model. Enabeling this will also allow the model to output the auxiliary loss. See [here]() for more details

• router_aux_loss_coef (float, optional, defaults to 0.0) – The aux loss factor for the total loss.

• router_jitter_noise (float, optional, defaults to 0.01) – Amount of noise to add to the router.

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

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

attach_custom_arguments(bits: Optional[int] = None, embd_pdrop: float = 0.0, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, **kwargs)

Attaches custom arguments to the configuration object.

This method allows dynamically adding or overriding configuration attributes. It primarily sets attributes related to quantization, dropout, and gradient checkpointing. Any additional keyword arguments are also set as attributes if they don’t already exist.

Parameters

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

• embd_pdrop (float, optional) – Dropout probability for embeddings. Defaults to 0.0.

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

• **kwargs – Additional keyword arguments to attach.

get_partition_rules(fully_sharded_data_parallel: bool = True)

Get the partition rules for the model.

Parameters

fully_sharded_data_parallel (bool, optional, defaults to True) – Whether to use fully sharded data parallelism.

Returns

The partition rules.

Return type

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

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 = 'phimoe'

class easydel.modules.phimoe.__init__.PhiMoeForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

PhiMoE model with a Causal Language Modeling head.

This model consists of the base PhiMoE transformer (PhiMoeModel) followed by a linear layer (lm_head) that projects the transformer’s output hidden states to the vocabulary size, producing logits for next token prediction. Optionally, the input token embeddings can be tied to the output projection layer.

config

Configuration object for the model.

Type

PhiMoeConfig

dtype

Data type for computation.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for JAX operations.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs

model

The core PhiMoE transformer model.

Type

PhiMoeModel

lm_head

The linear layer for projecting hidden states to vocabulary logits.

Type

ParallelLinear

class easydel.modules.phimoe.__init__.PhiMoeModel(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

The base PhiMoE model transformer.

This class represents the core transformer architecture of the PhiMoE model, consisting of an embedding layer, multiple PhiMoeDecoderLayer layers (which include Sparse Mixture of Experts blocks), and a final RMS normalization layer.

config

Configuration object for the model.

Type

PhiMoeConfig

dtype

Data type for computation.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for JAX operations.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs

embed_tokens

Embedding layer for input tokens.

Type

nn.Embed

layers

List of decoder layers.

Type

tp.List[PhiMoeDecoderLayer]

norm

Final layer normalization.

Type

RMSNorm

embed_dropout

Dropout layer applied after embeddings.

Type

nn.Dropout

gradient_checkpointing

Gradient checkpointing configuration.

Type

EasyDeLGradientCheckPointers