3 基础设施

本文为论文《DeepSeek-V3:强大高效的开源混合专家大模型》「中文全文翻译」板块的对应章节,与「论文原文」板块同一章节一一对应,层级保持一致。

3.1 计算集群

DeepSeek-V3 在一个配备 2048 块 NVIDIA H800 GPU 的集群上训练。H800 集群中的每个节点包含 8 块 GPU,节点内通过 NVLink 和 NVSwitch 互联。跨不同节点则使用 InfiniBand(IB)互联进行通信。

3.2 训练框架

DeepSeek-V3 的训练由 HAI-LLM 框架支持,这是一个由我们的工程师从零开始打造的高效轻量级训练框架。总体而言,DeepSeek-V3 应用了 16 路流水线并行(PP)(Qi et al., 2023a)、跨 8 个节点的 64 路专家并行(EP)(Lepikhin et al., 2021),以及 ZeRO-1 数据并行(DP)(Rajbhandari et al., 2020)。

为了实现 DeepSeek-V3 的高效训练,我们实施了精细的工程优化。首先,我们设计了 DualPipe 算法以实现高效的流水线并行。与现有的 PP 方法相比,DualPipe 具有更少的流水线气泡。更重要的是,它在前向和后向过程之间重叠了计算和通信阶段,从而解决了跨节点专家并行所带来的沉重通信开销这一挑战。其次,我们开发了高效的跨节点全互联通信内核,以充分利用 IB 和 NVLink 的带宽,并节省专用于通信的流式多处理器(SM)。最后,我们精细地优化了训练期间的内存占用,从而使我们能够在无需使用昂贵的张量并行(TP)的情况下训练 DeepSeek-V3。

3.2.1 DualPipe 与计算-通信重叠

图

图 4:一对独立前向与后向块(chunk)的重叠策略(Transformer 块的边界未对齐)。橙色表示前向,绿色表示“对输入的后向”,蓝色表示“对权重的后向”,紫色表示 PP 通信,红色表示屏障。全互联和 PP 通信均可被完全隐藏。

对于 DeepSeek-V3,跨节点专家并行带来的通信开销导致计算-通信比约为 1:1,效率低下。为应对这一挑战,我们设计了一种创新的流水线并行算法 DualPipe,它不仅通过有效地重叠前向和后向的计算-通信阶段来加速模型训练,还减少了流水线气泡。

DualPipe 的核心思想是在一对独立的前向块和后向块内重叠计算与通信。具体而言,我们将每个块划分为四个组件:注意力、全互联分发(dispatch)、MLP 和全互联组合(combine)。特别地,对于后向块,注意力和 MLP 都像 ZeroBubble(Qi et al., 2023b)中那样进一步拆分为两部分:对输入的后向和对权重的后向。此外,我们还有一个 PP 通信组件。如图 4 所示,对于一对前向和后向块,我们重新排列这些组件,并手动调整专用于通信与计算的 GPU SM 比例。在这种重叠策略下,我们可以确保在执行过程中全互联通信和 PP 通信都能被完全隐藏。鉴于高效的重叠策略,完整的 DualPipe 调度如图 5 所示。它采用了一种双向流水线调度,从流水线的两端同时送入微批次,并且很大一部分通信可以被完全重叠。这种重叠还确保了随着模型的进一步扩展,只要我们保持恒定的计算-通信比,就可以跨节点采用细粒度专家,同时实现近乎零的全互联通信开销。

图

图 5:8 个 PP rank 和 20 个微批次在两个方向上的 DualPipe 调度示例。反向方向的微批次与正向方向对称,为简化图示我们省略了它们的批次 ID。被共享黑色边框包围的两个单元格具有相互重叠的计算和通信。

