PyTorch 性能与显存优化手册

前一阵子在得道 APP 上读完了这本《PyTorch 性能与显存优化手册》，感觉是一本很不错的 PyTorch Training 入门读物，很适合刚刚接触这个领域的新手小白来读；同时整本书也提供了一个 PyTorch Training 的优化大纲，可以作为一个引子来扩展去学习更加底层的知识和技术。

这篇帖子主要是我读的时候标记下来的一些知识点，并没有包含书籍里面的全部内容。强烈推荐大家去阅读原书，会非常有收获！

基础性概念

• 在实际的应用中，我们一般会先优化显存然后再优化训练速度
• 在数据加载的过程中，可以通过将加载任务和模型计算容易进行重叠来进行优化，比如预加载技术；数据预处理也可以采用离线预处理技术或者优化 CPU 预处理代码的效率
• 随机读写模式的效率以硬盘的 IOPS 来衡量；而连续读写模式则是以 MB/s 来进行衡量
• 在 GPU 中，张量计算核心，标量计算核心决定了 GPU 的整体计算效率; L1 缓存，显存决定了 GPU 存取数据的效率；线程束调度器和 CUDA 编程模型中的线程束概念直接对应，负责线程间的通信和调度
• CUDA 编程模型的核心是要求程序员将算法的实现代码，拆分成多个可以独立执行的软件任务

PyTorch Tensor

• 在 PyTorch 中，张量和张量的存储是两个不同层级的概念，张量可以理解为对其底层数据存储的一种特定的访问和解释方式，诸如 shape, stride, offset 等额外属性来帮助对底层数据的访问
• 张量的不连续性可能在实际应用中造成一些问题，比如可能会导致算法的内存访问模式不理想，可能降低整体的计算效率；其次许多 PyTorch 算法在设计时就预设了张量在内存中是连续存储的，如果遇到不连续的张量，可能会抛出错误提示甚至得到错误的结果
• 基于高级索引的读取操作会尝尽新的内存存储，比如使用布尔或者整数张量作为索引
• 算子调用的开销恰恰是最主要的性能杀手之一
• 在 PyTorch 中，计算图是一个有向无环图，其中的节点代表各种算子操作，比如加法，乘法或者更复杂的操作入卷积等；而边则代表数据（指张量数据）的流动，这些边的方向描述了数据流行的路径和操作的执行顺序
• 在运算过程中对于同一个张量的梯度是累加而不是覆写的，而且这里累加的是具体数值而非符号表达式，这一点至关重要，它使得自动微分系统能够自然地兼容程序中出现的逻辑判断，如分支，循环和递归等，而这对符号微分系统是非常困哪的

常用优化流程

• 对于绝大多数场景，可以通过下面的标准流程来查找性能瓶颈
  • 观察 GPU 队列，如果 GPU 队列整体都非常稀疏，那么性能瓶颈在 CPU 上
  • 观察 GPU 队列，如果任务队列密集，而没有显著空白区域，说明 GPU 满载，那么性能瓶颈在 GPU 算子
  • 观察 GPU 队列，如果任务队列密集，同时存在 GPU 空闲区域，则需要放大空闲区域进行进一步观察
  • 观察 GPU 空闲区域，查看空闲前后 GPU 任务以及 CPU 任务详情，并以此推断导致 GPU 队列阻塞的原因
• 具体到训练过程来说，一般只有在优化 CUDA 算子时才会考虑使用 Nsight Compute 用于定位算子内部的性能瓶颈

数据读取

• 对于其他无直接接口的库如 Pandas 等，建议先转换为 Numpy ndarray 再导入到 PyTorch
• 简单来讲，Dataset 描述了读取单个数据的方法以及必要的预处理，输出的是单个张量; Dataloader 则定义了批量读取数据的方法，包括 batchsize, 预读取，多进程读取等，输出的结果是一批张量
• 读者可能会担心非阻塞模式会导致 GPU 数据错误，但 GPU 内部有自己的任务队列 CUDA stream 系统；在没有特别指定 CUDA 计算流的情况下，所有任务默认进入同一个队列，并且会按照任务提交的顺序串行执行
• 我们可以采取类似 CPU 预加载数据的策略，在当前训练轮次进行的同时，预先把下一轮训练所需的数据从 CPU 复制到 GPU 上。要做到这一点，我们需要通过配置不同的 GPU 计算流 (CUDA stream) 来创建一个并行的数据拷贝任务。为了实现数据传输与 GPU 计算的并行，我们将数据传输任务和模型计算任务交替提交到两个不同的 GPU 队列中，为了能够正确的更新参数，我们还要保证两个 GPU 队列的重叠部分仅限于数据传输，而计算部分不发生重叠，也也是为什么引入了 sumbit_stream.wait_stream(running_stream) 来进行 GPU 队列见间的同步和等待。对于推理部署熟悉的朋友可能会注意到，这一技巧与推理中常用的双重缓冲 (double buffering) 优化有些类似，这个方法主要加快数据的拷贝速度，因此在数据量较小，数据传输用时短的场景中，效果可能不明显

