moe/qwen3_moe_unstructured.ipynb

Qwen 3 from Alibaba is currently the best open-source model available, offering state-of-the-art performance across a wide range of tasks including reasoning, coding, math, and multilingual understanding. Its flagship version, Qwen3‑235B‑A22B, ranks at or near the top of key benchmarks like MMLU-Pro, LiveCodeBench, and AIME, often matching or exceeding the performance of leading proprietary models.

Built using a Mixture-of-Experts (MoE) architecture, Qwen 3 activates only a subset of its 235B parameters per query resulting in high efficiency without sacrificing quality. It also supports up to 128K-token contexts, handles 119 languages, and introduces dual “thinking” vs. “non-thinking” modes to balance deep reasoning with faster inference.

In this blog post, we'll walk through building a Qwen 3 Mixture-of-Experts (MoE) model from scratch. Our Qwen model features 0.6 billion parameters and includes two expert layers.

Table of Contents

Prerequisites

The good part is we won't be using object-oriented programming (OOP) coding, just plain Python programming. However, you should have a basic understanding of neural networks and Transformer architecture. These are the only two prerequisites needed to follow along with the blog.

Topic	Link
Transformer Theory	Video
Neural Networks Theory	Video
Python basics	Video

Understanding Qwen 3 MoE Architecture

Raschka, S. (2024). Build A Large Language Model (From Scratch), Published by Manning, ISBN 978-1633437166 [Computer software]. https://github.com/rasbt/LLMs-from-scratch

Let’s first understand the Qwen MoE architecture as an intermediate techy person, and then use an example “the cat sat” to see how it goes through the architecture to get a clear understanding of it.

Imagine you have a really tough job. Instead of hiring one person who kinda knows everything, you hire a team of specialists, each amazing at one particular thing (like an electrician, a plumber, a painter). You also hire a manager who looks at the current task and sends it to the right specialist(s).

MoE in AI models is kinda like that. Instead of one gigantic neural network trying to learn everything, an MoE layer has:

1. A Team of “Experts”: These are smaller, specialized neural networks (usually simple Feed-Forward Networks or MLPs). Each expert might get good at handling certain types of information or patterns.
2. A “Router” (The Manager): This is another small network. Its job is to look at the input data (like a word or part of a word) and decide which expert(s) are the best fit to handle it right now.

Imagine our model is processing the sentence: The cat sat.

1. Tokens: First, we break it into pieces (tokens): “The” “cat” “sat”
2. Router Gets a Token: The MoE layer receives the token cat (represented as a bunch of numbers, an embedding vector). The Router looks at this cat vector.
3. Router Chooses: Let’s say we have 4 experts (E1, E2, E3, E4). The Router decides which ones are best suited for cat.
4. Maybe it thinks E2 (perhaps good with nouns?) and E4 (perhaps good with animal concepts?) are the top choices. It gives scores or "weights" to these choices (e.g., 70% for E2, 30% for E4).

The cat vector is sent only to Expert 2 and Expert 4. Experts 1 and 3 don't do any work for this token, saving computation! E2 processes cat and generates its result (Output_E2). E4 processes cat and generates its result (Output_E4).

We now combine the results from the chosen experts using the router weights: Final_Output = (0.7 * Output_E2) + (0.3 * Output_E4).

This Final_Output is what the MoE layer passes on for the token cat. This happens for every token in the sequence! Different tokens might get routed to different experts.

So, when our model processes text like "The cat sat.", the overall journey looks like this:

Input Text goes into the Tokenizer.Tokenizer creates numerical Token IDs. Embedding Layer turns IDs into meaningful number vectors (Embeddings) and adds Positional Info (using RoPE later in attention).

These vectors go through multiple Transformer Blocks. Each block has:

• Self-Attention (where tokens look at each other, enhanced by RoPE).
• MoE Layer (where the router sends tokens to specific experts).
• Normalization (RMSNorm) and Residual connections help learning.

The output from the last block goes to a Final Layer. This layer produces Logits (scores) for every possible next token in our vocabulary.

We convert logits to Probabilities and Predict the Next Token.

Now that we have a feel for how MoE fits into the picture, let’s dive into the code and build these components step-by-step! We’ll start by setting up our coding environment.

Downloading Weights

The first step is to download the weights for the Qwen 3 model. We are targeting the Qwen 3.5 0.6B model, which is a smaller version of the Qwen 3 model with two layers of MoE (Mixture of Experts). This model is suitable for running on consumer hardware.

In [ ]:

# Import tqdm for progress bars and snapshot_download for downloading model files
from tqdm import tqdm
from huggingface_hub import snapshot_download

# Define the Hugging Face repository ID and the local directory to save the files
repo_id = "huihui-ai/Huihui-MoE-0.8B-2E"
local_dir = "Huihui-MoE-0.8B-2E"

# Download the model snapshot from Hugging Face, excluding .bin files
# This will fetch config, tokenizer, and safetensors weights only
snapshot_download(
    repo_id=repo_id,
    local_dir=local_dir,
    ignore_patterns=["*.bin"],      # Skip large .bin files, only get safetensors
    tqdm_class=tqdm                # Use standard tqdm for progress bar
)

In [3]:

from pathlib import Path
from tokenizers import Tokenizer
import json
import torch


