easydel.modules.qwen2_moe.__init__

class easydel.modules.qwen2_moe.__init__.Qwen2MoeConfig(vocab_size=151936, hidden_size=2048, intermediate_size=5632, num_hidden_layers=24, num_attention_heads=16, num_key_value_heads=16, hidden_act='silu', max_position_embeddings=32768, initializer_range=0.02, rms_norm_eps=1e-06, use_cache=True, tie_word_embeddings=False, qkv_bias=False, rope_theta=10000.0, use_sliding_window=False, sliding_window=4096, max_window_layers=28, attention_dropout=0.0, decoder_sparse_step=1, moe_intermediate_size=1408, shared_expert_intermediate_size=5632, num_experts_per_tok=4, num_experts=60, norm_topk_prob=False, output_router_logits=False, router_aux_loss_coef=0.001, mlp_only_layers=None, 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 151936) – Vocabulary size of the Qwen-2 MoE 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 2048) – Dimensionality of the encoder layers and the pooler layer.

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

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

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

• num_key_value_heads (int, optional, defaults to 16) – 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 32768) – 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-6) – 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.

• 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 10000.0) – The theta value to use for rotary position embeddings.

• use_sliding_window (bool, optional, defaults to False) – Whether to use a sliding window attention.

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

• max_window_layers (int, optional, defaults to 28) – The maximum number of layers to use for the sliding window attention.

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

• decoder_sparse_step (int, optional, defaults to 1) – The sparse step for the decoder.

• moe_intermediate_size (int, optional, defaults to 1408) – The intermediate size of the MoE layer.

• shared_expert_intermediate_size (int, optional, defaults to 5632) – The intermediate size of the shared expert.

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

• num_experts (int, optional, defaults to 60) – The number of experts.

• norm_topk_prob (bool, optional, defaults to False) – Whether to normalize the top-k probabilities.

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

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

• mlp_only_layers (list of int, optional) – The layers that should only contain an MLP.

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

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

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

static get_weight_decay_exclusions()

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

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

static rng_keys()

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

class easydel.modules.qwen2_moe.__init__.Qwen2MoeForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

Qwen2 MoE model with a Causal Language Modeling (CLM) head.

This class wraps the base Qwen2MoeModel and adds a linear layer (language model head) to predict the next token logits.

config

Configuration object for the model.

Type

Qwen2MoeConfig

model

The base Qwen2 MoE model.

Type

Qwen2MoeModel

lm_head

The language model head (linear layer).

Type

ParallelLinear

dtype

Data type for computations.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for matrix multiplications.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs

class easydel.modules.qwen2_moe.__init__.Qwen2MoeForSequenceClassification(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

Qwen2 MoE model with a sequence classification head.

This class wraps the base Qwen2MoeModel and adds a linear layer on top to perform sequence classification tasks.

config

Configuration object for the model.

Type

Qwen2MoeConfig

model

The base Qwen2 MoE model.

Type

Qwen2MoeModel

score

The sequence classification head (linear layer).

Type

ParallelLinear

dtype

Data type for computations.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for matrix multiplications.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs

class easydel.modules.qwen2_moe.__init__.Qwen2MoeModel(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

The base Qwen2 MoE transformer model.

This class implements the core transformer architecture, including embedding layers, decoder layers, and final normalization.

config

Configuration object for the model.

Type

Qwen2MoeConfig

embed_tokens

Embedding layer for input tokens.

Type

nn.Embed

layers

List of decoder layers.

Type

nn.List[Qwen2MoeDecoderLayer]

norm

Final layer normalization.

Type

RMSNorm

gradient_checkpointing

Gradient checkpointing strategy.

Type

str

dtype

Data type for computations.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for matrix multiplications.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs