qwen3 transformer decoder

Complete educational pure-PyTorch implementation of the Qwen Transformer (Qwen2 / Qwen3 style).

This is a clean, readable, from-scratch version focused on understanding (not production HF code). It includes:

• RMSNorm
• Rotary Positional Embeddings (RoPE)
• Grouped-Query Attention (GQA)
• Optional QK-Norm
• SwiGLU Feed-Forward (the standard Qwen MLP)
• Pre-Norm Transformer blocks
• Causal language model head

pythonCopy
1import math
2import torch
3import torch.nn as nn
4import torch.nn.functional as F
5
6# ============================================================
7# 1. RMSNorm
8# ============================================================
9class RMSNorm(nn.Module):
10    def __init__(self, dim: int, eps: float = 1e-6):
11        super().__init__()
12        self.eps = eps
13        self.weight = nn.Parameter(torch.ones(dim))
14
15    def forward(self, x: torch.Tensor) -> torch.Tensor:
16        # x: (batch, seq, dim)
17        variance = x.pow(2).mean(dim=-1, keepdim=True)
18        x_norm = x * torch.rsqrt(variance + self.eps)
19        return self.weight * x_norm
20
21
22# ============================================================
23# 2. Rotary Positional Embeddings (RoPE)
24# ============================================================
25def precompute_rope_freqs(head_dim: int, max_seq_len: int, theta: float = 10000.0, device="cpu"):
26    """Precompute cos and sin for RoPE."""
27    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
28    t = torch.arange(max_seq_len, device=device).float()
29    freqs = torch.outer(t, inv_freq)                    # (seq, head_dim//2)
30    emb = torch.cat((freqs, freqs), dim=-1)             # (seq, head_dim)
31    return emb.cos(), emb.sin()
32
33
34def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
35    """
36    x:   (batch, n_heads, seq_len, head_dim)
37    cos/sin: (seq_len, head_dim)
38    """
39    # Rotate half
40    x1 = x[..., : x.shape[-1] // 2]
41    x2 = x[..., x.shape[-1] // 2 :]
42    rotated = torch.cat((-x2, x1), dim=-1)
43
44    cos = cos[: x.shape[2]].unsqueeze(0).unsqueeze(0)   # (1,1,seq,head_dim)
45    sin = sin[: x.shape[2]].unsqueeze(0).unsqueeze(0)
46
47    return (x * cos) + (rotated * sin)
48
49
50# ============================================================
51# 3. Grouped-Query Attention (GQA)
52# ============================================================
53class GroupedQueryAttention(nn.Module):
54    def __init__(
55        self,
56        dim: int,
57        n_heads: int,
58        n_kv_heads: int,
59        head_dim: int | None = None,
60        qk_norm: bool = False,
61        bias: bool = False,          # Qwen2 uses bias=True for QKV
62    ):
63        super().__init__()
64        assert n_heads % n_kv_heads == 0, "n_heads must be divisible by n_kv_heads"
65
66        self.n_heads = n_heads
67        self.n_kv_heads = n_kv_heads
68        self.n_rep = n_heads // n_kv_heads          # how many times to repeat KV
69
70        self.head_dim = head_dim if head_dim is not None else dim // n_heads
71        self.scale = self.head_dim ** -0.5
72
73        self.q_proj = nn.Linear(dim, n_heads * self.head_dim, bias=bias)
74        self.k_proj = nn.Linear(dim, n_kv_heads * self.head_dim, bias=bias)
75        self.v_proj = nn.Linear(dim, n_kv_heads * self.head_dim, bias=bias)
76        self.o_proj = nn.Linear(n_heads * self.head_dim, dim, bias=False)
77
78        self.q_norm = RMSNorm(self.head_dim) if qk_norm else None
79        self.k_norm = RMSNorm(self.head_dim) if qk_norm else None
80
81    def forward(
82        self,
83        x: torch.Tensor,
84        cos: torch.Tensor,
85        sin: torch.Tensor,
86        attention_mask: torch.Tensor | None = None,
87    ) -> torch.Tensor:
88        B, T, _ = x.shape
89
90        # Projections
91        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
92        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
93        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
94
95        # Optional QK-Norm (used in some Qwen3 variants)
96        if self.q_norm is not None:
97            q = self.q_norm(q)
98        if self.k_norm is not None:
99            k = self.k_norm(k)
100
101        # Apply RoPE
102        q = apply_rope(q, cos, sin)
103        k = apply_rope(k, cos, sin)
104
105        # Expand KV heads for GQA
106        k = k.repeat_interleave(self.n_rep, dim=1)
107        v = v.repeat_interleave(self.n_rep, dim=1)
108
109        # Attention scores
110        scores = (q @ k.transpose(-2, -1)) * self.scale   # (B, n_heads, T, T)
111
112        if attention_mask is not None:
113            scores = scores + attention_mask
114
115        attn = F.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype)
116        out = attn @ v                                    # (B, n_heads, T, head_dim)
117
118        out = out.transpose(1, 2).contiguous().view(B, T, -1)
119        return self.o_proj(out)
120
121
122# ============================================================
123# 4. SwiGLU Feed-Forward (standard Qwen MLP)
124# ============================================================
125class SwiGLU(nn.Module):
126    def __init__(self, dim: int, hidden_dim: int, bias: bool = False):
127        super().__init__()
128        # Qwen uses intermediate_size ≈ (8/3)*dim, rounded to multiple of 256
129        self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
130        self.up_proj   = nn.Linear(dim, hidden_dim, bias=bias)
131        self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
132
133    def forward(self, x: torch.Tensor) -> torch.Tensor:
134        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
135
136
137# ============================================================
138# 5. Transformer Block (Pre-Norm)
139# ============================================================
140class TransformerBlock(nn.Module):
141    def __init__(
142        self,
143        dim: int,
144        n_heads: int,
145        n_kv_heads: int,
146        hidden_dim: int,
147        head_dim: int | None = None,
148        qk_norm: bool = False,
149        eps: float = 1e-6,
150    ):
151        super().__init__()
152        self.attn_norm = RMSNorm(dim, eps=eps)
153        self.attn = GroupedQueryAttention(
154            dim=dim,
155            n_heads=n_heads,
156            n_kv_heads=n_kv_heads,
157            head_dim=head_dim,
158            qk_norm=qk_norm,
159            bias=True,                 # Qwen2 uses bias on QKV
160        )
161        self.ffn_norm = RMSNorm(dim, eps=eps)
162        self.ffn = SwiGLU(dim, hidden_dim, bias=False)
163
164    def forward(
165        self,
166        x: torch.Tensor,
167        cos: torch.Tensor,
168        sin: torch.Tensor,
169        attention_mask: torch.Tensor | None = None,
170    ) -> torch.Tensor:
171        # Pre-Norm + residual
172        x = x + self.attn(self.attn_norm(x), cos, sin, attention_mask)
173        x = x + self.ffn(self.ffn_norm(x))
174        return x
175
176
177# ============================================================
178# 6. Full Qwen Model
179# ============================================================
180class QwenModel(nn.Module):
181    def __init__(
182        self,
183        vocab_size: int = 151936,
184        dim: int = 1024,                 # emb_dim / hidden_size
185        n_layers: int = 28,
186        n_heads: int = 16,
187        n_kv_heads: int = 8,             # GQA
188        head_dim: int | None = 128,      # often fixed in modern Qwen
189        intermediate_size: int = 3072,   # hidden_dim of MLP
190        max_seq_len: int = 40960,
191        rope_theta: float = 1_000_000.0,
192        qk_norm: bool = True,            # Qwen3 often uses this
193        rms_norm_eps: float = 1e-6,
194        tie_word_embeddings: bool = False,
195    ):
196        super().__init__()
197        self.dim = dim
198        self.max_seq_len = max_seq_len
199        self.n_heads = n_heads
200        self.head_dim = head_dim if head_dim is not None else dim // n_heads
201
202        # Token embeddings
203        self.embed_tokens = nn.Embedding(vocab_size, dim)
204
205        # Transformer layers
206        self.layers = nn.ModuleList([
207            TransformerBlock(
208                dim=dim,
209                n_heads=n_heads,
210                n_kv_heads=n_kv_heads,
211                hidden_dim=intermediate_size,
212                head_dim=self.head_dim,
213                qk_norm=qk_norm,
214                eps=rms_norm_eps,
215            )
216            for _ in range(n_layers)
217        ])
218
219        self.norm = RMSNorm(dim, eps=rms_norm_eps)
220
221        # Output head
222        self.lm_head = nn.Linear(dim, vocab_size, bias=False)
223        if tie_word_embeddings:
224            self.lm_head.weight = self.embed_tokens.weight
225
226        # Precompute RoPE
227        cos, sin = precompute_rope_freqs(
228            self.head_dim, max_seq_len, theta=rope_theta
229        )
230        self.register_buffer("cos", cos, persistent=False)
231        self.register_buffer("sin", sin, persistent=False)
232
233        # Initialize weights
234        self.apply(self._init_weights)
235
236    def _init_weights(self, module):
237        if isinstance(module, nn.Linear):
238            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
239            if module.bias is not None:
240                torch.nn.init.zeros_(module.bias)
241        elif isinstance(module, nn.Embedding):
242            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
243
244    def forward(
245        self,
246        input_ids: torch.Tensor,
247        attention_mask: torch.Tensor | None = None,
248    ) -> torch.Tensor:
249        """
250        input_ids: (batch, seq_len)
251        returns logits: (batch, seq_len, vocab_size)
252        """
253        B, T = input_ids.shape
254        assert T <= self.max_seq_len, f"Sequence length {T} > max {self.max_seq_len}"
255
256        # Causal mask
257        if attention_mask is None:
258            # (1, 1, T, T) lower-triangular causal mask with -inf above diagonal
259            causal_mask = torch.triu(
260                torch.full((T, T), float("-inf"), device=input_ids.device),
261                diagonal=1,
262            )
263            attention_mask = causal_mask.unsqueeze(0).unsqueeze(0)
264
265        # Embedding
266        x = self.embed_tokens(input_ids)          # (B, T, dim)
267
268        # Pass through all layers
269        for layer in self.layers:
270            x = layer(x, self.cos, self.sin, attention_mask)
271
272        x = self.norm(x)
273        logits = self.lm_head(x)
274        return logits
275
276
277# ============================================================
278# 7. Example: create a small Qwen-like model (≈0.6B style)
279# ============================================================
280if __name__ == "__main__":
281    # Config roughly matching Qwen3-0.6B
282    config = {
283        "vocab_size": 151936,
284        "dim": 1024,
285        "n_layers": 28,
286        "n_heads": 16,
287        "n_kv_heads": 8,
288        "head_dim": 128,
289        "intermediate_size": 3072,
290        "max_seq_len": 4096,          # smaller for demo
291        "rope_theta": 1_000_000.0,
292        "qk_norm": True,
293    }
294
295    model = QwenModel(**config)
296    print(f"Number of parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.1f} M")
297
298    # Dummy forward pass
299    batch_size, seq_len = 2, 32
300    input_ids = torch.randint(0, config["vocab_size"], (batch_size, seq_len))
301    logits = model(input_ids)
302    print("Logits shape:", logits.shape)   # (2, 32, 151936)

Key Qwen design choices explained

This code is intentionally minimal and readable so you can experiment, add KV-cache, flash-attention, MoE, etc.

For the official production implementation see Hugging Face transformers (modeling_qwen2.py / modeling_qwen3.py). For an even more complete from-scratch version with weight loading, check Sebastian Raschka’s excellent notebooks in the LLMs-from-scratch repository.