Implementing Modern Transformer Architecture from Scratch

2.2 Building Block 1: RMSNorm (Automatic Volume Balancer)

2.2.1 Why Must Deep Networks Be Normalized?

Picture a real-world audio setup: 12 loudspeakers connected in series, passing an audio signal along the chain. If each speaker amplifies the sound slightly (say, $1.2\times$ gain), by the 12th speaker $(1.2^{12} \approx 8.91)$, deafening howling feedback and distortion will shatter the speakers! Conversely, if each speaker slightly attenuates the sound ($0.8\times$), by the 12th speaker $(0.8^{12} \approx 0.068)$, the signal vanishes into inaudibility.

The essence of Normalization is placing an agile Automatic Gain Controller (AGC) before and after each layer: no matter how wild the input signal is, it is immediately normalized to standard energy, guaranteeing stable numerical propagation across dozens of layers.

2.2.2 RMSNorm vs. Traditional LayerNorm: Discarding the Redundant Mean

Traditional LayerNorm involves two distinct operations:

• First, compute the feature mean: $\mu = \frac{1}{d}\sum_{i=1}^d x_i$;
• Second, center each element $(x_i - \mu)$ and divide by standard deviation $\sigma$

However, empirical research reveals: in deep networks, what stabilizes gradients is scaling by variance energy; subtracting the mean offers virtually zero mathematical value! Subtracting the mean requires an extra memory access pass across GPU SRAM, dragging down kernel speeds.

RMSNorm (Root Mean Square Normalization) makes an elegant simplification: compute the root mean square directly, and divide!

Where:

• $\epsilon$ (Epsilon) is a tiny protection constant (e.g., $10^{-6}$) preventing division by zero;
• $\gamma$ (Gamma) is a learnable scaling parameter of dimension d_model (initialized to 1.0), allowing the network to adjust feature gains during training.

📐 Mathematical Deep Dive: Scale Invariance Proof of RMSNorm

Why does dividing directly by root mean square stabilize deep networks? Let us examine its Scale Invariance:

Suppose preceding layers unintentionally amplify the input signal by a factor of $\alpha$ (turning $x$ into $\alpha x$):

Substitute $\alpha x$ into the RMSNorm formula:

Mathematical Takeaway: No matter how wildly previous layers amplify the signal, RMSNorm snaps its magnitude back to standard energy instantly!

2.2.3 Manual Numerical Walkthrough

To dispel any lingering doubt, let us calculate a concrete numerical example with 4 numbers: suppose a token's hidden vector is $x = [2.0, -2.0, 2.0, -2.0]$:

• Step 1 (Squaring):

$[2.0^2, (-2.0)^2, 2.0^2, (-2.0)^2] = [4.0, 4.0, 4.0, 4.0]$

• Step 2 (Mean):

$\text{Mean} = (4.0 + 4.0 + 4.0 + 4.0) / 4 = 4.0$

• Step 3 (Square Root for RMS):

$\text{RMS} = \sqrt{4.0} = 2.0$ (this is its root mean square energy!)

• Step 4 (Energy Normalization):

Divide each value by 2.0:

Behold! Numbers that could have been arbitrarily large are instantly tamed within standard unit energy bounds.

2.2.4 Line-by-Line Code Implementation of RMSNorm

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps # 防止除以 0 的极小保护值
        # 可学习参数 gamma: 长度为 dim 的一维向量，初始全为 1.0
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        # 1. x.pow(2).mean(-1, keepdim=True): 沿最后一个维度求平方的均值
        # 2. torch.rsqrt: 即 1.0 / sqrt(...)，求倒数平方根，GPU 上有专用硬件指令，一步完成！
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 💡 导师经验提示: 均方根计算很容易在半精度 (float16) 下发生微小数值下溢。
        # 因此业界标准实践是: 强转为 float32 计算均方根，算完后再转回张量原本的数据类型！
        output = self._norm(x.float()).type_as(x)
        # 乘以每个通道的可学习缩放系数 weight
        return output * self.weight

2.3 Building Block 2: RoPE Rotary Position Embedding (The Clock Hand Method)

This is the most geometrically elegant and intellectually satisfying operator in modern LLMs, and a frequent core topic in machine learning interviews.

2.3.1 Why Do LLMs Need Positional Awareness?

In human language, word order is decisive: "I do not like you" and "Do you not like me?" contain identical words but opposite meanings. Without position encodings, self-attention dot products cannot distinguish them.

