Skip to content

Qwen3 MoE

About

The d9d.module.model.qwen3_moe package implements the Qwen3 Mixture-of-Experts model architecture.

The d9d.module.parallelism.model.qwen3_moe package implements default horizontal parallelism strategies for this model.

HuggingFace Compatibility

d9d provides out-of-the-box support for streaming and converting HuggingFace checkpoints into the optimized d9d runtime format (and vice versa).

These operations utilize the graph-based State Mapping engine. You may use the model state mappers provided for the Model Provider implementation.

d9d.module.model.qwen3_moe

Qwen3MoEExpertsFormat

Bases: StrEnum

Specifies the underlying layout of the experts parameters in Hugging Face.

Attributes:

Name Type Description
MODULE_LIST

Legacy (v4.x) nn.ModuleList with individual down, up, and gate nn.Linear layers.
FUSED

New (v5.x) format with 3D expert representations fused into single tensors.

Qwen3MoELayer

Bases: Module, ModuleLateInit

Implements a single Qwen3 Mixture-of-Experts (MoE) transformer layer.

This layer consists of a Grouped Query Attention mechanism followed by an MoE MLP block, with pre-RMSNorm applied before each sub-layer.

__init__(params)

Constructs a Qwen3MoELayer object.

Parameters:

Name Type Description Default
params Qwen3MoELayerParameters

Configuration parameters for the layer.

 required

forward(hidden_states, position_embeddings)

Performs the forward pass of the MoE layer.

Parameters:

Name Type Description Default
hidden_states Tensor

Input tensor of shape (batch, seq_len, hidden_dim).

 required
position_embeddings tuple[Tensor, Tensor]

Tuple containing RoPE precomputed embeddings (cos, sin).

 required

Returns:

Type Description
Tensor

Output tensor after attention and MoE blocks, shape (batch, seq_len, hidden_dim).

reset_parameters()

Resets module parameters.

Qwen3MoELayerParameters

Bases: BaseModel

Configuration parameters for a single Qwen3 MoE layer.

Attributes:

Name Type Description
hidden_size int

Dimension of the model's hidden states.
intermediate_size int

Dimension of the feed-forward hidden state.
num_experts int

Total number of experts in the MoE layer.
experts_top_k int

Number of experts to route tokens to.
num_attention_heads int

Number of attention heads for the query.
num_key_value_heads int

Number of attention heads for key and value.
rms_norm_eps float

Epsilon value found in the RMSNorm layers.
head_dim int

Dimension of a single attention head.

Qwen3MoEModel

Bases: Module, ModuleLateInit, ModuleSupportsPipelining[SequenceInput, SequenceTransfer[Tensor], SequenceShared, SequenceTransfer[Tensor]]

The Qwen3 Mixture-of-Experts (MoE) Transformer Decoder backbone.

It is designed to be split across multiple pipeline stages.

hidden_size property

Dimensionality of the backbone hidden states.

split_vocab_order property

The order in which vocabulary segments are concatenated.

split_vocab_size property

Mapping of vocabulary segment names to their sizes.

__init__(params, stage, hidden_states_snapshot_mode, enable_checkpointing)

Constructs the Qwen3MoEModel object.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Configuration parameters for the full model.

 required
stage PipelineStageInfo

Information about the pipeline stage this instance belongs to.

 required
hidden_states_snapshot_mode HiddenStatesAggregationMode

Configures intermediate hidden state aggregation & snapshotting mode

 required
enable_checkpointing bool

If True, enables activation checkpointing for transformer layers to save memory.

 required

forward(inputs, shared)

Executes the backbone forward pass for the current pipeline stage.

Parameters:

Name Type Description Default
inputs SequenceInput | SequenceTransfer[Tensor]

SequenceInput (token ids) on the first stage; the incoming SequenceTransfer otherwise.

 required
shared SequenceShared

The shared input broadcast to every stage (position ids and, if snapshotting is enabled, the aggregation mask).

 required

Returns:

Type Description
SequenceTransfer[Tensor]

The produced SequenceTransfer (hidden states and, optionally, the updated snapshot).

output_dtype()

Returns the data type of the model output hidden states.

Returns:

Type Description
dtype

The output hidden states data type.

reset_parameters()

Resets module parameters.

stage_transfer_spec(pipeline_input, boundary)

Describes the SequenceTransfer crossing the given boundary of this stage.

Parameters:

Name Type Description Default
pipeline_input SequenceInput

A representative SequenceInput microbatch; only shapes are read.

 required
boundary StageBoundary

Which inter-stage edge to describe.

 required

Returns:

Type Description
SequenceTransfer[TensorSpec]

