CUDA Profiling and Debugging Tools for LLM

Min Xu (徐敏), NVIDIA GPU应用研发, 资深工程师

LLM 开发中的常见挑战
Megatron Core 的性能最佳实践
性能优化工具
Nsight Systems 多报告视图与“掉队者”分析
Nsight Systems 多报告分析与 Recipes
NVIDIA Resiliency Extension: 慢节点检测
内存优化工具
PyTorch Memory Profiler (PyTorch 内存分析器
CUPTI (NVIDIA CUDA Profiling Tools Interface
性能数据收集方法
CUPTI API 概览
CUPTI Activity API
CUPTI Callback API
案例研究：构建自定义内存分析器
PyTorch Profiler
总结
未来工作展望

LLM 开发中的常见挑战

大型语言模型（LLM）的开发过程中通常面临三大挑战：配置选择、次优性能和内存不足（OOM）。

1. 配置选择 (Config Selection)

问题: 如何为特定用例选择最佳的性能配置。
解决方案: 遵循 Megatron Core 的性能最佳实践。

2. 次优性能 (Suboptimal Performance)

问题:
- 需要分析特定或所有GPU的性能信息。
- 需要探索定制化的分析工具以满足特定需求。
- 某些GPU运行较慢（成为“掉队者”），需要找出原因。
解决方案:
- 使用 Nsight Systems 和 Nsight Profile 进行性能分析。
- 使用 DCGM (Data Center GPU Manager) 进行定制化分析。
- 利用工具分析“掉队者”现象。

3. 内存不足 (Out-of-Memory, OOM)

问题: 需要可视化分析内存分配随时间变化的情况。
解决方案: 使用 PyTorch Memory Profiler 等工具。

Megatron Core 的性能最佳实践

为了解决配置选择的挑战，NVIDIA 提供了针对 Megatron Core 的性能最佳实践。

工作流

优化工作流旨在通过迭代式的评估和分析，找到最优配置。该流程始于模型和硬件参数，通过选择配置、性能评估和分析，并借助分析工具和内存估算，最终得到经过验证的优化配置。

配置建议

并行策略配置

类别	关键建议
模型并行 (Model Parallelism)	`intra_layer_model_parallel_size`
混合精度训练 (Mid-Autocast Training)	`dtype: bf16` 或 `fp16`；`bf16_O2=true` 或 `fp16_O2=true`
流水线并行 (Pipeline Parallelism)	启用该功能；流水线并行度（PP）应低于张量并行度（TP）或数据并行度（DP）。
专家并行 (Expert Parallelism)	针对 MoE 模型启用。
张量并行 (Tensor Parallelism)	TP ≥ 节点间网络带宽 / (2 * 单卡 HBM 带宽)
虚拟流水线并行 (Virtual PP)	VPP 值应为 `num_layers` 的公约数；当流水线并行度（PP）较小时使用。
上下文并行 (Context Parallelism)	序列长度（Sequence length） ≥ 4k

其他配置

类别	建议/选项
CUDA Graph	`capture_cudagraph_iters >= 10`；`cuda-graph-syn` (cudagraphs)
通信重叠 (Communication Overlap)	`async_p2p_allgather`；`overlap_p2p_comm`；`overlap_grad_reduce`；`overlap_param_gather`
分组GEMM (Grouped GEMM)	`use_grouped_gemm`
即时编译权重融合 (Fused JIT weights)	`jit_fusion`

内存估算

为避免内存不足（OOM）问题，可以使用 Megatron-LM 提供的内存计算器脚本来估算不同配置下的内存占用。通过在工具中启用或禁用特定模块，可以预估其对内存的影响，从而辅助配置选择。

脚本路径: https://github.com/NVIDIA/Megatron-LM/blob/main/tools/memory_calculator.py

性能优化工具

Nsight Systems 多报告视图与“掉队者”分析

NVIDIA Nsight Systems 是一款系统级的性能分析工具。其多报告视图功能对于分布式训练场景的性能调试至关重要。该功能允许用户在单个视图中同时打开和分析多个节点的报告（例如，一个分布式任务中所有GPU的报告），从而方便地进行跨节点性能对比和问题定位。

Page 14: Nsight Systems 多报告视图界面，可在一个窗口中查看8个报告

应用案例：定位分布式训练中的“掉队者”（Straggler）

在分布式训练中，某些节点（rank）的性能可能低于其他节点，成为系统瓶颈。Nsight Systems 的多报告视图可以有效地帮助分析此类问题。

问题现象: 通过对比不同 rank 的时间线，可以发现某个 rank 的计算或通信耗时明显长于其他 rank。
根本原因分析: 深入分析“掉队者” rank 的详细时间线，可以定位到具体的瓶颈。如下图所示，右侧视图揭示了原因是其中一个 rank 的 NCCL 通信流量缓慢，阻塞了其他 rank 的执行。Nsight Systems 能够详细展示分布式训练中此类“掉队者”现象的成因。

Page 15: 使用Nsight Systems多报告视图分析“掉队者”rank

Nsight Systems 多报告分析与 Recipes

Nsight Systems 提供了基于 Python 的脚本（称为 "Recipes"）用于自定义报告分析。用户可以运行这些脚本来从 .nsys-rep 文件中提取和可视化特定信息。

Nsight Systems Multi Report Analysis 列表 (Page 16)

上表展示了部分可用的 Recipes 及其功能，按类型分类：

Summary (摘要):
- cuda_api_sum: CUDA API 摘要
- cuda_gpu_kern_sum: CUDA GPU Kernel 摘要（按调用）
- cuda_gpu_mem_sum: CUDA GPU 内存操作摘要（按调用）
- nvtx_sum: NVTX 事件摘要
- osrt_api_sum: 操作系统运行时 API 摘要
Pace (步速):
- cuda_gpu_kern_pace: CUDA GPU Kernel 执行步速
- nvtx_pace: NVTX 执行步速
Heatmap (热力图):
- cuda_gpu_kern_sum_heatmap: CUDA GPU Kernel 摘要热力图
- os_mem_util_heatmap: OS 内存利用率热力图
Track (轨迹):
- nvtx_gpu_proj_track: NVTX GPU 投影轨迹
- nvtx_gpu_proj_sum: NVTX GPU 投影摘要
SanityCheck (健全性检查):
- cuda_sync_summary: CUDA 同步策略摘要
- cuda_api_stat: CUDA API 统计

将 Jupyter Notebooks 作为数据呈现工具

用户可以在 Jupyter Lab 环境中运行这些 Recipes，以便进行交互式分析和数据可视化。

示例代码:

# Set python path if nsys-recipe is not installed
# import sys; sys.path.insert(0, &quot;/path/to/nsight-systems/target-linux-x64/python/packages&quot;)

# Import module
from nsys_recipe.recipes import run_recipe

# Run recipe
df = run_recipe(
    'nvtx_sum',
    [nsys_rep_path],
    report_file_name_suffix=None
)

下图对比了在 Nsight Systems 原生界面和在 Jupyter Lab 中呈现分析结果的差异。

Nsight Systems 与 Jupyter Lab 对比 (Page 17)

Recipe 示例

摘要 (Summary)

通过运行 cuda_api_sum 和 nvtx_sum Recipes，可以生成关于 CUDA API 调用和 NVTX 事件的统计摘要。输出包括按名称排序的 CUDA Kernel 调用数据表格，以及所有 Kernels 的调用频率柱状图。

> nsys-recipe run --name cuda_api_sum,nvtx_sum <nsys-rep-file>

步速 (Pace)

Pace Recipe 用于分析任务随时间变化的执行情况。例如，nvtx_gpu_proj_trace 可以追踪 NVTX 范围在 GPU 上的投影，帮助评估工作负载的执行节奏和性能波动。

自定义 Recipe

当预置的 Recipes 无法满足特定的分析需求时，可以创建自定义 Recipe。例如：
* 现有 Recipe 功能不足。
* 图表不够清晰。
* 需要查看原始数据。
* 指标过多，但洞察不足。

自定义 Recipe 可以生成更具针对性的可视化图表，例如下图所示的饼图，以提供更直接的性能洞察。

NVIDIA Resiliency Extension: 慢节点检测

这是一个专注于分布式训练中慢节点（straggler）检测和处理的功能。

NVRx-Straggler Detection (慢节点检测)

目标与优势
- 识别表现不佳的节点：识别出执行缓慢的节点，确保单个表现不佳的节点不会显著拖慢整个分布式训练的进程。
- 最大化吞吐量：通过确保并可能终止表现不佳的节点，来保障整体的吞吐效率，从而维持各个周期的性能一致性。

重要性
- 可扩展性 (Scalability)：随着任务（rank）数量的增加，维持可扩展性变得至关重要，慢节点检测在其中扮演了关键角色。
- 可靠性 (Reliability)：通过确保进程不会因为表现不佳的任务而失败或被拖慢，增强了分布式系统的可靠性。

下图展示了在一次使用80个节点的 Lama 6B 模型训练中，吞吐量（Throughput）随训练步数（Step Number）的变化。在大约第30步时，吞吐量出现了急剧下降，这很可能是由于一个或多个慢节点造成的。

NVRx-Straggler Detection 与 PyTorch 的集成

检测流程
1. 比较分数 (Compare Scores)：计算相对和个体的性能分数。
2. 识别慢节点 (Identify Stragglers)：检查是否有节点的性能分数超出了指定的阈值。

性能分数类型
- 相对分数 (Relative)：将当前批次（rank）作为参考基准。
- 历史分数 (Historical)：使用当前批次的历史最佳性能作为参考。

下图展示了将慢节点检测逻辑集成到 PyTorch 训练代码中的示例。通过在训练循环中插入检查点，可以实时监控并处理慢节点问题。

内存优化工具

PyTorch Memory Profiler (PyTorch 内存分析器)

可视化显存分配与时间关系

PyTorch 提供了一系列 API 用于内存分析：
- PyTorch API:
- torch.cuda.memory_summary(device=None, abbreviated=False): 返回一个关于当前设备上所有张量内存使用情况的可读摘要。
- torch.cuda.memory_snapshot(): 返回一个内存分配器的快照。
- torch.cuda.memory_history(device=None): 返回一个内存分配/释放事件的历史记录。

与 Megatron 集成:
- --memory-profiler: 在 Megatron 框架中启用此参数，可以记录每个张量的内存历史。

使用内存分析器解决 OOM (Out of Memory) 问题

以下是一个使用内存分析器调试 CUDA graph 导致的 OOM 错误的案例：
1. 问题报告：有报告称，在使用 CUDA graph 时，aten::cat 操作会导致 OOM 错误，并且有大量的内存被保留（reserved）。
2. 缩小模型规模：为了在使用 CUDA graph 时不触发 OOM，首先将模型规模缩小。
3. 捕获内存使用情况：使用 record_memory_history 来捕获第一个迭代（iteration）的内存使用情况，并与不使用 CUDA graph 的情况进行对比。

代码示例：

if self.memory_profiler and self.iteration == 0:
    torch.cuda.memory._record_memory_history(
        max_entries=100000
    )

下图展示了在不使用 CUDA graph 和使用 CUDA graph 两种情况下，模型训练初期的内存占用情况。可以观察到，使用 CUDA graph 的情况下（右图），内存占用呈现出一种奇特的、不均衡的锯齿状模式。

进一步分析发现，CUDA graph 会为每个微批次（micro-batch）隐式地预分配一个输入/输出缓冲区，这导致了额外的内存开销。下图将有问题的内存模式（OOM-B）与正常情况（OOM-A）进行了对比，清晰地揭示了由于预分配缓冲区而产生的锯齿状内存增长模式，最终导致了 OOM。

Page 39: OOM 问题分析：CUDA graph 预分配缓冲区导致内存额外开销

CUPTI (NVIDIA CUDA Profiling Tools Interface)

本节探讨使用 CUPTI 创建自定义性能分析工具。

性能数据收集方法

性能数据的收集主要有三种方法：注入 (Instrumentation)、采样 (Sampling) 和混合模式 (Hybrid)。

注入 (Instrumentation / Event Tracing)
- 记录每一个被关注的事件，提供精确的时间和顺序信息。
- 对于高频事件，可能会引入较大的开销。
- 实现方式:
  - 源码级包装器和宏: 手动调用 CUDA/OSRT API。
  - CUPTI API: 官方的事件追踪和性能分析接口，需要专业的领域知识和固件支持。
  - LD_PRELOAD 拦截: 覆盖 CUDA 运行时库函数，无需修改或重新编译源码，但对静态链接的二进制文件无效，且通常需要 GPU 刷写（flush）。
采样 (Sampling / Statistical Profiling)
- 周期性地记录程序状态，而非记录每个事件。
- 优点: 开销低，即使在高采样率下对性能影响也较小；能自动过滤掉持续时间很短的事件，聚焦于主要瓶颈。
- 缺点: 精度较低，可能会遗漏某些事件；不适用于需要精确事件顺序的分析，更适合统计分析。
混合模式 (Tracing + Profiling)
- 结合注入和采样。典型用法是为 GPU 内核等轻量级函数进行注入，同时对 CPU 活动进行采样。
- 优点: 在控制总体开销的同时，能够精确地计时关键事件。

CUPTI API 概览

CUPTI 提供了丰富的 API 用于追踪和性能分析。

分类	CUPTI API	功能描述	模块
Tracing	Activity	异步记录 CPU/GPU 活动（如 CUDA API、内核、内存操作）。	`cupti_activity.h`
	Callback	向 CUPTI 注册回调函数，以在特定事件（如 API 调用、驱动事件）发生时接收通知。	`cupti_callbacks.h`
Profiling	Event/Metric	查询可用的性能事件和指标，并收集其值。	`cupti_events.h`
	Range Profiling	在一系列操作范围内收集性能指标的值。	`cupti_range_profiler.h`
	PC Sampling	对 GPU 上的程序计数器 (PC) 进行采样，以识别代码中的热点。	`cupti_pcsampling.h`
	Serialization	收集硬件状态信息，通过对 GPU 性能计数器执行一次读取操作来实现。	`cupti_serialization.h`
	Chip	提供对特定于芯片的性能分析功能的访问。	`cupti_chip.h`
	Overhead	提供对自动化性能分析开销和监控 CUDA 设备功能状态的支持。	`cupti_overhead.h`

CUPTI Activity API

CUPTI Activity API 用于异步收集 CPU 和 GPU 的活动记录。
* 客户端提供空缓冲区，CUPTI 线程将活动记录填充到缓冲区中。记录处理的顺序不保证。
* 使用 cuptiActivityEnable 可以启用特定类型的活动记录，并将主机 API 调用与内核活动记录关联起来。

缓冲管理

缓冲管理通过回调函数实现，每个线程独立管理。
buffer_requested 和 buffer_completed 回调用于优化多线程环境下的缓冲。
当缓冲区满或手动调用 flush 时，缓冲区将被递交处理。

CUPTI Callback API

CUPTI Callback API 允许用户订阅特定事件并注册回调函数。
* 订阅 (Subscription): 客户端订阅一个回调函数集合。
* 回调函数 (Callback Function): 用户定义的函数，在事件发生时被调用。
* 回调域 (Callback Domain): 定义回调的范围，如 Runtime、Driver、JIT/Linker、Library 等。
* 回调 ID (Callback ID): 标识特定的回调点。

案例研究：构建自定义内存分析器

本案例展示如何使用 CUPTI 构建一个自定义内存分析器。

用法: 通过 LD_PRELOAD 钩子 malloc 和 free 来追踪主机内存分配。
使用 CUPTI Callbacks 来追踪所有的 CUDA 运行时内存分配和释放 API。
使用 LD_PRELOAD 和 __attribute__((constructor)) 在 main 函数执行前初始化 CUPTI。
生成按峰值内存使用量排序的报告，并显示对应的 NVTX 范围。

分析器输出

分析器运行后，可以通过 Python 脚本对收集到的数据进行后处理，并生成详细的内存使用报告。报告中包含每个线程的内存操作、大小、设备以及关联的 NVTX 信息。最终的摘要表显示了 GPU 和 CPU 的峰值内存使用情况。

PyTorch Profiler

PyTorch Profiler 是一款集成在 PyTorch 中的性能分析工具，它底层利用了 CUPTI 和 NVTX。它能够提供详细的算子（operator）级别性能数据，包括 CPU 和 CUDA 上的执行时间。

在 DGX-A100 上进行 MobileNet-v2 训练时，启用 PyTorch Profiler 会带来约 7.9% 的性能下降。

使用 PyTorch Profiler 需要 import torch.cuda.nvtx as nvtx，并使用 nvtx.range_push/pop 或 torch.profiler.profile 上下文管理器来标记代码区域。

代码示例:

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(
        wait=1,
        warmup=1,
        active=2),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./result', worker_name='worker0'),
    record_shapes=True,
    with_stack=True
) as prof:
    for step, batch_data in enumerate(data_loader):
        if step >= (1 + 1 + 2) * 2:
            break
        train_step(batch_data)
        prof.step()

下图展示了使用 PyTorch Profiler 后在 Nsight Systems 中看到的时间轴视图。

通过可视化工具，开发者可以直观地分析模型训练或推理过程中的性能瓶颈。下图展示了 PyTorch Profiler 的一个可视化输出示例。左侧是类似于火焰图的时间轴视图，显示了不同操作的嵌套调用和执行时长。右侧则提供了聚合的算子性能统计表，详细列出了每个算子的名称、CPU 自我执行时间（Self CPU）、CPU 总时间（CPU total）、CUDA 自我执行时间（Self CUDA）和 CUDA 总时间（CUDA total）等关键指标。

Page 31: PyTorch Profiler 的可视化界面和性能统计表示例

总结

本节对所介绍的各类工具进行了总结和比较。

工具 (Tool)	真实用例 (Real Use Case)	优点与比较 (Pros & Cons Comparison)
NVIDIA Resiliency Extension	在性能分析之前，检查是否存在拖慢训练的慢节点。	易于使用，只需在代码中编辑几行，或现有 Megatron 选项。可应用于整个训练和推理生命周期，并评估其对性能的影响。
Nsight Systems Multi-Report Analysis	比较和聚合多节点分析报告。	适用于多节点训练/推理，可进行回归测试跟踪。支持 GPU 核心利用率关联。
Nsight Systems	选择单个 GPU/CPU 配置文件。	提供 GPU 核心利用率和详细分析功能。允许按单元格显示模型信息。支持 Python 模型装饰器。
PyTorch Profiler	模型/性能和代码跟踪以及内存分析。	显示节点位置。显示形状和模型信息（需要展开才能查看）。需要确认是主节点报告。支持通过点击查看算子。
DLProf	针对旧版 TensorFlow 的集合式查看器。	低开销的事件跟踪。需要额外开发。无模型上下文。

未来工作展望

本节讨论了未来希望在用户友好工具方面进行的功能改进。

特性 (Feature)	描述 (Description)
链接组件级指标 (Link Component-level metrics)	将 Nsys、cuBLAS/CUTLASS 内部指标和图表链接到 NVML。专注于模型可执行性、理论与实际 FLOPs。
时间轴上的模型信息 (Model information on Timeline)	时间轴显示特定于 CUDA 事件的模型信息，例如张量形状、数据类型、核函数等。
多节点通信关联 (Multi Rank Communication Correlation)	将跨多个计算节点的通信/传输链接在一个视图中，作为所有相关的操作。
稳健的内存分析 (Robust Memory Profiling)	防止文件损坏，完全验证与 OOM 之前保存的数据相同的保证。
官方支持与维护 (Official Support & Maintenance)	由 NVIDIA 进行专业审查和维护，提供向后兼容性和长期维护。