easydel.modules.qwen3_moe.modeling_qwen3_moe_flax#
- class easydel.modules.qwen3_moe.modeling_qwen3_moe_flax.Qwen3MoeAttention(*args: Any, **kwargs: Any)[source]#
Bases:
AttentionModuleQwen3Moe Attention module.
This module implements the multi-head attention mechanism used in the Qwen3Moe model. It supports Grouped Query Attention (GQA) and Rotary Position Embeddings (RoPE).
- config#
Configuration object for the model.
- Type
- dtype#
Data type for computations.
- 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
Dimensionality of the hidden states.
- Type
int
- head_dim#
Dimensionality of each attention head.
- Type
int
- num_key_value_groups#
Number of query head groups for each key/value head.
- Type
int
- q_proj#
Linear layer for query projection.
- Type
- k_proj#
Linear layer for key projection.
- Type
- v_proj#
Linear layer for value projection.
- Type
- o_proj#
Linear layer for the output projection.
- Type
- attention_performer#
Module to perform the core attention computation.
- rotary#
Rotary position embedding module.
- Type
RoPE
- class easydel.modules.qwen3_moe.modeling_qwen3_moe_flax.Qwen3MoeDecoderLayer(*args: Any, **kwargs: Any)[source]#
Bases:
ModuleQwen3Moe Transformer Decoder Layer.
This module represents a single decoder layer in the Qwen3Moe model, combining self-attention and MLP sub-layers with residual connections and RMS normalization.
- config#
Configuration object for the model.
- Type
- layer_idx#
The index of the layer in the model.
- Type
int
- dtype#
Data type for computations.
- 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
- self_attn#
The self-attention module.
- Type
- mlp#
The feed-forward (MLP) module.
- Type
- class easydel.modules.qwen3_moe.modeling_qwen3_moe_flax.Qwen3MoeForCausalLM(*args: Any, **kwargs: Any)[source]#
Bases:
EasyDeLBaseModuleQwen3Moe model with a Causal Language Modeling head.
This model consists of the base Qwen3Moe transformer (Qwen3MoeModel) 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
- 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 Qwen3Moe transformer model.
- Type
- lm_head#
The linear layer for projecting hidden states to vocabulary logits.
- Type
- class easydel.modules.qwen3_moe.modeling_qwen3_moe_flax.Qwen3MoeForSequenceClassification(*args: Any, **kwargs: Any)[source]#
Bases:
EasyDeLBaseModuleQwen3Moe model with a Sequence Classification head.
This model consists of the base Qwen3Moe transformer (Qwen3MoeModel) followed by a linear layer (score) that projects the transformer’s output hidden states (typically the hidden state of the last token or a pooled representation) to the number of classes for classification.
- config#
Configuration object for the model.
- Type
- 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 Qwen3Moe transformer model.
- Type
- score#
The linear layer for classification.
- Type
- class easydel.modules.qwen3_moe.modeling_qwen3_moe_flax.Qwen3MoeMLP(*args: Any, **kwargs: Any)[source]#
Bases:
ModuleQwen3Moe MLP module.
This module implements the feed-forward network (MLP) used in the Qwen3Moe model. It uses a Gated Linear Unit (GLU) structure with SiLU activation.
- config#
Configuration object for the model.
- Type
- dtype#
Data type for computations.
- Type
jnp.dtype
- param_dtype#
Data type for parameters.
- Type
jnp.dtype
- precision#
Precision setting for JAX operations.
- Type
jax.lax.PrecisionLike
- gate_proj#
Linear layer for the GLU gate.
- Type
- down_proj#
Linear layer for the down projection.
- Type
- up_proj#
Linear layer for the GLU value.
- Type
- act_fn#
Activation function (SiLU).
- Type
callable
- class easydel.modules.qwen3_moe.modeling_qwen3_moe_flax.Qwen3MoeModel(*args: Any, **kwargs: Any)[source]#
Bases:
EasyDeLBaseModuleThe base Qwen3Moe model transformer.
This class represents the core transformer architecture of the Qwen3Moe model, consisting of an embedding layer, multiple Qwen3MoeDecoderLayer layers, and a final RMS normalization layer.
- config#
Configuration object for the model.
- Type
- 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[Qwen3MoeDecoderLayer]
- gradient_checkpointing#
Gradient checkpointing configuration.
- class easydel.modules.qwen3_moe.modeling_qwen3_moe_flax.Qwen3MoeSparseMoeBlock(*args: Any, **kwargs: Any)[source]#
Bases:
ModuleSparse Mixture of Experts (MoE) block for Qwen3 MoE.
This block routes input hidden states to a selected subset of experts and combines their outputs.
- config#
Configuration object for the model.
- Type
- gate#
Linear layer for the gating network.
- Type
- experts#
List of expert MLP modules.
- Type
nn.List[Qwen3MoeMLP]
- 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