优化器

• 本质上来说，BatchSize 是通过增加计算并行度的方式来提高算子计算效率
• PyTorch 针对每种优化器，提供了三种不同的梯度更新实现，for-loop, for-each, fused，对于性能和显存的侧重不同

内存优化

• PyTorch 使用显存池机制来进行内存管理，不会只申请张量所需要的大小的缓存，而是一次性的申请一块儿更大的显存空间，本质上是对分配出来的若看显存进行二次管理，会有显存碎片化的问题，实用的解决方法是设置 max_split_size_mb
• PyTorch 的分布式函数库在实现的时候，调用了 Nvidia 提供的集合通信库 NCCL 来完成 GPU 节点间的通信，而 NCCL 库会自己管理进程间通信所占用的显存
• 由于动态分配的显存冗余程度打，优化这部分显存对于模型训练收敛性能影响风险较低，是优化的首选目标。同时，训练过程中显存占用的峰值通常书现在反向传播过程的某个反向算子的计算中，因此当遇到内存溢出问题而导致模型无法训练时，动态分配的显存是首先要排查的点
• 首先，通过自动微分生成的反向算子通常不是原位的，这是因为在反向传播过程中使用原位操作容易导致数值错误
• 如果受限于硬件资源，无法达到理想的 BatchSize，可以通过牺牲一些训练速度来增加 BatchSize，即使用跨批次梯度累加（cross-batch gradient accumulation）。这种方法的核心是降低优化器梯度更新的频率——通过累积多轮训练的梯度，最后再一起更新，从而实现增大 BatchSize 的效果
• 为了节省宝贵的 GPU 显存，一种常用的策略是将数据默认存放在内存里，而只在有需要时才临时加载到显存中。这本质上是牺牲了内存和性能来换取显存
• 循环垃圾回收机制实现了检测循环依赖的算法，能够打破循环依赖对引用计数造成的破坏；除了循环依赖以外，全局作用域中的张量也是导致显存资源不能被释放的重要原因

分布式训练

• 因此分布式训练的效率优化也主要围绕着如何尽可能降低和隐藏通信的时间开销这个核心思路展开；数据并行策略主要通信梯度张量，模型并行则会根据策略不同对模型参数、梯度、优化器状态和激活张量的通信都有可能涉及
• DDP 自动实现了两个优化
  • 分组传输: 为了减少每个参数独立进行 allreduce 操作所带来的通信开销，可以利用分组传输（Bucketing）技术。该技术自动将模型中的所有算子参数分成几个组，每组的参数梯度被合并成一个较大的张量后再执行通信。这样，每组内的梯度计算完成后只需执行一次通信操作，从而大幅降低了通信次数，提高了训练效率
  • 重叠计算和通信: 梯度计算较早完成的组会优先启动通信操作，这一过程与后续层的梯度计算重叠，从而大幅减少了由通信引起的延迟
• 其他的可以考虑的优化
  • 降低通信量: 通过梯度压缩技术减少传输数据量，例如量化、稀疏化或低秩近似，从而降低通信成本
