Skip to content

Qwen3 Dense

About

The d9d.module.model.qwen3_dense package implements the Qwen3 Dense model architecture.

The d9d.module.parallelism.model.qwen3_dense 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_dense

Qwen3DenseLayer

Bases: Module, ModuleLateInit

Implements a single Qwen3 Dense transformer layer.

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

__init__(params)

Constructs a Qwen3DenseLayer object.

Parameters:

Name Type Description Default
params Qwen3DenseLayerParameters

Configuration parameters for the layer.

 required

forward(hidden_states, position_embeddings)

Performs the forward pass of the dense 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 MLP blocks, shape (batch, seq_len, hidden_dim).

reset_parameters()

Resets module parameters.

Qwen3DenseLayerParameters

Bases: BaseModel

Configuration parameters for a single Qwen3 Dense 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_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.

Qwen3DenseModel

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

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

Parameters:

Name Type Description Default
params Qwen3DenseParameters

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 backbone shared input (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.

Qwen3DenseParameters

Bases: BaseModel

Configuration parameters for the Qwen3 Dense model backbone.

Attributes:

Name Type Description
layer Qwen3DenseLayerParameters

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

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

Base model parameters.

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

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

Base model parameters.

 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_dense

parallelize_qwen3_dense_model(dist_context, model, stage)

Parallelizes the base Qwen3 Dense 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, MLP).

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 Qwen3DenseModel

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