# This file contains the vocabulary, merge rules, and configuration.
tokenizer_path = Path("Huihui-MoE-0.8B-2E/tokenizer.json")

# The Tokenizer.from_file() method is the standard way to load tokenizers
# from the Hugging Face ecosystem.
tokenizer = Tokenizer.from_file(str(tokenizer_path))

# We can also load them from the special_tokens_map.json for confirmation
with open("Huihui-MoE-0.8B-2E/special_tokens_map.json", "r") as f:
    special_tokens_map = json.load(f)
    print(f"Special tokens from file: {special_tokens_map}")


# We'll follow the encode -> decode pattern to ensure it works correctly.
prompt = "The only thing I know is that I know"

# .encode() returns an Encoding object, we access the token IDs via .ids
encoded = tokenizer.encode(prompt)
print(f"\nOriginal prompt: '{prompt}'")
print(f"Encoded token IDs: {encoded.ids}")

# .decode() converts the token IDs back to a string.
decoded = tokenizer.decode(encoded.ids)
print(f"Decoded back to text: '{decoded}'")

# Step 5: Verify the vocabulary size
vocab_size = tokenizer.get_vocab_size()
print(f"\nTokenizer vocabulary size: {vocab_size}")

Special tokens from file: {'additional_special_tokens': ['<|im_start|>', '<|im_end|>', '<|object_ref_start|>', '<|object_ref_end|>', '<|box_start|>', '<|box_end|>', '<|quad_start|>', '<|quad_end|>', '<|vision_start|>', '<|vision_end|>', '<|vision_pad|>', '<|image_pad|>', '<|video_pad|>'], 'eos_token': {'content': '<|im_end|>', 'lstrip': False, 'normalized': False, 'rstrip': False, 'single_word': False}, 'pad_token': {'content': '<|endoftext|>', 'lstrip': False, 'normalized': False, 'rstrip': False, 'single_word': False}}

Original prompt: 'The only thing I know is that I know'
Encoded token IDs: [785, 1172, 3166, 358, 1414, 374, 429, 358, 1414]
Decoded back to text: 'The only thing I know is that I know'

Tokenizer vocabulary size: 151669

In [4]:

# Get the vocabulary as a dictionary: {token_string: token_id}
vocab = tokenizer.get_vocab()

# Display a slice of the vocabulary for inspection (tokens 5600 to 5609)
sample_vocab_slice = list(vocab.items())[5600:5610]
sample_vocab_slice

Out [4]:

[('_PARAM', 12073),
 ('ĠTCHAR', 67069),
 ('ĠMunicip', 35703),
 ('Ġnguyên', 129031),
 ('Javascript', 64567),
 ('Ġric', 29950),
 ('ðŁİĢ', 145315),
 ('ë°ı', 144499),
 ('ĠTul', 42216),
 ('ä¼ļ议室', 112818)]

In [5]:

# Define the path to the configuration file.
config_path = Path("Huihui-MoE-0.8B-2E/config.json")

# Open and load the JSON file into a Python dictionary.
with open(config_path, "r") as f:
    config = json.load(f)

# Print the configuration to see all the parameters.
# This gives us a complete overview of the model we're about to build.
print(json.dumps(config, indent=4))

{
    "architectures": [
        "Qwen3MoeForCausalLM"
    ],
    "attention_bias": false,
    "attention_dropout": 0.0,
    "bos_token_id": 151643,
    "decoder_sparse_step": 1,
    "eos_token_id": 151645,
    "head_dim": 128,
    "hidden_act": "silu",
    "hidden_size": 1024,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "max_position_embeddings": 40960,
    "max_window_layers": 28,
    "mlp_only_layers": [],
    "model_type": "qwen3_moe",
    "moe_intermediate_size": 3072,
    "norm_topk_prob": true,
    "num_attention_heads": 16,
    "num_experts": 2,
    "num_experts_per_tok": 1,
    "num_hidden_layers": 28,
    "num_key_value_heads": 8,
    "output_router_logits": false,
    "rms_norm_eps": 1e-06,
    "rope_scaling": null,
    "rope_theta": 1000000,
    "router_aux_loss_coef": 0.001,
    "sliding_window": null,
    "tie_word_embeddings": true,
    "torch_dtype": "bfloat16",
    "transformers_version": "4.52.4",
    "use_cache": true,
    "use_sliding_window": false,
    "vocab_size": 151936
}

In [ ]:

# --- Main Architecture Parameters ---
# Extract model hyperparameters from the config dictionary.

# Embedding dimension (hidden size of the model)
dim = config["hidden_size"]

# Number of transformer layers
n_layers = config["num_hidden_layers"]

# Number of attention heads
n_heads = config["num_attention_heads"]

# Number of key/value heads (for grouped-query attention)
n_kv_heads = config["num_key_value_heads"]

# Vocabulary size
vocab_size = config["vocab_size"]

# RMSNorm epsilon value for numerical stability
norm_eps = config["rms_norm_eps"]

# Rotary positional embedding theta parameter
rope_theta = torch.tensor(config["rope_theta"])

# Dimension of each attention head
head_dim = config["head_dim"]  # <-- This is important for attention calculations

# --- Mixture-of-Experts (MoE) Specific Parameters ---
# Number of experts in the MoE layer
num_experts = config["num_experts"]

