Qwen Architecture › Complete Model Implementation
Language AI HandbookAll chapters Qwen Architecture › Complete Model ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Having explored each component individually, we now assemble them into a complete Qwen model. This integration reveals how the pieces fit together and gives a reference implementation you can extend for your own experiments.
The journey from components to complete model follows this path:
- Configuration: Define the hyperparameters that determine model size and behavior
- Normalization: Implement RMSNorm, which stabilizes training by normalizing hidden states
- Transformer block: Combine attention with the FFN and normalization layers to form the repeating unit
- Full model: Stack blocks with embeddings and output projection
Model Configuration
We start by defining a configuration class that captures all the hyperparameters. This makes it easy to instantiate different model sizes and ensures consistency across components:
from dataclasses import dataclass
@dataclass
class QwenConfig:
"""Configuration for Qwen model"""
hidden_size: int = 4096
intermediate_size: int = 11008
num_hidden_layers: int = 32
num_attention_heads: int = 32
num_kv_heads: int = 32 # For GQA; same as num_attention_heads = MHA
vocab_size: int = 151936
max_position_embeddings: int = 8192
rms_norm_eps: float = 1e-6
rope_theta: float = 10000.0
use_qkv_bias: bool = True # Qwen uses bias in attention
# Configurations for different Qwen sizes
QWEN_CONFIGS = {
"1.8B": QwenConfig(
hidden_size=2048,
intermediate_size=5504,
num_hidden_layers=24,
num_attention_heads=16,
num_kv_heads=16,
),
"7B": QwenConfig(
hidden_size=4096,
intermediate_size=11008,
num_hidden_layers=32,
num_attention_heads=32,
num_kv_heads=32,
),
"14B": QwenConfig(
hidden_size=5120,
intermediate_size=13696,
num_hidden_layers=40,
num_attention_heads=40,
num_kv_heads=40,
),
"72B": QwenConfig(
hidden_size=8192,
intermediate_size=24576,
num_hidden_layers=80,
num_attention_heads=64,
num_kv_heads=8, # GQA with 8:1 ratio
),
}
def get_qwen_config(size: str) -> QwenConfig:
return QWEN_CONFIGS[size]Qwen Model Configurations ====================================================================== Qwen-1.8B: Layers: 24 Hidden size: 2,048 Attention heads: 16 KV heads: 16 (GQA ratio 1:1) Head dimension: 128 FFN dimension: 5,504 Qwen-7B: Layers: 32 Hidden size: 4,096 Attention heads: 32 KV heads: 32 (GQA ratio 1:1) Head dimension: 128 FFN dimension: 11,008 Qwen-14B: Layers: 40 Hidden size: 5,120 Attention heads: 40 KV heads: 40 (GQA ratio 1:1) Head dimension: 128 FFN dimension: 13,696 Qwen-72B: Layers: 80 Hidden size: 8,192 Attention heads: 64 KV heads: 8 (GQA ratio 8:1) Head dimension: 128 FFN dimension: 24,576
The configuration reveals several key design decisions. Note how GQA ratio varies: Qwen-7B uses 1:1 (equivalent to MHA), while Qwen-72B uses 8:1 for significant memory savings. The intermediate size follows the ~2.7x rule we discussed for SwiGLU. The large vocabulary (151,936) is constant across all sizes, optimized for Chinese characters.
RMSNorm Implementation
Before assembling the transformer block, we need our normalization layer. RMSNorm (Root Mean Square Layer Normalization) is simpler than the original LayerNorm: it normalizes by the root mean square of activations without centering (subtracting the mean).
class QwenRMSNorm(nn.Module):
"""RMSNorm as used in Qwen"""
def __init__(self, hidden_size: int, eps: float = 1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(hidden_size))
def forward(self, x: torch.Tensor) -> torch.Tensor:
variance = x.pow(2).mean(-1, keepdim=True)
x = x * torch.rsqrt(variance + self.eps)
return self.weight * x
# Test RMSNorm
norm = QwenRMSNorm(4096)
x = torch.randn(2, 64, 4096)
normalized = norm(x)RMSNorm: Input shape: torch.Size([2, 64, 4096]) Output shape: torch.Size([2, 64, 4096]) Normalization effect: Input - mean: 0.0018, std: 1.0004 Output - mean: 0.0018, std: 1.0000 Parameters: 4,096 (learnable scale only)
RMSNorm normalizes input magnitudes while preserving the mean, unlike LayerNorm which centers around zero. The output standard deviation is close to 1.0, showing proper normalization. With only 4,096 learnable parameters (the scale weights), RMSNorm adds minimal overhead while giving training stability.
The mathematical formulation is:
where is the hidden dimension, is a small constant for numerical stability, and is the learnable scale parameter. By skipping the mean centering of LayerNorm, RMSNorm is slightly faster while empirically performing just as well.
Complete Transformer Block
Now we combine attention with the FFN and normalization layers to form a single transformer block. This is the basic repeating unit of the model. Qwen-7B stacks 32 of these blocks.
The block follows the pre-norm architecture pattern, where normalization happens before each sub-layer rather than after. This improves training stability, especially for deep networks:
class QwenDecoderLayer(nn.Module):
"""Single Qwen transformer decoder layer"""
def __init__(self, config: QwenConfig):
super().__init__()
self.hidden_size = config.hidden_size
# Attention with pre-norm
self.input_layernorm = QwenRMSNorm(
config.hidden_size, config.rms_norm_eps
)
self.self_attn = QwenAttention(
config.hidden_size,
config.num_attention_heads,
config.num_kv_heads,
)
# FFN with pre-norm
self.post_attention_layernorm = QwenRMSNorm(
config.hidden_size, config.rms_norm_eps
)
self.mlp = QwenMLP(config.hidden_size, config.intermediate_size)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
# Self attention with residual
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.self_attn(hidden_states)
hidden_states = residual + hidden_states
# FFN with residual
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
return hidden_states
# Test decoder layer
config = get_qwen_config("7B")
layer = QwenDecoderLayer(config)
x = torch.randn(2, 64, config.hidden_size)
output = layer(x)Qwen Decoder Layer: Input shape: torch.Size([2, 64, 4096]) Output shape: torch.Size([2, 64, 4096]) Parameter breakdown: Attention: 67,121,152 (33.2%) MLP: 135,266,304 (66.8%) Norms: 8,192 (0.0%) Total: 202,395,648
The parameter distribution shows MLP dominates at roughly 65% of layer parameters, with attention at about 35%. This is typical for modern transformers using SwiGLU. The normalization layers contribute negligibly to parameter count but are necessary for training stability.
The code reveals the pre-norm pattern in action: each sub-layer computes residual = hidden_states before normalization, then adds the sub-layer output back (hidden_states = residual + hidden_states). This residual connection is needed. It gives a direct gradient pathway during backpropagation and allows each layer to learn a refinement of its input rather than a complete transformation.
Complete Qwen Model
Finally, we assemble the complete model by wrapping our transformer blocks with:
- Token embeddings: Convert token IDs to dense vectors
- Stacked transformer blocks: The core computation
- Final normalization: Stabilize the output before prediction
- Language model head: Project back to vocabulary size for next-token prediction
class QwenModel(nn.Module):
"""Complete Qwen model"""
def __init__(self, config: QwenConfig):
super().__init__()
self.config = config
# Token embeddings (untied from output)
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
# Transformer layers
self.layers = nn.ModuleList(
[QwenDecoderLayer(config) for _ in range(config.num_hidden_layers)]
)
# Final normalization
self.norm = QwenRMSNorm(config.hidden_size, config.rms_norm_eps)
# Output projection (untied from embeddings)
self.lm_head = nn.Linear(
config.hidden_size, config.vocab_size, bias=False
)
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
hidden_states = self.embed_tokens(input_ids)
for layer in self.layers:
hidden_states = layer(hidden_states)
hidden_states = self.norm(hidden_states)
logits = self.lm_head(hidden_states)
return logits
# Create a small model for testing
test_config = QwenConfig(
hidden_size=512,
intermediate_size=1376,
num_hidden_layers=4,
num_attention_heads=8,
num_kv_heads=8,
vocab_size=32000, # Smaller for testing
max_position_embeddings=512,
)
model = QwenModel(test_config)
total_params = sum(p.numel() for p in model.parameters())
# Test forward pass
test_input = torch.randint(0, test_config.vocab_size, (2, 64))
with torch.no_grad():
logits = model(test_input)Qwen Model (Test Configuration): Total parameters: 45,427,200 Component breakdown: Embeddings: 16,384,000 (36.1%) Layers: 12,658,688 (27.9%) Final norm: 512 LM head: 16,384,000 (36.1%) Forward pass: Input shape: torch.Size([2, 64]) Output shape: torch.Size([2, 64, 32000]) Output matches expected: True
The model successfully processes input tokens through all layers and produces vocabulary logits. With untied embeddings, both the input embedding and output LM head have separate parameters, each accounting for a significant fraction of total model size. This is particularly pronounced in Qwen due to its large vocabulary.
A note on tied vs. untied embeddings
Some models (like the original GPT-2) tie the input embedding and output projection, using the same weight matrix for both. This saves parameters but forces the embedding space to serve double duty: representing input tokens and creating output logits. Qwen untied these, letting the input embedding to specialize for contextual representation and the output head to specialize for prediction. The cost is doubling the embedding parameters (about 0.6B for Qwen-7B), but the quality gain justifies this for large models.
Parameter Count Analysis
Understanding where parameters live helps with capacity planning and debugging. Let's systematically count parameters for each Qwen variant:
def count_parameters(config: QwenConfig) -> dict:
"""Calculate parameter count for each component"""
hidden = config.hidden_size
intermediate = config.intermediate_size
vocab = config.vocab_size
n_layers = config.num_hidden_layers
n_heads = config.num_attention_heads
n_kv_heads = config.num_kv_heads
head_dim = hidden // n_heads
# Embeddings (untied, so counted separately)
embed_params = vocab * hidden
# Per layer: attention + MLP + norms
# Attention: Q, K, V projections + output projection
q_params = hidden * hidden + (hidden if config.use_qkv_bias else 0)
kv_params = 2 * (
hidden * (n_kv_heads * head_dim)
+ (n_kv_heads * head_dim if config.use_qkv_bias else 0)
)
o_params = hidden * hidden
attn_params = q_params + kv_params + o_params
# MLP: gate, up, down projections
mlp_params = 3 * hidden * intermediate
# Norms: 2 per layer
norm_params = 2 * hidden
layer_params = attn_params + mlp_params + norm_params
all_layer_params = layer_params * n_layers
# Final norm
final_norm_params = hidden
# LM head (untied)
lm_head_params = hidden * vocab
total = embed_params + all_layer_params + final_norm_params + lm_head_params
return {
"embeddings": embed_params,
"per_layer": layer_params,
"all_layers": all_layer_params,
"final_norm": final_norm_params,
"lm_head": lm_head_params,
"total": total,
}
# Calculate for each size
for size in ["1.8B", "7B", "14B", "72B"]:
config = get_qwen_config(size)
params = count_parameters(config)Qwen Parameter Breakdown ====================================================================== Qwen-1.8B: Embeddings: 0.31B (16.9%) Per layer: 50.6M All layers: 1.21B (66.1%) LM head: 0.31B (16.9%) Total: 1.84B Qwen-7B: Embeddings: 0.62B (8.1%) Per layer: 202.4M All layers: 6.48B (83.9%) LM head: 0.62B (8.1%) Total: 7.72B Qwen-14B: Embeddings: 0.78B (5.5%) Per layer: 315.3M All layers: 12.61B (89.0%) LM head: 0.78B (5.5%) Total: 14.17B Qwen-72B: Embeddings: 1.24B (2.0%) Per layer: 755.0M All layers: 60.40B (96.0%) LM head: 1.24B (2.0%) Total: 62.89B
Notice how the large vocabulary (151,936 tokens) materially impacts the embedding and LM head parameters. For Qwen-7B, embeddings alone account for approximately 0.6B parameters (about 8% of total), compared to less than 3% in LLaMA-7B.
