AI训练加速优化技术_aten::matmul cpu耗时很长-优快云博客

1. 背景

随着深度学习模型参数量和计算量越来越大，对算力的要求越来越高，GPU算力也越来越高，如何用好GPU变得越来越重要，一个简单的GPU的训练流程如下图所示：

模型训练中，面对的主要问题可以分为四类：IO、计算、显存、通信，而大模型训练推理中凸显的显存需求高、算力需求大、以及越来越多的数据通信的问题，都需要相应的加速方案来提高模型的训练和推理速度。

本文全文近7000字，整理耗时近一周，如果对你有帮助欢迎点赞收藏，或扫码支持作者更新～

2. 训练加速

2.1 IO显存优化

在模型训练中，IO 操作常常成为性能瓶颈。传统的数据传输方式，如从 CPU 内存经 PCIe 总线传输到 GPU 显存，速度受限，频繁的数据交换会导致 I/O 阻塞。

为解决这一问题，可以采用数据预取技术优化IO显存，这是一种在计算机系统中用于优化数据访问性能的策略，这种策略主要应用于 IO 数据优化环节，旨在解决 CPU 和 GPU 在数据读取与训练迭代过程中的效率问题。

在未采用数据预取技术时，流程中存在 “空闲” 阶段，这意味着 CPU 或 GPU 在等待数据读取完成，造成了计算资源的浪费。而数据预取技术通过提前将数据从存储设备（如硬盘）读取到缓存中，使得 CPU 或 GPU 在需要数据时能够立即获取，避免了等待时间，从而提高了整体的处理效率。在图片右侧的优化后流程中，不再有 “空闲” 阶段，蓝色的 “读取数据” 方块和绿色的 “训练迭代” 方块紧密衔接，展示了数据预取技术如何让数据读取和训练迭代能够并行进行。

此外，还可采用 Pinned Memory（固定内存）与 Zero Copy 技术。Pinned Memory 可使数据直接传输给 GPU，减少额外拷贝时间；Zero Copy 则让 GPU 直接访问 CPU 的 Pinned Memory，适用于小批量数据传输，显著降低 PCIe 传输开销。

对于大模型权重加载，由于模型权重数据量巨大，单个 GPU 显存难以容纳，可采用权重分片（Sharded Weights）与按需加载（Lazy Loading）策略。将大模型参数拆分成多个小块分布到多个 GPU，每个 GPU 仅加载其负责部分，且只在推理需要时加载相应权重，提前缓存 KV Cache，减少内存占用与 I/O 瓶颈，加快推理启动速度。此外，预加载（Prefetching）与 Streaming 加载技术也可有效缓解磁盘 I/O 慢的问题。在后台预加载即将用到的数据，边加载边计算，减少 GPU 等待时间，提升吞吐量，尤其适用于长序列推理。

拓展阅读：模型训练和计算过程中的一些优化 - zhangkele - 博客园

2.2 计算相关优化

算子优化

在模型训练实践当中，往往会使用融合算子（Operator Fusion）优化效率。简单说，就是把原本需要多个独立算子依次执行的计算逻辑，合并成一个融合后的算子来执行。

我们以下图里的注意力（Attention）计算为例：

在融合前，计算注意力要拆分调用 aten::matmul（矩阵乘法）、aten::softmax（softmax 归一化）等多个零散算子，每个算子都要单独做 “读取输入→计算→写回结果” 的流程，中间会有显存读写、调度等额外开销。

在融合后，用 aten::scaled_dot_product_attention（SDPA，缩放点积注意力）这类融合算子，把注意力计算涉及的矩阵乘、归一化、缩放等逻辑 “打包”，让这些计算在一个算子里完成。

深度学习训练时，算子间的数据传递要频繁读写显存，还涉及计算资源调度。比如融合前，matmul 算完把结果写回显存，softmax 再从显存读数据，这会产生明显的时间、性能损耗。
融合后，多个计算逻辑在一个算子里完成，中间结果不用频繁写回显存，直接在算子内部流转，大幅减少显存访问次数，也省了多次调度算子的耗时。

拓展阅读：

编译优化技术