• Sentence A: "I do not like you"
• Sentence B: "Do you not like me?"

Both sentences use identical words, but word order produces completely different meanings. However, self-attention dot product summation is mathematically permutation invariant! Without positional encoding, the model cannot distinguish word order.

2.3.2 The Fatal Flaw of Traditional Absolute Position Embeddings

In early models, the approach was rigid: "If you are the 1st word, issue an absolute badge labeled '1'; if you are the 2nd word, issue badge '2'..."

• Inability to Extrapolate (Cannot read long documents): If trained on sequences up to 1024 tokens, the model only possesses 1024 badges. Encountering token 1025 results in total failure;
• Violates Human Linguistic Intuition: When reading, humans do not care whether a verb is the 3421st word of a book; we only care about its relative distance to preceding nouns and modifiers.

2.3.3 The Clock Hand Epiphany: Relative Distance via Rotary Angles

How can we ensure the model focuses only on relative distance while extrapolating naturally? RoPE (Rotary Position Embedding) offers a brilliant insight:

Imagine a clock face:

• Suppose each token is a clock hand. Its position $m$ determines its rotation angle $m \cdot \theta;
• For example, token 1 points toward 1, while token 4 points toward 4;
• The angular difference between token 1 and token 4 is: $4\theta - 1\theta = 3\theta;

• If the entire sentence shifts by 10 positions (from token 11 to token 14):
• Token 11 rotates to $11\theta$, while token 14 rotates to $14\theta;
• Their angular difference remains strictly invariant: $14\theta - 11\theta = 3\theta$!

===================================================================================
                       RoPE 时钟旋转相对位置直观图解
===================================================================================

       位置 m = 1 的词 (1点钟)                   位置 m = 4 的词 (4点钟)
              12                                        12
           11    1                                   11    1
        10    ▲     2                             10          2
       9      │      3                           9             3 ──►
        8     │     4                             8           4
           7     5                                   7     5
              6                                         6
              
       【两根表针之间的相对夹角】: Δ = 4θ - 1θ = 3θ (只取决于它们之间相隔了 3 个字！)
===================================================================================

📐 Mathematical Deep Dive: Euler's Formula and Relative Position Proof

Why does rotating clock hands cause inner product attention scores to depend strictly on relative distance? Let us prove this with higher mathematics:

1. Euler's Formula and Complex Multiplication

Leonhard Euler established one of the greatest formulas in human history:

In the complex plane, any 2D vector $(x_1, x_2)$ can be represented as a complex number $z = x_1 + i x_2$. Multiplying by $e^{i\theta}$ rotates the vector counterclockwise by angle $\theta$:

In real Cartesian coordinate form, the new rotated coordinates $(x_1', x_2')$ are:

In matrix notation, this forms the familiar 2D rotation matrix:

2. Relative Distance Emergence in Dot Products

Now consider query vector $q$ at position $m$ and key vector $k$ at position $n$. We apply positional rotations:

In complex space, the inner product of two vectors is $\langle u, v \rangle_{\mathbb{C}} = \text{Re}(u^* v)$:

Applying exponent addition rules:

Stop and gaze at this magnificent equation! Position indices $m$ and $n$, through complex inner product, emerge cleanly as a relative distance difference: $(m - n)$!

3. Multi-Frequency Waveband Spectral Analysis

In our model, per-head dimension is head_dim = 64, partitioned into 32 independent 2D rotation planes. The angular frequency for plane $i$ is:

The corresponding wavelength (number of tokens required to complete a full $2\pi$ cycle) is:

• High-Frequency Plane 1 ($i=1$): $\theta_1 = 10000^0 = 1.0$, wavelength $\lambda_1 = 2\pi \approx 6.28$ tokens. Highly sensitive to adjacent word syntax!
• Low-Frequency Plane 32 ($i=32$): $\theta_{32} = 10000^{-62/64} \approx 1.25 \times 10^{-4}$, wavelength $\lambda_{32} \approx 50,000$ tokens! Captures long-range document discourse context.

2.3.4 Dissecting the Core Operator: rotate_half

In PyTorch implementation, how do we efficiently compute complex rotation without explicit matrix multiplies? Notice the second coordinate term: (-x2, x1).

