把前面的数组、函数变换、PyTree、随机 key、状态和 JIT 组合起来，写出一个 decoder-only Transformer。重点是看清数据流，不追求生产级性能。

1. 模型结构

token ids [B,T]
  → token embedding + position embedding [B,T,D]
  → N × Transformer block
      → RMSNorm → causal self-attention → residual
      → RMSNorm → MLP                   → residual
  → final RMSNorm
  → lm_head [B,T,V]

每个位置输出 V 个 logits，用来预测下一个 token。

1.1 先看懂结构中的名字

这里的 B 是 batch 中的序列数，T 是序列长度，D 是每个 token 的隐藏向量维度，V 是词表大小。模型的任务，是把整数 token id 转成向量，结合上下文加工这些向量，最后输出词表中各个候选 token 的分数。

Embedding（嵌入）是把离散编号转换成可学习的向量。 token id 只是词表里的编号，例如编号 100 并不比编号 50 多一倍语义。我们准备一张 [V,D] 的参数表，每个 token 对应一行，输入哪个 id，就取出哪一行向量。这张表会随训练更新，使向量逐渐包含对预测有用的信息。本例还使用位置嵌入（position embedding），给不同位置分配可学习的向量，让模型区分 token 在序列中的位置。

Transformer 是一种以注意力机制组织序列计算的网络结构。 本文采用 decoder-only 形式：只堆叠带因果约束的 Transformer block，没有独立的 encoder，也没有读取 encoder 输出的交叉注意力。每个 block 主要包含两个计算分支：注意力负责从允许访问的上下文中汇集信息，MLP 负责进一步加工每个位置的特征。多个 block 堆叠后，后面的层可以继续利用前面已经加工过的表示。

RMSNorm 是按均方根缩放向量的归一化层。 它对一个 token 的 D 个特征计算均方根，用这个尺度去除原向量，再乘以可学习的缩放参数。作用是控制进入后续计算的数值尺度，帮助训练保持稳定。它与 LayerNorm 的区别之一，是不先减去特征均值。本文把归一化放在注意力和 MLP 之前，这种组织方式通常称为 Pre-Norm。RMSNorm 的具体定义可参见原论文。

Residual（残差连接）就是把一个分支的输出加回它的输入：x + F(x)。 可以把 x 看成已有表示，F(x) 看成这一层补充的修正。这样每层可以在原表示上逐步修改，同时为反向梯度提供更直接的传播路径。为了能逐元素相加，分支输出必须回到与输入一致的 [B,T,D] 形状。

MLP（多层感知机）是这里的前馈网络，也常写成 FFN。 本文使用两次线性变换，中间接 GELU 激活：先把特征从 D 维扩展到 4D，经过非线性变换后，再投影回 D 维。注意力在不同 token 之间交换信息，MLP 则对每个位置的向量单独做特征变换，各位置共享同一套 MLP 参数。

MTP（Multi-Token Prediction，多 token 预测）是另一项概念。 普通 next-token 训练让每个位置预测紧随其后的一个 token；MTP 则增加对多个未来 token 的预测目标，具体实现可能增加预测模块和辅助损失，例如 DeepSeek-V3 技术报告中的设计。本文结构图和代码使用的是 MLP，没有实现 MTP。训练时同时计算序列中所有位置的 next-token loss，也仍然属于普通 next-token 训练：每个位置只负责自己的下一个 token。

最后的 lm_head 是语言模型的输出投影，把每个位置的 D 维隐藏表示变成 V 个 logits。训练时把这些分数送入交叉熵损失；生成时，才根据解码策略选择下一个 token。

1.2 为什么迷你模型也这样搭

这里的“迷你”主要体现在词表、隐藏维度、层数和序列长度较小，但仍然保留一条完整的语言模型计算路径。Embedding 解决离散输入的表示问题，位置表示提供顺序信息，因果注意力结合历史上下文，MLP 提供非线性特征加工，归一化与残差连接支持多层计算，输出投影再把隐藏表示映射回词表。保留这套结构，后面才能用小规模实验理解大模型里的前向、反向和参数更新。