编译优化技术可在代码编译阶段对深度学习模型进行优化。通过对计算图的分析和优化，将连续的算子操作合并为单一内核（Kernel）执行，减少显存搬运次数，提升计算效率。

以本图为例，我们写的 torch 代码（比如模型前向传播、自定义算子逻辑），是编译优化的 “原材料”，包含各种算子（aten::abs/aten::pow 等）和计算逻辑。

那么代码是如何一步步转化成为单一内核的执行逻辑呢？我们一步步看——

dynamo 与 IR 转换：拆解与标准化
- dynamo：PyTorch 的编译前端组件，负责 “捕获”torch 代码的计算图，把 Python 动态执行逻辑，转成更易分析的 中间表示（IR，Intermediate Representation） 。
- Fx IR → Inductor IR：IR 是一种 “中间语言”，把 PyTorch 算子拆成标准化的计算节点。先转成 Fx IR（PyTorch 专属 IR ），再进一步优化、转成 Inductor IR（更贴近硬件执行的 IR ）。这一步会做 算子融合、冗余计算消除 等初步优化。
Graph lowering + CodeGen：生成硬件可执行代码
- Graph lowering：把高层的 IR 计算图，“下推” 成更接近硬件指令的底层表示，适配 GPU 等加速硬件的执行模型。
- CodeGen（代码生成）：根据 lower 后的结果，生成 Triton kernel 代码（Triton 是专为 GPU 设计的高性能编程语言，能高效控制 GPU 线程、内存）。
编译与调用：生成的 Triton kernel 会被编译成 GPU 可直接执行的二进制指令，最终通过 Module Call 接入训练流程，替代原来的零散算子调用。

混合精度

混合精度（Mixed Precision），可能是最常见的优化方法之一，简单说就是在训练过程中，同时用不同精度的数值（如 FP16、FP32 ）做计算 。从下图里能看到，流程里既有 FP16（半精度浮点）的权重、梯度，也有 FP32（单精度浮点）的转换和计算，核心是 “精度混合搭配”，而非只用一种精度硬算。

由于FP16 数据占的显存，只有 FP32 的一半（FP32 是 4 字节，FP16 是 2 字节）。所以显存省了，就能塞更多数据进 GPU（batch_size 翻倍），训练效率更高。

此外，由于GPU 硬件对 FP16 计算的支持更 “友好”，算得比 FP32 快（底层硬件电路设计，让半精度浮点运算流水线更高效）。所以训练时，模型里的矩阵乘、激活函数等计算，用 FP16 跑更快，整体训练速度就被加速。

拓展阅读：

https://zhuanlan.zhihu.com/p/84219777

2.3. 显存优化

重计算技术：

重计算技术（也常被称为梯度检查点技术），是深度学习训练中为平衡显存占用和计算效率设计的优化手段。

我们结合下面这张图理解：

我们可以看到，训练包含 “前向计算（-> 方向）” 和 “反向计算（<- 方向）” 两个阶段。前向会生成中间结果（图里的 “节点” ），反向要用到这些中间结果算梯度。

而重计算的核心逻辑是：前向计算时，只保留部分关键节点（图里的三个紫色 checkpoint 节点）的中间结果，其他节点的中间结果不保存 。等反向计算需要这些 “没保存的节点” 时，重新跑一遍前向计算来生成它们 ，以此换显存的节省。

由于深度学习训练中，模型的中间结果（前向计算产生的张量）超占显存。尤其大模型，中间结果能把显存撑爆。所以重计算通过 “选择性丢弃中间结果”，只存少量 checkpoint 节点，让显存占用直接砍半甚至更多。

此外，尽管重新计算前向会多花点计算时间，但相比 “显存不够训不了模型”，这点计算开销是值得的。相当于用 “少量计算资源”，换 “显存资源的释放”，让训练能继续。

拓展阅读：

4.1. 前向重计算 — PaddleFleetX 0.1.0.beta documentation

2.4. 通信优化

首先我们看看通信的耗时公式：

$t = \alpha \cdot P + \beta \cdot \frac{M}{B}$

其中的变量分别为：

