使用原生 PyTorch 在 AWS Trainium2 上训练和服务 i14 的 Laplace 3B
Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch
- 谁
- AWS Machine Learning Blog
- 引文
This post explains the model design, the changes i14 made for Trainium2, and the training and inference results.
发生了什么
i14 和 AWS 介绍了如何通过 TorchNeuron Native 使用原生 PyTorch，将 Laplace 3B 的线性注意力架构适配到 AWS Trainium2 上进行训练和服务。该模型为每个序列维持固定的 29.5 MB 循环状态，训练速度约为每秒 120,000 个 token，报告的结果覆盖约 440 亿个 token。在单个 trn2.48xlarge 实例上，同时处理 1,280 个序列时，推理速度达到每秒 11,544 个输出 token。
推荐理由
这篇工程实践指南展示了通信模式、算子调度和批次宽度变化如何主导性能，并提供了一套可迁移的分析方法，用于在加速器上部署自定义架构。
这件事到目前为止
i14 与 AWS 将 i14 的 Laplace 3B 线性注意力模型适配到 AWS Trainium2，使用 TorchNeuron Native 通过原生 PyTorch 进行训练和服务，并报告训练速度约为每秒 120,000 个 token，推理吞吐量为每秒 11,544 个输出 token。 AWS 表示，其 GenAIIC 团队在一台 trn2.48xlarge 实例上使用全部 64 个逻辑 NeuronCore 训练了该模型，训练结果报告覆盖约 440 亿个 token。该模型为每个序列保持固定的 29.5 MB 循环状态。推理方面，团队在一台 trn2.48xlarge 实例上使用未经修改的 vLLM 引擎；报告的吞吐量是在 1,280 个并发序列下测得的。
继续阅读
最新日报 · 第 27 期TorchNeuron 为 AWS Trainium 上的 Hugging Face 模型带来原生 PyTorch 支持正在热议
明早 08:00 把这样一份日报发到你邮箱
和上面这份一样：几件值得看的事。免费，点确认信之后才会开始发送，每封信都能一键退订。 看一份往期日报 · 2026-10-20