此外,即使在更一般的、没有沉重通信负担的场景中,DualPipe 仍然表现出效率优势。在表 2 中,我们总结了不同 PP 方法下的流水线气泡和内存使用量。如表中所示,与 ZB1P(Qi et al., 2023b)和 1F1B(Harlap et al., 2018)相比,DualPipe 显著减少了流水线气泡,同时仅将峰值激活内存增加了 $\frac{1}{PP}$ 倍。尽管 DualPipe 需要保留两份模型参数副本,但由于我们在训练中使用较大的 EP 规模,这并不会显著增加内存消耗。与 Chimera(Li and Hoefler, 2021)相比,DualPipe 仅要求流水线阶段和微批次能被 2 整除,而不要求微批次能被流水线阶段整除。此外,对于 DualPipe,气泡和激活内存都不会随着微批次数量的增长而增加。

表 2:不同流水线并行方法的流水线气泡和内存使用量比较。$F$ 表示一个前向块的执行时间,$B$ 表示一个完整后向块的执行时间,$W$ 表示“对权重的后向”块的执行时间,$F\&B$ 表示两个相互重叠的前向和后向块的执行时间。

3.2.2 跨节点全互联通信的高效实现

为了确保 DualPipe 拥有充足的计算性能,我们定制了高效的跨节点全互联通信内核(包括分发和组合),以节省专用于通信的 SM 数量。这些内核的实现与 MoE 门控算法以及我们集群的网络拓扑进行了协同设计。具体而言,在我们的集群中,跨节点 GPU 通过 IB 全互联,而节点内通信通过 NVLink 处理。NVLink 提供 160 GB/s 的带宽,约为 IB(50 GB/s)的 3.2 倍。为了有效利用 IB 和 NVLink 的不同带宽,我们将每个词元限制为最多分发到 4 个节点,从而减少 IB 流量。对于每个词元,当做出其路由决策后,它将首先通过 IB 传输到其目标节点上具有相同节点内索引的 GPU。一旦到达目标节点,我们将努力确保它通过 NVLink 被即时转发到托管其目标专家的特定 GPU,而不会受到随后到达的词元的阻塞。通过这种方式,通过 IB 和 NVLink 的通信被完全重叠,并且每个词元可以在不产生额外 NVLink 开销的情况下,平均每节点高效地选择 3.2 个专家。这意味着,尽管 DeepSeek-V3 在实践中仅选择了 8 个路由专家,但它可以将这个数字扩展到最多 13 个专家(4 个节点 $\times$ 每节点 3.2 个专家),同时保持相同的通信成本。总体而言,在这种通信策略下,仅需 20 个 SM 就足以充分利用 IB 和 NVLink 的带宽。

具体而言,我们采用 warp 专门化技术(Bauer et al., 2014),并将 20 个 SM 划分为 10 个通信通道。在分发过程中,(1)IB 发送、(2)IB 到 NVLink 的转发、以及(3)NVLink 接收由各自的 warp 处理。分配给每个通信任务的 warp 数量会根据所有 SM 上的实际工作负载动态调整。类似地,在组合过程中,(1)NVLink 发送、(2)NVLink 到 IB 的转发与累加、以及(3)IB 接收与累加也由动态调整的 warp 处理。此外,分发和组合内核都与计算流重叠,因此我们也考虑了它们对其他 SM 计算内核的影响。具体来说,我们采用了定制的 PTX(并行线程执行)指令,并自动调优通信块大小,这显著减少了 L2 缓存的使用以及对其他 SM 的干扰。

3.2.3 极小的开销实现极致的内存节省

为了减少训练期间的内存占用,我们采用了以下技术。

RMSNorm 与 MLA 上投影的重新计算。

我们在反向传播过程中重新计算所有的 RMSNorm 操作和 MLA 上投影,从而无需持久存储它们的输出激活。以微小的开销为代价,该策略显著降低了存储激活所需的内存需求。

CPU 上的指数移动平均。

在训练过程中,我们保留模型参数的指数移动平均(EMA),用于在学习率衰减后提前估计模型性能。EMA 参数存储在 CPU 内存中,并在每个训练步骤后异步更新。这种方法使我们能够维护 EMA 参数,而不会产生额外的内存或时间开销。

多词元预测的共享嵌入与输出头。

