工作笔记 · 模型训练 · Infra
OE Embedding 机制图解
主题 Over-Encoding 原理 · v2 patch · CPU 卸载模型 WeLM v4.5 · 80B-A3B日期 2026-06-13
● 面向非 infra 同学 · monkey_patch_welmv4_5_moe_v2.py · 附字节 OTT / DeepSeek Engram 对照
📌 一句话先讲清楚
WeLM v4.5 模型里有
4 张超大查表（共 65.6 GB） ，本来要塞进显卡，会爆显存；这版 patch 把它们
搬到内存（CPU） ，每次训练时去内存里查一下结果再送回显卡，
顺便修了一个跨样本污染的小 bug 。
1 什么是 OE Embedding？
OE 是 Over-Encoding（过度编码） 的缩写，是 WeLM v4.5 这个大模型自带的一个"加料"机制。
你可以把模型理解成一个"读句子的人" 。当前这个字 是什么（这叫 token embedding）。顺便偷瞄一下前面 1～4 个字组合起来是什么 ，然后把这个"上下文味道"也加到当前字上一起读。
怎么实现这种"偷瞄"呢？模型准备了 4 张巨大的查询表（hash 表） ：
每张表 1600 万行 × 512 列，用 bf16 半精度存，每张 16.4 GB
4 张加起来 65.6 GB
用法：把"前 n 个 token"做个 hash 算出一个行号，去表里把那一行（512 个数字）抠出来当作"短上下文向量"
📊 这 4 张表到底有多大？参数量算给你看
每张表的参数量 = 行数 × 列数：
16,000,000 行 × 512 列 = 8,192,000,000 参数 ≈ 81.92 亿（8.19B）
4 张表合计：
8.19B × 4 = 327.68 亿参数（32.77B）
对比对象 参数量 对比 OE 4 张表（32.77B） Llama-3-8B 整个模型 8 B OE 是它的 4 倍 WeLM v4.5 (80B-A3B) 单步激活参数 3 B OE 是它的 10 倍 GPT-3 175 B OE 约是它的 1/5 显存（bf16，每参数 2 字节） 32.77B × 2B = 65.54 GB
这 4 张表本身就是一个超大模型的体量。单 H100 总共才 80 GB 显存，光它们就占 65.6 GB——这就是为什么必须卸载到 CPU。
好在这些表本质是预训练阶段就训好的 hash 查询表 ，SFT 阶段冻结它们对效果影响很小；真正起作用的小投影层 oe_up_proj（参数量约几百万）仍然留在 GPU 上正常训练。
🔬 n-gram 是怎么算出来的？4 张表分别干嘛？
第一步：每张表对应一个 "n-gram 阶数"
WeLM v4.5 的配置里 oe_grams = [2, 3, 4, 5]，oe_vocab_sizes = [16M, 16M, 16M, 16M]，意思是：
表号 阶数 含义 类比（读到"果"字时） 表 1 2-gram 同时看 当前 + 前 1 个 字 看 "苹果" 这个组合 表 2 3-gram 同时看 当前 + 前 2 个 字 看 "吃苹果" 这个组合 表 3 4-gram 同时看 当前 + 前 3 个 字 看 "爱吃苹果" 这个组合 表 4 5-gram 同时看 当前 + 前 4 个 字 看 "我爱吃苹果" 这个组合
4 张表是从短到长的"上下文滤镜" ，每张抓不同长度的搭配信息。短的（2-gram）抓常见词组（"苹果"），长的（5-gram）抓固定短语（"我爱吃苹果"）。最后把 4 段信息拼在一起 ，让模型同时拥有"近距离 + 远距离"的搭配感。
第二步：怎么把"前 n 个 token"压成一个数字？
这是 OE 最巧妙的地方。代码里的做法分两小步：
① 把 n 个 token id 拼成一个大整数（V = 词表大小）：
combined = t0 + t1 ·V + t2 ·V² + t3 ·V³ + t4 ·V⁴
就像把多位数字拼起来：个位放 t0 ，"V 进制"的十位放 t1 ，百位放 t2 ……这样不同的 token 组合一定得到不同的 combined 值 ，互不冲突。
② 用 Knuth 黄金比例哈希把这个超大整数打散：
hashed = (combined × 2654435761 ) & 0xFFFFFFFF
（这个常数是 2³² × (√5−1)/2，黄金分割的整数版，让 hash 结果在 0~2³² 之间均匀分布）
③ 取模落到表的某一行（每张表 16M = 2²⁴ 行）：
row_id = hashed % 16,000,000
"取模"就像把无穷多种 token 组合抽签 分到 1600 万个抽屉里。同一个组合永远落到同一个抽屉（确定性），但不同组合可能撞到同一个抽屉（叫 hash 冲突）—— 这是用"压缩存储"换来的小代价，模型能学着容忍它。
第三步：4 张表查出 4 个向量，拼起来投影
每张表查出来的是一行 512 维向量。4 张表 → 4 个 512 维向量 → 拼接成 2048 维 → 过一个小线性层 oe_up_proj（2048 → hidden_size）→ 和普通 token embedding 取平均。
第四步：和普通 embedding 取平均
代码最后一行：
hidden_states = (普通 embedding + OE 投影结果) / 2.0
最终送进 Transformer 的"果"字向量 = "果"自己的语义 （普通 embedding） + "我爱吃苹果"这串组合的味道 （OE）—— 两碗汤一勺一勺各舀一半，混合后口味更丰富。
为什么要 4 张独立的表，而不是 1 张大表？
原因 解释 不同 n-gram 的"组合空间"差太多 2-gram 组合空间是 V²，5-gram 是 V⁵，量级差几十亿倍。如果共用一张表，hash 冲突会被长 n-gram 完全淹没短 n-gram。 语义维度不同 2-gram 学的是固定搭配 （"苹果"），5-gram 学的是固定句式 （"我爱吃苹果"），用不同表分别学，互不干扰。 分桶可控 每张表独立 16M 行，每张表的 hash 冲突率独立可控。
问题 A：65.6 GB 装不进显卡
单张 H100 显卡总共才 80 GB 显存，光这 4 张表就吃掉 65.6 GB ，再加上模型本身 80B 参数、激活值、优化器状态，根本装不下。
在多卡训练的 ZeRO-3 模式下，正常做法是把这些参数按卡切片 ，每张卡只存自己的一片。但是每次前向计算时还要临时把所有切片合起来（all-gather） ，那一瞬间整张表又会出现在每张卡上 → 瞬时显存爆炸 → 训练直接 OOM 崩溃。
v1 版的解法是直接把 OE 关掉 （粗暴但有效），相当于让模型不再看这 4 张表。保留 OE 的能力 ，必须想别的办法。
问题 B：训练样本之间会"串味"
大模型训练为了省算力，会把多条短样本拼成一条长序列 （叫 packing），比如：
样本 A：今天天气真好
样本 B：苹果是一种水果
样本 C：北京是中国的首都
↑ 三个完全不相关的样本被拼成一条长序列一起训练
原版 OE 的"偷瞄前 n 个字"是无脑往前数 n 位 的，不知道哪里是样本边界。跨样本串味了！ 苹 果是一种水果"，是真的语义错误，会拉低训练质量。
2 v2 patch 是怎么修的？v2 同时解决两个问题，思路非常清晰：
1
把 4 张大表搬到"内存条"（CPU 内存）里
显卡显存（GPU）容量只有 80 GB，但服务器的内存条（CPU 内存）通常有几百 GB甚至 1TB ，65.6 GB 放进去毫无压力。
就像你电脑的固态硬盘装不下大型游戏，但机械硬盘空间大——把不常用的素材包放到机械硬盘 ，需要时再读到内存里使用。这里把 GPU 当"小而快的固态"，把 CPU 内存当"大而稍慢的机械硬盘"。
对应代码：_detach_oe_embed_to_cpu() 函数
2
让 GPU 训练框架"看不见"这 4 张表
光搬走还不够，还要从模型结构里把它们的引用彻底删掉 （del inner_model.oe_embed），否则 ZeRO-3 框架还会自作主张去切片、去 all-gather，依然会炸。
相当于不光把素材包搬到机械硬盘，还要把游戏配置里的引用路径也改掉 ，免得游戏启动时还去固态盘上找它，找不到就崩溃。
3
每次训练时，去 CPU 上查表，结果搬回 GPU
替换掉模型原本的前向函数，改成：
把要查的"行号"从 GPU 拷到 CPU（很小，几十 KB）
在 CPU 上查表 得到结果（约 16 MB / 张 × 4 张 = 64 MB）
把结果通过 PCIe 总线搬回 GPU ，约 3 毫秒，可以接受
对应代码：_oe_cpu_lookup() 函数
4
修复"跨样本串味" bug
把"无脑往前数 n 位"改成"懂边界的版本" ：每个样本的开头 n 个位置，强制写 0（相当于"前面没有字可以偷瞄"），就把跨样本污染斩断了。
对应代码：_packing_aware_skip_firstn_token_ids() 函数
修复前（错误）：
样A末
↓串味
样B首
样B
B 开头偷瞄到 A 的尾巴 ❌
修复后（正确）：
样A末
置 0
样B首
样B
B 开头啥也偷瞄不到 ✅
3 训练时数据怎么流？一图看懂每一步训练的 5 个动作
GPU 拿到输入的 token ids，先查普通 embedding 表得到基础向量。
GPU 上算出"前 n 个 token 的 hash 行号"，这一步用 packing-aware 版本 ，把样本边界写 0。
把行号（很小）从 GPU 拷到 CPU。
CPU 上对 4 张表分别查表 → 得到 4 段 OE 向量 → 通过 PCIe 拷回 GPU（约 64 MB / 3 ms）。
GPU 上：4 段 OE 向量拼起来，过 oe_up_proj 投影，再和普通 embedding 取平均 ，作为最终输入送进 Transformer。
4 代价和收益
项目
v1（关掉 OE）
v2（CPU 卸载）
4 张大表是否参与训练
❌ 完全不用
✅ 参与前向 但冻结不更新
oe_up_proj 投影层不更新（前向不走它）
可训练 ，留在 GPU
显存占用
低
低（4 张表 在 CPU 内存 ）
每步耗时
基线
+约 3ms（PCIe 拷贝）
跨样本串味 bug
不触发（OE 关了）
✅ 已修复
OE 能力是否保留
❌ 没了
✅ 保留
结论： v2 用"4 张表搬到 CPU、SFT 阶段不更新它们"的代价，换来了 OE 能力的保留 + 跨样本污染的修复 。oe_up_proj 仍然留在 GPU 上正常训练。
🔑 三个核心函数对照表
函数
干什么
_detach_oe_embed_to_cpu启动时把 4 张大表搬去 CPU，并从模型结构里删掉引用
_oe_cpu_lookup每步训练时去 CPU 查表，把结果搬回 GPU
_packing_aware_skip_firstn_token_ids修跨样本串味 bug：样本开头 n 位强制清零
本文档基于 monkey_patch_welmv4_5_moe_v2.py 自动整理 · 仅作内部学习交流用途
5 OE 给工程带来的真实挑战前面讲的是「OE 是什么、v2 patch 怎么修」。但当你真的要把这套 65.6 GB 的查表机制跑在训练 / 推理集群上，会撞上一连串 infra 层面的硬骨头。下面把我们实际踩过、或预期会踩的坑系统地列一下。
5.1 显存墙：参数比模型本体还大
OE 的 4 张表合计 32.77B 参数 / 65.6 GB（bf16） ，而 WeLM v4.5（80B-A3B）单步只激活 3B 。也就是说——一个「辅助 embedding 机制」的体量是主干单步激活的 10 倍 。
单张 H100 才 80 GB 显存。如果把 OE 表和模型一起塞进卡里，光这 4 张表就吃掉 65.6 GB，剩下的连一层 attention 的激活都放不下，必然 OOM。
v2 patch 的核心动作就是把这 4 张表 卸载到 CPU 内存 ，GPU 只保留可训练的小投影层 oe_up_proj。代价是每个 step 都要做一次「GPU→CPU 取行号→CPU 查表→搬回 GPU」的往返。
5.2 通信瓶颈：PCIe 往返与 host-device 拷贝
表在 CPU、计算在 GPU，意味着每个 forward 都要跨 PCIe 搬数据。一个 batch 里每个 token 都要从 4 张表各查一行 512 维向量：
单 step 取回数据量 ≈ batch × seq_len × 4 表 × 512 维 × 2 字节(bf16)
以 batch=8、seq_len=4096 估算，单 step 仅 OE 查表回搬就有 8×4096×4×512×2 ≈ 268 MB 。PCIe 4.0 单向带宽约 32 GB/s，看似不大，但它夹在前向关键路径上 ，且 host→device 拷贝难以和计算完全 overlap，TTFT / 单步耗时都会被拉高。
就像图书馆把最常翻的 4 套大百科全书放在隔壁楼（CPU 内存），你（GPU）每算一句话都得跑过去抄一段再跑回来。书是省地方了，但你的腿（PCIe）成了新瓶颈。
5.3 随机访存：hash 决定了「没有局部性」
OE 用 Knuth 黄金比例 hash 故意 把相邻 token 组合打散到 1600 万行的不同位置——这对「均匀分布、减少冲突」是好事，但对内存子系统是灾难 ：
Cache 几乎全 miss ：连续两个 token 的 row_id 天差地别，CPU L1/L2/L3 缓存命中率趋近于 0，每次查表都是一次主存随机访问（~100ns 级延迟）。无法预取（prefetch） ：row_id 要等当前 token 算完才知道，硬件预取器猜不到下一个地址。NUMA 不友好 ：多路 CPU 下，表可能落在远端 NUMA node，跨 node 访存再加一层延迟。
对照 DeepSeek Engram 的解法 Engram 论文专门强调 deterministic addressing（确定性寻址）+ host memory runtime prefetching ——因为 N-gram 是确定的，可以在需要某行之前就异步把它从 host memory 预取过来，把随机访存延迟藏在计算后面。这正是 OE 朴素实现所缺的优化方向。
5.4 跨样本污染：状态泄漏 bug 的工程根源
v2 patch 修的那个「跨样本污染」bug，本质是有状态的 n-gram 缓冲区在 batch 内 / 跨 micro-batch 之间没有正确按样本边界清零 。这类 bug 在普通无状态 embedding 里根本不存在，是 OE「要回看前 n 个 token」这一跨时间步依赖 带来的工程复杂度。
在序列并行 / 流水并行 / packing（多条短样本拼成一条长序列）等场景下，「前 n 个 token」可能跨越了真正的样本边界，把上一条样本的尾巴当成了当前样本的上文，污染 hash。debug 时它表现为「loss 抖动、复现困难」，极其隐蔽。
5.5 与并行策略的耦合
并行维度 OE 带来的麻烦 TP（张量并行） oe_up_proj 是小层好切，但 65.6 GB 的表若要切分，hash row_id 与分片的映射要重新设计，否则每张卡都要存全表。 SP（序列并行 / Ulysses） n-gram 需要「前 n 个 token」，而 SP 把序列切到了不同卡上——边界处的上文 token 可能在邻卡，需要额外通信换边界。 PP（流水并行） OE 在 embedding 层（第一段 pipeline stage），但 CPU 卸载让这一 stage 的耗时不确定，容易成为 pipeline bubble 的源头。 Packing 样本边界处 n-gram 必须 reset，与 5.4 的污染 bug 同源。
6 延伸：OE 思路用到多模态训练的挑战OE 的本质是「给输入端挂一个超大的、按内容寻址的静态记忆表 」。这个思路天然诱人推广到多模态——但难度会再上一个台阶。这里把我对未来训练多模态大模型时引入 OE-style 机制的挑战做一个前瞻梳理。
6.1 「token 是离散的」这个前提在视觉/音频里不成立
文本 OE 能 work，靠的是 token id 是离散整数 ，可以拼成 combined 值再 hash。但图像 patch、音频帧的原始表征是连续向量 ：
文本里「苹果」永远是同样两个 id，hash 出来落同一行；但同一只猫的两张照片，像素从来不会一模一样，连续特征 hash 一下就分到了完全不同的行——相同语义却查不到同一段记忆 ，N-gram 记忆表瞬间失效。
要复用 OE，必须先把连续模态离散化 （VQ-VAE / 残差量化 codebook 把图像/音频映射成离散 code）。这一步本身就是开放难题：codebook 坍缩、码本利用率低、量化损失都会直接传染给上层的 N-gram 记忆。
6.2 跨模态 n-gram 的组合爆炸
文本只有一个序列，n-gram 是一维的。多模态是图文音交错的序列 ，「前 n 个 token」可能横跨模态边界：
2-gram 可能是 [图像patch, 文字]、[文字, 音频帧] 这种跨模态对 ——它真的有「上下文搭配」意义吗？还是纯噪声？
各模态的「有效词表」量级差异巨大（文本 ~10⁵，视觉 codebook ~10⁴，音频又不同），拼成 combined 时进制设计极易溢出或分布严重不均。
2D 图像的「n-gram」该取空间邻域（上下左右 patch）还是光栅扫描序？这让一维 hash 公式直接失配。
6.3 显存 / 访存压力指数级放大
文本 OE 已经是 65.6 GB。多模态意味着每个模态都要自己的记忆表 ，甚至跨模态联合表：
若视觉、音频、文本各配一套 4 张 16 GB 的表，再加跨模态联合表，总量轻松突破数百 GB。6.2/6.3 讲的 PCIe 往返与随机访存问题会按模态数量倍增，CPU 卸载也未必兜得住。
可能的出路 (1) 分层记忆：高频跨模态组合留 GPU，长尾放 CPU/host；(2) 借鉴 Engram 的 host prefetch + 确定性寻址；(3) 用学习式寻址（可微检索）替代纯 hash，让相似的连续特征映射到相近的记忆槽，解决 7.1 的「同语义不同 code」问题。
6.4 训练稳定性：多模态本就难，再叠一层静态记忆
多模态训练本身就有模态不平衡、梯度尺度差异、对齐难等问题。再挂一个冻结的大记忆表，会引入新的耦合：记忆表在预训练阶段「学到」的搭配，在 SFT / 跨模态对齐阶段可能与新数据分布冲突，且 5.4 的样本边界污染在「图文交错 packing」下更难定位。
小结 OE 把「静态记忆」作为一条独立于「神经计算」的稀疏性轴——这条轴在纯文本上已被验证有效（见下文字节 / DeepSeek 两项工作），但要迁到多模态，离散化、跨模态组合语义、访存工程 三道关卡缺一不可，是值得长期投入的方向。
WeLM v4.5 的 OE 并非凭空设计，它属于近一两年「给 LLM 加一条静态记忆 / 超大 embedding 稀疏轴 」这一研究脉络。这里详细介绍工业界两项最有代表性的工作——字节跳动的 Over-Tokenized Transformer 与 DeepSeek 的 Engram，帮助理解 OE 的来龙去脉。
① 字节跳动 · Over-Tokenized TransformerICML 2025 · arXiv 2501.16975
Hongzhi Huang 等（ByteDance Seed）·
arXiv:2501.16975 · 标题《Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling》
这篇是 WeLM「OE」命名与思想的直接学术源头 。核心主张：把输入词表（input vocabulary）和输出词表（output vocabulary）解耦，并大幅放大输入词表 。
关键设计
Over-Encoding（OE） ：输入端不再只用单 token，而是把 多元 n-gram（multi-gram）token 也作为输入词表的一部分——这正是 WeLM 那 4 张 2/3/4/5-gram 表的来历。输入词表从几万膨胀到上亿级别，用 hash embedding 压缩存储。Over-Decoding（OD） ：与 OE 对称，在输出端预测更细粒度的多 token 目标，提供更丰富的训练信号。OE/OD 合起来构成 Over-Tokenized 框架。log-linear scaling law ：论文最重要的实证——训练 loss 与输入词表大小呈 log-linear 关系 ，且不论模型多大都成立 。即放大输入词表是一条几乎「免费」的提升路径。
最亮眼的结论：用足够大的输入词表，可以让模型达到「两倍大小 baseline 的性能，而几乎不增加计算成本 」——因为 embedding 查表是 O(1)，不进 FLOPs 密集的 transformer 主干。
对我们的意义：它从 scaling law 层面证明了「OE 这条路值得走」，WeLM v4.5 的 4 张 hash 表就是这套理论的工程落地。本文前半讲的所有 infra 痛点（65.6 GB、CPU 卸载、hash 随机访存），本质都是为这条 scaling law 买的工程单。
② DeepSeek · Engram2026-01 · arXiv 2601.07372
Xin Cheng, Damai Dai, … Wenfeng Liang 等（DeepSeek-AI）·
arXiv:2601.07372 · 标题《Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models》
如果说字节证明了「放大输入词表有用」，DeepSeek Engram 则把它系统化、理论化 ，并直面了本文第 5 节列的所有工程挑战。它把这类机制重新定义为一条与 MoE 并列的新稀疏性轴 ：MoE 是「条件计算（conditional computation）」，Engram 是「条件记忆（conditional memory） 」。
关键贡献
Engram 模块 ：把经典 N-gram embedding 现代化为 $\mathcal{O}(1)$ 查找 的可扩展记忆，retrieve 出静态 N-gram 记忆后与动态 hidden state 融合——结构上与 WeLM OE 高度同构。Sparsity Allocation（稀疏性分配） ：形式化「神经计算（MoE）」与「静态记忆（Engram）」之间的权衡，发现一条 U 型 scaling law ——存在一个最优的「算力 vs 记忆」分配点，两头都不是最优。iso-param / iso-FLOPs 严格对照 ：在参数量、FLOPs 都对齐的前提下，Engram-27B 全面超过 MoE baseline ：MMLU +3.4、BBH +5.0、HumanEval +3.0 等，覆盖知识 / 推理 / 代码 / 数学。机制分析 ：Engram 把「重建静态模式」的负担从浅层 transformer 卸下来 ，从而为复杂推理保留了「有效深度（effective depth）」——给出了「为什么有用」的解释，而不只是「有用」。系统效率 ：采用 deterministic addressing（确定性寻址） ，使得巨大的 embedding 表可以 offload 到 host memory 且推理开销极小——正面解决了本文 5.2 / 5.3 的 PCIe 往返与随机访存难题。
和 WeLM OE 的关系 Engram = OE 的「理论完成态」。WeLM 的 v2 patch 在工程上做的事（CPU 卸载、确定性 hash 寻址），Engram 用 scaling law 证明了它不仅是省显存的妥协，更是一条独立且最优的稀疏性维度 。下次设计 OE 表大小 / 卸载策略时，Engram 的 U 型曲线与 host prefetch 方案是直接可借鉴的。
7.1 三者一图对照
维度 字节 OTT DeepSeek Engram WeLM v4.5 OE 核心主张 放大输入词表，OE/OD 解耦 条件记忆=新稀疏轴 4 张 n-gram hash 表加料 n-gram 阶数 multi-gram N-gram（可配） 2/3/4/5-gram 寻址方式 hash embedding 确定性寻址 O(1) Knuth 黄金比例 hash scaling law log-linear（越大越好） U 型（有最优分配点） —（工程落地） 卸载方案 — host memory + prefetch CPU 卸载（v2 patch） 规模 多尺度验证 27B，超 MoE baseline 4×8.19B=32.77B 表
一句话收尾 OE 不是 WeLM 独有的奇技淫巧，而是「条件记忆 」这条稀疏性轴的一个工程实例。字节证明它有用 ，DeepSeek 证明它最优且可解释 ，WeLM 把它跑在了真实训练管线上 ——而代价，就是本文第 5 节那一长串 infra 硬骨头。