【大模型】embed_tokens.bin 与 lm_head简介

声明：本文主要参考开源资料进行学习整理，如有错漏，欢迎评论交流~

1. 引言

当我们把一个训练好的大语言模型（以 Qwen3-VL-2B-Instruct 为例，vocab_size=151936，hidden_size=2048）部署到端侧时，会看到一个有意思的现象：

模型里有一个 embed_tokens 模块，负责把 token id 转成向量；
模型里还有一个 lm_head 模块，负责把最后的隐藏状态转成词表上的 logits；
配置里写着 tie_word_embeddings: true，说这俩"共享权重"。

可部署产物里，embed_tokens 被单独导出成一个 xxx_embed_tokens.bin 文件（float32），而 lm_head 却被编进了量化后的端侧模型（权重量化成 fp16/int8）。

矛盾点来了：

embed_tokens 是 gather 查表，lm_head 是 matmul/linear，根本不是同一个算子，怎么还能叫"同一份权重"？
既然"同一份权重"，部署时为什么又要拆成两种格式、两个地方存放？

接下来一起分析下~

2. 两个算子定义

先把两个模块的定义摆出来：

它们在模型里干的事完全不同：

	embed_tokens	lm_head
类型	nn.Embedding	nn.Linear（无 bias）
算子本质	gather（按索引取行）	matmul（x @ Wᵀ）
位置	输入端	输出端
输入	token id i（整数）	隐藏状态 x（H 维向量）
输出	一个 H 维向量	V 个 logits
干的事	"这个 token 长什么样"	"当前状态该预测哪个 token"

算子确实不同。但注意最后一行——它们的权重形状都是 (151936, 2048)。这不是巧合，而是能共享的前提。

3. 算子共享权重

输入端（embedding，gather）：

取出一行，得到 token i 的向量表示。

输出端（lm_head，matmul）：

让 hidden x 和每一行做点积，得到 V 个分数——即"x 跟 token 0/1/…/V-1 各有多像"，作为预测每个 token 的 logit。

有对称性：

gather 读W 的某一行 → "token i 的向量"
matmul 读 W 的所有行做点积 → "x 最像哪个 token"

两种算子，读的是同一个矩阵 W 的同一批行。gather 一次取一行，matmul 一次把所有行当匹配目标。这就是"算子不同、但同一份权重"的数学真相。

小贴士：因为 gather ≡ one-hot @ W，所以 gather 其实是 matmul 的一个特例（输入是 one-hot）。二者本就是同一族操作，只是实现和访问模式不同。

4. 权重绑定Weight Tying

为什么"同一份权重"是合理的？这叫 Weight Tying（权重绑定），核心是一个很强的归纳偏置：

"一个 token 进入网络时学到的向量表示，应当和它作为预测目标时的代表向量，是同一个。"

这俩描述的是同一个语言符号 i，理应共用一份 embedding。让它们共享，等于强制"词的输入语义"和"词的输出语义"对齐，这是一种很好的先验，效果往往更好，还能省一半参数（vocab × hidden 只存一份，对大词表模型动辄省 1+ GiB）。

代码配置中的 tie_word_embeddings: true 就是开关。

4.1 PyTorch 层面是"同一份"吗？

是字面意义的同一块内存。打开 tying 后，框架直接执行：

改一个，另一个跟着变。不是"形状相同、数值恰好一样"的两个独立张量，而是同一个张量。

一个实用 gotcha：因为 tied，HuggingFace checkpoint 里根本不会单独存 lm_head.weight——它和 embed_tokens.weight 指向同一块内存，加载时由框架的 tying 机制自动接上。自己解析权重文件发现少了一个 key，别慌。

4.2 形状为什么恰好能对上？

这是能 tying 的前提：

nn.Embedding(vocab, hidden) 的权重形状是 (vocab, hidden) = (V, H)；
nn.Linear(in, out, bias=False) 的权重形状是 (out_features, in_features) = (vocab, hidden) = (V, H)。

两者形状一致，所以可以指向同一块存储。nn.Linear 把权重按 (out, in) 布局而不是 (in, out)，本是为了前向 x @ Wᵀ 高效，却恰好和 embedding 表的 (vocab, hidden) 对齐——这种"巧合"让 weight tying 在工程上几乎零成本。

