Skip to content

Qwen3 Dense

About

The d9d.module.model.qwen3_dense package implements the dense Qwen3 backbone. The d9d.module.parallelism.model.qwen3_dense package applies the default horizontal parallelism to it: HSDP for all modules. 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). An embedding model cannot be exported to Hugging Face if its embedding head has a projection.

Usage

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
from d9d.module.block.hidden_states_aggregator import HiddenStatesAggregationMode
from d9d.module.model import DecoderForCausalLM
from d9d.module.model.qwen3_dense import Qwen3DenseModel, mapper_from_huggingface_qwen3_dense_for_causal_lm
from d9d.module.parallelism.model import parallelize_causal_lm_head, parallelize_qwen3_dense_model

# In ModelProvider.initialize_model_stage.
backbone = Qwen3DenseModel(
    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_dense_for_causal_lm(params)

# In ModelProvider.parallelize_model_stage.
parallelize_qwen3_dense_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_dense

Qwen3 Dense model and its Hugging Face state mappers.

Qwen3DenseLayer

Bases: Module, ModuleLateInit

A single Qwen3 Dense transformer layer.

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

__init__(params)

Constructs the Qwen3DenseLayer object.

Parameters:

Name Type Description Default
params Qwen3DenseLayerParameters

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.

Qwen3DenseLayerParameters

Bases: BaseModel

Configuration for a single Qwen3 Dense layer.

Attributes:

Name Type Description
hidden_size int

Size of the hidden states.
intermediate_size int

Size of the feed-forward hidden layer.
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.

Qwen3DenseModel

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

The Qwen3 Dense 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 Qwen3DenseModel object.

Parameters:

Name Type Description Default
params Qwen3DenseParameters

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.

Qwen3DenseParameters

Bases: BaseModel

Configuration for the Qwen3 Dense model backbone.

Attributes:

Name Type Description
layer Qwen3DenseLayerParameters

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_dense(params)

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

Parameters:

Name Type Description Default
params Qwen3DenseParameters

Base model parameters.

 required

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_from_huggingface_qwen3_dense_for_causal_lm(params, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 Dense 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 Qwen3DenseParameters

Base model parameters.

 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_dense_for_classification(params, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 Dense classification Hugging Face keys into the d9d format.

Parameters:

Name Type Description Default
params Qwen3DenseParameters

Base model parameters.

 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_dense_for_embedding(params)

Creates a state mapper translating Qwen3 Dense 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 Qwen3DenseParameters

Base model parameters.

 required

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_to_huggingface_qwen3_dense(params)

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

Parameters:

Name Type Description Default
params Qwen3DenseParameters

Base model parameters.

 required

Returns:

Type Description
ModelStateMapper

A composite state mapper.

mapper_to_huggingface_qwen3_dense_for_causal_lm(params, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 Dense causal LM d9d keys back into the Hugging Face format.

Parameters:

Name Type Description Default
params Qwen3DenseParameters

Base model parameters.

 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_dense_for_classification(params, *, head_prefix=SINGLE_HEAD_PREFIX)

Creates a state mapper translating Qwen3 Dense classification d9d keys back into the Hugging Face format.

Parameters:

Name Type Description Default
params Qwen3DenseParameters

Base model parameters.

 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_dense_for_embedding(params, *, embedding_dim=None)

Creates a state mapper translating Qwen3 Dense 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 Qwen3DenseParameters

Base model parameters.

 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_dense

Horizontal parallelism for the Qwen3 Dense model.

parallelize_qwen3_dense_model(dist_context, model, stage)

Parallelizes a Qwen3 Dense backbone within one pipeline stage.

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

Parameters:

Name Type Description Default
dist_context DistributedContext

The distributed context.

 required
model Qwen3DenseModel

The Qwen3 Dense backbone to parallelize.

 required
stage PipelineStageInfo

The current pipeline stage.

 required

Raises:

Type Description
ValueError

If tensor parallelism or context parallelism is enabled.