A SequenceTransfer of TensorSpec.

Qwen3MoEParameters

Bases: BaseModel

Configuration parameters for the Qwen3 Mixture-of-Experts model backbone.

Attributes:

Name Type Description
layer Qwen3MoELayerParameters

Configuration shared across all transformer layers.
num_hidden_layers int

The total number of transformer layers.
rope_base int

Base value for RoPE frequency calculation.
max_position_ids int

Maximum sequence length.
split_vocab_size dict[str, int]

A dictionary mapping vocabulary segment names to their sizes.
split_vocab_order list[str]

The sequence in which vocabulary splits are correctly ordered.
pipeline_num_virtual_layers_pre int

The number of 'virtual' layers representing the computational cost of modules on the first stage, before the main layers (e.g., token and positional embeddings).
pipeline_num_virtual_layers_post int

The number of 'virtual' layers representing the computational cost of modules on the last stage, after the main layers (e.g., the final layer normalization and LM head).

mapper_from_huggingface_qwen3_moe(params, experts_format)

Creates a state mapper translating base Qwen3 MoE HuggingFace keys into the d9d format.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_from_huggingface_qwen3_moe_for_causal_lm(params, experts_format, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 MoE Causal LM HuggingFace keys into the d9d format.

HuggingFace models carry exactly one head, so the mapper needs to know where in the composed model it lands. The default targets a single-head decoder; loading the same checkpoint into a multi-head model is a matter of passing that head's prefix (f"heads.{name}.") instead, and the remaining heads keep their initialization.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required
head_prefix str

FQN prefix of the head that receives the HuggingFace head in the target d9d model.

 SINGLE_HEAD_PREFIX

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_from_huggingface_qwen3_moe_for_classification(params, experts_format, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 MoE classification HuggingFace keys into the d9d format.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required
head_prefix str

FQN prefix of the head that receives the HuggingFace head in the target d9d model.

 SINGLE_HEAD_PREFIX

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_from_huggingface_qwen3_moe_for_embedding(params, experts_format)

Creates a state mapper translating Qwen3 MoE embedding HuggingFace keys into the d9d format.

The HuggingFace reference for an embedding model is the bare backbone with no head weights, so no head is named here: the embedding head has nothing to load and keeps its initialization.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_to_huggingface_qwen3_moe(params, experts_format)

Creates a state mapper translating base Qwen3 MoE d9d keys back into the HuggingFace format.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_to_huggingface_qwen3_moe_for_causal_lm(params, experts_format, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 MoE Causal LM d9d keys back into the HuggingFace format.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required
head_prefix str

FQN prefix of the head holding the causal LM weights in the source d9d model.

 SINGLE_HEAD_PREFIX

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_to_huggingface_qwen3_moe_for_classification(params, experts_format, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 MoE classification d9d keys back into the HuggingFace format.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required
head_prefix str

FQN prefix of the head holding the classification weights in the source d9d model.

 SINGLE_HEAD_PREFIX

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_to_huggingface_qwen3_moe_for_embedding(params, experts_format, *, embedding_dim=None)

Creates a state mapper translating Qwen3 MoE embedding d9d keys back into the HuggingFace format.

The HuggingFace reference for an embedding model is the bare backbone with no head weights, so no head is named here and a trained projection has nowhere to go.

Parameters:

Name Type Description Default
params Qwen3MoEParameters

Base model parameters.

 required
experts_format Qwen3MoEExpertsFormat

Format of the MoE experts storage.

 required
embedding_dim int | None

The embedding head's projection dimensionality, or None if it has no projection.

 None

Returns:

Type Description
ModelStateMapper

A composite state mapper.

Raises:

Type Description
ValueError

If the head has a trained embedding projection, which has no HuggingFace counterpart.

d9d.module.parallelism.model.qwen3_moe

parallelize_qwen3_moe_model(dist_context, model, stage)

Parallelizes the base Qwen3 MoE model components.

This function configures the model layers for distributed execution within a pipeline stage. It applies Hybrid Sharded Data Parallelism (HSDP) to dense components (embeddings, norms, attention) and Expert Parallelism (EP) to the Mixture-of-Experts (MLP) layers.

Current usage constraints: * Tensor Parallelism is not supported (we may implement it later). * Context Parallelism is not supported (we will implement it later).

Parameters:

Name Type Description Default
dist_context DistributedContext

The distributed context.

 required
model Qwen3MoEModel

The Qwen3 MoE base model to parallelize.

 required
stage PipelineStageInfo

Information about the current pipeline stage.

 required

Raises:

Type Description
ValueError

If Tensor Parallel or Context Parallel is enabled in the context.