搭建这类教学模型时，常见做法是先选一套明确的结构配方，并保持各层接口一致。本文采用以下约定：

• 主干中的隐藏状态统一为 [B,T,D]，注意力头数 H 能整除 D，每个头的维度为 D/H。
• 每个 block 保持输入输出形状不变，便于残差相加和重复堆叠；MLP 只在分支内部扩宽。
• 使用因果 mask，并把输入与目标错开一个 token，让训练目标与自回归生成一致。
• 使用 Pre-Norm，并在最后一个 block 后增加一次归一化，再接输出投影。
• 先采用容易核对的普通多头注意力、GELU MLP 和独立输出权重，跑通小 batch，再逐步引入更复杂的结构与性能优化。

这些是本文选择的常见搭法。不同模型会采用 LayerNorm 或 RMSNorm、学习式位置嵌入或 RoPE、GELU 或门控 MLP 等不同配方。做 AI Infra 时，应先明确参考模型的配方，再做数值对齐；这些差异都会改变计算，不能在性能优化时随意替换。原始 Transformer 的注意力、前馈网络和残差结构可参考 Attention Is All You Need，本文采用的是其中一些组件与后续常见设计的组合。

2. 基础组件

import math
import jax
import jax.numpy as jnp

def linear(w, x):
    return x @ w

def rms_norm(scale, x, eps=1e-6):
    x_f32 = x.astype(jnp.float32)
    rms = jnp.sqrt(jnp.mean(x_f32 ** 2, axis=-1, keepdims=True) + eps)
    return (x_f32 / rms).astype(x.dtype) * scale

RMSNorm 对最后一维归一化。scale 的形状是 [D]，广播到 [B,T,D]。

linear 中的 x @ w 只变换最后一维特征。例如 x 是 [B,T,D]、w 是 [D,Dout]，输出就是 [B,T,Dout]，batch 和序列维保持不变。本例省去 bias，便于把注意力放在矩阵形状和数据流上。

rms_norm 会为每个 token 独立计算一个尺度，不跨 batch 或序列位置求平均。keepdims=True 保留最后一维，方便广播；eps 则避免全零或很小的向量导致除零、放大数值误差。平方和均值先用 FP32 计算，有助于降低低精度归约的风险。不过最终输出 dtype 还会受到 scale 的 dtype 和类型提升规则影响，后续引入混合精度时，应明确参数、归约和输出各自使用什么精度。

3. 因果多头自注意力

def causal_self_attention(params, x, num_heads):
    # x: [B, T, D]
    B, T, D = x.shape
    assert D % num_heads == 0
    Dh = D // num_heads

    q = linear(params["wq"], x)
    k = linear(params["wk"], x)
    v = linear(params["wv"], x)

    # [B, T, D] -> [B, H, T, Dh]
    def split_heads(t):
        return t.reshape(B, T, num_heads, Dh).transpose(0, 2, 1, 3)

    q, k, v = map(split_heads, (q, k, v))

    # [B,H,T,Dh] @ [B,H,Dh,T] -> [B,H,T,T]
    scores = (q @ k.swapaxes(-1, -2)) / math.sqrt(Dh)

    # 第 i 个位置只能看见 <= i 的位置
    mask = jnp.tril(jnp.ones((T, T), dtype=bool))
    scores = jnp.where(mask, scores, -jnp.inf)
    weights = jax.nn.softmax(scores, axis=-1)

    context = weights @ v                    # [B,H,T,Dh]
    context = context.transpose(0, 2, 1, 3)  # [B,T,H,Dh]
    context = context.reshape(B, T, D)        # [B,T,D]
    return linear(params["wo"], context)

因果 mask 是语言模型与双向 encoder 的重要区别：训练时虽然一次输入整个序列，但位置 t 不能偷看未来 token。