# Number of experts selected per token by the router
num_experts_per_tok = config["num_experts_per_tok"]

# Intermediate size of the MoE feed-forward network
moe_intermediate_size = config["moe_intermediate_size"]

# --- Print out all key architecture parameters for verification ---
print(f"Embedding Dimension (dim): {dim}")
print(f"Number of Layers (n_layers): {n_layers}")
print(f"Number of Heads (n_heads): {n_heads}")
print(f"Number of KV Heads (n_kv_heads): {n_kv_heads}")
print(f"Head Dimension (head_dim): {head_dim}")  # <-- Print head_dim for clarity
print(f"Vocabulary Size (vocab_size): {vocab_size}")
print(f"RoPE Theta: {rope_theta}")
print(f"--- MoE Params ---")
print(f"Total Experts (num_experts): {num_experts}")
print(f"Experts per Token (num_experts_per_tok): {num_experts_per_tok}")
print(f"MoE Intermediate Size: {moe_intermediate_size}")

Embedding Dimension (dim): 1024
Number of Layers (n_layers): 28
Number of Heads (n_heads): 16
Number of KV Heads (n_kv_heads): 8
Head Dimension (head_dim): 128
Vocabulary Size (vocab_size): 151936
RoPE Theta: 1000000
--- MoE Params ---
Total Experts (num_experts): 2
Experts per Token (num_experts_per_tok): 1
MoE Intermediate Size: 3072

In [ ]:

# Import the safetensors loader
from safetensors.torch import load_file

# Define the path to the model weights file
model_weights_path = Path("Huihui-MoE-0.8B-2E/model.safetensors")

# Load the weights into a dictionary; each key is a layer/parameter name, and each value is a torch tensor
model_weights = load_file(model_weights_path)

# Inspect the loaded weights: print the first 20 layer names to confirm successful loading
print("First 20 keys in model_weights:")
print(json.dumps(list(model_weights.keys())[:20], indent=4))

[
    "model.embed_tokens.weight",
    "model.layers.0.input_layernorm.weight",
    "model.layers.0.mlp.experts.0.down_proj.weight",
    "model.layers.0.mlp.experts.0.gate_proj.weight",
    "model.layers.0.mlp.experts.0.up_proj.weight",
    "model.layers.0.mlp.experts.1.down_proj.weight",
    "model.layers.0.mlp.experts.1.gate_proj.weight",
    "model.layers.0.mlp.experts.1.up_proj.weight",
    "model.layers.0.mlp.gate.weight",
    "model.layers.0.post_attention_layernorm.weight",
    "model.layers.0.self_attn.k_norm.weight",
    "model.layers.0.self_attn.k_proj.weight",
    "model.layers.0.self_attn.o_proj.weight",
    "model.layers.0.self_attn.q_norm.weight",
    "model.layers.0.self_attn.q_proj.weight",
    "model.layers.0.self_attn.v_proj.weight",
    "model.layers.1.input_layernorm.weight",
    "model.layers.1.mlp.experts.0.down_proj.weight",
    "model.layers.1.mlp.experts.0.gate_proj.weight",
    "model.layers.1.mlp.experts.0.up_proj.weight"
]

In [ ]:

import torch.nn as nn

# Get special token IDs using the tokenizer
im_start_id = tokenizer.token_to_id("<|im_start|>")
im_end_id = tokenizer.token_to_id("<|im_end|>")
newline_id = tokenizer.encode("\n").ids[0]  # ID for newline character

# Get token IDs for 'user' and 'assistant'
user_ids = tokenizer.encode("user").ids
assistant_ids = tokenizer.encode("assistant").ids

# Print token IDs for verification
print(f"ID for '<|im_start|>': {im_start_id}")
print(f"ID for '<|im_end|>': {im_end_id}")
print(f"ID for '\\n': {newline_id}")
print(f"IDs for 'user': {user_ids}")
print(f"IDs for 'assistant': {assistant_ids}\n")

# Encode the main prompt
prompt_ids = tokenizer.encode(prompt).ids

# Build the chat template token sequence
# Template: <|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n
prefix_ids = [im_start_id] + user_ids + [newline_id]
suffix_ids = [im_end_id, newline_id, im_start_id] + assistant_ids + [newline_id]

tokens_list = prefix_ids + prompt_ids + suffix_ids
tokens = torch.tensor(tokens_list)

print(f"Final combined token IDs: {tokens}")
print(f"Number of tokens: {len(tokens)}")

# Decode each token individually for inspection
prompt_split_as_tokens = [tokenizer.decode([token.item()]) for token in tokens]
print(f"\nPrompt split into tokens: {prompt_split_as_tokens}")

ID for '<|im_start|>': 151644
ID for '<|im_end|>': 151645
ID for '\n': 198
IDs for 'user': [872]
IDs for 'assistant': [77091]

Final combined token IDs: tensor([151644,    872,    198,    785,   1172,   3166,    358,   1414,    374,
           429,    358,   1414, 151645,    198, 151644,  77091,    198])
Number of tokens: 17

Prompt split into tokens: ['', 'user', '\n', 'The', ' only', ' thing', ' I', ' know', ' is', ' that', ' I', ' know', '', '\n', '', 'assistant', '\n']

In [7]:

dim

Out [7]:

1024

In [8]:

import torch.nn as nn

embedding_layer = nn.Embedding(vocab_size, dim)
embedding_layer.weight.data.copy_(model_weights["model.embed_tokens.weight"])

token_embeddings_unnormalized = embedding_layer(tokens).to(torch.bfloat16)

print("\nShape of the token embeddings:", token_embeddings_unnormalized.shape)

Shape of the token embeddings: torch.Size([17, 1024])

In [9]:

# First, let's define the RMSNorm function.
# This function will be reused throughout the model.

def rms_norm(tensor, norm_weights):
    """
    Apply Root Mean Square Normalization.
    
    The formula is: (tensor / sqrt(mean(tensor^2) + epsilon)) * weights
    Using rsqrt (reciprocal square root) is a common optimization.
    """
    # Save the original data type of the input tensor
    input_dtype = tensor.dtype
    
    # Perform the calculation in float32 for numerical stability
    tensor_float = tensor.to(torch.float32)
    variance = tensor_float.pow(2).mean(-1, keepdim=True)
    normalized_tensor = tensor_float * torch.rsqrt(variance + norm_eps)
    
    # Apply the weights and cast the result back to the original data type
    return (normalized_tensor * norm_weights).to(input_dtype)

In [10]:

# --- Building the first part of the first Transformer layer ---

# Step 1: Normalize the token embeddings.
# This is the "input_layernorm" for the first attention block. We retrieve
# its specific weights from our loaded model_weights dictionary.
token_embeddings_normalized = rms_norm(token_embeddings_unnormalized, model_weights["model.layers.0.input_layernorm.weight"])

# Step 2: Verify the shape.
# Normalization doesn't change the shape of the tensor, only the values within it.
# The shape should still be [17, 1024].
print("Shape of the normalized token embeddings:", token_embeddings_normalized.shape)

Shape of the normalized token embeddings: torch.Size([17, 1024])

In [11]:

# --- Unwrapping the Query weights ---
# Shape: [n_heads * head_dim, dim] -> [16 * 128, 1024] -> [2048, 1024]
# New Shape: [n_heads, head_dim, dim] -> [16, 128, 1024]
q_layer0 = model_weights["model.layers.0.self_attn.q_proj.weight"]
q_layer0 = q_layer0.view(n_heads, head_dim, dim)

# --- Unwrapping the Key and Value weights ---
# Shape: [n_kv_heads * head_dim, dim] -> [8 * 128, 1024] -> [1024, 1024]
# New Shape: [n_kv_heads, head_dim, dim] -> [8, 128, 1024]
k_layer0 = model_weights["model.layers.0.self_attn.k_proj.weight"]
k_layer0 = k_layer0.view(n_kv_heads, head_dim, dim)

v_layer0 = model_weights["model.layers.0.self_attn.v_proj.weight"]
v_layer0 = v_layer0.view(n_kv_heads, head_dim, dim)

print(f"Reshaped Query weights shape: {q_layer0.shape}")
print(f"Reshaped Key weights shape: {k_layer0.shape}")
print(f"Reshaped Value weights shape: {v_layer0.shape}")

Reshaped Query weights shape: torch.Size([16, 128, 1024])
Reshaped Key weights shape: torch.Size([8, 128, 1024])
Reshaped Value weights shape: torch.Size([8, 128, 1024])

In [12]:

# --- Calculating Q, K, and V for the first head of the first layer ---

# Step 1: Get the weight matrix for the first head (head 0).
# Shape of each is [head_dim, dim] -> [128, 1024]
q_layer0_head0 = q_layer0[0]
k_layer0_head0 = k_layer0[0] # For the first head, we use the first KV head
v_layer0_head0 = v_layer0[0]

# Step 2: Calculate the Query, Key, and Value vectors for each token.
# We matrix-multiply our normalized embeddings with the head's weight matrices.
# Input: [17, 1024] @ [1024, 128] -> Output: [17, 128]
q_per_token = torch.matmul(token_embeddings_normalized, q_layer0_head0.T)
k_per_token = torch.matmul(token_embeddings_normalized, k_layer0_head0.T)
v_per_token = torch.matmul(token_embeddings_normalized, v_layer0_head0.T)

# --- Step 3: Verify the shapes ---
# We now have a 128-dimensional Q, K, and V vector for each of our 17 input tokens.
print("\n--- Output for Head 0 ---")
print("Shape of Query vectors per token:", q_per_token.shape)
print("Shape of Key vectors per token:", k_per_token.shape)
print("Shape of Value vectors per token:", v_per_token.shape)

--- Output for Head 0 ---
Shape of Query vectors per token: torch.Size([17, 128])
Shape of Key vectors per token: torch.Size([17, 128])
Shape of Value vectors per token: torch.Size([17, 128])

In [13]:

import matplotlib.pyplot as plt

# --- Step 1: Pre-compute the RoPE frequencies ---
# We will pre-compute the rotation matrices (in complex number form) for all positions.

# The `max_position_embeddings` from the config tells us the maximum sequence length.
max_seq_len = config["max_position_embeddings"]

# Calculate the frequencies (theta_i) for each of the 64 pairs.
# The formula is theta_i = 1.0 / (rope_theta^(2i / head_dim))
freqs = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2) / head_dim))

