将 Gemma-4（2B / 4B / 12B）移植到 AWS Inferentia2

将 Gemma-4（2B / 4B / 12B）移植到 AWS Inferentia2

一份关于在AWS Inferentia2（inf2）上运行Google Gemma-4系列模型的实地报告，涵盖了破坏供应商技术栈的三个架构障碍——混合注意力头、vLLM / optimum-neuron / NxD的死胡同，以及Neuron编译器（neuronx-cc）的限制——以及让所有三种模型规模都能连贯运行的解决方案。

1. 背景：为何如此困难

Gemma-4并非普通的解码器。该系列模型结合了多个特性，每个特性在TPU/XLA（模型设计平台）上都能完美运行，但单独来看都会破坏AWS推理路径：

• 逐层嵌入（PLE）和MatFormer嵌套（E2B/E4B上的“有效2B/4B”技巧）。
• 跨层KV共享：在E2B/E4B上，许多层不计算自己的Key/Value——它们重用相邻层的KV投影。
• 分组查询注意力（GQA），KV头数量较少。
• 混合注意力类型：交错的滑动窗口和全局注意力层，在12B上每种类型具有不同的KV头数量。
• Logit软上限（tanh上限为30），词汇表大小为262,144个token。

AWS供应商技术栈（optimum-neuron + neuronx-distributed + Neuron vLLM后端）根本没有Gemma-4模型类，其图构建器无法表达KV共享。我们最初使用的公共Neuron vLLM端点加载了某些内容，并生成了看似流畅但实际无意义的乱码。以下所有内容都是为了从那里走向正确、快速、廉价的推理。

1.1 三种模型概览

三种模型规模，三种不同的供应商路径失败原因，以及——事实证明——三种不同的“混合头”含义。

2. “混合头”的三种面貌

“混合头”是本次移植中最昂贵的一类错误。随着模型规模的增大，它以三种完全不同的形式出现。

2.1 跨层混合KV共享（E2B、E4B）

在E2B/E4B上，层的注意力被标记为self_attn.is_kv_shared_layer。共享层不运行自己的k_proj/v_proj；它们针对由较早的“所有者”层产生的KV张量进行注意力计算。在TPU上，这是一个自由的图视图。在AWS的neuronx-distributed（NxD）ModelBuilder中，它期望一个静态的、逐层的权重→缓冲区映射，无法表示——没有地方可以说明“这一层的K/V是那一层的K/V，计算一次，实时使用。”

修复方法（“选项B”，见§3）是停止与框架对抗，直接追踪Hugging Face的前向传播。KV共享随后作为追踪图中一个普通的实时数据依赖关系自然出现。移植只需要枚举哪些层实际拥有缓存：

# 发现写入KV的层；共享层从不接触缓存
NONSHARED = []
for i, lyr in enumerate(lang.layers[:cfg.num_hidden_layers]):
    a = lyr.self_attn
    if not a.is_kv_shared_layer:                 #  4/rank, keep groups
    a.k_proj = col(a.k_proj); a.v_proj = col(a.v_proj)