Taking the second half of the vector, negating it, and prepending it is the famous operator rotate_half:

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """
    将向量的前后两半拆开，后半截取负号并与前半截调换位置:
    输入: [x1, x2]
    输出: [-x2, x1]
    """
    half_dim = x.shape[-1] // 2
    x1 = x[..., :half_dim] # 前半截
    x2 = x[..., half_dim:] # 后半截
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """
    按照公式: x * cos(θ) + rotate_half(x) * sin(θ) 执行高速并行旋转
    """
    cos = cos.unsqueeze(0).unsqueeze(2) # 调整广播维度以匹配 [Batch, Seq_Len, Heads, Head_Dim]
    sin = sin.unsqueeze(0).unsqueeze(2)
    return (x * cos) + (rotate_half(x) * sin)

2.4 Building Block 3: SwiGLU Gated Feed-Forward Network (The Private Study)

2.4.1 Work Division: Seminars vs. Private Study

If a model only contains attention mechanisms, it acts like a room of people whispering rumors to each other without anyone having read actual books. In Transformers:

• Self-Attention handles "Information Gathering": Tokens ask peers what contextual intelligence they possess;
• Feed-Forward Network (FFN) handles "Knowledge Digestion": Tokens retire to their "Private Study", accessing deep long-term memory to process findings!

In modern LLMs, over 60% of all model parameters reside in the FFN; it is the true factual knowledge vault of the model!

2.4.2 Qualitative Leap: Traditional FFN vs. Modern SwiGLU

In legacy Transformers, the FFN structure was single-channeled and rigid:

All tokens pass through the exact same non-linear activation regardless of relevance.

Modern frontier architectures (LLaMA-3, Mistral, DeepSeek) exclusively use SwiGLU (Swish Gated Linear Unit):

• Content Pipe ($W_{up}$): Carries expanded candidate representations;
• Gating Pipe ($W_{gate} + \text{SiLU}$): Acts as an intelligent electromagnetic valve, generating continuous flow values between $0.0 \sim 1.0$;

Element-wise multiplication allows the network to: "close the valve against noise (multiply by 0), and open wide for vital clues (multiply by 1)"!

📐 Mathematical Deep Dive: Derivative and Self-Gating of SiLU

Why choose SiLU (Swish) over classic ReLU for the gate? Let us compare mathematical derivatives:

Differentiating $\text{SiLU}(x) = x \cdot \sigma(x)$:

Since the Sigmoid derivative satisfies $\sigma'(x) = \sigma(x)(1 - \sigma(x))$, substituting and simplifying yields:

| Activation Function | Mathematical Formulation | Negative Half-Axis Property | Derivative Continuity & Dying Neuron Risk | | ReLU | $\max(0, x)$ | Constant 0, hard cutoff | Discontinuous at $x=0$; negative inputs permanently "die" | | GELU | $x \cdot \Phi(x)$ | Smooth slight negative dip | Everywhere smooth, but computation involves costly error functions | | SiLU (Swish)<br>[ Our Choice ] | $x \cdot \sigma(x)$ | Smooth slight negative dip | Everywhere differentiable, self-gated, and native hardware friendly! |

This smooth differentiability ensures gradients flow uninterrupted during deep backpropagation, permanently resolving dying neuron hazards.

2.4.3 Line-by-Line Code Implementation of SwiGLU

import torch.nn.functional as F

class SwiGLUFeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        # 升维与门控投影矩阵 (无 bias，追求纯粹的高速线性特征)
        self.w_gate = nn.Linear(config.d_model, config.d_ffn, bias=False) # 512 -> 1408
        self.w_up = nn.Linear(config.d_model, config.d_ffn, bias=False)   # 512 -> 1408
        # 降维输出矩阵: 将知识浓缩回原本的 512 维空间
        self.w_down = nn.Linear(config.d_ffn, config.d_model, bias=False) # 1408 -> 512

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. gate 管道: 计算 0~1 的软阀门开度
        gate_out = F.silu(self.w_gate(x))
        # 2. up 管道: 抽取宽阔的知识特征
        content = self.w_up(x)
        # 3. 逐元素相乘 (阀门放行) 并降维输出
        return self.w_down(gate_out * content)

2.5 Building Block 4: GQA Grouped-Query Attention (The Detective Bureau)

Self-attention is the central engine enabling the model to comprehend context and exhibit intelligent behavior.

2.5.1 What Are $Q, K, V$? The Detective Bureau Metaphor

Rather than reciting matrix formulas, let us visualize $Q, K, V$ with a vivid detective scenario:

• $Q$ (Query): The detective's missing-person bulletin ("Seeking words describing bright moonlight");
• $K$ (Key): Identity badges worn by individuals in the room ("I am a nighttime celestial body");
• $V$ (Value): Dossier evidence carried in each person's satchel ("Here is the descriptive prose of Tang poetry").

Self-attention is an orderly investigative retrieval process:

• Scoring ($Q \cdot K^T$): The detective compares bulletin $Q$ against badge $K$ via dot products, assessing relevance;
• Normalization (Softmax): Converts raw scores into percentage probabilities (e.g., 80% focus on Word A, 20% on Word B);
• Retrieval ($\dots \times V$): Blends satchel contents $V$ according to attention weights into a fused representation!

📐 Mathematical Deep Dive: Why Divide Attention by $\sqrt{d_k}$? Variance Proof

This is one of the classic core interview questions in deep learning:

What disaster does that unassuming $\sqrt{d_k}$ prevent? What happens if we omit it?

1. Mathematical Derivation of Inner Product Variance

Suppose query vector $q$ and key vector $k$ are $d_k$-dimensional vectors ($d_k = 64$ in our model). Assume their elements are independent identically distributed with mean 0 and variance 1:

Their dot product is:

Let us derive the expectation and variance of inner product $S$:

• Expectation:

• Single Term Variance:

• Total Variance (Sum of independent variables):

• Standard Deviation:

2. The Fatal Disaster of Omitting $\sqrt{d_k}$: Softmax Gradient Saturation

When $d_k = 64$, the standard deviation of raw dot products escalates to $\sigma_S = \sqrt{64} = 8$! This means dot products frequently reach values like $+24$ or $-24$.

The maximum exponent $e^{24}$ overwhelms all other terms in the Softmax denominator, driving maximum probability to $p_{\max} \approx 1.0$ while all other probabilities vanish to $0.0$!

Let us examine the Softmax backward derivative:

• When $p_i \to 1.0$: $p_i(1 - p_i) \to 1.0 \times 0 = \mathbf{0.0}$!
• When $p_i \to 0.0$: $p_i(1 - p_i) \to 0.0 \times 1 = \mathbf{0.0}$!

Catastrophic Freeze: Gradients across the entire attention mechanism collapse to absolute zero, plunging backpropagation into total brain death!

3. The Life-Saving Antidote: Variance Normalization

Therefore, we must divide dot products by $\sqrt{d_k}$:

Mathematical Takeaway: Dividing by $\sqrt{d_k}$ resets inner product variance back to $1.0$, keeping inputs centered in the most sensitive gradient zone of Softmax!

2.5.2 Mathematical Limit of the Causal Mask

LLMs are autoregressive: token $t$ is strictly forbidden from peeking at future token $t+1$. In the attention score matrix, future positions are masked with $-\infty$:

According to calculus limits of the exponential function:

After Softmax, attention weights on future tokens vanish to $0.0\%$, completely barring temporal information leakage!

2.5.3 Why GQA? Detective Corps and Shared Secretaries

In Chapter 05 we will analyze autoregressive generation: at every step, past Key and Value tensors must reside in GPU VRAM (KV-Cache).

• Multi-Head Attention (MHA): 12 Query detective heads require 12 KV secretary pairs. For long sequences, KV-Cache memory consumes gigabytes of VRAM!
• Multi-Query Attention (MQA): All 12 detective heads share a single secretary. VRAM is slashed, but all heads see identical context, severely crippling model IQ.
• Grouped-Query Attention (GQA, Our Golden Choice): 12 Query heads share 4 KV secretary pairs (3 grouping)! Every 3 detectives share 1 dedicated secretary!
• VRAM footprint is instantly slashed by 66.7%!
• Semantic diversity and reasoning performance remain virtually unharmed!

| Attention Scheme | Query to KV Head Ratio | Per-Step KV Cache VRAM | Semantic Expressiveness & Speed Trade-off | | Multi-Head Attention (MHA) | $1:1$ (e.g., 12 Q heads : 12 KV heads) | $100\%$ (full memory) | Highest expressiveness, but severe inference VRAM bottleneck | | Multi-Query Attention (MQA) | $N:1$ (e.g., 12 Q heads : 1 KV head) | Drops to $\approx 8.3\%$ | Lowest memory, but noticeable quality degradation | | Grouped-Query Attention (GQA)<br>[ Our Choice ] | $G:1$ (e.g., 12 Q heads : 4 KV heads, 3) | Slashes VRAM by 66.7% | The Industry Benchmark: Matches MHA quality with MQA speed! |