3.1 Q、K、V 各负责什么

同一份输入 x 经过三个不同的投影得到 Q、K、V。可以把 Q（query）理解为“当前位置要寻找什么信息”，K（key）理解为“各位置提供什么线索”，V（value）则是“匹配后实际取回的内容”。这些含义是理解计算的类比，投影矩阵本身都是训练学出来的。

q @ kᵀ 为每个查询位置计算它与各个候选位置的匹配分数；沿最后一维做 softmax，将分数转换成权重；weights @ v 再按权重汇总内容。比如预测句尾时，某个头可以对前面相关的主体信息分配较大权重，但不会人为指定某个头必须学习某种语法关系。

多头的做法，是把 D 维投影拆成 H 组 Dh 维表示，让各头用各自的分数和权重汇总上下文，再拼回 [B,T,D]，经 wo 混合各头的结果。拆头和合头时，reshape、transpose 的顺序很关键；shape 对上了，也要确认维度的语义没有被打乱。

3.2 缩放与 mask 为什么放在这里

点积会累加 Dh 个分量，在常见的初始化假设下，分数的波动随维度增大。除以 sqrt(Dh) 是为了控制分数尺度，避免 softmax 过早变得过于尖锐。这是注意力中的维度缩放，与生成时人为调节采样分布的温度是不同的用途。

mask 中保留对角线，是因为当前位置的输入 token 已经给定，可以用它预测下一个 token；对角线右边才是未来输入。被屏蔽的位置设成负无穷，softmax 后权重为 0。本例没有 padding，每一行至少能看到自身；以后合并 padding mask 时，要留意整行都被屏蔽会导致 softmax 出现无效数值。

生产代码通常使用融合的 scaled-dot-product attention/FlashAttention 实现，以降低中间 [T,T] 矩阵的内存开销。

这里显式生成 [B,H,T,T] 的分数和权重，容易读懂，但内存占用会随序列长度平方增长。替换融合实现时，重点核对缩放系数、因果 mask 的方向、维度布局和计算精度，并一起验证前向结果与反向梯度。

4. MLP 与 Transformer block

def mlp(params, x):
    h = jax.nn.gelu(linear(params["w1"], x))
    return linear(params["w2"], h)

def block(params, x, num_heads):
    x = x + causal_self_attention(
        params["attn"],
        rms_norm(params["norm1"], x),
        num_heads,
    )
    x = x + mlp(params["mlp"], rms_norm(params["norm2"], x))
    return x

残差连接让每一层学习对已有表示的修正，也为深层网络提供更直接的梯度路径。

MLP 的两个权重形状分别是 [D,4D] 和 [4D,D]。中间的 GELU 提供非线性；如果去掉激活，连续两个无 bias 的线性变换可以合并成一次矩阵乘法，表达能力就不同了。4D 是传统 FFN 中常见的扩展宽度，也是本例的选择；采用 SwiGLU 等门控结构时，中间维度通常需要结合参数量和计算量重新设置。

block 中第一次赋值完成注意力分支的残差更新，第二次 MLP 使用的是已经更新后的 x。两处分支都先做 RMSNorm，再计算分支输出，最后加回未经过该次归一化的主干输入。这个顺序定义了本文的 Pre-Norm block；做算子融合或重排代码时，应保持这条依赖关系。

5. 参数初始化

def normal(key, shape, scale=0.02):
    return scale * jax.random.normal(key, shape)

def init_block(key, dim):
    keys = jax.random.split(key, 6)
    return {
        "norm1": jnp.ones((dim,)),
        "attn": {
            "wq": normal(keys[0], (dim, dim)),
            "wk": normal(keys[1], (dim, dim)),
            "wv": normal(keys[2], (dim, dim)),
            "wo": normal(keys[3], (dim, dim)),
        },
        "norm2": jnp.ones((dim,)),
        "mlp": {
            "w1": normal(keys[4], (dim, 4 * dim)),
            "w2": normal(keys[5], (4 * dim, dim)),
        },
    }