t：总通信耗时（比如分布式训练中，节点间同步梯度的时间）。
α：通信延迟系数（和硬件、网络拓扑有关，比如网卡握手、建立连接的固定耗时）。
P：通信次数（每次通信都要先经历 α∗P 的延迟）。
β：带宽相关系数（和网络带宽上限、硬件传输效率有关）。
M：总通信数据量（要传输的梯度、参数等数据的总大小）。
B：实际网络带宽（单位时间能传的数据量，受硬件和网络负载影响）

通信耗时由两部分组成：

固定延迟α∗P：每次通信都有 “建立连接、协议握手” 等固定开销，通信次数 P 越多，这部分耗时越大。
传输耗时 $\beta*\frac{M}{B}$ *：和总数据量 M 成正比，和实际带宽 B 成反比 —— 数据越多、带宽越低，传输越慢。

融合通信：

主要从减少固定延迟的角度减少通信耗时，简单说就是 把多个 “数据小包” 合并成 “数据大包” 再通信 。

如上图所示，反向计算会产生多个零散的梯度（“梯度小包” ），原本要多次通信传输。但是融合通信先把这些梯度 “合并” 成一个大的梯度包，再进行一次（或少数几次）通信传输。

这其中有两个超参数需要注意：

融合数据大小的选择：要平衡 “大包减少次数” 和 “大包可能超带宽、反而变慢” 的矛盾，需选最优大小。
梯度 tiling 策略：把大梯度拆成合理的块（tiling ），再融合通信，进一步适配硬件和网络，最大化优化效果。

分层通信：

分层通信是依据硬件拓扑和通信需求，将分布式训练中的通信过程拆分为多个层级的优化策略。

结合上图来看，它把原本 “大一统” 的通信，拆成机器内、机器间等不同层次，每一层用适配的通信算子（Allreduce、Reducescatter 等）和策略，让通信更贴合硬件能力。

那么它是如何提升训练效率的呢？

利用机器内通信并行性：
机器内的 GPU 间，带宽高（比如图里 Reducescatter 能到 300GB/s ）、延迟低。分层通信先在机器内并行通信（同一机器的多个 GPU 同时传数据），把机器内的高带宽用满，减少整体耗时。
降低机器间通信拓扑复杂度：
机器间网络带宽通常比机器内低（图里机器间 Allreduce 是 3GB/s ）。分层通信通过 “先聚合机器内数据，再少量机器间交互”，大幅减少机器间通信量。比如图里 16 台 8 卡机器，通信拓扑从 256 降到 48 ，机器间传输压力直接减轻。
GDR 技术加速 GPU 直连：
传统通信要经 “GPU→CPU→网络”，GDR（GPU Direct RDMA ）让 GPU 跳过 CPU，直接通过网络卡通信，省去 GPU - CPU 数据拷贝的耗时 ，尤其大模型训练时，能省超多时间。

1-bit 通信技术

把梯度等数据压缩成 1 比特（0 或 1 ）传输。比如原本用 32 位浮点存梯度，现在用 1 位表示（比如符号位），数据量直接砍到 1/32 。

这能够极大减少通信数据量，适合机器间带宽紧张的场景。但会损失精度，需配合 “梯度恢复” 策略（比如传符号 + 缩放因子），尽量减少对训练的影响。

TopK 通信技术

只传输梯度中 “最大的 K 个值”（TopK 梯度），其他小梯度暂时不传（或用近似值代替）。因为训练中，大梯度对参数更新影响更大，小梯度的 “存在感” 低。

这种方法能减少通信数据量（只传 K 个大值，而非全部梯度），同时保留对训练影响关键的信息。适合大模型训练，能在加速通信的同时，尽量保证训练效果。

2.5 加速收敛

这是一类通过优化 “参数更新策略”，让模型训练更快找到最优解的算法。

典型代表是 LARS（Layer-wise Adaptive Rate Scaling） 和 LAMB（Layer-wise Adaptive Momentum Batch） ：

由伪代码可见，它们针对神经网络的 “层” 做自适应优化（Layer-wise ），不是整个模型用同一套学习率。通过更智能的 “梯度调整、学习率缩放”，让模型参数更新更高效，从而加速收敛。

3. 大模型技术

大模型训练加速，面临的两个主要问题是算力和显存问题

对于算力问题，通过模型并行的方法将算力需求分摊到更多的GPU资源上；对于显存问题通过ZerO的方法以及cpu offload策略大幅度减少了显存的需求。

