Level 08 | Architecture Tricks

Mission 1 TODO 1

Gemma 的 (1+w) 缩放：一个 +1 为何能稳住训练

🎯 先猜一猜（下注解锁）

Gemma 把 RMSNorm 的缩放从 x·w 改成 x·(1+w)，而且 weight 还是初始化成 0。看着只是个 +1。

先判断：这个 +1 最主要解决什么问题？

让 weight=0 的初始时刻等价于“纯归一化、原样通过”，使训练早期梯度平滑、不被清零让 RMSNorm 的计算量翻倍，从而更精确把输出强行限制在 0~1 之间

先押一个假设，下面的讲解就会解锁。猜错完全没关系——带着疑问读，记得最牢。

先补的知识

回忆 RMSNorm：先把向量按它自己的均方根归一化到稳定尺度，再乘一个可学习的缩放向量 weight，让模型微调每一维的强弱。
PyTorch 里 weight 常初始化成 0（或极小值）。标准写法 y = x_norm · w，此时 w=0 → 输出被乘成 0，等于把这一层“掐断”了。
Gemma 改成 y = x_norm · (1 + w)。w=0 时 (1+0)=1，相当于“先什么都不缩放、原样通过”。
工程细节：统计量要在 FP32 下算，最后 type_as(x) 转回原精度，避免半精度数值不稳。

图解原理

把缩放因子的“默认值”从 0 挪到 1。标准写法里参数学的是“我要把这维放大到多少”（从 0 起步很尴尬）；Gemma 写法里参数学的是“我要在默认 1 的基础上，多一点还是少一点”。训练刚开始 w≈0，整层就是一个老老实实的纯归一化，梯度平滑、不会一上来就把信号清零。

0. 先看坑点：标准写法会把信号清零

x_norm * ww=0 时 scale=0，输出被掐断 x_norm * (1+w)w=0 时 scale=1，纯归一化原样通过

Gemma 的参数语义

weight 不再表示“绝对缩放值”，而表示“相对默认 1 的偏移”。初始化为 0 时就是健康默认。

w=0 scale=1 训练平滑启动

TODO 主线：FP32 归一化已经算好，只补最后的 (1 + weight) 缩放

x输入 dtype x.float()稳定算统计量 x_normRMS 归一化 scale1 + self.weight outputtype_as(x)

notebook 里前面的归一化已经帮你铺好；TODO 只需要把缩放写对，并转回输入精度。

1. 为什么不直接初始化成 1

写成 1+w 后，参数仍能从 0 开始学习偏移，weight decay 拉向 0 时也代表“回到默认 scale=1”，不会把层拉向清零。

2. 测试应该怎么反查

当 self.weight=0，输出应该等于纯 RMSNorm 结果；如果输出接近 0，说明你写成了 x_norm * weight。

行业视角：训练稳定性是大模型的隐形成本

Gemma 的 1+w、LLaMA 的 Pre-Norm、各种初始化技巧，目标都一样：让训练早期别突然发散。一个小小的 +1，背后是大模型训练不要崩在前几千步的工程要求。

语法热身：把基准 1 写进公式

gain = nn.Parameter(torch.zeros(4))
scale = 1.0 + gain
y = x_norm * scale
y = y.type_as(x)

举一反三：搬回 TODO 1

gainself.weight x_norm前面 FP32 算出的归一化结果 scale1 + self.weight 返回 dtype.type_as(x)

迁移口诀：Gemma 不是乘 weight，而是乘 1 + weight；最后把 dtype 还给输入。

巩固一下

判断：Gemma RMSNorm 中 weight 初始化为 0 时，实际作用在归一化结果上的缩放是多少？

0（输出被清零） hidden_size 1（等价于不额外缩放，原样通过）

回到 notebook 的作业

这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。

TODO 1：在已算好的 x_norm 上实现 output = x_norm * (1 + self.weight)。记得 .type_as(x) 把结果转回输入精度（前面是 FP32 计算的）。跑测试：weight=0 时输出应与“无缩放纯归一化”结果一致。

Mission 2 TODO 2

