i14在AWS Trainium2上用原生PyTorch训练并服务Laplace 3B线性注意力模型
AWS称，i14的Laplace 3B采用门控线性注意力，每序列循环状态固定为 29.5 MB，不随上下文长度增长；在单台trn2.48xlarge上训练吞吐约 12 万token/秒，推理在 1280 并发序列下达到 11544 输出token/秒。
AI解读：Laplace 3B要解决的是标准Transformer的KV缓存随上下文增长的问题。AWS的表述是，i14用门控线性注意力（GLA）把历史压缩进固定大小的循环状态，每序列 29.5 MB，与上下文长度无关；代价是历史被压缩而非逐条保留，因此来源明确提示需要针对具体负载评估检索质量。
训练侧的关键改动有 3 项：用标准PyTorch算子重写GLA内核，按 128 token分块以对齐NeuronCore片上缓冲的 128 个分区；Muon优化器改用两次all_to_all交换，使每卡通信量降到 0.45 GB、减少 24 倍，吞吐从 73378 提升到 119934 token/秒（1.63 倍）；将 16 步循环反向传播静态展开以适配编译器。
推理结果的口径需要注意：11544 输出token/秒是在单台trn2.48xlarge、1280 并发序列、未修改的vLLM引擎下测得，且来自编译解码步骤，相比eager模式的 6772 token/秒提升 1.7 倍。来源还指出，固定批次基准会掩盖真实流量下新批次宽度首次编译约 11 秒的开销，该开销占总解码时间 47.5%。
来源给出的训练数据截至第 84210 步（约 440 亿token），完整配方为 1000 亿FineWeb-Edu token、190735 步；loss从约 11.4 降至 2.3729，作者称未出现发散或NaN/Inf。这是训练进行中的结果，不是完整训练后的最终指标。
这些数字均出自AWS博客中AWS与i14合作项目的自述，属于厂商公布的单项实现与基准，不能据此推断该架构在通用任务上的质量或与其他模型的竞争胜负。
AWS Machine Learning博客发文介绍，AI研究与技术公司i14在AWS Trainium2上用原生PyTorch（TorchNeuron Native）训练并服务其Laplace 3B线性注意力模型。文章由AWS方面作者撰写，描述的是双方合作项目的实现与测试结果。
Laplace 3B是 24 层解码器模型，每层用Simple GLA替代softmax注意力。AWS称其循环状态每序列固定为 29.5 MB，不随上下文长度增长，而标准Transformer的KV缓存会随对话或文档变长而增大。模型总参数 3.028B，训练数据为 1000 亿FineWeb-Edu token，序列长度 2048。
训练侧的 3 项适配改动
i14原实现使用GPU优化的Triton内核。AWS Generative AI Innovation Center团队先用标准PyTorch算子重写GLA内核，并在比较后选择TorchNeuron Native的beta版本。
具体改动包括：用 16 个 128 token的分块匹配NeuronCore片上缓冲的 128 个分区，AWS称这一对齐使实测吞吐提升 14%；Muon优化器在FSDP2下改用两次all_to_all集合通信，每卡通信量从 10.87 GB降至 0.45 GB；将GLA循环的 16 个分块步骤静态展开，以适配编译器的图规模与片上内存限制。
- 分块数固定为 16、每块 128 token，是编译器图规模、片上中间结果和分区对齐三者共同约束下的选择
- Muon改动后训练吞吐从 73378 提升到 119934 token/秒，AWS称loss曲线逐位一致
- Newton-Schulz只占单步FLOPs约 2%，AWS判断通信是瓶颈
训练与推理的实测数字及其范围
在单台trn2.48xlarge（16 颗Trainium2芯片、64 个逻辑NeuronCore）上，AWS称训练吞吐稳定在约 12 万token/秒。公布的结果截至第 84210 步，约 440 亿token，loss从约 11.4 降到 2.3729，梯度范数降至约 0.05，作者称未出现训练后的尖峰、发散或NaN/Inf。完整配方为 190735 步。
推理方面，通过未修改的vLLM引擎，在 1280 个并发序列下达到 11544 输出token/秒。AWS称 29.5 MB的循环状态约为 8192 token上下文下可比KV缓存的 82 分之 1；提示从 2048 增至 8192 token，每输出token成本仅增加 0.5%。
来源还给出两点限制：编译解码与eager解码在BF16下输出略有差异，需要精确复现时应使用eager；固定批次基准会掩盖真实流量下新批次宽度首次编译的开销，该开销占总解码时间 47.5%，按固定宽度填充可消除。以上均为AWS与i14合作项目的自述结果。
- 推理吞吐 11544 token/秒为编译解码下的整机结果，eager模式为 6772 token/秒
- 无张量并行时通过复制引擎扩展，每逻辑NeuronCore 20个序列
- 线性注意力压缩历史，来源建议按具体负载评估检索质量