借助 DualPipe 策略,我们将模型最浅层(包括嵌入层)和 deepest 层(包括输出头)部署在同一个 PP rank 上。这种安排使得共享嵌入和输出头的参数与梯度能够在 MTP 模块和主模型之间实现物理共享。这种物理共享机制进一步提升了我们的内存效率。

3.3 FP8 训练

图

图 6:采用 FP8 数据格式的整体混合精度框架。为清晰起见,仅展示了 Linear 算子。

受近期低精度训练进展(Peng et al., 2023b; Dettmers et al., 2022; Noune et al., 2022)的启发,我们提出了一种细粒度混合精度框架,利用 FP8 数据格式来训练 DeepSeek-V3。虽然低精度训练前景广阔,但它常常受到激活、权重和梯度中离群点存在的限制(Sun et al., 2024; He et al., ; Fishman et al., 2024)。尽管推理量化(Xiao et al., 2023; Frantar et al., 2022)已取得显著进展,但展示低精度技术在大规模语言模型预训练中的成功应用的研究相对较少(Fishman et al., 2024)。为了应对这一挑战并有效扩展 FP8 格式的动态范围,我们引入了细粒度量化策略:以 $1\times N_{c}$ 个元素为单位的分块(tile-wise)分组,或以 $N_{c}\times N_{c}$ 个元素为单位的块状(block-wise)分组。相关的反量化开销在我们的高精度累加过程下得到了大幅缓解,这是实现精确 FP8 通用矩阵乘法(GEMM)的关键方面。此外,为了进一步降低 MoE 训练中的内存和通信开销,我们以 FP8 格式缓存和分发激活,同时以 BF16 格式存储低精度优化器状态。我们在两个与 DeepSeek-V2-Lite 和 DeepSeek-V2 规模相当的模型上验证了所提出的 FP8 混合精度框架,训练约 1 万亿词元(更多细节见附录 B.1)。值得注意的是,与 BF16 基线相比,我们的 FP8 训练模型的相对损失误差始终保持在 0.25% 以下,这一水平完全处于训练随机性的可接受范围内。

3.3.1 混合精度框架

基于低精度训练中广泛采用的技术(Kalamkar et al., 2019; Narang et al., 2017),我们提出了一种用于 FP8 训练的混合精度框架。在该框架中,大多数计算密集型操作以 FP8 执行,而少数关键操作则被策略性地保留在其原始数据格式中,以平衡训练效率和数值稳定性。整体框架如图 6 所示。

首先,为了加速模型训练,大多数核心计算内核(即 GEMM 操作)以 FP8 精度实现。这些 GEMM 操作接受 FP8 张量作为输入,并以 BF16 或 FP32 产生输出。如图 6 所示,与 Linear 算子相关的全部三个 GEMM,即 Fprop(前向传播)、Dgrad(激活后向传播)和 Wgrad(权重后向传播),都以 FP8 执行。与原始 BF16 方法相比,这种设计在理论上将计算速度提升了一倍。此外,FP8 Wgrad GEMM 允许以 FP8 格式存储激活,用于后向传播。这显著降低了内存消耗。

尽管 FP8 格式具有效率优势,某些算子由于其对于低精度计算的敏感性,仍需要更高的精度。此外,一些低成本算子也可以利用更高精度,而对整体训练成本的开销可以忽略不计。因此,经过仔细研究,我们对以下组件保留了原始精度(例如 BF16 或 FP32):嵌入模块、输出头、MoE 门控模块、归一化算子以及注意力算子。这些有针对性的高精度保留确保了 DeepSeek-V3 稳定的训练动态。为了进一步保证数值稳定性,我们将主权重、权重梯度和优化器状态存储在更高精度中。虽然这些高精度组件会带来一些内存开销,但其影响可以通过我们分布式训练系统中跨多个 DP rank 的高效分片来最小化。

图

图 7:(a) 我们提出了一种细粒度量化方法,以缓解由特征离群点引起的量化误差;为简化图示,仅展示了 Fprop。(b) 结合我们的量化策略,我们通过在以 $N_{C}=128$ 个元素的间隔将 MMA 提升(promote)到 CUDA Cores 进行高精度累加,从而提升了 FP8 GEMM 的精度。

