i14 在 AWS Trainium2 上用原生 PyTorch 训练并部署 Laplace 3B 线性注意力模型

Laplace 3B 是 i14 的 24 层解码器模型，每层用 Simple GLA 线性注意力替换 softmax 注意力，把历史上下文压缩进每头 96×192 的矩阵状态；16 个头按几何衰减调度分配不同记忆长度，半衰期从 1 个 token 到 2048 个 token。训练注意力从 O(T²) 降为 O(T)，推理每 token 为 O(1)，单序列循环状态固定为 29.5 MB，不随上下文长度增长。 为适配 Trainium2，团队做了三处改动：用标准 PyTorch 算子重写 GLA 内核，采用 16 个 128 token 分块对齐 NeuronCore 片上缓冲的 128 个分区，吞吐提升 14%；在 FSDP2 下把 Muon 正交化改为两次 all_to_all 交换，单卡通信量从 0.45 GB 降至原来的 1/24，吞吐从 73378 提升到 119934 token/s；静态展开 16 步 GLA 递推以适配编译器图规模限制。 对一个 trn2.48xlarge 实例（16 颗 Trainium2、64 个 NeuronCore），Laplace 3B 训练吞吐约 12 万 token/s，推理在 1280 条并发序列下达到 11544 输出 token/s。使用该架构的团队需要注意：线性注意力压缩而非保留全部键值，需针对自身负载评估检索质量；在 Trainium2 上部署自定义架构的开发者则要面对软件管理的片上内存和静态编译器带来的调优成本。