Week 1 Day 5: The Qwen3 Model
On Day 5, we will combine the components from the previous chapters into the complete Qwen3 model.
Model-level tests require the corresponding model files. Start with the default 0.6B model; download the larger models only if you want to test them as well:
hf download Qwen/Qwen3-0.6B-MLX-4bit
# Optional larger models:
hf download Qwen/Qwen3-1.7B-MLX-4bit
hf download Qwen/Qwen3-4B-MLX-4bit
The deterministic checkpoint tests do not need downloaded model files. The default 0.6B model is required for the final Day 5 integration check; only checks for the optional larger models are skipped when those files are unavailable. A skip means that the model-backed check did not run, not that it passed.
Task 1: Implement Qwen3TransformerBlock
src/tiny_llm/qwen3_week1.py
📚 Readings
Qwen3 uses the following Transformer block structure:
input
/ |
| input_layernorm (RMSNorm)
| |
| Qwen3MultiHeadAttention
\ |
Add (residual)
/ |
| post_attention_layernorm (RMSNorm)
| |
| MLP
\ |
Add (residual)
|
output
Read head_dim independently from the model configuration rather than deriving it from hidden_size and the number of
attention heads. Match the reference block’s output dtype.
Run the tests for this task with:
pdm run test --week 1 --day 5 -- -k task_1
Task 2: Implement Embedding
src/tiny_llm/embedding.py
📚 Readings
The embedding layer maps token IDs (integers) to vectors of length embedding_dim. In this task, you will implement
that lookup operation.
Embedding::__call__
weight: vocab_size x embedding_dim
Input: N.. (tokens)
Output: N.. x embedding_dim (vectors)
This can be implemented with array indexing.
Token IDs may have any number of leading dimensions. The lookup appends embedding_dim to that shape and preserves the
BF16 model-facing dtype.
When input and output embeddings are tied, Qwen3 also uses the embedding weight as a linear projection from hidden vectors back to vocabulary logits.
Embedding::as_linear
weight: vocab_size x embedding_dim
Input: N.. x embedding_dim
Output: N.. x vocab_size
The projection accepts the same arbitrary leading dimensions, replacing only the final embedding_dim axis with
vocab_size and preserving BF16.
Run the tests for this task with:
# Deterministic lookup and projection checks always run. The downloaded model adds integration parity.
pdm run test --week 1 --day 5 -- -k task_2
Task 3: Implement Qwen3ModelWeek1
Now that we have built all the Qwen3 components, we can implement Qwen3ModelWeek1.
src/tiny_llm/qwen3_week1.py
You will not implement the process of reading model parameters from tensor files. Instead, load the model with mlx_lm,
then transfer its parameters into our implementation. The Qwen3ModelWeek1 constructor therefore accepts an MLX model.
The Qwen3 model has the following layers:
input
| (tokens: N..)
Embedding
| (N.. x hidden_size); note that hidden_size == embedding_dim
Qwen3TransformerBlock
| (N.. x hidden_size)
Qwen3TransformerBlock
| (N.. x hidden_size)
...
|
RMSNorm
| (N.. x hidden_size)
Embedding.as_linear OR linear (lm_head)
| (N.. x vocab_size)
output
Read the number of layers, hidden size, head dimension, and other configuration values from mlx_model.args, whose type
is defined by ModelArgs. The loaded weights
are available through mlx_model.model; use the Qwen3 implementation and model metadata to identify the corresponding
layer names.
Instantiate every configured Transformer block in model order, run all of them, and apply the final RMSNorm before the
vocabulary projection. head_dim, the query and key/value head counts, intermediate size, RMS epsilon, maximum positions,
and RoPE theta are separate configuration fields; map each one directly.
By this point, you have implemented RMSNorm. Replace the temporary Day 3 calls to mx.fast.rms_norm with
RMSNorm(head_dim, q_norm, eps=...) and RMSNorm(head_dim, k_norm, eps=...). They implement the same formula; the built-in
calls existed only to keep the GQA chapter focused on attention.
Different Qwen3 model variants map hidden vectors back to vocabulary logits in different ways. Some tie the input and
output embeddings and use Embedding.as_linear; others have a separate lm_head linear layer. Select the strategy with
mlx_model.args.tie_word_embeddings: if it is True, use Embedding.as_linear; otherwise, load and use lm_head.
The model takes a sequence of token IDs and returns unnormalized logits for every sequence position. On Day 6, we will use the final position’s logits to select the next token and generate a response.
The MLX models used in this course have quantized weights. Dequantize each
linear or embedding layer before loading it into tiny-llm by using the provided
quantize.dequantize_linear function, then store the readable Week 1 weight as
BF16. Model activations and layer outputs should remain BF16. A readable
attention or normalization expression may compute in FP32 for stability, but it
must cast its model-facing result back to BF16.
Pass mask="causal" to every Transformer block. For a one-token sequence the mask has no effect; for longer sequences,
it prevents each position from attending to future tokens.
Run the tests for this task with:
# The required 0.6B integration check exercises Task 3; download it if it is not already cached.
hf download Qwen/Qwen3-0.6B-MLX-4bit
pdm run test --week 1 --day 5 -- -k task_3
You may also download the optional 1.7B and 4B models from the commands at the top of this chapter. Their checks skip when the corresponding files are absent.
At the end of the day, you should be able to pass all tests of this day:
pdm run test --week 1 --day 5
Before treating this as a pass, confirm that the deterministic Task 2 checks ran and that the required 0.6B Task 3 integration check passed rather than skipped.
Your feedback is greatly appreciated. Join our Discord community.
Found an issue? Open an issue or pull request at github.com/skyzh/tiny-llm.
tiny-llm-book © 2025 by Alex Chi Z is licensed under CC BY-NC-SA 4.0.