# Create a tensor of all possible positions [0, 1, ..., max_seq_len-1]
t = torch.arange(max_seq_len)

# Create the full frequency matrix by taking the outer product of positions and frequencies.
# Shape: [max_seq_len, head_dim / 2] -> [40960, 64]
freqs_for_each_token = torch.outer(t, freqs)

# Convert this to a matrix of complex numbers using Euler's formula: e^(i*theta) = cos(theta) + i*sin(theta).
# torch.polar is perfect for this. The magnitude is 1, the angle is the frequency.
# Shape: [max_seq_len, head_dim / 2] -> [40960, 64]
freqs_cis = torch.polar(torch.ones_like(freqs_for_each_token), freqs_for_each_token)

# --- Step 2: Visualize the rotations for one position (optional but insightful) ---
# This shows how the rotation angle changes for each of the 64 pairs for a single token.
value = freqs_cis[3] # Let's look at the rotations for the 4th token (index 3)
plt.figure()
for i, element in enumerate(value):
    plt.plot([0, element.real], [0, element.imag], color='blue', linewidth=1)
    plt.annotate(f"{i}", xy=(element.real, element.imag), color='red')
plt.xlabel('Real')
plt.ylabel('Imaginary')
plt.title('RoPE Rotations for each 2D-pair at a single token position')
plt.grid(True)
plt.show()

In [14]:

# --- Step 3: Apply RoPE to Query and Key vectors of Head 0 ---

# Take the Q vectors for our 17 tokens and view them as 64 pairs of real numbers.
# Shape: [17, 128] -> [17, 64, 2]
q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)

# Now, view these pairs as complex numbers.
# Shape: [17, 64, 2] -> [17, 64]
q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)

# Get the rotation frequencies for just our 17 token positions.
# Shape of freqs_cis is [40960, 64], we slice the first 17 rows.
freqs_cis_for_tokens = freqs_cis[:len(tokens)]

# Apply the rotation by performing element-wise complex number multiplication.
# Shape: [17, 64] * [17, 64] -> [17, 64]
q_per_token_as_complex_numbers_rotated = q_per_token_as_complex_numbers * freqs_cis_for_tokens

# Convert the rotated complex numbers back into pairs of real numbers.
# Shape: [17, 64] -> [17, 64, 2]
q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers_rotated)

# Finally, merge the pairs back into our original vector shape.
# Shape: [17, 64, 2] -> [17, 128]
q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)

# --- Repeat the exact same process for the Key vectors ---
k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
k_per_token_as_complex_numbers_rotated = k_per_token_as_complex_numbers * freqs_cis_for_tokens
k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers_rotated)
k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)


# --- Step 4: Verify the shapes ---
# The shapes are unchanged, but the vectors now contain positional information.
print("\n--- Output for Head 0 after RoPE ---")
print("Shape of rotated Query vectors:", q_per_token_rotated.shape)
print("Shape of rotated Key vectors:", k_per_token_rotated.shape)

--- Output for Head 0 after RoPE ---
Shape of rotated Query vectors: torch.Size([17, 128])
Shape of rotated Key vectors: torch.Size([17, 128])

In [ ]:

import matplotlib.pyplot as plt

# Calculate the raw attention scores by multiplying rotated Q and K vectors
qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)

# Scale the attention scores by the square root of the head dimension
qk_per_token_scaled = qk_per_token / (head_dim**0.5)

# Visualize the raw attention scores before masking
print("--- Raw Attention Scores (Before Masking) ---")

def display_qk_heatmap(qk_matrix, title="Attention Heatmap"):
    _, ax = plt.subplots()
    im = ax.imshow(qk_matrix.to(torch.float32).detach(), cmap='viridis')
    ax.set_xticks(range(len(prompt_split_as_tokens)))
    ax.set_yticks(range(len(prompt_split_as_tokens)))
    ax.set_xticklabels(prompt_split_as_tokens, rotation=90)
    ax.set_yticklabels(prompt_split_as_tokens)
    ax.figure.colorbar(im, ax=ax)
    plt.title(title)
    plt.show()

display_qk_heatmap(qk_per_token_scaled, title="Raw Attention Scores (Before Masking)")

--- Raw Attention Scores (Before Masking) ---

In [16]:

# --- Step 1: Create the causal mask ---
mask = torch.full((len(tokens), len(tokens)), float("-inf"))
mask = torch.triu(mask, diagonal=1)

# --- Step 2: Apply the mask to the scaled scores ---
qk_per_token_masked = qk_per_token_scaled + mask

# --- Step 3: Visualize the masked scores ---
print("\n--- Attention Scores After Causal Masking ---")
display_qk_heatmap(qk_per_token_masked, title="Attention Scores After Masking")

--- Attention Scores After Causal Masking ---

In [17]:

mask

Out [17]:

tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

In [18]:

# --- Step 1: Apply softmax ---
qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_masked.float(), dim=1).to(torch.bfloat16)

# --- Step 2: Visualize the final attention weights ---
print("\n--- Final Attention Weights (After Softmax) ---")
display_qk_heatmap(qk_per_token_after_masking_after_softmax, title="Final Attention Weights")

# --- Step 3: Verify the shape ---
print("Shape of the final attention weights matrix:", qk_per_token_after_masking_after_softmax.shape)