3.3.2 通过量化与乘法提升精度

基于我们的 FP8 混合精度框架,我们引入了若干策略来提升低精度训练的准确性,重点同时放在量化方法和乘法过程上。

细粒度量化。

在低精度训练框架中,由于 FP8 格式受限的动态范围(受其减少的指数位数所限),溢出和下溢是常见的挑战。作为标准做法,输入分布通过将输入张量的最大绝对值缩放到 FP8 的最大可表示值来对齐到 FP8 格式的可表示范围(Narang et al., 2017)。这种方法使得低精度训练对激活离群点高度敏感,这会严重降低量化精度。为了解决这个问题,我们提出了一种细粒度量化方法,在更细粒度级别上应用缩放。如图 7 (a) 所示,(1) 对于激活,我们以 1x128 的分块(即每个词元每 128 个通道)为基础进行元素分组和缩放;(2) 对于权重,我们以 128x128 的块(即每 128 个输入通道每 128 个输出通道)为基础进行元素分组和缩放。这种方法确保量化过程能够通过根据更小的元素组调整缩放比例来更好地适应离群点。在附录 B.2 中,我们进一步讨论了当我们以与权重量化相同的方式按块对激活进行分组和缩放时出现的训练不稳定性。

我们方法中的一个关键修改是引入了沿 GEMM 操作内维的每组的缩放因子。这一功能在标准 FP8 GEMM 中并不直接支持。然而,结合我们精确的 FP32 累加策略,它可以高效地实现。

值得注意的是,我们的细粒度量化策略与微缩放(microscaling)格式的思想高度一致(Rouhani et al., 2023b),而 NVIDIA 下一代 GPU(Blackwell 系列)已宣布支持具有更小量化粒度的微缩放格式(NVIDIA, 2024a)。我们希望我们的设计能够为未来的工作提供参考,使其与最新的 GPU 架构保持同步。

提升累加精度。

低精度 GEMM 操作经常受到下溢问题的困扰,其精度在很大程度上依赖于高精度累加,这通常以 FP32 精度执行(Kalamkar et al., 2019; Narang et al., 2017)。然而,我们观察到 NVIDIA H800 GPU 上 FP8 GEMM 的累加精度仅能保留约 14 位,远低于 FP32 累加精度。当内维 K 较大时(Wortsman et al., 2023),这个问题会变得更加明显,这是大规模模型训练中批量大小和模型宽度都增大的典型场景。以两个 K = 4096 的随机矩阵的 GEMM 操作为例,在我们的初步测试中,Tensor Cores 中受限的累加精度导致的最大相对误差接近 2%。尽管存在这些问题,受限的累加精度仍是少数 FP8 框架中的默认选项(NVIDIA, 2024b),严重制约了训练精度。

为了解决这个问题,我们采用了提升到 CUDA Cores 以获得更高精度的策略(Thakkar et al., 2023)。该过程如图 7 (b) 所示。具体而言,在 Tensor Cores 上执行 MMA(矩阵乘累加)时,中间结果使用受限位宽进行累加。一旦达到 $N_{C}$ 的间隔,这些部分结果将被复制到 CUDA Cores 上的 FP32 寄存器中,在那里执行全精度 FP32 累加。如前所述,我们的细粒度量化沿内维 K 应用每组的缩放因子。这些缩放因子可以在 CUDA Cores 上作为反量化过程高效地相乘,而额外的计算成本极小。

值得注意的是,这一修改降低了单个 warpgroup 的 WGMMA(Warpgroup 级矩阵乘累加)指令发射速率。然而,在 H800 架构上,通常两个 WGMMA 会并发持续存在:当一个 warpgroup 执行提升操作时,另一个能够执行 MMA 操作。这种设计使两个操作能够重叠,保持 Tensor Cores 的高利用率。根据我们的实验,设置 $N_{C}=128$ 个元素(相当于 4 个 WGMMAs)是能够显著提升精度而不引入大量开销的最小累加间隔。

尾数优先于指数。

