Qwen3-0.6B From Scratch

Hi, I'm Sidhant. This was my first technical blog post: a way to strengthen my own understanding while sharing the path with anyone following along.

After working with RAG and AI agents, I wanted to refresh my PyTorch fundamentals. Sebastian Raschka's Qwen3 notebook was the prompt to dig in. The implementation for this post is available in the companion repository.

Qwen is Alibaba Cloud's open-weight LLM family. This post focuses on the 0.6B dense model and its use of Rotary Positional Embeddings, RMSNorm, and Grouped Query Attention. Its small size makes it approachable to implement on consumer hardware while still exposing the architecture used by much larger models.

• Grouped Query Attention
• Rotary Positional Embeddings (RoPE)
• SiLU activations
• A Mixture-of-Experts variant, alongside the dense feed-forward model covered here.

Config

Self-attention is permutation-invariant: changing input order changes output order, but does not itself give the model a sense of sequence. Positional encoding supplies that missing information.

Absolute positional encodings emphasize each token's absolute location. Relative distance is often more useful, and adding a positional vector directly to an embedding can blur its semantic representation.

RoPE

Rotary Positional Embeddings, introduced in RoFormer, encode relative distance through the angle between token representations. Only Query and Key are rotated, so Values remain untouched and the embedding semantics are preserved.

Application

Take a Query matrix with shape [5, 8]: five positions, each with an eight-dimensional head. We pair adjacent dimensions, then rotate every pair according to that token position and the pair's frequency.

Frequencies and angles

For an eight-dimensional head, the four dimension pairs use indices 0, 2, 4, and 6. With a base of 100,000, their frequencies are [1.0000, 0.1000, 0.0100, 0.0010].

Rotation

Each position receives one 2×2 rotation matrix per pair. Position zero has angle zero, so it stays unchanged. At position one, the four matrices use angles 1, 0.1, 0.01, and 0.001. Applying the block-diagonal rotation to every row produces the rotated Query matrix.

The same transformation applies to Keys after Q and K have been split into heads. Sine and cosine values are precomputed for the maximum context length and stored as PyTorch buffers: persistent model state, but not trainable parameters.

def calculate_sin_cos():
    i = torch.arange(0, config.head_dim, 2).float()
    freqs = 1.0 / (config.base ** (i / config.head_dim))
    positions = torch.arange(config.max_ctx_length).float()
    angles = positions[:, None] * freqs[None, :]
    angles = torch.cat([angles, angles], dim=1)
    return torch.sin(angles), torch.cos(angles)

def rotate(x, sin, cos):
    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    rotation = torch.cat((-x2, x1), dim=-1)
    return (x * cos[:x.shape[-2]].unsqueeze(0).unsqueeze(0)) + (rotation * sin[:x.shape[-2]].unsqueeze(0).unsqueeze(0))

Root Mean Square Normalization stabilizes and accelerates training by normalizing inputs using their RMS value. Normalization helps control exploding gradients and vanishing gradients in deep networks.

• γ is a learned scale parameter.
• β is a learned shift parameter.
• ε prevents division by zero when the RMS is near zero.

class RMSNORM(nn.Module):
    def __init__(self, embed_dim, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(embed_dim))
        self.beta = nn.Parameter(torch.zeros(embed_dim))
        self.eps = eps

    def forward(self, x):
        variance = x.pow(2).mean(dim=-1, keepdim=True)
        rms = torch.sqrt(variance + self.eps)
        return self.gamma * (x / rms) + self.beta

SiLU, or Sigmoid Linear Unit, is also known as Swish when parameterised. It is non-monotonic: it can decrease across part of its input range before increasing again.

SiLU is Swish when β = 1 in fβ(x) = x · σ(βx).

SiLU vs ReLU

ReLU sets all negative inputs to zero, which can contribute to dying neurons. SiLU permits small negative values and yields smoother gradients, helping preserve learning signals.

SiLU vs GELU

Both GELU and SiLU are smooth, non-monotonic activations. SiLU is cheaper to compute, and GELU can be approximated by a scaled SiLU.

Grouped Query Attention shares a Key-Value pair across multiple Query heads. It reduces redundant projections and memory traffic while retaining more flexibility than sharing a single KV pair across every Query head.

Multi-Head Attention

Standard MHA assigns every head its own Query, Key, and Value matrices. It is expressive, but the number of K and V matrices raises memory costs during training and inference.

Multi-Query Attention

MQA keeps one Query per head but makes all heads share a single Key-Value pair. It significantly cuts parameters and memory, at some cost to accuracy.

GQA

GQA is the middle ground: groups of Query heads share KV projections. It reduces memory, parameters, and matrix operations, improving training and inference throughput with relatively small quality loss compared with full MHA.

Implementation

class GQA(nn.Module):
    def __init__(self):
        super().__init__()
        self.w_q = nn.Linear(config.embed_dim, config.n_heads * config.head_dim, bias=False)
        self.w_k = nn.Linear(config.embed_dim, config.kv_heads * config.head_dim, bias=False)
        self.w_v = nn.Linear(config.embed_dim, config.kv_heads * config.head_dim, bias=False)
        self.out_proj = nn.Linear(config.n_heads * config.head_dim, config.embed_dim, bias=False)
        self.q_norm = RMSNORM(config.n_heads * config.head_dim)
        self.k_norm = RMSNORM(config.kv_heads * config.head_dim)

    def forward(self, x, sin, cos):
        b, seq_len, _ = x.size()
        Q, K, V = self.q_norm(self.w_q(x)), self.k_norm(self.w_k(x)), self.w_v(x)
        Q = Q.view(b, seq_len, config.n_heads, config.head_dim).transpose(1, 2)
        K = K.view(b, seq_len, config.kv_heads, config.head_dim).transpose(1, 2)
        V = V.view(b, seq_len, config.kv_heads, config.head_dim).transpose(1, 2)
        Q, K = rotate(Q, sin, cos), rotate(K, sin, cos)
        K = K.repeat_interleave(config.n_heads // config.kv_heads, dim=1)
        V = V.repeat_interleave(config.n_heads // config.kv_heads, dim=1)
        scores = Q @ K.transpose(-2, -1) / (config.head_dim ** 0.5)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
        weights = torch.softmax(scores.masked_fill(mask.unsqueeze(0).unsqueeze(0), float('-inf')), dim=-1)
        output = weights @ V
        return self.out_proj(output.transpose(1, 2).contiguous().view(b, seq_len, -1))

QWEN(
  (embeddings): Embedding(151936, 1024)
  (blocks): ModuleList(
    (0-27): 28 x Block(
      (rms_1): RMSNORM()
      (mgqa): GQA(w_q: 1024→2048, w_k: 1024→1024, w_v: 1024→1024)
      (rms_2): RMSNORM()
      (ffn): FFN(1024→3072→1024)
    )
  )
  (rms_3): RMSNORM()
  (lm_head): Linear(1024, 151936, bias=False)
)

1. Qwen3 Technical Report
2. Sebastian Raschka's Qwen3 Notebook
3. Designing Positional Encoding, Hugging Face
4. RoFormer paper
5. Grouped Query Attention paper
6. Grouped Query Attention, KLU.ai