--- Final Attention Weights (After Softmax) ---

Shape of the final attention weights matrix: torch.Size([17, 17])

In [19]:

# --- Step 1: Calculate the output of the attention head ---
# Matrix multiply the final attention weights with the Value vectors.
# Input: [17, 17] @ [17, 128] -> Output: [17, 128]
qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)

# --- Step 2: Verify the shape ---
# The output tensor has the same shape as the Q, K, and V vectors. Each of the 17
# tokens now has a new 128-dimensional vector that contains context from other tokens.
print("Shape of the final attention output for Head 0:", qkv_attention.shape)

Shape of the final attention output for Head 0: torch.Size([17, 128])

In [20]:

# Create an empty list to store the attention output of each head.
qkv_attention_store = []

# Iterate over each of the 16 attention heads.
for head in range(n_heads):
    # Get the Q, K, and V weights for the current head.
    # Note the use of `head // 4` for K and V due to Grouped-Query Attention.
    # Every 4 query heads share the same key and value heads.
    q_layer0_head = q_layer0[head]
    k_layer0_head = k_layer0[head // (n_heads // n_kv_heads)]
    v_layer0_head = v_layer0[head // (n_heads // n_kv_heads)]

    # Project the normalized embeddings into Q, K, and V vectors for this head.
    q_per_token = torch.matmul(token_embeddings_normalized, q_layer0_head.T)
    k_per_token = torch.matmul(token_embeddings_normalized, k_layer0_head.T)
    v_per_token = torch.matmul(token_embeddings_normalized, v_layer0_head.T)

    # Apply RoPE to the Q and K vectors for this head.
    q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
    q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
    q_per_token_as_complex_numbers_rotated = q_per_token_as_complex_numbers * freqs_cis_for_tokens
    q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers_rotated)
    q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)

    k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
    k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
    k_per_token_as_complex_numbers_rotated = k_per_token_as_complex_numbers * freqs_cis_for_tokens
    k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers_rotated)
    k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)

    # Calculate and scale the attention scores.
    qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T) / (head_dim**0.5)

    # Apply the causal mask.
    qk_per_token_masked = qk_per_token + mask

    # Apply softmax to get the attention weights.
    qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_masked.float(), dim=1).to(torch.bfloat16)

    # Aggregate the Value vectors using the attention weights.
    qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
    
    # Append the result of this head to our list.
    qkv_attention_store.append(qkv_attention)

# --- Verify the result ---
# We should have a list containing 16 tensors, each of shape [17, 128].
print("Number of heads processed:", len(qkv_attention_store))
print("Shape of the output for a single head:", qkv_attention_store[0].shape)

Number of heads processed: 16
Shape of the output for a single head: torch.Size([17, 128])

In [21]:

# --- Step 1: Concatenate the outputs of all heads ---
# We combine the 16 tensors of shape [17, 128] along the last dimension.
# The result is a single tensor of shape [17, 16 * 128], which is [17, 2048].
stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)

# --- Step 2: Pass through the output projection layer ---
# This is a final linear layer that mixes the information from all heads.
# Its weight matrix 'o_proj.weight' has a shape of [dim, dim] -> [1024, 2048]
# which is incorrect for this model. Qwen3's o_proj is [dim, n_heads*head_dim] -> [1024, 2048]
# but since hidden_size (dim) is 1024, the effective output dimension should match.
# Let's check the weight shape.
w_layer0 = model_weights["model.layers.0.self_attn.o_proj.weight"]

# The expected shape is [dim, n_heads * head_dim] -> [1024, 2048] for this specific model,
# which has a different hidden size vs. attention output size.
# Actually, the attention output size IS the hidden size. n_heads*head_dim = 16*128 = 2048.
# The `o_proj` weight must project this 2048-dim vector back to the model's `dim` of 1024.
# So, the weight shape must be [dim, n_heads * head_dim] -> [1024, 2048].

# Let's re-verify this logic.
# Ah, I see a mismatch in my previous comments vs. the config.
# `hidden_size` is 1024, `n_heads` is 16, `head_dim` is 128.
# Total attention output dimension = 16 * 128 = 2048.
# This 2048-dim vector must be projected back down to `hidden_size` (1024).
# So `w_layer0` should be [1024, 2048]. Let's check.
print("Shape of the output projection weights (o_proj):", w_layer0.shape)

# Perform the matrix multiplication.
# Input: [17, 2048] @ [2048, 1024] -> Output: [17, 1024]
embedding_delta = torch.matmul(stacked_qkv_attention, w_layer0.T)

# --- Step 3: Verify the shape ---
# The output is now back to the model's main embedding dimension.
# This `embedding_delta` is the final result of the entire attention block.
print("\nShape of the final attention output (embedding_delta):", embedding_delta.shape)

Shape of the output projection weights (o_proj): torch.Size([1024, 2048])

Shape of the final attention output (embedding_delta): torch.Size([17, 1024])

In [22]:

# --- Step 1: Add the attention output to the original embeddings ---
# We add the output of the attention block back to the input of the layer.
# This creates the output of the first sub-layer.
embedding_after_attention = token_embeddings_unnormalized + embedding_delta

# --- Step 2: Verify the shape ---
# The shape remains the same: [17, 1024].
print("Shape of the embeddings after the first residual connection:", embedding_after_attention.shape)