5. 拆分部署

理解了"同一份权重"，再看部署产物就不矛盾了：同一个 tied 张量 W，在部署时被读取两次，以两种格式分别落地，给两种算子用。

用途	存储形式	算子	精度
输入端（embed_tokens）	embed_tokens.bin 独立文件	gather 查表	float32
输出端（lm_head）	编进编译后的端侧模型	matmul x @ Wᵀ	量化（fp16/int8）

算子不同、精度不同、存放位置不同，但底层是同一批数。

5.1 为什么 embed_tokens 要单独导出成一个 .bin？

这是整件事里最工程化、也最值得讲的部分。四个原因：

gather 不适合端侧的连续计算流水。 端侧推理芯片优化的是连续的大矩阵乘（GEMM）。而 nn.Embedding 是按索引离散取行，访问模式不规律，塞进连续 GEMM 流水里效率很低。所以让 gather 留在软件/CPU 侧做，模型主体只保留擅长 GEMM 的 transformer 层和 lm_head。
软件侧需要灵活处理输入。 真实推理时，输入端要干很多"非固定图"的活：可变长度序列、batching、以及（多模态场景下）把图像/视频 embedding 通过 masked_scatter 散布到文本 embedding 的对应位置。这些逻辑放在固定编译图里做不到，留在软件侧最灵活。
避免量化损失。 embedding 表很大（vocab × hidden × 4 ≈ 1.16 GiB），单独存成 float32 能保留输入端精度；而模型主体（含 lm_head）可以放心量化到 fp16/int8。同一份权重，输入端要精度、输出端可量化，所以各存一份。
多实例共享。 embedding 表巨大，runtime 通常用一个带锁的全局缓存把它共享给多个推理实例/线程，避免重复加载这一 GiB 级的表。

5.2 embed_tokens文件大小

embed_tokens.bin 是原始 float32 二进制转储，形状 (vocab_size, hidden_size)：

换模型或改词表，大小随之变化；导出格式始终是 float32。

6. 推理链路pipeline

把上面拼起来，端侧推理的数据流是这样的：

一个关键细节：编译后模型的第一个输入是 input_embeddings（浮点张量），不是 input_ids（整数）。也就是说，"token id → 向量"这一步在进模型之前就由 runtime 完成了，模型本身只接收已经查好的向量。

7. 总结

问题	答案
embed_tokens.bin 是什么	embed_tokens 层权重的 float32 原始二进制转储
作用	runtime 侧完成 token id→向量查表，把结果作为浮点输入喂给模型
大小	vocab_size × hidden_size × 4
和 lm_head 的关系	tie_word_embeddings：共享同一份 (vocab, hidden) 权重张量
算子不同为何能共享	gather 和 matmul 读的是同一个矩阵 W 的同一批行；nn.Linear 权重 (out,in) 天然等于 embedding (vocab,hidden)，形状匹配
为什么共享	Weight Tying：token 输入向量 = token 输出目标向量，是同一语言符号的两种角色
为什么部署时拆开	gather 不适合 GEMM 流水 / 软件侧需灵活处理输入 / 输入端要 float32 保精度而输出端可量化 / 多实例共享大表

"同一份权重"这个词容易让人误解成"做同一件事"。其实正相反——embed_tokens 和 lm_head 是两个完全不同的算子（gather vs matmul），干的是输入端查表和输出端打分两件不同的事。它们能共享权重，靠的是：

形状对齐：nn.Linear(out,in) 的权重布局恰好等于 nn.Embedding(vocab,hidden)；
数学对称：同一个矩阵 W，gather 读一行、matmul 读全部行，本就是同一族操作；
语义同源：token 进入时的向量和它作为预测目标的向量，描述的是同一个词。

而部署时把它们拆成 embed_tokens.bin（float32，软件侧 gather）和编进模型的 lm_head（量化，端侧 matmul），是同一个 tied 张量按"输入端要精度+灵活、输出端可量化+走 GEMM"的工程权衡各取一份。

理解了这层"算子不同、权重同源、按用途分存"，再看到 tie_word_embeddings 和那个单独的 .bin 文件，就不会觉得矛盾了。