与先前工作(NVIDIA, 2024b; Peng et al., 2023b; Sun et al., 2019b)采用的混合 FP8 格式不同——后者在 Fprop 中使用 E4M3(4 位指数和 3 位尾数),在 Dgrad 和 Wgrad 中使用 E5M2(5 位指数和 2 位尾数)——我们在所有张量上采用 E4M3 格式以获得更高精度。我们将这种方法的可行性归因于我们的细粒度量化策略,即分块和块状缩放。通过在更小的元素组上操作,我们的方法有效地在这些分组元素之间共享指数位,从而减轻了受限动态范围的影响。

在线量化。

张量级量化框架(NVIDIA, 2024b; Peng et al., 2023b)采用了延迟量化,它维护先前迭代中最大绝对值的历史以推断当前值。为了确保精确的缩放比例并简化框架,我们对每个 1x128 的激活分块或 128x128 的权重块在线计算最大绝对值。基于此,我们推导出缩放因子,然后将激活或权重在线上量化为 FP8 格式。

3.3.3 低精度存储与通信

结合我们的 FP8 训练框架,我们通过将缓存的激活和优化器状态压缩为更低精度的格式,进一步降低了内存消耗和通信开销。

低精度优化器状态。

我们采用 BF16 数据格式而非 FP32 来跟踪 AdamW(Loshchilov and Hutter, 2017)优化器中的一阶和二阶矩,而不会引起可观察到的性能下降。然而,主权重(由优化器存储)和梯度(用于批量大小累积)仍保留在 FP32 中,以确保整个训练过程中的数值稳定性。

低精度激活。

如图 6 所示,Wgrad 操作以 FP8 执行。为了减少内存消耗,缓存 FP8 格式的激活以用于 Linear 算子的后向传播是自然的选择。然而,对于若干算子,我们采取了特殊考虑以实现低成本的高精度训练:

(1) 注意力算子之后的 Linear 的输入。这些激活也用于注意力算子的后向传播,这使其对精度敏感。我们专门为这些激活采用了定制的 E5M6 数据格式。此外,这些激活在后向传播中会从一个 1x128 的量化分块转换为 128x1 的分块。为了避免引入额外的量化误差,所有缩放因子都采用取整缩放,即 2 的整数次幂。

(2) MoE 中 SwiGLU 算子的输入。为了进一步降低内存成本,我们缓存 SwiGLU 算子的输入,并在后向传播中重新计算其输出。这些激活也以我们细粒度量化方法的 FP8 格式存储,在内存效率和计算精度之间取得了平衡。

低精度通信。

通信带宽是 MoE 模型训练中的关键瓶颈。为了缓解这一挑战,我们将 MoE 上投影之前的激活量化为 FP8,然后应用分发(dispatch)组件,这与 MoE 上投影中的 FP8 Fprop 兼容。与注意力算子之后 Linear 的输入类似,该激活的缩放因子也是 2 的整数次幂。类似的策略也应用于 MoE 下投影之前的激活梯度。对于前向和反向的组合(combine)组件,我们将其保留在 BF16 中,以在训练流程的关键部分保持训练精度。

3.4 推理与部署

我们将 DeepSeek-V3 部署在 H800 集群上,其中节点内的 GPU 通过 NVLink 互联,集群内所有 GPU 通过 IB 全互联。为了同时确保在线服务的服务级别目标(SLO)和高吞吐率,我们采用以下将预填充和解码阶段分离的部署策略。

3.4.1 预填充

预填充阶段的最小部署单元由 4 个节点共 32 块 GPU 组成。注意力部分采用 4 路张量并行(TP4)配合序列并行(SP),并结合 8 路数据并行(DP8)。其较小的 TP 规模 4 限制了 TP 通信的开销。对于 MoE 部分,我们使用 32 路专家并行(EP32),这确保每个专家处理足够大的批量,从而提升计算效率。对于 MoE 全互联通信,我们采用与训练相同的方法:首先通过 IB 在节点间传输词元,然后通过 NVLink 在节点内 GPU 之间转发。特别地,我们对浅层的稠密 MLP 使用 1 路张量并行以节省 TP 通信。