权重绑定：输入查表和输出预测，共用同一张词表

🎯 先猜一猜（下注解锁）

embed_tokens（读词）和 lm_head（猜词）的权重形状一模一样，都是 [vocab, hidden]。Qwen 干脆让它们共用同一块参数。

先判断：权重绑定（weight tying）带来的最直接收益是？

省掉一整张 [vocab, hidden] 的参数，词表越大省得越多；且输入/输出梯度共同更新词向量让模型推理速度翻倍把 vocab_size 自动缩小一半

先押一个假设，下面的讲解就会解锁。猜错完全没关系——带着疑问读，记得最牢。

先补的知识

embed_tokens：把 token id 查成向量，权重形状 [vocab_size, hidden]。
lm_head：把 hidden 向量映射回每个词的分数(logits)，权重形状 [vocab_size, hidden]——和 embedding 一模一样！
正因为两者形状相同、语义相关（都是“词 ↔ 向量”的对应表），可以让它们共用同一块参数。
关键区分：相等(==，值一样) ≠ 共享(is，同一个对象)。绑定要的是“同一块内存”。

图解原理

模型开头“读词”（id→向量）和结尾“猜词”（向量→打分）其实在用同一套“词 ↔ 向量”的对应关系。与其各学一张大表，不如共用一张：参数省一半，而且训练时输入端和输出端的梯度都更新到同一张表，让词向量学得更充分。实现只有一行——让 lm_head.weight 直接指向 embed_tokens.weight。

0. 先看为什么能绑

embed_tokens 负责 id 到向量，lm_head 负责 hidden 到 vocab logits。两者的权重形状都是 [vocab, hidden]，都在描述“词和向量的对应关系”。

输入表embed_tokens.weight [vocab, hidden] 输出表lm_head.weight [vocab, hidden] 目标两者指向同一个 Parameter

绑定不是复制

TODO 2 不是让两张表数值相等，而是让两个属性指向同一块内存。复制一份权重，测试仍会判错。

要写赋值，不要写 clone()、copy_() 或重新创建 Parameter。

主线：两个名字，指向同一个 Parameter

before两张独立表 assignlm_head.weight = embed.weight after同一块参数

从此输入端和输出端的梯度都会更新同一张词表；词表越大，省下的参数越明显。

1. 怎么验真共享

对象级lm_head.weight is embed_tokens.weight 内存级lm_head.weight.data_ptr() == embed_tokens.weight.data_ptr() 不可靠shape 相同或数值相等

2. 参数账

词表 V、隐藏维 H，一张表就是 V*H 个参数。绑定后两张变一张，直接省下 V*H；大词表模型能省下好几 GB 显存。

不绑定
2*V*H 绑定
1*V*H

行业视角：绑不绑，是容量与成本的权衡

Qwen、GPT-2 绑定权重，省参数且让词向量获得输入/输出两端监督；超大模型有时解绑，让输入表征和输出预测各自特化。能讲清这个权衡，比只会写一行赋值更重要。

语法热身：让两个属性指向同一块权重

emb = nn.Embedding(100, 16)
head = nn.Linear(16, 100, bias=False)
head.weight = emb.weight
assert head.weight is emb.weight
assert head.weight.data_ptr() == emb.weight.data_ptr()

举一反三：搬回 TODO 2

embself.embed_tokens headself.lm_head 正确动作self.lm_head.weight = self.embed_tokens.weight 检查is 或 data_ptr()

迁移口诀：绑定是“同一块内存”，不是“复制一份一样的值”。

巩固一下

判断权重是否真的绑定（共享同一块内存），最可靠的检查是？

lm_head.weight is embed_tokens.weight（或 data_ptr 相等）两个模块的变量名相似两者 .shape 相等

回到 notebook 的作业

这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。

TODO 2：在 __init__ 里写 self.lm_head.weight = self.embed_tokens.weight（直接赋值，不要 clone/copy）。用 is 或 data_ptr() 验证两者共享内存；更新 embedding 后确认 lm_head 同步变化。确认 forward 的 logits 形状仍是 [..., vocab_size]，没被绑定影响。