08. Architecture Tricks | 模型配置与参数归属

难度： Easy | 环境： CPU-first | 标签： 模型结构, 架构技巧, Transformer | 目标人群： 模型结构学习者

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行：

(国内推荐：魔搭社区免费实例)

---

本节导读

本节聚焦配置层面的两个局部选择：输入词嵌入与语言模型输出头（LM Head）是否共享参数，以及 RMSNorm 的缩放参数如何初始化。它们不改变解码器块的主路径，却会影响参数账本、梯度归属和训练初期的数值行为。

关键词： Qwen, Gemma, Tie Embeddings

---

前置阅读

导语： 先理解 RMSNorm 在解码器块中的缩放位置，以及词嵌入、LM Head 与残差主路径在解码器层中各自承担什么角色；随后再把局部配置对应到参数归属和审计字段。

• 01. RMSNorm Tutorial | RMSNorm 教程
• 05. LLaMA3 Block Tutorial | LLaMA3 Block 教程
• 可选回看：P0: 09. PyTorch nn.Module Basics | PyTorch nn.Module 基础（理解 Parameter 与模块属性时使用）

Step 1：配置选择如何落在模型结构上

模型配置不仅记录层数和隐藏宽度，还决定参数由谁持有、残差如何缩放、哪些层使用不同组件，以及检查点能否按预期加载。读取真实模型时，应把配置声明映射到模块对象、参数账本和状态接口，而不是只摘录一个字段名。

配置类别	代表选择	需要核对的实际对象	直接后果
参数归属	词嵌入 / LM Head 是否共享	参数身份与梯度归属	参数量、更新路径与检查点语义
数值约定	归一化类型、eps、缩放初始化	每个块的归一化模块	初始尺度与加载兼容性
残差与深度	残差缩放、层级模式	块更新和层索引	深层激活与梯度稳定性
注意力 / MLP	KV 头、潜在状态、稠密/MoE	投影、路由器与专家模块	状态容量、路由和执行路径

Step 2：权重绑定的参数账本与梯度路径

权重绑定（weight tying）让输入词嵌入和输出 LM Head 共享同一份 $W\in {\mathbb{R}}^{V\times d}$。输入侧通过 $W$ 查表，输出侧使用同一参数把主状态投影回词表；反向传播会把两条路径的梯度累积到同一份参数。数值相同的两份拷贝不等于参数共享。

检查项	绑定（tied）	不绑定（untied）
参数存储	两个模块指向同一矩阵	两个独立矩阵
参数账本	一份 V × D	两份 V × D
梯度归属	两条路径更新同一参数	两条路径分别更新
身份关系	两个位置引用同一份参数	两个位置各自持有参数

Step 3：从局部字段到可复查的结构审计

结构审计要回答三件事：配置声称采用什么结构，模型实际采用什么结构，这个选择会改变哪些参数、状态或执行条件。局部字段只有进入这条证据链后，才能用于比较模型架构。本节选择权重共享和归一化缩放作为两个可运行例子，其他结构沿用同一审计方法进入对应机制页和项目页。

审计对象	配置声明	结构事实	结构后果与后续验证
权重共享	输入与输出是否共享	参数身份与梯度归属	词表参数账本、加载语义
归一化变体	类、eps、初始化	块中的归一化模块	尺度路径与深层稳定性
残差策略	缩放、层模式	块更新公式	激活/梯度随深度变化
注意力表示	查询头、KV 头、潜在状态字段	Q/K/V 投影与缓存接口	表示容量、重建和执行后端条件
稠密 / MoE	专家、Top-K、共享专家	路由器与专家模块	激活容量、容量溢出与通信
层级混合	层类型排列	每层实际模块类型	注意力 / SSM 等记忆职责分配

Step 4：代码设计与结构配置审计

本题用小型 RMSNorm、词嵌入和 LM Head 复现配置如何改变真实模块。骨架已提供模块接口、输入检查和配置审计；学习者补全缩放、参数绑定与参数账本。测试分别覆盖初值与数据类型、参数和梯度身份、合法与非法账本，以及配置声明不一致。

TODO / 函数	学习者完成的机制	必须满足的约束	测试证据
TODO 1：应用 Gemma 缩放 GemmaRMSNorm.forward	将归一化状态乘以可学习的 1 + weight	输出形状和 dtype 与输入一致；零初始化不改变归一化结果	test_gemma_rmsnorm_scale 覆盖零/非零缩放与 FP16 输出
TODO 2：绑定共享参数 QwenTieEmbeddings.__init__	让 LM Head 与词嵌入引用同一参数	两个模块共享同一 Parameter，更新与梯度同步	test_tied_parameter_identity_and_gradient 覆盖参数、更新和梯度身份
TODO 3：计算词表参数账本 summarize_parameter_layout	计算绑定与不绑定配置的参数数量	词嵌入为 V × D；不绑定时 LM Head 额外持有一份	test_parameter_layout 覆盖两种配置与非法维度

import torch
import torch.nn as nn