Shape of the embeddings after the first residual connection: torch.Size([17, 1024])

In [23]:

# --- Step 1: Apply RMSNorm before the MoE block ---
# This is the 'post_attention_layernorm' in the model's weight names.
embedding_after_attention_normalized = rms_norm(embedding_after_attention, model_weights["model.layers.0.post_attention_layernorm.weight"])

# --- Step 2: Verify the shape ---
print("Shape after MoE input normalization:", embedding_after_attention_normalized.shape)

Shape after MoE input normalization: torch.Size([17, 1024])

In [24]:

# --- Step 2: The MoE Router ---
# The router is a simple linear layer that determines which expert to send each token to.
# It projects our [17, 1024] tensor to a [17, num_experts] tensor of scores (logits).
gate = model_weights["model.layers.0.mlp.gate.weight"]
router_logits = torch.matmul(embedding_after_attention_normalized, gate.T)

# We apply softmax to the logits to get probabilities, and then find the expert with the
# highest probability for each token.
routing_weights = torch.nn.functional.softmax(router_logits.float(), dim=1).to(torch.bfloat16)
routing_expert_indices = torch.argmax(routing_weights, dim=1)

print("Router logits shape:", router_logits.shape)
print("Expert chosen for each of the 17 tokens:", routing_expert_indices)

# --- Step 3: The Expert Layers ---
# Each expert is a SwiGLU-style Feed-Forward Network, similar to Llama 3's FFN,
# but smaller. We'll load the weights for both experts.
expert0_w1 = model_weights["model.layers.0.mlp.experts.0.gate_proj.weight"]
expert0_w2 = model_weights["model.layers.0.mlp.experts.0.down_proj.weight"]
expert0_w3 = model_weights["model.layers.0.mlp.experts.0.up_proj.weight"]

expert1_w1 = model_weights["model.layers.0.mlp.experts.1.gate_proj.weight"]
expert1_w2 = model_weights["model.layers.0.mlp.experts.1.down_proj.weight"]
expert1_w3 = model_weights["model.layers.0.mlp.experts.1.up_proj.weight"]

# --- Step 4: Process tokens with their chosen expert ---
# We'll create a final output tensor and iterate through each token, sending it
# to its assigned expert.

final_expert_output = torch.zeros_like(embedding_after_attention_normalized)

for i, token_embedding in enumerate(embedding_after_attention_normalized):
    chosen_expert_index = routing_expert_indices[i]
    
    # Get the weights for the chosen expert
    if chosen_expert_index == 0:
        w1, w2, w3 = expert0_w1, expert0_w2, expert0_w3
    else:
        w1, w2, w3 = expert1_w1, expert1_w2, expert1_w3
    
    # Apply the SwiGLU activation for this token's chosen expert
    # This is the same FFN logic as in the Llama 3 notebook.
    silu_output = torch.nn.functional.silu(torch.matmul(token_embedding, w1.T))
    gated_output = silu_output * torch.matmul(token_embedding, w3.T)
    expert_output = torch.matmul(gated_output, w2.T)
    
    # Weight the expert's output by its routing probability
    final_expert_output[i] = expert_output * routing_weights[i, chosen_expert_index]

# --- Step 5: Verify the shape ---
# The output of the MoE block has the same shape as its input.
print("\nShape of the final MoE output:", final_expert_output.shape)

Router logits shape: torch.Size([17, 2])
Expert chosen for each of the 17 tokens: tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

Shape of the final MoE output: torch.Size([17, 1024])

In [25]:

# --- Step 1: Add the MoE output to the output from the attention block ---
# This is the second and final residual connection in the transformer block.
layer_0_embedding = embedding_after_attention + final_expert_output

# --- Step 2: Verify the shape ---
# The shape is still [17, 1024]. This tensor is now the final output of the entire first layer
# and will serve as the input to the second layer.
print("Shape of the final output for Layer 0:", layer_0_embedding.shape)

Shape of the final output for Layer 0: torch.Size([17, 1024])

In [26]:

# The final embedding starts as the output from the token embedding layer.
# We will update this tensor in-place as it passes through the layers.
final_embedding = token_embeddings_unnormalized

