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的报告），从而方便地进行跨节点性能对比和问题定位。

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

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

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

Nsight Systems 多报告分析与 Recipes

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

上表展示了部分可用的 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, "/path/to/nsight-systems/target-linux-x64/python/packages")

# 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 中呈现分析结果的差异。

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。

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）等关键指标。

总结

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

工具 (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 进行专业审查和维护，提供向后兼容性和长期维护。