2.5.4 Line-by-Line Implementation of Attention

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """将 [Batch, Seq_Len, 4, 64] 沿着头维度重复复制 3 次，扩展为 [Batch, Seq_Len, 12, 64]"""
    if n_rep == 1:
        return x
    batch, slen, n_kv, head_dim = x.shape
    return (
        x[:, :, :, None, :]
        .expand(batch, slen, n_kv, n_rep, head_dim)
        .reshape(batch, slen, n_kv * n_rep, head_dim)
    )

class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_q_heads = config.n_q_heads   # 8 个 Query 侦探头
        self.n_kv_heads = config.n_kv_heads # 2 个 KV 秘书头
        self.n_rep = config.n_q_heads // config.n_kv_heads # 组大小 = 4
        self.head_dim = config.head_dim     # 每个头的维度 = 64 (8 * 64 = 512)

        # 线性映射矩阵 (全无 bias，保持干净纯粹)
        self.wq = nn.Linear(config.d_model, config.n_q_heads * config.head_dim, bias=False)
        self.wk = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
        self.wv = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
        self.wo = nn.Linear(config.n_q_heads * config.head_dim, config.d_model, bias=False)

    def forward(self, x, cos, sin, mask=None, kv_cache=None):
        b_sz, seq_len, _ = x.shape
        # 1. 将 512 维特征分别投影出 Q, K, V，并拆分为多头
        xq = self.wq(x).view(b_sz, seq_len, self.n_q_heads, self.head_dim)
        xk = self.wk(x).view(b_sz, seq_len, self.n_kv_heads, self.head_dim)
        xv = self.wv(x).view(b_sz, seq_len, self.n_kv_heads, self.head_dim)

        # 2. 为 Q 和 K 注入时钟旋转位置编码 (RoPE)
        xq = apply_rotary_emb(xq, cos, sin)
        xk = apply_rotary_emb(xk, cos, sin)

        # 3. 维护增量推理 KV-Cache (第 05 章详细讲解)
        new_kv = None
        if kv_cache is not None:
            prev_k, prev_v = kv_cache
            xk = torch.cat([prev_k, xk], dim=1) if prev_k is not None else xk
            xv = torch.cat([prev_v, xv], dim=1) if prev_v is not None else xv
            new_kv = (xk, xv)

        # 4. GQA 广播: 把 2 个 KV 头复制 4 遍，对齐 8 个 Query 头
        xk_expanded = repeat_kv(xk, self.n_rep)
        xv_expanded = repeat_kv(xv, self.n_rep)

        # 5. 调整维度以符合 PyTorch 官方底层硬件加速格式: [Batch, Heads, Seq_Len, Head_Dim]
        xq = xq.transpose(1, 2)
        xk_expanded = xk_expanded.transpose(1, 2)
        xv_expanded = xv_expanded.transpose(1, 2)

        # 6. 调用超强底层硬件融合内核 (FlashAttention / SDPA)
        is_causal = (mask is None and seq_len > 1 and kv_cache is None)
        output = F.scaled_dot_product_attention(
            xq, xk_expanded, xv_expanded,
            attn_mask=mask,
            dropout_p=0.0,
            is_causal=is_causal # 自动构建严密的下三角因果防偷看掩码！
        )

        # 7. 多头汇报汇总: 还原回 [Batch, Seq_Len, 512]
        output = output.transpose(1, 2).contiguous().view(b_sz, seq_len, -1)
        return self.wo(output), new_kv

2.6 Two Crowning Touches: Residual Connections and Weight Tying

Before assembling the building blocks into a skyscraper, we must appreciate two brilliant innovations in modern deep learning.

2.6.1 Residual Connections: The Non-Stop Express Elevator

Inside each layer, you encounter this signature line: h = x + attn_out

• Why not simply write h = attn_out? Why add the original x?
• The Broken Telephone Metaphor: Imagine 12 people whispering a message down a line. By the 12th person, the original meaning is completely unrecognizable.
• The Express Elevator: The residual addition x + ... builds a non-stop elevator alongside the 12-story building. The raw input signal $x$ travels straight to upper floors without distortion, while backpropagation gradients flow losslessly downward!

2.6.2 Weight Tying: The Mirror Pact of Input and Output Vocabularies