为了实现 MoE 部分不同专家之间的负载均衡,我们需要确保每个 GPU 处理大约相同数量的词元。为此,我们引入了一种冗余专家部署策略,对高负载专家进行复制并冗余部署。高负载专家是根据在线部署期间收集的统计信息检测出来的,并定期(例如每 10 分钟)进行调整。在确定冗余专家集合后,我们根据观察到的负载,在节点内 GPU 之间小心地重新安排专家,力求在不增加跨节点全互联通信开销的情况下尽可能平衡 GPU 之间的负载。对于 DeepSeek-V3 的部署,我们为预填充阶段设置了 32 个冗余专家。对于每个 GPU,除了其原本托管的 8 个专家外,它还将额外托管一个冗余专家。

此外,在预填充阶段,为了提高吞吐率并隐藏全互联和 TP 通信的开销,我们同时处理两个计算负载相似微批次,将一个微批次的注意力与 MoE 与另一个微批次的分发和组合重叠。

最后,我们正在探索一种动态的专家冗余策略,其中每个 GPU 托管更多专家(例如 16 个专家),但在每个推理步骤中仅激活 9 个。在每一层全互联操作开始之前,我们即时计算全局最优的路由方案。鉴于预填充阶段涉及的大量计算,计算该路由方案的开销几乎可以忽略不计。

3.4.2 解码

在解码过程中,我们将共享专家视为一个路由专家。从这个角度看,每个词元在路由过程中将选择 9 个专家,其中共享专家被视为一个始终会被选中的高负载专家。解码阶段的最小部署单元由 40 个节点共 320 块 GPU 组成。注意力部分采用 TP4 配合 SP,并结合 DP80,而 MoE 部分使用 EP320。对于 MoE 部分,每个 GPU 仅托管一个专家,64 块 GPU 负责托管冗余专家和共享专家。分发和组合部分的全互联通信通过 IB 上的直接点对点传输执行,以实现低延迟。此外,我们利用 IBGDA(NVIDIA, 2022)技术进一步最小化延迟并提升通信效率。

与预填充类似,我们根据在线服务的统计专家负载,在一定时间间隔内定期确定冗余专家集合。然而,由于每个 GPU 仅托管一个专家,我们无需重新安排专家。我们也在探索解码的动态冗余策略。但这需要对计算全局最优路由方案的算法以及与分发内核的融合进行更仔细的优化,以减少开销。

此外,为了提升吞吐率并隐藏全互联通信的开销,我们也在探索在解码阶段同时处理两个计算负载相似的微批次。与预填充不同,注意力在解码阶段消耗了更大比例的时间。因此,我们将一个微批次的注意力与另一个微批次的分发+MoE+组合重叠。在解码阶段,每个专家的批量大小相对较小(通常在 256 个词元以内),瓶颈在于内存访问而非计算。由于 MoE 部分只需加载一个专家的参数,内存访问开销最小,因此使用较少的 SM 不会显著影响整体性能。因此,为了避免影响注意力部分的计算速度,我们可以仅分配一小部分 SM 给分发+MoE+组合。

3.5 关于硬件设计的建议

基于我们对全互联通信和 FP8 训练方案的实现,我们向 AI 硬件厂商提出以下关于芯片设计的建议。

3.5.1 通信硬件

在 DeepSeek-V3 中,我们实现了计算与通信之间的重叠,以隐藏计算过程中的通信延迟。与串行计算和通信相比,这显著降低了对通信带宽的依赖。然而,当前的通信实现依赖于昂贵的 SM(例如,我们从 H800 GPU 可用的 132 个 SM 中分配了 20 个用于此目的),这将限制计算吞吐率。此外,使用 SM 进行通信会导致显著的效率低下,因为张量核心完全未被利用。

目前,SM 主要为全互联通信执行以下任务:

  • •

在 IB(InfiniBand)和 NVLink 域之间转发数据,同时将从单个 GPU 发往同一节点内多个 GPU 的 IB 流量聚合起来。

  • •