def init_model(key, vocab_size, max_seq_len, dim, num_layers):
    keys = jax.random.split(key, num_layers + 3)
    return {
        "token_embedding": normal(keys[0], (vocab_size, dim)),
        "position_embedding": normal(keys[1], (max_seq_len, dim)),
        "blocks": [
            init_block(keys[i + 2], dim) for i in range(num_layers)
        ],
        "final_norm": jnp.ones((dim,)),
        "lm_head": normal(keys[-1], (dim, vocab_size)),
    }

这里为教学清晰度保留独立的 lm_head。许多语言模型会让输出投影与 token embedding 共享权重（weight tying）。

共享时，可以用 embedding 表的转置完成输出投影：x @ token_embedding.T，从而复用 [V,D] 的参数。是否共享属于模型配置，需要训练和加载 checkpoint 时保持一致。

每个随机权重使用独立拆出的 key，避免相同形状的矩阵因复用 key 而初始化成相同数值。归一化的 scale 初始化为全 1，使初始归一化不再叠加额外的逐通道缩放。

固定标准差 0.02 是这里方便演示的初始化方案。模型扩宽、加深后，需要结合层宽、残差分支和参考训练配方重新检查初始化尺度，观察激活与梯度是否稳定。参数表中的形状也可以用于估算参数量：本例每个 block 的四个注意力矩阵共 4D²，两个 MLP 矩阵共 8D²，再加两组 RMSNorm 的 2D 个参数。

6. 完整前向传播

def model(params, token_ids, num_heads):
    # token_ids: [B,T]
    _, T = token_ids.shape
    token_x = params["token_embedding"][token_ids]    # [B,T,D]
    pos_x = params["position_embedding"][:T]         # [T,D]
    x = token_x + pos_x

    for layer_params in params["blocks"]:
        x = block(layer_params, x, num_heads)

    x = rms_norm(params["final_norm"], x)
    return linear(params["lm_head"], x)               # [B,T,V]

层数是 Python 列表长度，对一个模型配置来说是静态的，因此这个循环可在 tracing 时展开。层数很大或需要更紧凑的编译表示时，可以进一步学习 jax.lax.scan。

token_embedding[token_ids] 是按 id 查表，输入应当是合法范围内的整数；position_embedding[:T] 给当前序列的每个位置取出一个向量，并沿 batch 维广播相加。这里要求 T <= max_seq_len，因为位置表只初始化了这么多行。后续换成 RoPE 时，位置处理会进入 Q、K 的计算，不能简单沿用这里的位置表相加逻辑。

输出 [B,T,V] 同时包含所有位置对其下一个 token 的预测分数。训练时可以并行计算所有位置的 loss；生成时则通常取最后一个位置的 logits，选出新 token 后继续生成。本例每次处理完整输入前缀，适合验证模型计算，实际解码会使用 KV cache 复用历史 K、V。

7. 损失与训练步骤

def cross_entropy(logits, targets):
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    chosen = jnp.take_along_axis(
        log_probs, targets[..., None], axis=-1
    )[..., 0]
    return -chosen.mean()

def loss_fn(params, batch, num_heads):
    logits = model(params, batch["inputs"], num_heads)
    return cross_entropy(logits, batch["targets"])

@jax.jit
def train_step(params, batch):
    loss, grads = jax.value_and_grad(loss_fn)(params, batch, 4)
    new_params = jax.tree.map(
        lambda p, g: p - 1e-3 * g,
        params,
        grads,
    )
    return new_params, loss

为避免 num_heads 的动态/静态问题，教学示例直接在 jitted 函数中固定为 4。更整洁的课程版本可以把模型配置放在闭包中，或使用 static_argnames。

7.1 输入和标签如何对齐