• 拓扑感知策略: 根据计算节点的网络拓扑结构设计更高效的通信操作，优化数据传输过程
• 随着节点数量的增加，节点间的通信量和通信次数通常也会相应增加，这导致通信延迟逐渐增大; 其次，随着节点数量的增加，系统中任何不稳定的组件的负面影响也会相应放大; 除了通信开销以外，分布式系统还会带来额外的显存占用，这与通信部分的实现细节有关。例如，分组传输技术本质上就是一种用显存来换取性能的优化方法，它通过将小的通信请求合并，减少通信次数，但也就不可避免地需要占用额外的显存来存储这些合并后的数据
• 这种方法涉及使用 PyTorch 的 DistributedDataParallel 模块中提供的 register_comm_hook() 接口，通过这个接口，我们可以将 DDP 中默认的节点通信函数替换为自定义函数，从而获取更精确的通信开销数据
• 将静态显存分块存储在不同的GPU上，本质上是将持续占用显存的静态数据转变为动态的“按需分配”，从而有效降低单GPU显存占用的峰值
• 动态显存切分的关键在于将模型的不同部分分配到多个 GPU 卡上，每张卡负责处理模型的一部分，这样静态显存和动态显存都会被切分，这类方法被称为模型并行
• 虽然流水线并行的方法能有效地分割模型计算，但它也有很大的局限性，尤其是当模型中显存占用最大的层无法在单个GPU上运行时。例如，在大型语言模型中，一些大规模的矩阵乘法操作可能会超出单个GPU的显存容量，这时流水线并行就无法解决问题。在这种情况下，我们需要采用张量并行（tensor parallel）策略。这种策略主要是通过矩阵乘法的分块计算，实现单卡无法容纳的大型矩阵乘法操作。比如要实现下矩阵乘法 $Y=X×W$ ，在参数矩阵 $W$ 非常大甚至超出单张卡的显存容量时，我们可以把它在特定维度上切分到多张卡上，并通过all-gather集合通信汇集结果，保证最终结果在数学计算上等价于单卡计算结果
• 与流水线并行能在多台机器协同训练不同，目前张量并行的应用范围存在较大局限性。这主要是因为张量并行需要传输大量数据，当这种传输需要通过网络设备跨机器进行时，受限的网络带宽会严重阻碍张量并行训练的效率。因此，张量并行通常只在配备了NVLink的单机多卡范围中使用

混合精度训练

• 数值范围主要由指数部分决定, 而精度主要由尾数的位数决定
• FP16 训练所带来的问题
  • 首先，训练初期数值波动往往较大，这容易导致使用 float16 时发生数据溢出，从而产生 NaN 或 Inf 等问题。需要采取措施来处理不同训练轮次间数值范围的差异
  • 其次，训练中一个普遍的问题是前向传播中的张量与反向传播中的梯度在数值范围上有显著差异——前向张量的数值通常较大，而反向梯度的数值较小。在更换为 float16 后会加剧数值溢出的风险，因此我们需要平衡前向张量和反向梯度的数值范围
  • 最后，使用 float16 本质上是在性能和精度之间进行取舍，不同的算子对数值精度的需求不一，因此受益于 float16 加速的程度也会有所不同。我们希望能够自动判断哪些算子适合使用 float16 加速，并自动对这些算子的输入输出张量进行类型转换
• PyTorch 提供的核心 API: torch.autocast 和 torch.cuda.amp.GradScaler

图优化

• torch.compile 的常用参数
  • 启用 fullgraph 模式获取完整的计算图
  • 支持动态形状输入的编译 (dynamic=True)
  • 调整编译和执行模式, 用户可通过激活r educe-overhead 模式来启用 CUDA graph 功能
• 基于Python运行时的跟踪（tracing）方法本质上是在模型执行过程中动态捕捉计算图，也就是通过监视 PyTorch 操作的执行，来实时记录这些操作及其上下游之间的依赖关系
• 基于源码分析（source code analysis）的方法本质上是通过分析模型的源代码结构来构建计算图。这种方法不需要实际运行模型，而是直接解析代码中的静态结构
• 通过CPython提供了内部接口，torch.compile 和 Torch Pynamo 技术在 Python 运行时捕获 PyTorch 的张量操作，并将这些操作转化成计算图。这个动态生成的计算图随后可以被进一步优化，并用于生成更高效的执行代码，这些代码在执行时将取代原来Python解释器中的函数调用

其他

• accelerator的基本思路与DDP一样，但是增加了更多的优化而且用户接口非常友好，详细使用方法可以参考其官方文档，或者直接参考当前示例的代码
• 尽管本书详细讨论了许多性能优化的策略，但有一个非常重要的技巧尚未提及，那就是“始终从小处开始（start small）”这个原则