在 RDMA 缓冲区(已注册的 GPU 内存区域)与输入/输出缓冲区之间传输数据。

  • •

为全互联组合(combine)执行归约(reduce)操作。

  • •

在通过 IB 和 NVLink 域向多个专家传输分块数据的过程中,管理细粒度的内存布局。

我们希望看到未来的厂商开发出能够将这些通信任务从宝贵的计算单元 SM 上卸载下来的硬件,充当类似于 NVIDIA SHARP(Graham et al., 2016)的 GPU 协处理器或网络协处理器。此外,为了降低应用程序编程的复杂性,我们希望这种硬件从计算单元的角度统一 IB(横向扩展)和 NVLink(纵向扩展)网络。借助这种统一的接口,计算单元可以基于简单原语提交通信请求,轻松地在整个 IB-NVLink 统一域内完成读、写、多播和归约等操作。

3.5.2 计算硬件

Tensor Cores 中更高的 FP8 GEMM 累加精度。

在 NVIDIA Hopper 架构当前的 Tensor Core 实现中,FP8 GEMM(通用矩阵乘法)采用定点累加,在加法之前基于最大指数右移对齐尾数乘积。我们的实验表明,在符号填充右移之后,它仅使用每个尾数乘积的最高 14 位,并截断超出此范围的位。然而,例如,要从 32 个 FP8$\times$FP8 乘法的累加中获得精确的 FP32 结果,至少需要 34 位精度。因此,我们建议未来的芯片设计提高 Tensor Cores 中的累加精度,以支持全精度累加,或者根据训练和推理算法的精度要求选择合适的累加位宽。这种方法确保误差保持在可接受的范围内,同时保持计算效率。

对分块和块状量化的支持。

当前 GPU 仅支持按张量的量化,缺乏对像我们的分块和块状量化这样的细粒度量化的原生支持。在当前实现中,当达到 $N_{C}$ 间隔时,部分结果将从 Tensor Cores 复制到 CUDA cores,乘以缩放因子,并加到 CUDA cores 上的 FP32 寄存器中。尽管结合我们精确的 FP32 累加策略,反量化开销得到了显著缓解,但 Tensor Cores 和 CUDA cores 之间频繁的数据移动仍然限制了计算效率。因此,我们建议未来的芯片通过使 Tensor Cores 能够接收缩放因子并实现带组缩放的 MMA 来支持细粒度量化。这样,整个部分和累加和反量化都可以直接在 Tensor Cores 内部完成,直到产生最终结果,避免了频繁的数据移动。

对在线量化的支持。

尽管我们的研究已经证明了在线量化的有效性,当前的实现却难以有效支持它。在现有流程中,我们需要从 HBM(高带宽内存)读取 128 个 BF16 激活值(前一次计算的输出)进行量化,然后将量化后的 FP8 值写回 HBM,仅仅是为了再次读取以进行 MMA。为了解决这种低效,我们建议未来的芯片将 FP8 转换和 TMA(张量内存加速器)访问整合为单一融合操作,以便量化可以在激活从全局内存传输到共享内存的过程中完成,避免频繁的内存读写。我们还建议支持 warp 级转换指令以加速,这进一步促进了层归一化与 FP8 转换的更好融合。或者,可以采用近内存计算(near-memory computing)方法,将计算逻辑放置在 HBM 附近。在这种情况下,BF16 元素在从 HBM 读入 GPU 时可以直接转换为 FP8,将片外内存访问减少约 50%。

对转置 GEMM 操作的支持。

当前的架构使得将矩阵转置与 GEMM 操作融合变得繁琐。在我们的工作流程中,前向传播过程中的激活被量化为 1x128 的 FP8 分块并存储。在后向传播过程中,需要对矩阵进行读出、反量化、转置、重新量化为 128x1 的分块,并存储在 HBM 中。为了减少内存操作,我们建议未来的芯片在 MMA 操作之前,支持从共享内存直接转置读取矩阵,用于那些在训练和推理中都需要用到的精度。结合 FP8 格式转换与 TMA 访问的融合,这一增强将显著简化量化工作流程。