对于一段 token 序列 [a,b,c,d]，通常构造 inputs = [a,b,c]、targets = [b,c,d]。位置 0 看见 a，预测 b；位置 1 看见 a,b，预测 c。因果 mask 防止模型读取后续输入，标签右移则定义了每个位置到底预测什么。两者需要配合，否则可能出现训练 loss 很低、实际生成却不正常的情况。

log_softmax 直接得到各个 token 的对数概率，通常比先 softmax 再取 log 更稳定。take_along_axis 按正确目标 id，从词表维选出对应的对数概率，再取负数、求平均，就是交叉熵。本例默认所有位置都有效；如果加入 padding 或只训练回答部分，需要改为按有效位置的 loss 求和，再除以有效 token 数，并处理有效数量为 0 的情况。

7.2 一次训练步骤做了什么

value_and_grad 对参数树求导，返回这一步更新前的 loss 和梯度；jax.tree.map 再按相同的树结构逐项执行 SGD 更新。返回的 new_params 是更新后的参数，而 loss 对应更新之前的那次前向计算。如果要观察本次更新后的 loss，需要用新参数再计算一次。

这里使用固定学习率的普通 SGD，便于理解完整链路。实际语言模型训练通常还会组织 AdamW 状态、学习率调度、梯度裁剪等逻辑。先用这个简单版本确认前向、反向和更新正确，再接入更复杂的训练状态，排查问题会更容易。

8. 最小运行检查

V, T, D, H, L = 32, 16, 64, 4, 2
key = jax.random.key(0)
params = init_model(key, V, T, D, L)

inputs = jnp.arange(2 * T, dtype=jnp.int32).reshape(2, T) % V
targets = (inputs + 1) % V
batch = {"inputs": inputs, "targets": targets}

logits = model(params, inputs, H)
print(logits.shape)  # [2,16,32]

params, loss = train_step(params, batch)
print(float(loss))

这个例子先用人造数据检查模型是否能运行。

这里的 token id 按顺序循环，targets = (inputs + 1) % V 对应这个人造序列的下一个 token，因此构成一个容易学习的玩具任务。真实文本应该先分词，再从同一段序列构造错开一位的输入和标签，不能直接把 token id 加 1 当作真实的下一个词。V=32、D=64、L=2 等配置主要用于快速运行和检查计算路径。

运行一次能输出形状和 loss，说明基本调用链路可以执行；接下来还应检查：

• logits shape 是否为 [B,T,V]；
• 所有参数和梯度是否为有限值；
• 对一个很小的固定 batch 反复训练，模型能否过拟合、loss 能否明显下降；
• 改变未来 token 后，较早位置的 logits 是否保持不变（验证因果 mask）；
• 参数量是否与手算数量级一致。

本例没有 bias，也没有共享 embedding 与输出权重，总参数量为 2VD + max_seq_len × D + L × (12D² + 2D) + D，代入上面的配置是 103,744。可以把参数树各叶子的元素数相加，与这个数字比较，帮助发现漏初始化或多注册参数的问题。

验证因果性时，应固定参数和其他条件，只替换位置 k 及之后的输入，比较 k 之前的 logits。做小 batch 过拟合检查时，固定数据和初始化，观察多步 loss 是否明显下降；如果无法下降，先核对标签偏移、mask、梯度和实际参数变化，再扩大模型规模。

小结

这个模型已经包含大语言模型的核心数学结构，但离生产训练还有很长距离：

• tokenizer、数据管线、padding mask；
• AdamW、学习率调度、梯度累积与裁剪；
• dropout、checkpoint、评估和断点续训；
• RoPE、SwiGLU、分组查询注意力等现代结构；
• bfloat16/混合精度和稳定性策略；
• 多设备 sharding、张量/数据并行；
• FlashAttention、KV cache 和高性能解码。

回过头看，大型框架就在替我们组织这些数组、函数变换、训练状态和并行计算。这个迷你模型可以作为后续实验的参考版本：每引入一项结构变化或性能优化，先检查对应模块的形状、前向数值和梯度，再比较小规模训练的 loss 走势，逐步建立对实现的信心。