Laplace 3B 线性注意力模型登陆 Trainium2:单实例每秒 12 万 token 训练
AWS ML Blog · rss · 2026-10-21
AWS ML Blog 详解 i14 公司的 Laplace 3B——一个用门控线性注意力(GLA)替代 softmax attention 的 24 层 3B 模型——如何在 AWS Trainium2 上以原生 PyTorch 完成训练与服务。
架构要点
- GLA 将历史压缩为每序列固定 29.5 MB 的循环状态,不随上下文增长;训练注意力从 O(T²) 降到 O(T)
- 16 个头按几何衰减调度分配记忆长度(半衰期 1 到 2048 token),配 Muon 优化器
Trainium2 适配三项改动
- 用标准 PyTorch 算子重写 GPU 专用 Triton GLA kernel,切成 16 个 128-token chunk 以匹配 NeuronCore 片上缓存的 128 个分区
- 重设计 Muon 通信:MuonFSDP2 中两个 alltoall 集合通信将每 rank 流量削减 24 倍,吞吐提升 1.63 倍
- 静态展开 16 步 GLA 循环以适配编译器图规模与片上内存限制
性能:单台 trn2.48xlarge(16 颗 Trainium2)训练约 12 万 token/秒,推理在 1280 并发序列下吞吐 11544 输出 token/秒。
「Infra」频道最新
- NVIDIA 推出 FP4 本地化检查点:Qwen 与 GLM 模型省显存跑 agent — NVIDIAAI · 2026-10-21
- 曝小团队后训练 Qwen3 27B,GQA 转 MLA 省 8.4 倍 KV 缓存 — ChrisGPT · 2026-10-21
- EAGLE-4 投机解码实测:agentic 编码场景提速最高 5.21 倍 — zhyncs42 · 2026-10-21
- EAGLE-4 发布:投机解码冲上 961.6 tokens/s,代码与仓库开源 — zhyncs42 · 2026-10-21
- REAX 引擎竞赛开放一晚,矿工将决策模型提速近 3 倍 — airesearch12 · 2026-10-21
- 马斯克称算力明年或六个月回本,分析师质疑 neocloud 融资难 — firstadopter · 2026-10-21