# Loop through each of the 28 layers of the transformer.
for layer in range(n_layers):
    
    # --- Attention Sub-Layer ---
    
    # 1. RMS Normalization before attention
    attention_input = rms_norm(final_embedding, model_weights[f"model.layers.{layer}.input_layernorm.weight"])
    
    # 2. Multi-Head Attention
    q_layer = model_weights[f"model.layers.{layer}.self_attn.q_proj.weight"].view(n_heads, head_dim, dim)
    k_layer = model_weights[f"model.layers.{layer}.self_attn.k_proj.weight"].view(n_kv_heads, head_dim, dim)
    v_layer = model_weights[f"model.layers.{layer}.self_attn.v_proj.weight"].view(n_kv_heads, head_dim, dim)
    w_layer = model_weights[f"model.layers.{layer}.self_attn.o_proj.weight"]
    
    qkv_attention_store = []
    for head in range(n_heads):
        q_layer_head = q_layer[head]
        k_layer_head = k_layer[head // (n_heads // n_kv_heads)]
        v_layer_head = v_layer[head // (n_heads // n_kv_heads)]

        q_per_token = torch.matmul(attention_input, q_layer_head.T)
        k_per_token = torch.matmul(attention_input, k_layer_head.T)
        v_per_token = torch.matmul(attention_input, v_layer_head.T)

        q_per_token_rotated = torch.view_as_real(torch.view_as_complex(q_per_token.float().view(q_per_token.shape[0], -1, 2)) * freqs_cis_for_tokens).view(q_per_token.shape)
        k_per_token_rotated = torch.view_as_real(torch.view_as_complex(k_per_token.float().view(k_per_token.shape[0], -1, 2)) * freqs_cis_for_tokens).view(k_per_token.shape)
        
        qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T) / (head_dim**0.5)
        qk_per_token_masked = qk_per_token + mask
        qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_masked.float(), dim=1).to(torch.bfloat16)
        
        qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
        qkv_attention_store.append(qkv_attention)
        
    stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
    embedding_delta = torch.matmul(stacked_qkv_attention, w_layer.T)
    
    # 3. First Residual Connection
    embedding_after_attention = final_embedding + embedding_delta
    
    # --- Mixture-of-Experts Sub-Layer ---
    
    # 1. RMS Normalization before MoE
    moe_input = rms_norm(embedding_after_attention, model_weights[f"model.layers.{layer}.post_attention_layernorm.weight"])
    
    # 2. Router
    gate = model_weights[f"model.layers.{layer}.mlp.gate.weight"]
    router_logits = torch.matmul(moe_input, gate.T)
    routing_weights = torch.nn.functional.softmax(router_logits.float(), dim=1).to(torch.bfloat16)
    routing_expert_indices = torch.argmax(routing_weights, dim=1)
    
    # 3. Expert Layers
    final_expert_output = torch.zeros_like(moe_input)
    
    expert0_w1 = model_weights[f"model.layers.{layer}.mlp.experts.0.gate_proj.weight"]
    expert0_w2 = model_weights[f"model.layers.{layer}.mlp.experts.0.down_proj.weight"]
    expert0_w3 = model_weights[f"model.layers.{layer}.mlp.experts.0.up_proj.weight"]

    expert1_w1 = model_weights[f"model.layers.{layer}.mlp.experts.1.gate_proj.weight"]
    expert1_w2 = model_weights[f"model.layers.{layer}.mlp.experts.1.down_proj.weight"]
    expert1_w3 = model_weights[f"model.layers.{layer}.mlp.experts.1.up_proj.weight"]
    
    for i, token_embedding in enumerate(moe_input):
        chosen_expert_index = routing_expert_indices[i]
        
        if chosen_expert_index == 0:
            w1, w2, w3 = expert0_w1, expert0_w2, expert0_w3
        else:
            w1, w2, w3 = expert1_w1, expert1_w2, expert1_w3
        
        silu_output = torch.nn.functional.silu(torch.matmul(token_embedding, w1.T))
        gated_output = silu_output * torch.matmul(token_embedding, w3.T)
        expert_output = torch.matmul(gated_output, w2.T)
        
        final_expert_output[i] = expert_output * routing_weights[i, chosen_expert_index]
        
    # 4. Second Residual Connection
    final_embedding = embedding_after_attention + final_expert_output

# --- Verify the final shape ---
print("Shape of the final embeddings after all layers:", final_embedding.shape)

Shape of the final embeddings after all layers: torch.Size([17, 1024])

In [27]:

# --- Step 1: Apply the final RMS Normalization ---
# This uses the 'model.norm.weight' tensor.
final_embedding_normalized = rms_norm(final_embedding, model_weights["model.norm.weight"])

# --- Step 2: Get the weights for the LM Head ---
# Due to weight tying, the LM head weights are the same as the token embedding weights.
lm_head_weights = model_weights["model.embed_tokens.weight"]

# --- Step 3: Select the embedding for the last token ---
# To predict the next token, we only need the output corresponding to the
# last token of our input sequence (which was 'assistant\n').
last_token_embedding = final_embedding_normalized[-1]

# --- Step 4: Calculate the Logits ---
# Perform the final matrix multiplication to get the scores for each word in the vocabulary.
# Input: [1024] @ [151936, 1024].T -> Output: [151936]
logits = torch.matmul(last_token_embedding, lm_head_weights.T)

# --- Step 5: Verify the shapes ---
print("Shape of the final normalized embeddings:", final_embedding_normalized.shape)
print("Shape of the last token's embedding:", last_token_embedding.shape)
print("Shape of the final logits:", logits.shape)

Shape of the final normalized embeddings: torch.Size([17, 1024])
Shape of the last token's embedding: torch.Size([1024])
Shape of the final logits: torch.Size([151936])

In [28]:

# --- Step 1: Find the token with the highest logit score ---
# torch.argmax() finds the index of the maximum value in the logits tensor.
# This index is the predicted token ID.
next_token_id = torch.argmax(logits, dim=-1)

print(f"Predicted Token ID: {next_token_id.item()}")

# --- Step 2: Decode the token ID to get the word ---
# We use our tokenizer's .decode() method to turn the ID back into text.
predicted_word = tokenizer.decode([next_token_id.item()])

print(f"\nPredicted Word: '{predicted_word}'")

Predicted Token ID: 64088

Predicted Word: 'stem'