TT Lab
开始
学习 学习路径 课程

MiniMind — 亲手从头到尾训练一个小型语言模型

缩小 MiniMind 配置构建模型，用数字验证每层的设计

在 TT Lab 中继续学习

目标

把 MiniMindConfig 缩小到 128 维、4 层、Q 头 4、KV 头 2、词表 1024 来创建模型，并让参数量在公式和模型之间对上。用 MiniMind 的代码亲自测出 GQA 节省的 KV 缓存、RMSNorm 的输出、RoPE 的相对位置性质、训练前的损失，以及 MoE 的总参数和激活参数。

为什么重要

一份配置文件就决定了模型的大小和服务成本。如果能用公式数出参数量，就能立刻知道有没有漏掉共享权重或归一化；如果能把 KV 缓存按每个令牌的字节数算出来，就能算出一张卡能接待多少人。 而且如今的模型所用的办法（RMSNorm、GQA、RoPE、SwiGLU、嵌入共享）各有道理。这个实验不是相信说明，而是用 MiniMind 的真实代码，以数字确认这些性质——不减去平均值的归一化、减半的 K·V、只留下距离的旋转。

步骤

  1. 把小配置保存到 /root/mm/arch/config.json（hidden_size 128，num_hidden_layers 4，vocab_size 1024，num_attention_heads 4，num_key_value_heads 2，max_position_embeddings 512）。
  2. 用公式计算参数量，并在模型中也数一遍，以 intermediate_size、per_layer、formula、model 写入 /root/mm/arch/count.json。
  3. 把第一层注意力的 q_proj、k_proj 输出大小，以及用 fp32 缓存一个令牌时的字节数（GQA 原样的情况 / KV 头数与 Q 头数相同的情况），写入 /root/mm/arch/gqa.json。
  4. 在 torch.manual_seed(0) 之后，把 randn(4, 128)*5+3 放入第一层的 input_layernorm，把输入、输出的 RMS 以及输出的平均值写入 /root/mm/arch/rmsnorm.json。
  5. 在 torch.manual_seed(0) 之后各取出一个 q 和 k，用 MiniMind 的 apply_rotary_pos_emb 把它们放在 (3,7)、(103,107)、(3,50) 的位置上，把内积写入 /root/mm/arch/rope.json。
  6. 用 mmkit.seed_all(0) 创建训练前的模型，用验证数据最前面的 16×128 个令牌测出损失，以 loss、ln_vocab 写入 /root/mm/arch/init_loss.json。
  7. 把在同样配置上加了 use_moe=True 的模型的总参数和激活参数，写入 /root/mm/arch/moe.json。
  8. 在 /root/mm/arch/report.md 中写 ## 파라미터는 어디에、## GQA 와 KV 캐시、## RoPE 와 RMSNorm 三节（三个标题为韩文，依次意为“参数在哪里”“GQA 与 KV 缓存”“RoPE 与 RMSNorm”），并放入第 2 步的参数量和第 3 步 GQA 每个令牌的缓存字节数。

参考

一份小配置

在 /root/mm/arch/config.json 中写入 hidden_size 128、num_hidden_layers 4、vocab_size 1024、num_attention_heads 4、num_key_value_heads 2、max_position_embeddings 512。其余的值使用 MiniMindConfig 的默认值。

MiniMind-3 是 768、8、6400、8、4。缩小层数和维度，参数量按维度的平方缩小，用 CPU 几分钟内就能训练完。评分器会用这个文件真正创建一个模型。

用公式数参数

用公式计算参数量（formula），并在模型中也数一遍（model），以 intermediate_size、per_layer、formula、model 写入 /root/mm/arch/count.json。两个值必须相同。

一层 = q、k、v、o 投影 + q_norm、k_norm（各 head_dim）+ gate、up、down（各 hidden×intermediate）+ 两个 RMSNorm（各 hidden）。再加上嵌入（词表×hidden，与输出层共享，所以只算一次）和最后一个 RMSNorm。intermediate_size 是 ceil(hidden·π/64)·64。