else:                                        # 全局层：nkv=1，不可分割
    # 保持k/v复制；将组缩小到每个rank分片的q头数量
    a.num_key_value_groups = (nq // TP) // nkv

两个陷阱使这变成了一场持续数小时的战斗（提交d70dc94）：

1. 顺序： 你必须在替换q_proj为ColumnParallelLinear之前读取nq = q_proj.out_features // head_dim，因为替换后的.out_features报告的是分片后的宽度——否则你会从已经减半的数字计算组数，得到AttributeError/错误的组。
2. 全局层的else分支从未被E4B执行过（所有E4B层都是nkv=2，可整除），所以它是死代码，直到12B遇到它才被测试。

要点： “混合头”不是一个问题。它是（a）跨层混合KV共享，（b）层内混合查询/KV数量（GQA），以及（c）混合注意力类型，每种类型具有不同的KV数量。每种都需要不同的分片规则，而天真的“只需将所有内容除以TP”在三种情况下都是错误的。

3. 为什么vLLM / optimum-neuron / NxD无法做到

本次移植故意不使用AWS供应商推理路径。以下是每个供应商层失败的原因，以及替代方案。

4. Neuron 编译器（neuronx-cc）的瓶颈

当计算图可追踪后，neuronx-cc 会施加其自身的硬性限制。反复出现的问题是 SBUF（片上状态缓冲区，即 SRAM），其上限为 每个分区 196 608 B。Gemma-4 的多个操作会直接超出此限制，每个操作都有不同的修复方法。

4.1 PLE 表无法驻留在设备上（E2B/E4B）

逐层嵌入表的大小为 262144 × hidden。若在设备上实例化，会触发编译器错误，并且其本身就会超出 16 GB 核心的容量。修复方法： 将 PLE（和词）嵌入保留在主机上，在 CPU 上执行 gather 操作，并将结果作为激活值输入。设备计算图永远不会看到该表。

4.2 Logit 软上限导致 SBUF 溢出（12B）

Gemma-4 在整个 262 144 词表上应用 softcap * tanh(logits / softcap)。在 fp32 下，tanh 是一个自定义调用，其工作集大小为 128 × 524288 字节——524 288 > 196 608 B/分区：

[NCC_INLA001] Allocated memory out of bound (128x524288) 524288 vs 196608 B/partition

修复方法（提交 328ca88）： 不要在设备上应用软上限。 软上限是单调的，因此 argmax(softcap(x)) == argmax(x)——贪婪解码完全不受影响。设备返回原始 logits；主机服务器仅在采样时（temperature > 0）重新应用上限。正确且零开销。（在 E2B/E4B 上，由于隐藏层大小较小，显式的 tanh 上限可以容纳，因此保留在计算图中。）

4.3 在 sliding_window=1024 时，融合/SDPA 注意力导致 SBUF 溢出（12B）

移除 softcap 是必要的，但还不够——12B 模型上真正的溢出是融合注意力自定义调用，因为 12B 的 sliding_window = 1024 是 E4B 的 512 的 2 倍，导致注意力块大小翻倍。修复方法（提交 a0a2a75）： 在设备模型上强制使用 eager 注意力（在配置及其 text_config 上设置 cfg._attn_implementation = "eager"）。Eager 模式将注意力实现为显式的矩阵乘法，这些矩阵乘法可以适配 SBUF。（E2B/E4B 本来就使用 eager 模式；在 sliding_window=512 时，融合路径也可以容纳，但 eager 模式是更具可移植性的选择。）

4.4 fp32 常量超出 16 GB 核心容量 → 使用 bf16 以及 张量并行

单个 Option-B neff 文件为 E4B 实例化了 约 15.4 GB 的 fp32 模型常量；同一核心上的第二个 neff 文件（prefill + decode）会超出 16 GB 的 NeuronCore 预算 → NRT_RESOURCE / status=4 Allocation Failure。需要两件事：

• bf16 权重（MB_WDTYPE=bf16）：NxD 的 shard_children 将检查点转换为层的 dtype，使 neff 文件大小减半。但是 bf16 缩小的是磁盘上的 neff 文件，而不是设备上的常量计数——15.4 GB 的常驻常量仍然存在。
• TP=2：实际适配 E4B/12B 的唯一方法是将模型分片到两个核心上，这样每个 rank 持有约一半的权重（E4B 上约 7.7 GB/核心；256 上下文时 12B 上约 12 GB/核心）。

E2B 是个例外：在 bf16 下，它可以适配单个核心，这就是为什么它保持单核心且速度约为 44 tok/s。

4.5 保持计算图为基本操作（所有尺寸）

neuronx-cc 对纯算术运算最友好。以下几项是手动展开的，而不是依赖库辅助函数或动态控制流：

• tanh-GELU 展开为 0.5*x*(1+tanh(0.7978845608*(x + 0.044715*x^3)))。
• 在追踪时使用 DynamicCache，但在解码时使用固定的静态 KV 缓冲区；每步写入是一个独热掩码散射操作——buf*(1-oh) + k*oh——纯算术运算，追踪安全，无散射操作或数据相关索引。
• 语言模型和 LM head 注册为真正的子模块，以便它们编译到计算图中，而不是被不透明地调用。

4.6 一个无声的、非编译器的陷阱：layer_scalar 是一个缓冲区

每个Gemma-4层都会执行hidden_states *= self.layer_scalar。layer_scalar是一个注册的buffer（默认值为1.0，实际值约为0.06）。NxD的权重分片（shard_children / get_sharded_checkpoint）仅加载参数，从不加载buffer——因此如果不加干预，每一层都会过度缩放约16倍，并在所有层中累积，最终导致cos ≈ 0的垃圾输出。修复方法（提交e785f6d）： 从检查点读取每层的layer_scalar，并在跟踪前手动将其复制到模块中。这个看似简单的一行物理bug在E4B和12B上都浪费了大量时间，是整个移植过程中最重要的"陷阱"。

4.7 编译可用性

neuronx-cc在CPU上运行——编译无需NeuronCore，因此可以在任何机器上构建neff文件。但TP=2编译会同时为两个rank运行neuronx-cc，峰值超过128 GB主机RAM；编译前需添加≥55 GB的交换空间（生成的neff无需交换空间即可正常运行）。

5. 设备驻留KV缓存与预填充问题

正确性和SBUF只是故事的一半；吞吐量来自于将KV缓存保留在设备上，并且永远不通过主机进行往返。设计了三种方案：

5.1 双图静态KV（E2B，单核）

两个neff共享一个静态KV缓冲区：预填充（填充后的prompt ≤ KV_BUCKET，返回15个非共享层的K/V）和解码（针对固定KV_MAX缓冲区的单token前向传播，one-hot掩码写入）。KV张量作为图输入/输出。约44 tok/s，一次性约100秒的neff加载后，预填充约0.06秒。

5.2 TP + 设备上别名KV（E4B/12B）

在TP下，一个neff跨越两个核，因此不能将预填充放在核0上，解码放在核1上。两种子设计处理预填充：

• tp_alias：仅解码neff驻留；预填充在主机CPU上运行一次以填充缓存。KV缓冲区是设备驻留的nn.Parameter，别名为图I/O（input_output_aliases），每一步就地更新，因此缓存永远不会离开核心。每个rank的KV使用头r进行填充。第一个token约1.4–1.6秒（主机预填充），解码约33 tok/s。
• 设备预填充（ModelBuilder）：预填充和解码是一个权重共享模型的两个桶，都在设备上，共享别名KV。第一个token降至约0.1–0.16秒；解码约39 tok/s（E4B）/ 约15 tok/s（12B）。这是推荐的构建方式。

# KV参数别名为图I/O，因此缓存永远不会通过主机往返
aliases = {}
for j in range(NK): aliases[w.kbuf[j]] = 1 + j
for j in range(NK): aliases[w.vbuf[j]] = 1 + NK + j

每个token的成本在上下文长度上基本持平，因为缓存是设备驻留的。

5.3 保存/重新加载

跟踪的执行器使用torch.jit.save(model, path)保存，并使用torch.jit.load + nxd_model.initialize_with_saved_weights(torch.tensor(start_rank))重新加载。（执行器自身的.save是TorchScript的，而非NxD的——这是一个微妙的陷阱。）

6. 结果

所有三个移植版本在贪心解码（SEQ_MATCH True）上与CPU浮点参考逐token匹配，例如"The capital of France is Paris."

12B较低的解码速率是固有的（密集12B，每个token所有权重都激活），而非移植缺陷；每个token的成本与8xlarge完整服务器匹配。值得注意的是，所有三个版本都可以在单个inf2.xlarge上运行（价格为8xlarge的¼，相同的2核加速器），通过"精简版"服务器将主机端嵌入保留在CPU上。

7. 值得单独说明的操作陷阱

• 缺少tokenizer.json会伪装成设备bug。 一次挂起的hf download曾导致tokenizer.json缺失；GemmaTokenizer加载时没有词汇表，将每个prompt映射为单个。模型随后忠实地输出了垃圾（未使用的高ID token），这看起来完全像是设备/重新加载/精度故障，耗费了数小时。决定性的诊断是在同一台机器上运行CPU参考：真实CPU模型产生了相同的垃圾，排除了加速器的嫌疑，并将问题指向输入。在怀疑编译器之前，始终检查tok("hello world").input_ids的合理性。
• inf2.xlarge（16 GB）需要交换空间。 服务运行只需约3.6–8 GB，但一次性neff加载峰值约14.5 GB；没有交换空间的默认DLAMI会OOM杀死容器以及SSM代理。请先添加交换文件。
• Docker镜像超过70 GB，无法在inf2.xlarge上构建。 基础镜像加上模型文件再加上交换空间，会在解压过程中溢出根磁盘。请在inf2.8xlarge（700 GB根磁盘）上构建精简标签；推送时仅上传微小的服务器层，因为基础镜像已在Hub上。

8. 跨领域发现

上述三个障碍是切入点。移植还产生了一组超越Gemma-4的发现，总体而言，其价值超过任何单个修复。

8.1 "垃圾输出"几乎从不意味着"加速器坏了"

[第 4/5 部分]

每当这个端口产生垃圾输出时，NeuronCore 都是无辜的。按出现频率排序，原因依次是：损坏/缺失的 tokenizer（§7）、错误恢复的 权重重载（§8.2）、错误的 buffer（layer_scalar，§4.6），以及不匹配的 driver/SDK 版本（预编译的 neff 在错误的运行时上可能错误执行产生垃圾而非报错）。这四种情况排查成本都比设备 bug 更低，而且它们看起来都像是精度/硬件故障。

最快的判断方法是在同一台机器上运行 CPU 参考：使用 from_pretrained 以 bf16 加载模型（权重下载后约 5 秒从页面缓存加载），然后运行一次前向传播。如果 CPU 参考产生与设备相同的垃圾输出，则加速器被排除嫌疑，bug 在上游（tokenizer、输入、权重）。这一单一技术将耗时数小时的"设备 bug"排查压缩为两分钟的 tokenizer 修复。在分析 neff 之前，请先使用 CPU 参考。

8.2 验证服务路径，而非trace

整个项目中最危险的幻觉是绿色的 SEQ_MATCH。正确性验证是在进程内对新 trace 的模型进行的（ModelBuilder.trace(initialize_model_weights=True)）——但服务器加载的是已保存的模型，在新进程中调用 initialize_with_saved_weights()。这是不同的代码路径，进程内的验证从未测试过服务器实际使用的路径。（torch.jit.load 没有 init 调用会直接失败——"This model is not initialized"——因此 init 步骤是承载性的，而非装饰性的。）

结合自动端口对剥离了 PLE 的 golden 产生的虚假的"100% PASS"（§3），教训很直接：针对错误的 oracle 或在错误的执行路径上通过的测试，比失败的测试更糟糕。 验证（a）你将实际交付的确切工件，（b）以你实际加载它的确切方式，（c）针对独立的浮点参考——而不是 trace，也不是源自相同错误假设的 golden。

8.3 "有效参数"是容量的谎言，bf16 无法拯救核心

E2B 被宣传为"2B"，E4B 被宣传为"4B"——这是 MatFormer/PLE 有效参数计数。但设备占用的是完整参数计数（约 5B 和约 8B），这才是必须适配 16 GB 的内容。这正是为什么 E4B（"4B"模型）无法适配一个核心，而 E2B（"2B"模型）可以。根据实际参数规划容量，永远不要看有效的宣传数字。

陷阱的第二部分：bf16 将磁盘上的 neff 减半，但不会减少设备上的常量数量。 "只需使用更低精度来适配"的直觉在这里是错误的——E4B 约 15.4 GB 的常驻 fp32 常量无论边界处的 dtype 技巧如何，都保持约 15.4 GB 的插槽，而第二个 neff 仍然超过 16 GB。张量并行，而非精度，才是真正将模型适配到核心上的杠杆（§4.4）。

8.4 KV 共享也是礼物，而不仅仅是障碍

NxD 无法表示的相同跨层 KV 共享（§2.1，§3）正是 E2B 能够适配一个核心的原因。只有15 个非共享层分配 KV 缓冲区；其余层通过读取共享。因此设备驻留缓存远小于天真的每层一个缓冲区的实现，这是内存预算得以闭合的重要原因。一个打破供应商抽象层的架构特性，一旦你停止与之对抗，仍然可以成为净收益。

8.5 利用数值不变性适配 SBUF

两个编译器优化来自数学，而非编译器：

• 单调性 ⇒ argmax 不变性。 Logit 软裁剪是单调的，因此它不能改变 argmax——贪心解码无论是否运行它结果都相同。将其移至主机端（仅采样时需要）释放了它溢出的 SBUF（§4.2）。推广：任何单调的、逐元素的输出变换都可以在贪心解码时从设备图中移除。
• 参数 vs 缓冲区。 "加载检查点"的框架通常指参数，会静默跳过缓冲区，如 layer_scalar（§4.6）。当模型乘以一个位于缓冲区中的学习标量时，该标量在任何仅参数加载后都会出错——这是一个没有错误消息的 bug，只会导致输出质量下降。

8.6 廉价盒子是相同的加速器

inf2.xlarge 和 inf2.8xlarge 搭载相同的 2-NeuronCore / 32 GB-HBM 加速器；它们仅在主机 vCPU（4 vs 32）和 RAM（16 vs 128 GB）上有所不同。由于 Gemma-4 的 transformer 完全在核心上运行，唯一阻碍约 4 倍更便宜的盒子实现完整性能的是主机内存——通过"精简"服务器解决，这些服务器仅将嵌入表保留在 CPU 上，丢弃 transformer 层（它们位于 neff 中）。结果：所有三个模型都在单个 inf2.xlarge 上提供服务，并且由于吞吐量几乎不下降，廉价盒子不仅每小时更便宜，而且每个 token 更便宜。（一个不太明显的需求是交换空间：neff 加载时短暂峰值约 14.5 GB，在 16 GB 主机上——§7。）

8.7 预填充是工作负载决策，而非固定设计

两种预填充策略共存且权衡清晰，因此"正确"的策略取决于工作负载：

由于 KV 缓存在两种策略中都是设备驻留的，每个 token 的延迟在上下文长度上基本平坦——没有随序列增长的主机往返。需要调整的旋钮是首个 token 延迟与持续解码速度之间的权衡，根据用例而非模型选择。

9. 工件

Hugging Face（配方 + 编译后的 neff + Dockerfile + 模型卡片）：

• xbill9/gemma-4-E2B-it-inferentia2
• xbill9/gemma-4-E4B-it-inferentia2
• xbill9/gemma-4-12B-it-inferentia2

Docker Hub（预构建，运行时需添加 --device /dev/neuron0 --ipc=host）：

• xbill9/gemma4-optb — E2B 版本：latest/512-128、slim、tp2-slim、tp2-2048
• xbill9/gemma4-optb-e4b — E4B 版本：latest/512-128、tp