超越KV缓存：在AWS Trainium2上使用原生PyTorch训练和部署i14的Laplace 3B线性注意力模型
aws ·
本文详细介绍了 i14 如何利用 TorchNeuron Native，通过原生 PyTorch 在 AWS Trainium2 上训练并部署其 Laplace 3B 线性注意力模型。与传统依赖不断增长的键值缓存的 Transformer 不同，Lapace 3B 采用门控线性注意力（GLA），无论上下文长度如何，每个序列仅维持 29.5 MB 的固定循环状态。该实现涉及三项重大工程调整：使用标准 PyTorch 算子重写 GLA 内核、为 Muon 优化器优化分布式通信，以及为编译器静态展开反向传播过程。在单个 trn2.48xlarge 实例上的训练实现了每秒约 120,000 个 token 的处理速度，且损失收敛稳定。通过将解码步骤编译以减少算子调度开销，推理性能在 1,280 个并发序列下达到了每秒 11,544 个输出 token。这些结果表明，自定义架构能够在 AWS 硬件上实现高效率，并为长上下文应用提供可预测的内存使用。
aws