class GemmaRMSNorm(nn.Module):
    """用 `1 + weight` 作为通道缩放的最小 RMSNorm。"""

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.zeros(hidden_size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """保持输入形状与 dtype；内部以 FP32 计算 RMS。"""
        x_f32 = x.float()
        variance = x_f32.pow(2).mean(-1, keepdim=True)
        x_norm = x_f32 * torch.rsqrt(variance + self.eps)

        # ==========================================
        # TODO 1：应用 Gemma 缩放
        # output = ???  # 与 x_norm 同形；应用逐通道 `1 + weight` 后保持有限
        # ==========================================
        raise NotImplementedError("TODO 1：请应用 Gemma 缩放")
        return output.type_as(x)


class QwenTieEmbeddings(nn.Module):
    """构造词表 embedding 与 LM Head，并在 TODO 中绑定同一参数。"""

    def __init__(self, vocab_size: int, hidden_size: int):
        super().__init__()
        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

        # ==========================================
        # TODO 2：绑定共享参数
        # shared_weight = ???  # 必须是 embed_tokens 当前持有的 Parameter，而不是数值副本
        # ==========================================
        raise NotImplementedError("TODO 2：请绑定共享参数")
        self.lm_head.weight = shared_weight

    def forward_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

    def forward_lm_head(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.lm_head(hidden_states)


def summarize_parameter_layout(vocab_size: int, hidden_size: int, tied: bool) -> dict:
    """返回 embedding、LM Head、总计与节省的词表参数数量。"""
    if vocab_size <= 0 or hidden_size <= 0:
        raise ValueError("vocab_size 和 hidden_size 必须为正数")

    # ==========================================
    # TODO 3：计算词表参数账本
    # embedding_params = ???  # 一份 [vocab_size, hidden_size] 词表矩阵的参数数
    # lm_head_params = ???  # 绑定时为 0；不绑定时为一份独立词表矩阵
    # ==========================================
    raise NotImplementedError("TODO 3：请计算词表参数账本")
    return {
        'embedding_params': embedding_params,
        'lm_head_params': lm_head_params,
        'total_params': embedding_params + lm_head_params,
        'saved_params': embedding_params if tied else 0,
    }


def audit_model_config(model: nn.Module, config: dict) -> dict:
    """由固定骨架比较配置声明与模块实际状态。"""
    declared_tied = config.get('tie_word_embeddings')
    actual_tied = model.embed_tokens.weight.data_ptr() == model.lm_head.weight.data_ptr()
    norm_scale_init = config.get('norm_scale_init')
    if declared_tied is None or norm_scale_init is None:
        status = 'needs_checkpoint_review'
    elif bool(declared_tied) != actual_tied:
        status = 'config_mismatch'
    else:
        status = 'consistent'
    return {
        'status': status,
        'declared_tied': declared_tied,
        'actual_tied': actual_tied,
        'norm_scale_init': norm_scale_init,
    }

# 测试设计：分别验证缩放语义、共享 Parameter、参数账本和配置审计。
# 每个测试只对应一项结构选择，便于定位未完成或错误的 TODO。


def test_gemma_rmsnorm_scale():
    """验证 Gemma 风格缩放的初值、非零权重和输出数据类型。"""
    hidden_size = 16
    norm = GemmaRMSNorm(hidden_size)
    x = torch.randn(2, 3, hidden_size)
    expected = (x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + norm.eps)).to(x.dtype)
    assert torch.allclose(norm(x), expected, atol=1e-4, rtol=1e-3)

    with torch.no_grad():
        norm.weight.fill_(0.1)
    assert not torch.allclose(norm(x), expected, atol=1e-4)
    assert norm(x.half()).dtype == torch.float16


def test_tied_parameter_identity_and_gradient():
    """验证词嵌入与 LM Head 共享参数存储和梯度。"""
    model = QwenTieEmbeddings(vocab_size=32, hidden_size=8)
    assert model.embed_tokens.weight.data_ptr() == model.lm_head.weight.data_ptr()

    before = model.embed_tokens.weight.detach().clone()
    with torch.no_grad():
        model.embed_tokens.weight.add_(1.0)
    assert torch.allclose(model.lm_head.weight, before + 1.0)

    input_ids = torch.randint(0, 32, (2, 4))
    loss = model.forward_lm_head(model.forward_embed(input_ids)).sum()
    loss.backward()
    assert model.embed_tokens.weight.grad is not None
    assert model.embed_tokens.weight.grad.data_ptr() == model.lm_head.weight.grad.data_ptr()


def test_parameter_layout():
    """验证绑定与不绑定配置的词表参数账本。"""
    tied = summarize_parameter_layout(vocab_size=1000, hidden_size=64, tied=True)
    untied = summarize_parameter_layout(vocab_size=1000, hidden_size=64, tied=False)
    assert tied == {
        'embedding_params': 64000, 'lm_head_params': 0,
        'total_params': 64000, 'saved_params': 64000,
    }
    assert untied == {
        'embedding_params': 64000, 'lm_head_params': 64000,
        'total_params': 128000, 'saved_params': 0,
    }
    for vocab_size, hidden_size in [(0, 64), (1000, 0)]:
        try:
            summarize_parameter_layout(vocab_size, hidden_size, tied=True)
        except ValueError:
            continue
        raise AssertionError('非法词表或 hidden size 应被拒绝')


def test_config_audit():
    """验证配置声明与模块实际共享状态的一致性分类。"""
    model = QwenTieEmbeddings(vocab_size=32, hidden_size=8)
    consistent = audit_model_config(model, {'tie_word_embeddings': True, 'norm_scale_init': 0.0})
    mismatch = audit_model_config(model, {'tie_word_embeddings': False, 'norm_scale_init': 0.0})
    incomplete = audit_model_config(model, {'tie_word_embeddings': True})
    assert consistent['status'] == 'consistent'
    assert mismatch['status'] == 'config_mismatch'
    assert incomplete['status'] == 'needs_checkpoint_review'


def run_architecture_trick_tests():
    """依次运行三个 TODO 与固定配置审计的机制测试。"""
    test_gemma_rmsnorm_scale()
    test_tied_parameter_identity_and_gradient()
    test_parameter_layout()
    test_config_audit()
    print('✅ 结构配置：缩放、参数共享、账本和审计测试通过')


try:
    run_architecture_trick_tests()
except NotImplementedError:
    print('请先完成 TODO 1–3，再运行测试。')
    raise

---

🛑 STOP HERE 🛑

请先完成题目区并通过测试，再继续查看参考代码与解析。

---

参考代码与解析

代码

class GemmaRMSNorm(nn.Module):
    """用 `1 + weight` 作为通道缩放的最小 RMSNorm。"""

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.zeros(hidden_size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """保持输入形状与 dtype；内部以 FP32 计算 RMS。"""
        x_f32 = x.float()
        variance = x_f32.pow(2).mean(-1, keepdim=True)
        x_norm = x_f32 * torch.rsqrt(variance + self.eps)

        # TODO 1：应用 Gemma 缩放
        output = x_norm * (1 + self.weight)
        return output.type_as(x)


class QwenTieEmbeddings(nn.Module):
    """构造词表 embedding 与 LM Head，并在 TODO 中绑定同一参数。"""

    def __init__(self, vocab_size: int, hidden_size: int):
        super().__init__()
        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

        # TODO 2：绑定共享参数
        shared_weight = self.embed_tokens.weight
        self.lm_head.weight = shared_weight

    def forward_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

    def forward_lm_head(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.lm_head(hidden_states)


def summarize_parameter_layout(vocab_size: int, hidden_size: int, tied: bool) -> dict:
    """返回 embedding、LM Head、总计与节省的词表参数数量。"""
    if vocab_size <= 0 or hidden_size <= 0:
        raise ValueError("vocab_size 和 hidden_size 必须为正数")

    # TODO 3：计算词表参数账本
    embedding_params = vocab_size * hidden_size
    lm_head_params = 0 if tied else embedding_params
    return {
        'embedding_params': embedding_params,
        'lm_head_params': lm_head_params,
        'total_params': embedding_params + lm_head_params,
        'saved_params': embedding_params if tied else 0,
    }


def audit_model_config(model: nn.Module, config: dict) -> dict:
    """由固定骨架比较配置声明与模块实际状态。"""
    declared_tied = config.get('tie_word_embeddings')
    actual_tied = model.embed_tokens.weight.data_ptr() == model.lm_head.weight.data_ptr()
    norm_scale_init = config.get('norm_scale_init')
    if declared_tied is None or norm_scale_init is None:
        status = 'needs_checkpoint_review'
    elif bool(declared_tied) != actual_tied:
        status = 'config_mismatch'
    else:
        status = 'consistent'
    return {
        'status': status,
        'declared_tied': declared_tied,
        'actual_tied': actual_tied,
        'norm_scale_init': norm_scale_init,
    }

解析

TODO 1：应用 Gemma 缩放

• 先以 FP32 计算 RMS，再让 x_norm 乘以 1 + weight，最后恢复输入 dtype。
• 因此 weight=0 时输出就是未额外缩放的 RMS 归一化结果；非零权重才产生逐通道调整。

TODO 2：绑定共享参数

• 两个模块必须引用同一 Parameter，而不是复制一份数值相同的张量。
• data_ptr() 与梯度指针共同验证存储和反向传播确实共享。

TODO 3：计算词表参数账本

• 词嵌入始终持有一份 V × D 词表矩阵；不绑定时 LM Head 再增加一份同样大小的矩阵。
• 因而绑定方案的 saved_params 正好是一份 V × D；参数账本测试同时核对两种配置与非法维度。

固定审计函数：对照声明与实际对象

• audit_model_config 同时返回配置声明、参数指针关系和审计状态。
• 当配置缺字段时返回 needs_checkpoint_review；声明与实际共享关系相反时返回 config_mismatch。

相关阅读

完成本节后，可对照真实模型的配置与模块对象，再把这些字段带入架构审计项目。

• Transformers 中的 LLaMA 模型实现
• Qwen3 开源仓库
• Gemma 开源仓库
• 61. 模型架构探索
• 归一化演进：架构专题
• 代表模型：架构专题