跳到正文
AWS Machine Learning Blog· Kevin Shabahang·· 3 天前精选AI 评分46

i14 将 Laplace 3B 线性注意力模型移植至 AWS Trainium2 训练与推理

Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch

AI 导读

i14 与 AWS 使用原生 PyTorch 支持，将 Laplace 3B 线性注意力模型部署到 AWS Trainium2，保留 Gated Linear Attention（GLA）架构和 Muon 优化器。

正文 · AI 翻译

大多数大型语言模型（LLM）采用带有 softmax 注意力机制的 Transformer 架构。在推理过程中，它们的键值（KV）缓存会随着对话或文档上下文的增长而增大。这会降低并发能力、增加服务成本，还会迫使智能体应用丢弃有用的历史信息。

i14 将 Laplace 3B 设计为一种线性注意力替代方案。其循环状态在每个序列中保持为 29.5 MB，不受上下文长度影响。i14 与 AWS 合作，借助 通过 TorchNeuron Native 提供的原生 PyTorch 支持，将这一自定义架构部署到 AWS Trainium2 上。此方案保留了其门控线性注意力（GLA）设计和 Muon 优化器。

在一台 trn2.48xlarge 实例上，Laplace 3B 的训练速度约为每秒 120,000 个 token。该模型在 1,280 个并发序列上实现了每秒 11,544 个输出 token。本文将介绍模型设计、i14 为 Trainium2 所做的改动，以及训练和推理结果。

关于 i14

i14 是一家 AI 研究与技术公司，致力于开发模型架构、组件和运行时。其工作借鉴了 AI、计算认知科学，以及关于生物系统如何随时间学习的研究。i14 将这些理念应用于模型，使其能够保留长距离上下文，同时避免服务状态随每个 token 增长。

为什么采用线性注意力模型？

Transformer 的 softmax 注意力会比较上下文中的各个 token。在训练过程中，这项计算会随序列长度呈二次方增长。在推理过程中，KV 缓存会随着模型处理更多 token 而增大。

线性注意力会围绕一个逐 token 更新的循环状态重新组织计算。Laplace 3B 使用 GLA 将先前的上下文压缩为固定大小的矩阵状态。这使训练注意力的复杂度从 O(T2) 降为 O(T)，并将每个序列的推理状态固定为 29.5 MB。

考量因素 Transformer softmax 注意力 采用 GLA 的 Laplace 3B
训练注意力 构建成对的 token 分数，复杂度为 O(T2) 更新分块循环状态，复杂度为 O(T)
推理内存 KV 缓存随上下文增长 循环状态大小保持不变
更长的上下文 每个序列会占用更多加速器内存 上下文长度不会增加循环状态的内存分配
实现 使用经过广泛优化的注意力内核 需要针对特定架构的内核和编译器调优

线性注意力会压缩历史信息，而不是保留每个键和值。因此，必须针对自己的工作负载评估检索质量。固定状态有利于长时间运行的智能体、大型文档分析，以及需要持久记忆和可预测服务需求的应用。

i14 如何将 Laplace 3B 适配到 AWS Trainium2

Laplace 3B 的原始实现使用了针对 GPU 优化的 Triton 内核。AWS 生成式 AI 创新中心（GenAIIC）团队首先使用标准 PyTorch 算子重写了适用于 Trainium2 的 GLA 内核。比较两种方案后，团队选择了测试版 TorchNeuron Native（AWS Trainium 的原生 PyTorch 后端），因为其性能更好。

该实现需要进行三项改动：

  1. 使用原生 PyTorch 实现 GLA 注意力内核。团队将面向 GPU 的融合式 fused_chunk_simple_gla Triton 内核替换为标准 PyTorch 算子，并使用 16 个每块 128 个 token 的分块，以匹配 NeuronCore 片上缓冲区的 128 个分区。
  2. 重新设计 Muon 的通信方式。 MuonFSDP2 中的两个 all_to_all 集合通信将每个 rank 的流量减少了 24 倍，吞吐量提升了 1.63 倍。
  3. 编译顺序反向传播过程。 团队对 GLA 递归中的 16 个分块步骤进行静态展开，以满足编译器对图规模和片上内存的限制。