GQA 减少了什么

把第一层注意力的 q_proj.out_features、k_proj.out_features、n_rep，以及用 fp32 缓存一个令牌的字节数（K·V 两份 × 层数 × KV 头数 × head_dim × 4）——GQA 原样的情况（kv_bytes_per_token_gqa）和 KV 头数与 Q 头数相同的情况（kv_bytes_per_token_mha）——写入 /root/mm/arch/gqa.json。

MiniMind 的 Attention 只生成 num_key_value_heads 个 K·V，计算时用 repeat_kv 复制 n_rep 次。缓存（past_kv）里放的是复制 之前 的 K·V。

RMSNorm 不减去平均值

在 torch.manual_seed(0) 之后生成 x = torch.randn(4, 128) * 5 + 3，放入第一层的 input_layernorm，把输入和输出的 RMS（每行 sqrt(mean(x²)) 的平均值）以及整个输出的平均值，以 rms_before、rms_after、mean_after 写入 /root/mm/arch/rmsnorm.json。

MiniMind 的 RMSNorm 是 x * rsqrt(mean(x²) + eps) 再乘以权重（一开始为 1）。RMS 会回到 1，但由于不减去平均值，输入的 +3 会在输出中留下痕迹——如果是 LayerNorm，平均值就会变成 0。

RoPE 只留下距离

在 torch.manual_seed(0) 之后取出 q = torch.randn(1,1,1,head_dim)、k = torch.randn(1,1,1,head_dim)，用模型的 freqs_cos、freqs_sin 和 apply_rotary_pos_emb 把 q 转到第一个位置、把 k 转到第二个位置，把内积的值针对 (3,7)、(103,107)、(3,50)，以 dot_3_7、dot_103_107、dot_3_50 写入 /root/mm/arch/rope.json。

apply_rotary_pos_emb(q, k, cos[p:p+1], sin[p:p+1]) 会把 q 和 k 转到同一个位置，所以要分别调用 q 和 k，把它们各自转到自己的位置。距离同为 4 的两对内积，应该一直到小数点后第五位都相同。

训练前的损失是 ln(词表)

用 mmkit.seed_all(0) 固定种子之后重新创建模型，把 /opt/mm/ref/val.npy 最前面 16×128 个令牌（view(16, 128)）的损失（model(x, labels=x).loss），以 loss、ln_vocab 写入 /root/mm/arch/init_loss.json。

随机初始化的模型会给所有令牌几乎相同的概率（1/1024）。它的交叉熵是 ln 1024 ≈ 6.93。如果离这个值很远，就是初始化或标签处理出了问题。

MoE 很大，但一个令牌用得很少

把在第 1 步配置上加了 use_moe=True 的模型（默认 4 个专家、每个令牌 1 个）的总参数，以及一个令牌实际用到的参数（总数 − 一个专家的大小×专家数 + 一个专家的大小×每个令牌的专家数），以 num_experts、top_k、total_params、active_params 写入 /root/mm/arch/moe.json。

一个专家的大小，把名称里含 mlp.experts.0. 的参数加起来即可（每层各有一个，所以是所有层的 0 号专家之和）。MiniMind 的 trainer_utils.get_model_params 就是用同样的算法打印出“198M-A64M”的。

说明结构的报告

在 /root/mm/arch/report.md 中写 ## 파라미터는 어디에、## GQA 와 KV 캐시、## RoPE 와 RMSNorm 三节（三个标题为韩文，依次意为“参数在哪里”“GQA 与 KV 缓存”“RoPE 与 RMSNorm”），并以数字放入第 2 步的参数量（model）和第 3 步的 kv_bytes_per_token_gqa。

第一节写 FFN、注意力、嵌入各占百分之几，第二节写 GQA 把缓存减少了几倍，第三节各用一行写下第 4、5 步里看到的性质即可。