3.1. 张量并行

以图片里的 MLP（多层感知机）为例，原本一个 GPU 要完成全部矩阵乘、激活函数等计算，

$Y = \text{GeLU}(XA), \quad Z = \text{Dropout}(YB).$

现在把矩阵参数（如 A、B ）拆分到多个 GPU，每个 GPU 算一部分。

比如按列拆分第一个矩阵乘的weight值为［A1，A2］，则Y可以表达为：

$[Y_1, Y_2] = \left[ \text{GeLU}(XA_1), \text{GeLU}(XA_2) \right]$

然后按行拆分第二个矩阵乘的B矩阵：

$B = \begin{bmatrix} B_1 \\ B_2 \end{bmatrix}, \quad Y = [Y_1, Y_2]$

最后汇总结果：

3.2. Pipeline并行

Pipeline 并行（Pipeline Parallelism）是大模型训练中优化多 GPU 协同效率的核心策略，本质是将模型的层（如 Transformer 层）拆分到多个 GPU，让不同 GPU 依次处理模型的不同阶段，像流水线一样并行执行前向和反向传播 。

恰如上图里的 Transformer 层拆分：每个 Pipeline 分区（Pipeline MP partition）包含若干层，不同 GPU 负责不同分区，前一个 GPU 处理完一层，把中间结果传给下一个 GPU 继续处理，形成流水线。

但是Pipeline 并行中会出现一个问题：由于前向（Forward）和反向（Backward）传播的依赖关系（反向必须等前向完成），会出现 “设备空闲（Bubble）” —— 部分 GPU 要等前一个 GPU 传数据，导致算力浪费。

Bubble 耗时公式：

$t_{pb} = (p - 1) \cdot (t_f + t_b)$

其中：

p 是 Pipeline 阶段数（拆分的 GPU 数量）；
$t_f$ 是单阶段前向传播耗时；
$t_b$ 是单阶段反向传播耗时。

实际计算耗时公式：

$t_{id} = m \cdot (t_f + t_b)$

其中 m 是训练的 micro-batch 数量。

也即：实际有效计算的总耗时，由 m 个 batch 的前向 + 反向组成。

所以pipeline采用多阶段分配（Assign multiple stages to each device）解决问题。其核心逻辑是让每个 GPU 负责多个 Pipeline 阶段，减少阶段间的空闲。

比如上面图片里，原本每个 GPU 只负责 1 个阶段，优化后负责多个阶段，让前向和反向的 “等待” 重叠，压缩 Bubble 耗时。

此时的耗时公式为：

$t_{pb}^{\text{int.}} = \frac{(p - 1) \cdot (t_f + t_b)}{v}$

此时的耗时比为：

$\frac{t_{pb}^{\text{int.}}}{t_{id}} = \frac{1}{v} \cdot \frac{p - 1}{m}$

拓展阅读：

Pipeline并行bubble优化-进阶调优-大模型算法优化-性能调优-大模型训练迁移和优化-CANN商用版7.0.0开发文档-昇腾社区

3.3. ZerO优化技术

Zero（ZeRO：Memory Optimizations Toward Training Trillion Parameter Models ）是大模型训练中针对显存优化的核心技术，全称 “零冗余优化器” 。其本质是 对模型训练时的显存占用进行 “精细化拆分” ，把原本在单个 GPU 上存储的 “参数、梯度、优化器状态”，按照不同策略分散到多个 GPU 上，从而大幅减少单 GPU 的显存压力，支持训练更大模型。

3.4. Flash attention

FlashAttention 是大模型训练中针对注意力（Attention）计算效率优化的核心技术，本质是 通过 “内存层次优化 + 数学公式拆解”，减少注意力计算中对高延迟显存（HBM）的访问，从而加速训练 。

传统注意力计算（如 PyTorch 原生实现）会频繁读写 GPU 高带宽内存（HBM），导致大量时间浪费在 “数据搬运” 上。FlashAttention 则利用 GPU 内存层级（SRAM→HBM→DRAM） 的带宽差异，结合数学上的 softmax 可拆分性质，让计算更贴近高速内存（SRAM），减少 HBM 访问次数。