In our 0.04B model:

• The input layer has a token_embeddings matrix mapping 4096 tokens to 512 dimensions (parameters: $4096 \times 512 = 2,097,152$);
• The output layer has an lm_head projection mapping 512 dimensions back to 4096 vocabulary scores (parameters: $2,097,152$).

Together, these two matrices consume 4.19M parameters—a massive 11.0% of our entire 0.04B model! Since one maps words to vectors and the other maps vectors back to words, why keep two duplicate matrices?

self.lm_head.weight = self.token_embeddings.weight

We force them to share the exact same physical memory matrix! The model instantly saves 2.10 million redundant parameters, shrinking VRAM and tightly aligning token representation spaces!

2.7 Final Assembly: Building the Complete MiniLLaMAForCausalLM

All foundations, beams, and windows are ready. We now assemble them into a full 35.93M parameter language model!

2.7.1 The Single-Layer Round Table: TransformerBlock

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        # 注意力层及其门前音量平衡器
        self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.attn = Attention(config)
        # 前馈网络层及其门前音量平衡器
        self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.ffn = SwiGLUFeedForward(config)

    def forward(self, x, cos, sin, mask=None, kv_cache=None):
        # 1. 第一阶段: 前置归一化 + 注意力情报交互 + 残差直连
        attn_out, new_kv = self.attn(self.attn_norm(x), cos, sin, mask=mask, kv_cache=kv_cache)
        h = x + attn_out

        # 2. 第二阶段: 前置归一化 + 前馈书房独立思考 + 残差直连
        out = h + self.ffn(self.ffn_norm(h))
        return out, new_kv

2.7.2 The Complete Causal Language Model: MiniLLaMAForCausalLM

class MiniLLaMAForCausalLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        # 1. 词嵌入映射
        self.token_embeddings = nn.Embedding(config.vocab_size, config.d_model)

        # 2. 预先计算 RoPE 旋转角度查找表，注册为非求导缓存 (Buffer)
        cos, sin = precompute_freqs_cis(config.head_dim, config.max_seq_len, config.rope_theta)
        self.register_buffer("cos_cached", cos, persistent=False)
        self.register_buffer("sin_cached", sin, persistent=False)

        # 3. 堆叠 12 层 TransformerBlock 圆桌研讨层
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])

        # 4. 终层归一化平滑
        self.norm = RMSNorm(config.d_model, eps=config.norm_eps)

        # 5. 语言模型预测输出头
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # 6. 开启权重绑定: 输入输出共享词典矩阵
        if config.tie_word_embeddings:
            self.lm_head.weight = self.token_embeddings.weight

    def forward(self, input_ids, labels=None, past_key_values=None, start_pos=0):
        b_sz, seq_len = input_ids.shape

        # 第一步: 查表获取初始 512 维特工档案袋
        h = self.token_embeddings(input_ids)

        # 第二步: 切出当前窗口对应的 RoPE 旋转角度
        cos = self.cos_cached[start_pos : start_pos + seq_len]
        sin = self.sin_cached[start_pos : start_pos + seq_len]

        new_past_key_values = [] if past_key_values is not None else None

        # 第三步: 顺次走过 12 层圆桌会议的深度研讨
        for idx, layer in enumerate(self.layers):
            layer_kv = past_key_values[idx] if past_key_values is not None else None
            h, updated_kv = layer(h, cos, sin, kv_cache=layer_kv)
            if new_past_key_values is not None:
                new_past_key_values.append(updated_kv)

        # 第四步: 终层音量稳定
        h = self.norm(h)
        # 第五步: 投影回 4096 个备选词的未归一化打分
        logits = self.lm_head(h)

        # 第六步: 如果给定了真实标签 labels，计算自回归交叉熵损失
        loss = None
        if labels is not None:
            # 错位 1 个 Token 计算自回归目标: 用当前词的打分去预测下一个真实的词！
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1)
            )

        return logits, loss, new_past_key_values

📐 Mathematical Deep Dive: From Maximum Likelihood to Autoregressive Cross-Entropy Loss

Why do lines 385~390 in the code shift logits and labels by 1 position before computing cross-entropy? Let us unveil the information theory roots:

1. Joint Probability Distribution of Language Models

Given a text sequence of length $T$, $W = (w_1, w_2, \dots, w_T)$, the core objective of an LLM is learning the joint probability distribution:

2. Maximum Likelihood Estimation (MLE) and Negative Log-Likelihood (NLL)

In statistics, we seek parameters $\theta$ that maximize the probability of real training corpora (Maximum Likelihood Estimation):

To avoid floating-point underflow from multiplying thousands of tiny probabilities, we take the natural logarithm $\ln$ and negate it to form Negative Log-Likelihood (NLL) minimization:

3. Why Is It Identical to Cross-Entropy?

At step $t$, the true next token $w_t$ is a one-hot probability vector $y_t \in \{0, 1\}^V$. Cross-entropy between predicted distribution $P$ and ground-truth $y$ is:

Since $y_t$ is 1 only at $v = w_t$ and 0 everywhere else, the formula collapses to:

Mathematical Conclusion: Taking the mean loss over the sequence is precisely what our PyTorch code executes: F.cross_entropy(shift_logits.view(-1, V), shift_labels.view(-1))!

2.8 Mentor Summary and Hands-on Lab

Through this systematic assembly, the inner workings of modern LLMs are completely demystified:

• Pre-RMSNorm: Like an agile mixer console, eliminates mean calculations with strict scale invariance to stabilize deep signals;
• RoPE: Uses Euler complex multiplications to naturally encode relative distance differences $(m - n)$ in inner products;
• SwiGLU: Smoothly differentiable SiLU gating delivers high factual knowledge storage density;
• GQA: Resets inner product variance to 1.0 with $\sqrt{d_k}$ and saves 75% KV-Cache VRAM via 4 grouping;
• Residual Connections & Weight Tying: Serve as express elevators and mirrored vocabularies, making the model deep yet lean.

💡 Hands-on Lab: Awakening Your 0.04B Model

Step 1: Verify Environment Dependencies (Install PyTorch)

Our model construction relies on native PyTorch tensor operations. If your terminal reports No module named 'torch', install PyTorch according to your hardware environment.

pip install torch

Step 2: Run in Project Root Directory (Troubleshooting Clinic ⚠️)

Before running tests, ensure your terminal working directory is the project root (llm-start), not the subdirectory src:

• Incorrect ❌: Navigating into llm-start\src and executing python -c "from src.model import ...";
• Correct ✅: Keeping terminal at project root d:\code\project\llm-start when running commands.

Method A: Recommended Automated Unit Test (Zero Friction)

Run the factory acceptance test script tests/test_model.py:

python tests/test_model.py

Expected Terminal Console Output:

============================================================
       Mini-LLaMA-0.04B 架构出厂体检与前向验证
============================================================
1. 正在初始化模型 (隐层维度: 512, 层数: 12, 词表: 4096)...
2. 模型总参数量核算: 35,926,528 个参数 (35.93 M)
3. 模拟输入张量形状: torch.Size([1, 5]) (Batch=1, Seq_Len=5)
4. 模型输出 Logits 形状: torch.Size([1, 5, 4096]) (符合预期: [1, 5, 4096])

🎉 恭喜！Mini-LLaMA-0.04B 神经网络架构前向推理与形状校验全部通过！✅
============================================================

Method B: Step-by-Step Experience in Python REPL

To inspect intermediate tensors interactively, launch python in your terminal and run:

# 1. 导入配置与模型定义
from src.model import MiniLLaMAForCausalLM, MiniLLaMAConfig
import torch

# 2. 实例化配置与模型
config = MiniLLaMAConfig()
model = MiniLLaMAForCausalLM(config)

# 3. 打印真实参数量
total_params = sum(p.numel() for p in model.parameters())
print(f"总参数量: {total_params:,} ({total_params/1e6:.2f} M)")
# 终端打印: 总参数量: 35,926,528 (35.93 M)

# 4. 模拟送入一句话 (Batch=1, 长度为 5 的假 Token 序列)
dummy_input = torch.tensor([[1, 232, 58, 209, 2]])
logits, loss, _ = model(dummy_input)
print(f"输出 Logits 形状: {logits.shape}") 
# 终端打印: 输出 Logits 形状: torch.Size([1, 5, 4096])

Observe the terminal printouts:

• Is the parameter count precisely the 35.93 M we calculated?
• Is the output tensor shape cleanly formatted as [1, 5, 4096]?

The thinking cortex is forged! In the next chapter, we construct the high-speed data pipeline that fuels this cortex with knowledge: Chapter 03 ｜ High-Performance Data Pipeline and Memory Mapping.