将 Google 的 Gemma-4 31B 模型移植到 AWS Inferentia2 inf2.24xlarge 后，其输出与 CPU 参考实现实现了完美的逐 token 匹配——然而，生成的每一句话都是乱码。这种“匹配”与“可用”之间的差距，为任何试图将大规模 LLM 挤进 Amazon 定制推理芯片的人敲响了警钟。

为什么逐 token 匹配并不足够

开发者将 Inferentia 设备的每个输出 token 与 CPU 运行该模型产生的 token 进行了对比。由于流完全一致，硬件似乎精确地复现了参考实现。但实际上，两个流都将一个格式错误的提示词（prompt）输入到了一个被剥离了聊天模板（chat template）且提供了错误轮次标记（turn markers）的模型中。缺失的模板导致模型进入了死循环，不断吐出无意义的内容。硬件完成了它的工作——它复现了参考代码中已存在的 bug。

教训很简单：SEQ_MATCH（顺序 token 相等）并不等于正确性。如果参考实现本身是损坏的，那么忠实的硬件副本也会继承同样的错误。验证必须超越 token 级别的对等性；它需要使用格式正确的输入进行端到端的函数检查。

伪装成参数的缓冲区 (Buffers)

在加载阶段，模型加载器跳过了一个名为 layer_scalar 的组件。在 PyTorch 模型定义中，代码将该对象注册为缓冲区 (buffer) 而非参数 (parameter)。缓冲区是训练过程中不会更新的静态张量，许多加载器在转换为 Neuron 兼容格式时会忽略它们。跳过该组件导致多个层的缩放因子（scaling factors）保持在默认值，从而扭曲了整个网络的数学计算。系统没有报错；模型成功编译，推理流水线也能运行，但数值结果是错误的。

对于任何将大型模型迁移到 Inferentia 的人来说，请审计每一个非参数张量。即使一个张量不用于学习，它对于正确的正向传播计算也可能至关重要。手动验证缓冲区的包含情况，可以防止难以诊断的隐性缩放错误。

Spot 实例的波动性与 39 分钟的编译时间

在 Spot 实例上运行 310 亿参数的模型看起来很便宜，但这种节省伴随着不可预测的实例回收事件。开发者的编译时间——将模型转换为 Neuron 兼容代码大约需要 39 分钟——在 AWS 回收实例时化为乌有。为了应对中断，他们构建了一个三层安全网：

  • ModelBuilder 将内存使用量控制在 384 GB 的主机限制内，避免了导致必须重启的崩溃。
  • 对原始权重文件和编译后的 “neffs”（Neuron 可执行文件）进行 S3 即时镜像，使新实例能够从前一个实例中断的地方精确接续。
  • 一个 多区域轮询器 (multi-region poller) 扫描 AWS 各个区域的可用 Spot 容量，并在一旦出现可用容量时立即启动新实例。

这些步骤将脆弱的、单点的编译过程转变为一个能够抵御 Spot 市场波动的韧性流水线。

混合注意力布局带来的分片 (Sharding) 陷阱

Gemma-4 31B 使用了两种注意力配置。某些层采用四个键值 (KV) 头，而其他层则采用不同的数量。当某层的 KV 头数量无法被 8 个并行 rank 整除时，将模型均匀拆分到 8 个并行 rank 的操作就会失败。尝试将一个 4 头的层分片到 8 个 rank 中，会迫使每个 rank 处理半个头——这在数学上是不可能的，会导致形状不匹配和运行时错误。

解决方案是：在所有 rank 中复制全局分片的层（即那些头数量兼容的层），并且仅对头数量允许均匀拆分的“滑动”层进行分片。这种混合策略在保持张量并行效率的同时，避免了对 KV 头的非法拆分，消除了困扰早期尝试的张量并行化错误。

总结

将巨型 LLM 移植到 Inferentia 不仅仅是一个“编译并运行”的过程。它需要超越 token 相等性的严格功能测试，需要细致验证每一个张量（无论是参数还是缓冲区）是否处理正确，还需要能够预见 Spot 实例回收的部署策略。最后，分片必须尊重模型的内部注意力几何结构；否则，原本承诺提速的并行化反而会成为隐性失败的根源。