Qwen3 MoE

About

The d9d.module.model.qwen3_moe package implements the Qwen3 Mixture-of-Experts (MoE) backbone. The d9d.module.parallelism.model.qwen3_moe package applies the default horizontal parallelism to it: HSDP for the dense modules and expert parallelism for the MoE layers. Tensor parallelism and context parallelism are not supported yet.

Hugging Face Compatibility

The package provides state mappers that convert Hugging Face checkpoints to the d9d layout and back. They use the graph-based State Mapping engine, so you can pass them to your Model Provider.

There is one mapper pair for the bare backbone and one for each head type (causal LM, classification and embedding). Pass Qwen3MoEExpertsFormat to choose the expert layout of the checkpoint: MODULE_LIST for Transformers v4.x or FUSED for Transformers v5.x.

An embedding model cannot be exported to Hugging Face if its embedding head has a projection.

Usage

from d9d.module.block.hidden_states_aggregator import HiddenStatesAggregationMode
from d9d.module.model import DecoderForCausalLM
from d9d.module.model.qwen3_moe import (
    Qwen3MoEExpertsFormat,
    Qwen3MoEModel,
    mapper_from_huggingface_qwen3_moe_for_causal_lm,
)
from d9d.module.parallelism.model import parallelize_causal_lm_head, parallelize_qwen3_moe_model

# In ModelProvider.initialize_model_stage.
backbone = Qwen3MoEModel(
    params=params,
    stage=context.stage,
    hidden_states_snapshot_mode=HiddenStatesAggregationMode.no,
    enable_checkpointing=True,
)
model = DecoderForCausalLM(backbone, context.stage)
state_mapper = mapper_from_huggingface_qwen3_moe_for_causal_lm(params, Qwen3MoEExpertsFormat.FUSED)

# In ModelProvider.parallelize_model_stage.
parallelize_qwen3_moe_model(context.dist_context, context.model.model, context.stage)
if context.stage.is_current_stage_last:
    parallelize_causal_lm_head(context.model.head, context.dist_context)

API Reference

`d9d.module.model.qwen3_moe`

Qwen3 MoE model and its Hugging Face state mappers.

`Qwen3MoEExpertsFormat`

Bases: StrEnum

Layout of the expert parameters in a Hugging Face Transformers checkpoint.

Attributes:

Name	Type	Description
`MODULE_LIST`		The Transformers v4.x layout: an `nn.ModuleList` of experts, each with separate `gate_proj`, `up_proj` and `down_proj` `nn.Linear` layers.
`FUSED`		The Transformers v5.x layout: 3D tensors that stack all experts, with `gate_proj` and `up_proj` fused into `gate_up_proj`.

`Qwen3MoELayer`

Bases: Module, ModuleLateInit

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

The layer applies Grouped Query Attention, then an MoE MLP block. Each sub-layer has a pre-RMSNorm and a residual connection.

`init(params)`

Constructs the Qwen3MoELayer object.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoELayerParameters`	Configuration parameters for the layer.	required

`forward(hidden_states, position_embeddings)`

Runs the forward pass of the layer.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input hidden states. Shape: `(batch, seq_len, hidden_size)`.	required
`position_embeddings`	`tuple[Tensor, Tensor]`	Precomputed RoPE embeddings as a `(cos, sin)` tuple.	required

Returns:

Type	Description
`Tensor`	Output hidden states. Shape: `(batch, seq_len, hidden_size)`.

`reset_parameters()`

Resets the module parameters.

`Qwen3MoELayerParameters`

Bases: BaseModel

Configuration for a single Qwen3 MoE layer.

Attributes:

Name	Type	Description
`hidden_size`	`int`	Size of the hidden states.
`intermediate_size`	`int`	Size of the hidden layer of each expert.
`num_experts`	`int`	Total number of experts in the MoE layer.
`experts_top_k`	`int`	Number of experts each token is routed to.
`num_attention_heads`	`int`	Number of query attention heads.
`num_key_value_heads`	`int`	Number of key and value attention heads.
`rms_norm_eps`	`float`	Epsilon of the RMSNorm layers.
`head_dim`	`int`	Size 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.

The backbone supports pipeline parallelism: each instance holds the layers of one stage.

`hidden_size` `property`

The size of the backbone hidden states.

`split_vocab_order` `property`

The order in which vocabulary segments are concatenated.

`split_vocab_size` `property`

The 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`	The pipeline stage this instance belongs to.	required
`hidden_states_snapshot_mode`	`HiddenStatesAggregationMode`	How to aggregate intermediate hidden states into snapshots.	required
`enable_checkpointing`	`bool`	Whether to apply activation checkpointing to each transformer layer.	required

`forward(inputs, shared)`

Executes the backbone forward pass for the current pipeline stage.

Parameters:

Name	Type	Description	Default
`inputs`	`SequenceInput \| SequenceTransfer[Tensor]`	`SequenceInput` with token ids on the first stage, or the incoming `SequenceTransfer` on other stages.	required
`shared`	`SequenceShared`	The backbone shared input: position ids and, if snapshotting is enabled, the aggregation mask.	required

Returns:

Type	Description
`SequenceTransfer[Tensor]`	The produced `SequenceTransfer` with the hidden states and, if enabled, the updated snapshot.

`output_dtype()`

Returns the dtype of the output hidden states.

Returns:

Type	Description
`dtype`	The dtype of the output hidden states.