工程挑战部分详细介绍了每项改动。

为什么选择 AWS Trainium 加速器？

AWS Trainium 是 AWS 机器学习 (ML) 加速器系列，专为模型训练和推理而设计。一个 trn2.48xlarge 实例结合了 16 颗 Trainium2 芯片、64 个逻辑 NeuronCore、高带宽内存和 NeuronLink-v3 互连。这些资源支持本文所述的分布式训练和推理配置。

Trainium2 的三个特性影响了实现方式：软件管理的片上内存（挑战 1）、高速芯片间互连（挑战 2）和静态编译器（挑战 3）。编译器将 GLA 递归跟踪为固定图。

软件管理的片上内存。每个 NeuronCore 都配有一个 28 MiB 的状态缓冲区 (SBUF)，包含 128 个分区，以及一个 2 MiB 的部分和缓冲区 (PSUM)。PSUM 与一个 128×128 脉动式 Tensor Engine 相连，此外还有独立的向量和标量单元。SBUF 不是缓存：编译器会提前安排每次 HBM-to-SBUF 传输，因此数据会就位，不会发生逐出或未命中（图 1）。

NeuronCore-v3 internal architecture

图 1：NeuronCore-v3 内部架构。软件管理的 SBUF（28 MiB，128 个分区）、Tensor Engine（128×128 脉动阵列）、Vector Engine、Scalar Engine 和 GPSIMD 核心。

i14 和 Laplace 3B 架构

Laplace 3B 是由 24 层组成的 i14 解码器，每一层都以 Simple GLA 取代 softmax 注意力。图 2 展示了一层的结构，从顶部的输入 xt 到底部的输出 xt + 1。本节将介绍各个模块：归一化和卷积、QKVG 投影、GLA 混合器以及前馈通道混合器。

归一化和卷积。每个 token 的残差流向量 xt ∈ ℝD（D = 3, 072）都会经过 LayerNorm 和宽度为 k = 4、带 SiLU 的深度因果 Conv1d。这样，在混合器处理之前，每个位置都会获得一个 4-gram 摘要。随后，GLA 混合器通过其衰减状态，对整个历史序列进行基于内容的检索。

QKVG 投影。卷积后的输出经过一个融合矩阵 Wqkvg ∈ ℝ9, 216 × 3, 072，其中 9, 216 = H × (2dk + 2dv)。它会拆分为每个 token 对应的四个流：查询 qt、键 kt、值 vt 和门控 gt（H = 16，dk = 96，dv = 192）。门控在混合器之后使用。

Laplace decoder layer architecture

图 2：Laplace 解码器层架构。绿色 = 固定（无参数），灰色 = 已学习，黄色 = 循环状态，蓝色 = 投影。

GLA 混合器。每个头 h 都维护一个大小为 St(h) 的矩阵状态 （96 × 192），它是对过去键值对的压缩衰减记录。在每个 token 处，状态更新如下：

GLA recurrent-state update equation

GLA 循环状态更新公式


这里的 λh 是每个头固定的衰减系数，ktvt⊤ 是秩为 1 的外积，用于将当前键值对写入状态。读取时，查询向量会查询该状态：

GLA state-readout equation

GLA 状态读取公式


接下来，经过自增强步骤，为每个头的输出添加一个门控值副本：

GLA self-augmentation equation

GLA 自增强公式


这里的 σ 是 sigmoid 函数，⊙ 表示逐元素乘积。门控值是查询向量与经过学习得到的逐通道缩放的键向量的点积，其中 a 是学习得到的缩放参数（aug_balance）。初始化时，a = 0，因此值向量会以一半强度跳跃连接。每个头的输出先在各个衰减层级内经过 Head RMSNorm，再进行拼接，随后由 SiLU(gt) 门控，最后通过 Wout ∈ ℝD × D 投影回残差流。

衰减计划。每个头的衰减率 λh 控制该头遗忘过去 token 的速度。它在初始化时固定（不可学习），并按照几何计划设定。该计划为 16 个头设定不同的记忆长度，并定义每个头的对数衰减率 γh：

GLA decay schedule equations

GLA 衰减计划公式


