第13章. PyTorch的性能分析 、调优与扩展
本作品已使用人工智能进行翻译。欢迎您提供反馈和意见:translation-feedback@oreilly.com
AI训练与推理管道可能在每个层级遭遇性能瓶颈,包括Python解释器开销、CPU主机端数据加载阻塞、CUDA内核利用率不足以及GPU设备内存争用等问题。要实现有效优化,需借助覆盖整个系统的多种工具,在堆栈的多层级进行性能分析。
本章重点探讨在现代NVIDIA GPU上运行PyTorch工作负载时的性能分析、调试及系统级优化.我们将探讨如何利用PyTorch内置分析器、NVIDIA Nsight工具、Linuxperf的CPU分析功能——以及PyTorch内存分析与内存分配器调优来识别并解决瓶颈。同时阐述PyTorch如何运用CUDA流实现并发处理,以及如何通过CUDA图结构降低内核启动开销。
接下来,我们将演示如何通过PyTorch分布式数据并行(DDP)、全分片数据并行(FSDP)及其他模型并行策略优化数据管道并扩展至多GPU环境。随后将展示如何对多GPU和多节点环境进行性能分析,包括整体跟踪分析(HTA)和Perfetto工具的应用。
本章将重点阐述性能权衡关系,并通过量化示例聚焦内核执行时长、硬件利用率指标、内存占用、数据加载效率及整体扩展成本效益。完成本章学习后,您将掌握如何在整个技术栈中实施高效、全面的PyTorch工作负载性能分析与调优方案。
NVTX马克笔与性能分析工具
要全面掌握性能状况,必须在多层级进行分析并使用覆盖整个系统的工具。实践者和性能工程师普遍采用一套通用工具和最佳实践,用于对系统堆栈所有层级进行整体分析。
在介绍工具前,需重点说明NVIDIA工具扩展(NVTX)及其马克笔机制。这些马克笔用于标注分析器时间轴视图中的时间区间,使不同分析器能关联相同阶段的事件。
例如,"forward" 的NVTX时间段既会出现在PyTorch分析器追踪记录中,也会显示在Nsight Systems的时间轴上。这极大简化了栈中不同层级的跨工具分析。NVTX马克笔得到多数现代AI框架和库的支持,包括PyTorch及所有与CUDA生态相关的工具。
NVTX马克笔可通过CUDA C++、PyTorch或任何支持NVIDIA GPU的C++/Python库(如OpenAI Triton、PyCUDA、CuPy、cuTile、cuTe、CUTLASS等)注入代码。多数库已为关键代码区域自动注入标记,例如:"train_step" 、"forward" 、"backward" 、"optimizer_step" 等。您也可通过PyTorch中的torch.profiler.record_function() 和torch.cuda.nvtx.range_push() 自行注入标记。
在阐述了如何使用 NVTX标记注释代码关键段落后,接下来将探讨能采集、对齐并可视化这些标记的工具。表13-1汇总了常用性能分析工具的适用范围、核心功能及典型应用场景,可帮助您在优化流程的各个阶段选择合适工具。
| 工具 | 范围 | 功能 | 典型用例 |
|---|---|---|---|
| PyTorch 性能分析器(Kineto) | PyTorch内部运算级性能分析(CPU/GPU) | 支持NVTX马克笔、形状记录、内存统计、跟踪导出、编译图中断识别 | 模型代码精细化分解;识别低效操作、GPU内核启动开销或正向/反向时间失衡 |
Nsight Systems(nsys )
|
全系统时间线(CPU、GPU、操作系统、I/O) ... |
Become an O’Reilly member and get unlimited access to this title plus top books and audiobooks from O’Reilly and nearly 200 top publishers, thousands of courses curated by job role, 150+ live events each month,
and much more.
Read now
Unlock full access