`reset_parameters()`

Resets the 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 its 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 for the Qwen3 Mixture of Experts (MoE) model backbone.

Attributes:

Name	Type	Description
`layer`	`Qwen3MoELayerParameters`	Configuration shared across all transformer layers.
`num_hidden_layers`	`int`	Total number of transformer layers.
`rope_base`	`int`	Base of the RoPE frequencies.
`max_position_ids`	`int`	Maximum sequence length.
`split_vocab_size`	`dict[str, int]`	Mapping of vocabulary segment names to their sizes.
`split_vocab_order`	`list[str]`	Order in which the vocabulary segments are concatenated.
`pipeline_num_virtual_layers_pre`	`int`	Number of virtual layers that stand for the compute cost of modules on the first stage before the main layers, such as token embeddings. Pipeline layer distribution uses it.
`pipeline_num_virtual_layers_post`	`int`	Number of virtual layers that stand for the compute cost of modules on the last stage after the main layers, such as the final norm and the LM head. Pipeline layer distribution uses it.

`mapper_from_huggingface_qwen3_moe(params, experts_format)`

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

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	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 Hugging Face keys into the d9d format.

A Hugging Face model has exactly one head, so the mapper must know where it goes in the d9d model. The default targets a single-head decoder. To load into a multi-head model, pass that head's prefix (f"heads.{name}."). The other heads keep their initialization.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	required
`head_prefix`	`str`	FQN prefix of the head that receives the Hugging Face 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 Hugging Face keys into the d9d format.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	required
`head_prefix`	`str`	FQN prefix of the head that receives the Hugging Face 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 Hugging Face keys into the d9d format.

A Hugging Face embedding model is the bare backbone without head weights. The d9d embedding head has nothing to load and keeps its initialization.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	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 Hugging Face format.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	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 Hugging Face format.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	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 Hugging Face format.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	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 Hugging Face format.

A Hugging Face embedding model is the bare backbone without head weights, so the mapper cannot export an embedding projection.

Parameters:

Name	Type	Description	Default
`params`	`Qwen3MoEParameters`	Base model parameters.	required
`experts_format`	`Qwen3MoEExpertsFormat`	Layout of the expert parameters in the Hugging Face checkpoint.	required
`embedding_dim`	`int \| None`	The output size of the embedding head projection, or `None` if the head has no projection.	`None`

Returns:

Type	Description
`ModelStateMapper`	A composite state mapper.

Raises:

Type	Description
`ValueError`	If `embedding_dim` is set, because Hugging Face has no counterpart for the projection.

`d9d.module.parallelism.model.qwen3_moe`

Horizontal parallelism for the Qwen3 MoE model.

`parallelize_qwen3_moe_model(dist_context, model, stage)`

Parallelizes a Qwen3 MoE backbone within one pipeline stage.

Applies Hybrid Sharded Data Parallel (HSDP) to the embeddings, norms and attention modules, and expert parallelism to the MoE layers. Tensor parallelism and context parallelism are not supported yet.

Parameters:

Name	Type	Description	Default
`dist_context`	`DistributedContext`	The distributed context.	required
`model`	`Qwen3MoEModel`	The Qwen3 MoE backbone to parallelize.	required
`stage`	`PipelineStageInfo`	The current pipeline stage.	required

Raises:

Type	Description
`ValueError`	If tensor parallelism or context parallelism is enabled.

Qwen3 MoE

About

Hugging Face Compatibility

Usage

API Reference

d9d.module.model.qwen3_moe

Qwen3MoEExpertsFormat

Qwen3MoELayer

__init__(params)

forward(hidden_states, position_embeddings)

reset_parameters()

Qwen3MoELayerParameters

Qwen3MoEModel

hidden_size property

split_vocab_order property

split_vocab_size property

__init__(params, stage, hidden_states_snapshot_mode, enable_checkpointing)

forward(inputs, shared)

output_dtype()

reset_parameters()

stage_transfer_spec(pipeline_input, boundary)

Qwen3MoEParameters

mapper_from_huggingface_qwen3_moe(params, experts_format)

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

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

mapper_from_huggingface_qwen3_moe_for_embedding(params, experts_format)

mapper_to_huggingface_qwen3_moe(params, experts_format)

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

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

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

d9d.module.parallelism.model.qwen3_moe

parallelize_qwen3_moe_model(dist_context, model, stage)

`d9d.module.model.qwen3_moe`

`Qwen3MoEExpertsFormat`

`Qwen3MoELayer`

`init(params)`

`forward(hidden_states, position_embeddings)`

`reset_parameters()`

`Qwen3MoELayerParameters`

`Qwen3MoEModel`

`hidden_size` `property`

`split_vocab_order` `property`

`split_vocab_size` `property`

`init(params, stage, hidden_states_snapshot_mode, enable_checkpointing)`

`forward(inputs, shared)`

`output_dtype()`

`reset_parameters()`

`stage_transfer_spec(pipeline_input, boundary)`

`Qwen3MoEParameters`

`mapper_from_huggingface_qwen3_moe(params, experts_format)`

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

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

`mapper_from_huggingface_qwen3_moe_for_embedding(params, experts_format)`

`mapper_to_huggingface_qwen3_moe(params, experts_format)`

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

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

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

`d9d.module.parallelism.model.qwen3_moe`

`parallelize_qwen3_moe_model(dist_context, model, stage)`