衰减计划的范围。对于此模型，Tmin = 1，Tmax = 2, 048。因此，头 0 的 λ0 = 0.5（半衰期为一个 token），头 15 的 λ15 ≈ 0.9997（半衰期为 2,048 个 token）。头按层级分组，每个层级共用一个 Head RMSNorm。

通道混合器（FFN）。注意力块及其残差之后，第二个 LayerNorm 会将输入送入一个内维度为 8,192 的 GLU（SiLU）FFN。其三个矩阵分别是两个升维投影 W1, W2 ∈ ℝ8, 192 × D 和一个降维投影 W3 ∈ ℝD × 8, 192。之后再接一个残差连接。

状态 S(h) 对每个头固定为 dk × dv = 96 × 192，不受序列长度影响。因此，推理时每个 token 的复杂度为 O(1)，训练复杂度为 O(T)，而非 O(T2)。

参数 值
参数量 总计 3.028B（2.718B Muon + 309.6M AdamW）
层数 24 个解码器层，隐藏维度 D = 3, 072
头数 16 个头，dk = 96（半宽 Q/K），dv = 192（全宽 V）
混合器 采用固定几何衰减的简单 GLA（Tmin = 1，Tmax = 2, 048）
卷积 深度方向因果 Conv1d，核宽度为 4，SiLU
FFN GLU（SiLU），内部维度为 8,192
归一化 注意力前 LayerNorm + Head RMSNorm（每个层级）+ FFN 前 LayerNorm
词表 50,304（填充至 64 的倍数）
序列 2,048
训练数据 100B tokens，FineWeb-Edu
优化器 二维权重矩阵使用 Muon；嵌入和归一化层使用 AdamW
批次大小 524,288 tokens/步，余弦学习率，预热 3,815 步

工程挑战

这三项改动分别带来了各自的工程挑战。

挑战 1：在 Trainium 上使用原生 PyTorch 实现 GLA

在 GPU 上，GLA 通过 flash-linear-attention 库中的融合 Triton 内核 fused_chunk_simple_gla 运行。GenAIIC 团队使用标准 PyTorch 运算符重新实现该功能，通过 TorchNeuron Native 运行，并以显式循环处理固定大小的块。块数在编译时固定，并决定反向传播展开的深度。团队使用 fla_num_chunks = 16，因此每个块包含 C = 128 个 token。

在每个块 c 内，块内输出采用掩码线性注意力，运行状态 S 则纳入先前上下文：

Chunked GLA intra-chunk output

分块 GLA 的块内输出

随后，状态推进到下一个块：

Chunked GLA state update

分块 GLA 的状态更新

其中，λC 是每个头在整个块上的累积衰减。16 个块内计算并行运行，只有状态更新是顺序执行的（图 3）。

GLA chunked recurrence decomposition

图 3：GLA 分块递归：块内注意力并行运行，状态链顺序执行。

由于 SBUF 有 128 个分区，每个包含 128 个 token 的块都能将一个 token 映射到一个分区，无需填充或跨分区移动。这种对齐使实测吞吐量提升了 14 percent（图 4）。

GLA chunk alignment with SBUF partitions

图 4：块大小 C = 128 与 128 个 SBUF 分区一一对应。

挑战 2：FSDP2 下的分布式 Muon 正交化

Muon 会在更新前对每个梯度进行正交化。对于 96 个二维权重矩阵（共 2.718 billion 个参数），它执行五次 Newton–Schulz 迭代（算法 1）。

The Newton-Schulz algorithm

Newton-Schulz 算法

分布式问题。Gram 矩阵 A = XX⊤ 会计算权重的每一行与其他所有行之间的相关性。在完全分片数据并行（FSDP2）下，64 个 rank 中的每个 rank 只持有 1/64 的行。因此，在分片本地计算 XX⊤ 只能得到一个对角块，并会丢弃所有跨分片项（图 5）。

Newton-Schulz communication patterns under FSDP2

图 5：FSDP2 下的 Newton–Schulz：分片本地 Gram 块，以及 gather/broadcast 与 all_to_all 的对比。

朴素方法：Gather 和 broadcast。每个 rank 都收集完整梯度（5.44 GB），由一个 rank 执行 Newton–Schulz，再广播结果。每步需传输 10.87 GB，其余 63 个 rank 则处于等待状态。

生产方案：all_to_all 转置。在 MuonFSDP2 中，一次 all_to_all 会让每个 rank 获得若干个完整矩阵，全部 64 个 rank 并行运行 Newton–Schulz，随后再通过第二次 all_to_all 返回结果。每个 rank 的通信量降至 0.45 GB，减少了 24 倍。吞吐量从每秒 73,378 个 token 提升至 119,934 个，增幅为 1.63 倍，且损失曲线逐位一致（图 6）。

Communication volume per training step

图 6：每步通信量：使用 all_gather 时为 10.87 GB，使用 all_to_all 时为 0.45 GB。

通信开销。Newton–Schulz 只占训练步骤 FLOPs 的约 2%（1.95 × 1014 / 9.53 × 1015），因此瓶颈在于通信。通过带宽为 1.28 TB/s 的 NeuronLink-v3，0.45 GB 的数据交换耗时不到 0.4 ms，而每步约需 4.4 秒。

挑战 3：顺序递归的反向传播展开

反向传播需要对全部 16 次顺序状态更新求导。在 GPU 上，autograd 会在运行时构建该计算图。在 Trainium2 上，编译器必须预先看到完整展开的 16 步计算图，这带来了三项限制。

计算图深度会随着分块数量增加：当有 32 个分块（C = 64）时，计算图超出了调度器的容量。每个分块的内存占用会随着分块大小增加：当有 8 个分块（C = 256）时，中间结果溢出到 HBM。由于所有形状都在跟踪时固定，因此每个循环中间结果都需要单独的静态分配。

16 个分块（C = 128）满足全部三项要求：展开后的计算图适配调度器，中间结果留在芯片上，且图块与 128 个 SBUF 分区对齐。

训练结果

GenAIIC 团队在一台 trn2.48xlarge 实例上使用全部 64 个逻辑 NeuronCore 和 all_to_all Muon 路径训练 Laplace 3B。完整训练方案为 190,735 步（FineWeb-Edu 的 1000 亿个 token）。以下结果截至第 84,210 步，约处理了 440 亿个 token。

吞吐量。编译预热后，吞吐量稳定在每秒约 120,000 个 token，并一直保持到第 84,210 步。

损失与梯度稳定性。损失（图 7a）从约 11.4 平稳下降至第 84,210 步时的 2.3729。梯度范数（图 7b）如预期般在早期出现峰值，随后稳步下降至约 0.05；预热后没有峰值、发散或 NaN/Inf 值。

Laplace 3B training loss and gradient norm

图 7：训练动态。(a) 损失从 11.4 降至 2.3729。(b) 来自 TensorBoard 数据的梯度范数。

在 Trainium2 上优化推理

GenAIIC 团队通过未经修改的 vLLM 引擎，在一台 trn2.48xlarge 实例上以 1,280 个并发序列实现了每秒 11,544 个输出 token。

恒定的服务状态。每个序列的循环状态为 29.5 MB，约为 8,192 个 token 时可比 KV 缓存的 1/82。将提示长度从 2,048 个 token 增至 8,192 个 token，单个输出 token 的成本仅增加了 0.5%。

未经修改的 vLLM。通过一个包装器替换每个解码器层的一个属性——token mixer，以便在 vLLM 池中读写循环状态。由于模型声明自身不使用注意力机制，vLLM 会为每个序列分配一个固定页面。输出与 i14 参考结果逐 token 匹配（BF16 下，3/3 个提示的 48/48 步均匹配）。

找出瓶颈。将解码时间对解码器层数 L 进行线性拟合，得到：

Decode-time linear fit equation

解码时间线性拟合方程

At L = 24, that is 168.1 ms per step, 96.9 percent of it inside the decoder layers. The large projections cost 13.27 ms, near the 12 ms HBM-streaming floor. Most time went to small operators dispatched one at a time under eager PyTorch.

扩展到完整实例。没有张量并行时，吞吐量通过副本扩展，每个逻辑 NeuronCore 配置一个引擎和 20 个序列。Eager 吞吐量从单个副本时每秒 133 个输出 token 增长到 64 个副本时的 6,772 个，此时主机 CPU 使用率达到 82 percent。编译解码步骤消除了逐算子调度，使吞吐量提高 1.7 倍，达到 11,544（图 8）。

Whole-instance inference throughput

图 8：整实例推理吞吐量。(a) Eager 和编译解码的副本扩展。(b) 编译解码将吞吐量从每秒 6,772 个输出 token 提升至 11,544 个。

固定批次基准测试掩盖的问题。在真实流量下，每种新的解码批次宽度都是一个单独编译的形状，首次使用时会耗费约 11 秒。这些步骤仅占总量的 1.1 percent，却消耗了 47.5 percent 的解码时间。填充到固定宽度消除了这项开销，且结果逐位精确。随后，在交错到达的情况下，编译解码使每个副本的吞吐量提高了 3.52 倍；相比之下，在固定批次下仅提高了 1.46 倍。由于 BF16 重新结合会导致编译输出与 eager 输出略有不同，因此在需要精确复现时应使用 eager 解码。

关键要点

该项目在训练和推理方面都积累了经验：

  • 原生 PyTorch 可支持 Trainium2 上的自定义架构。i14 保留了 GLA 架构，并使用 TorchNeuron Native 完成最终实现。
  • 通信设计对自定义优化器至关重要。Muon 的 all_to_all 设计将每个 rank 的流量减少了 24 倍，并使训练吞吐量提高了 1.63 倍。
  • 固定状态让推理内存保持可预测。无论上下文长度如何，每个序列都使用 29.5 MB。
  • 编译前先进行性能分析。解码时间的大部分消耗在算子调度，而非算术运算或 HBM 带宽。
  • 使用真实的到达模式进行基准测试。固定批次掩盖了为不断变化的批次宽度编译新形状所带来的开销。

结论

Laplace 3B 表明，AWS Trainium2 可以支持标准 Transformer 之外的模型架构。借助 TorchNeuron Native 和 AWS Neuron SDK 使用原生 PyTorch，i14 训练并部署了基于 GLA 的模型，同时保留了固定的循环状态。

对于智能体系统而言，这改变了一项基本的服务约束。更长的历史记录不需要更大的单序列状态，这有助于团队在保持内存使用可预测的同时保留更多上下文。

如果上下文增长正在推高服务成本、降低并发量或限制智能体的记忆容量，请试用 i14，或联系您的 AWS 客户团队，评估在 AWS Trainium2 上运行 Laplace 3B。

资源

作者简介

Kevin Shabahang

Kevin Shabahang

Kevin 是 i14 的 AI 研究副总裁，负责研究不同的模型架构。他的工作重点是将认知科学和计算神经科学的研究成果应用于开发更高效的大型语言模型。

Massimo Bardetti

Massimo Bardetti

Massimo 是 i14.ai 的首席技术官，该公司是一家 AI 研究初创公司。他曾领导 Amazon Web Services、Microsoft 以及 AI、神经技术和娱乐领域初创公司的工程团队。他还创作互动艺术装置，这些作品曾在世界各地的艺术节上展出。

Iman Abbasnejad

Iman Abbasnejad

Iman 是 AWS 的应用科学家，与客户合作设计、开发和部署生成式 AI 解决方案。他的工作连接 AI 研究与实际应用，重点关注大型语言模型、智能体 AI 系统、模型微调和性能优化，致力于在 AWS 上构建可靠、高效且经济实惠的 AI 解决方案。

Xiaodong Yu

Xiaodong Yu

Xiaodong 是 AWS 的机器学习工程师，专注于生成式 AI、大型语言模型和机器学习性能优化。他的工作涉及在 AWS 专用机器学习加速器（包括 AWS Trainium）上实现高效训练和推理。

Daniel Wirjo

Daniel Wirjo

Daniel 是 AWS 的解决方案架构师，专注于 AI 和 SaaS 初创公司。作为曾任初创公司 CTO 的从业者，他与创始人和工程负责人合作，帮助他们在 AWS 上构建业务并实现增长。工作之余，Daniel 喜欢一边喝咖啡一边散步、亲近自然，也喜欢学习新知。

来源：AWS Machine Learning Blog · aws.amazon.com

部署工程教程实践#教程/实践#部署/工程