CUDA Profiling and Debugging Tools for LLM

Min Yu (余愁), NVIDIA GPU加速计算专家团队, 高级工程师

目录

1. 议程 (Agenda
2. LLM 开发中的常见挑战
3. Megatron Core 性能最佳实践
4. 性能优化工具
5. CUPTI (NVIDIA CUDA Profiling Tools Interface
6. PyTorch Profiler
7. NVIDIA Resiliency Extension Straggler Detection (NVIDIA 弹性扩展：落后者检测
8. 内存优化工具 (Tools for Memory Optimization
9. 总结 (Summary
10. 未来工作讨论 (Future Work Discussion

1. 议程 (Agenda)

• 大型语言模型（LLM）开发中的常见挑战
• 性能优化工具
• 内存优化工具

2. LLM 开发中的常见挑战

大型语言模型（LLM）开发过程中通常面临三大挑战：

1. 配置选择 (CONFIG SELECTION)
2. 次优性能 (SUBOPTIMAL PERFORMANCE)
3. 内存不足 (OUT-OF-MEMORY, OOM)

2.1. 配置选择 (Configuration Selection)

面临的核心问题是如何为性能选择最佳配置。推荐参考 Megatron Core 的性能最佳实践来解决此问题。

2.2. 性能问题 (Performance Problems)

针对次优性能问题，可以采用以下工具进行分析与解决：

• 分析特定或所有 GPU 的性能信息:

  • Nsight System
  • PyTorch Profiler

• 探索满足特定需求的定制化分析工具:

  • CUPTI

• 解决部分 GPU 运行较慢的问题（掉队者检测）:

  • NVRx--Straggler Detection

2.3. 内存不足问题 (OOM Issues)

为了解决内存不足（OOM）问题，关键在于可视化分析所有内存分配随时间变化的情况。推荐使用以下工具：

• PyTorch Memory Profiler

3. Megatron Core 性能最佳实践

本节将介绍 Megatron Core 的性能最佳实践。

3.1. 工作流程 (Workflow)

优化 Megatron Core 性能的工作流程如下：

1. 输入: 从模型超参数和硬件规格开始。
2. 配置选择: 选择初始配置。

3. 迭代优化循环:

   • 性能评估与分析: 使用分析工具（Profiling Tools）评估当前配置的性能。
   • 内存估算: 估算内存使用情况，如果需要则启用重计算（Recompute）。
   • 配置优化: 根据分析结果调整和优化配置。

4. 输出: 获得最优或可接受的配置。

3.2. 选择并行配置 (Select Parallelism Configurations)

以下是针对不同并行策略的关键建议：

• 模型并行 (Model Parallelism): 保持并行规模最小化。
• MoE 并行折叠 (MoE Parallel Folding): 注意 tp-cp-dp-pp 和 tp-ep-dp-pp 的配置，MoE 的 EPxTP 组是 DPxCPxTP 的子集。
• 流水线并行 (Pipeline Parallelism): 优先于 EP 和 TP。
• 专家并行 (Expert Parallelism): 优先于 TP。
• 张量并行 (Tensor Parallelism): 对于密集部分，TP 应小于等于 NVLink 域的大小；对于 MoE 部分，TP = 1。
• 虚拟流水线 (Virtual PP): vpp_size 应该是 num_layers/pp_size 的公约数，中等值通常较好。
• 上下文并行 (Context Parallelism): 序列长度应大于等于 8k。

3.3. 选择其他配置 (Select Other Configurations)

以下是其他重要配置的推荐选项：

• CUDA Graph:

  • --external-cuda-graph (via TE)
  • --cuda-graph-scope [submodule1, submodule2]

• 通信重叠 (Communication Overlap):

  • DP 通信: --overlap-param-gather --overlap-grad-reduce
  • TP 通信: --tp-comm-overlap
  • EP 通信: --moe-token-dispatcher-type=flex --moe-enable-deepep

• 分组 GEMM (Grouped GEMM): --moe-grouped-gemm

• 原生 FP8 权重 (Native FP8 Weights): --fp8-param-gather

3.4. 内存估算 (Estimate Memory)

可以使用内存估算器（Memory Estimator）来预测内存使用情况。
* 工具地址: https://github.com/SEEKWAN/mbridge/tree/main/memory_estimator
* 如果发生 OOM，可以选择性地启用重计算（recompute）来释放内存，通过以下参数指定需要重计算的模块：
--recompute-granularity selective --recompute-modules [submodule1, submodule2, ...]

4. 性能优化工具

本节介绍用于性能优化的具体工具，重点介绍 Nsight Systems。

4.1. Nsight Systems 多报告视图 (Multi Report)

假设用户已熟悉 Nsight Systems 和 NVTX。Nsight Systems 提供了多报告视图功能，可以将多个分析报告（例如，来自不同 rank 的报告）在同一个视图中打开，便于对比分析。

4.2. 多报告视图的局限性

当前的 nsys multi-report view 存在跨 rank 关联性不足的问题。

• nsys multi-report view: 可以在同一 rank 内关联 CPU 和 GPU 的活动。此关联基于 CUPTI 的 correlationId。
• XPU-Timer: 能够链接所有 rank 中相同的计算或通信操作。它基于 Perfetto 的跨道事件关联和时间线对齐机制，利用事件名称、开始和结束时间以及跟踪中的上下文信息（如 kernel ID）来实现。

下图左侧展示了 Nsight Systems 的视图，右侧为 XPU-Timer 视图，后者可以清晰地对齐不同 rank 上的同类操作。

4.3. Nsight Systems 多报告分析 (Multi Report Analysis)

Nsight Systems 提供了“配方 (recipe)”功能，用于对多个分析报告进行聚合与分析。
命令行使用方法：nsys recipe [<args>] <recipe name> [<recipe args>]
更多信息请参阅官方文档：https://docs.nvidia.com/nsight-systems/2025.1/UserGuide/index.html#multi-report-analysis

可用的配方类型包括：
* 摘要 (Summary):
* cuda_api_sum: CUDA API 摘要
* cuda_gpu_kern_sum: CUDA GPU Kernel 摘要
* cuda_gpu_mem_size_sum: CUDA GPU 内存操作摘要 (按大小)
* cuda_gpu_mem_time_sum: CUDA GPU 内存操作摘要 (按时间)
* mpi_sum: MPI 摘要
* nccl_gpu_proj_sum: NCCL GPU 投影摘要
* nccl_sum: NCCL 摘要
* nvlink_sum: NVLINK 网络带宽摘要

• 步调 (Pace):

  • cuda_gpu_kern_pace: CUDA GPU Kernel 步调
  • nvtx_gpu_proj_pace: NVTX GPU 投影步调
  • nvtx_pace: NVTX 步调

• 热图 (Heatmap):

  • cuda_gpu_time_util_map: CUDA GPU 时间利用率热图
  • gpu_metric_util_map: GPU 指标利用率热图

• 追踪 (Trace):

  • nccl_gpu_overlap_trace: NCCL GPU 重叠追踪
  • nvtx_gpu_proj_trace: NVTX GPU 投影追踪

• 同步/异步 (Sync/Async):

  • cuda_memcpy_async: 使用页锁定内存的 CUDA 异步内存拷贝
  • cuda_api_sync: CUDA 同步 API

4.4. 使用 Jupyter Notebooks 进行数据呈现

Nsight Systems 的配方分析结果可以方便地在 Jupyter Notebooks 中加载和可视化。

工作流程:

1. 在集群上运行 nsys profile 收集数据：
   nsys profile -t nvtx,cuda -o /path/to/nsys-reps/xxx_%p.nsys-rep python train.py
2. 运行 nsys recipe 处理数据：
   nsys recipe <recipe_name> --input /path/to/nsys-reps/ --output /path/to/output/ <other_args>
3. 将输出文件夹复制到本地，以便在 Nsight Systems UI 或 Jupyter Lab 中打开：
   scp -r <username>@<cluster_ip>:/path/to/output/ ./

上图左侧为 Nsight Systems UI 中显示的 GPU 利用率摘要，右侧为 Jupyter Lab 中显示的热图，展示了数据呈现的灵活性。

4.5. 配方示例：摘要 (Summary)

此示例展示了如何使用 cuda_gpu_kern_sum 配方来生成 GPU Kernel 的时间成本摘要。

命令:
$ nsys recipe cuda_gpu_kern_sum --input /path/to/nsys-reps/ --output .

输出:
分析结果包含按名称排序的 Kernel 时间成本表和所有 Kernel 的时间成本曲线图。图表有助于快速识别耗时最长的 Kernel，例如图中的 ncclDevKernel_SendRecv (约 392.415M)。

4.6. 配方示例：步调 (Pace)

此示例展示了如何使用 nvtx_gpu_proj_pace 配方来分析 optimizer_step 的 NVTX GPU 投影步调，以发现性能异常。

命令:
$ nsys recipe nvtx_gpu_proj_pace --name optimizer_step --input /path/to/nsys-reps/ --output .

分析发现:
通过分析图表，可以发现 rank 11 上的优化器中增加了刻意的 CPU 开销。

• 进度图 (Progress): 按持续时间定义的迭代。图中显示 rank 11 的执行时间（约 12.9423M）明显高于其他 ranks。
• 一致性图 (Consistency): 按持续时间定义的迭代。图中显示在 rank 11 处有一个明显的尖峰，表明其行为与其他 ranks 不一致。

4.7. 自定义配方 (Customized Recipe)

如果现有的配方无法满足分析需求，例如图表不够清晰、数据繁杂或洞察不足，用户可以构建自己的配方。

思路: 如果轮子不好用，就造一个更好的。

用户可以参考或基于现有实现进行开发，例如：https://github.com/hyxcl/nsys_recipes

上图左侧是一个自定义配方生成的饼图，用于分析计算与通信的百分比。

5. CUPTI (NVIDIA CUDA Profiling Tools Interface)

探索自定义性能分析工具。

5.1. 性能数据收集：注入式测量 vs. 采样

注入式测量 (Instrumentation / Event Tracing)
* 通过插入的钩子 (hooks) 或探针 (probes) 记录每个事件，获取精确的计时和计数。
* 对于频繁发生的事件，开销较高。
* 实现方法:
* 源码级包装器和宏: 手动使用 CUDA Event API（如 cudaEventRecord/cudaEventElapsedTime）。
* CUPTI APIs: NVIDIA 官方的追踪与分析 API，提供精确的进入/退出时间戳和硬件计数器。
* LD_PRELOAD 拦截: 通过 custom.so 覆盖 CUDA 符号，可用于闭源二进制文件，部署简单。典型应用：XPU-Timer。

采样 (Sampling / Statistical Profiling)

• 周期性地收集程序状态数据，而不是记录每个事件。

• 优点:

  • 开销低，即使在高采样率下对程序影响也很小。
  • 自动过滤掉执行时间极短的事件，专注于热点。

• 缺点:

  • 精度较低，可能会遗漏罕见事件。
  • 不适用于精确的事件计时，更适合统计分析。

混合模式 (Tracing + Profiling)

• 结合注入式测量和采样。
• 典型应用: Nsight Systems 对 GPU Kernel 使用注入式测量，对 CPU 活动使用采样。
• 优点: 在控制整体开销的同时，精确追踪关键事件。

5.2. CUPTI APIs 概览

CUPTI 提供了一系列 API 用于追踪和分析 CUDA 应用。更多信息请参阅：https://docs.nvidia.com/cuda/cupti/index.html

API 分类:
* 追踪 (Tracing):
* Activity API: 异步记录 CUDA 活动（API 调用、Kernel 执行、内存拷贝）。
* Callback API: 当特定的 CUDA 事件发生时，通过回调函数通知订阅者。

• 分析 (Profiling):
  • Host Profiling: 用于枚举、配置和评估性能指标的主机端 API。
  • Range Profiling: 针对一段执行区间收集性能指标的目标 API。
  • PC Sampling: 采样 warp 程序计数器和 warp 调度器状态。
  • SASS Metrics: 使用 SASS 指令级插桩在源码级别收集 Kernel 性能指标。
  • PM Sampling: 通过周期性采样 GPU 性能监视器 (PM) 来收集硬件指标。
  • Profiling API (已弃用): 用于收集执行区间性能指标的目标 API，在 CUDA 13.0 中已弃用，推荐使用 Range Profiling API。
  • Checkpoint API: 支持自动保存和恢复 CUDA 设备的功能状态。

5.3. CUPTI Activity API

CUPTI Activity API 异步地收集 CPU 和 GPU 活动，并将其存储在活动记录结构中。

• 活动缓冲区用于传输记录；客户端需提供空缓冲区以防止数据丢失；记录的顺序不保证。
• 使用 correlationId 来关联主机端 API 调用和设备端 Kernel 活动。

上图展示了从主机调用 cudaLaunchKernel 到 CUPTI 线程记录活动，再到设备执行 Kernel 的流程，以及如何通过 correlationId 将它们关联起来。右侧是活动记录的数据结构 Cupti_ActivityKind。

CUPTI Activity API 代码示例
以下 Python 代码展示了如何使用 Activity API：

• 通过回调函数管理缓冲：为每个线程创建缓冲区以优化多线程性能。

  • buffer_requested: 请求缓冲区。
  • buffer_completed: 处理已填满的缓冲区。

• 当缓冲区满或被手动刷新时，数据将被递送。

代码流程包括：定义回调函数、注册回调、启用 KERNEL 和 MEMCPY 活动、刷新缓冲区以及禁用活动。

5.4. CUPTI Callback API

CUPTI Callback API 允许用户订阅 CUDA 事件并在事件发生时执行自定义函数。

• 订阅者 (Subscriber): 注册回调的实体。
• 回调函数 (Callback function): 事件发生时执行的函数。
• 回调域 (Callback Domain): 事件的范围，如驱动 (driver)、运行时 (runtime)、资源 (resource)、同步 (sync) 等。
• 回调 ID (Callback ID): 标识域内的具体事件。

上图展示了主机调用 API 时，在 API 进入 (API_ENTER) 和退出 (API_EXIT) 点触发已注册的回调函数，该函数可以记录日志或处理用户数据。

5.5. 案例研究：内存分析器 (Memory Profiler)

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

使用方法:
LD_PRELOAD=libmemprof.so NVTX_INJECTION64_PATH=/usr/local/cuda/lib64/libcupti.so python <py>

实现技术:
* 使用 CUPTI Callbacks 追踪所有内存分配和释放的 CUDA 运行时 API。
* 利用 LD_PRELOAD 和 __attribute__((constructor)) 在 main 函数之前插入 CUPTI 初始化代码，实现对目标应用的零代码侵入。
* 生成一个按峰值内存使用量降序排列的报告，并显示对应的 NVTX 范围。

上图代码展示了如何使用构造函数属性注册回调，并在回调函数中通过 switch 语句处理不同的 CUDA API 调用（如 cudaMalloc）和 NVTX 事件。

内存分析器工作流程示例

以下 Python 代码和追踪表演示了内存分析器如何工作。

• Python 代码: 定义了带有 NVTX 注解的嵌套函数，并在函数内部分配内存。
• 追踪表:
  • Step/Action: 显示程序的执行步骤。
  • Stack State: 追踪函数调用栈。
  • Memory Allocation: 记录每次内存分配的大小。
  • Statistics Update: 更新内存使用统计信息，包括 mem_self_peak (函数自身峰值)、mem_stack_peak (当前调用栈峰值) 和 mem_total (总内存)。

6. PyTorch Profiler

PyTorch Profiler 是一个用于分析 PyTorch 代码性能的工具。

• 性能开销: 在 H100 GPU 上对 Mixtral 8x7B 模型进行训练时，启用分析器会带来约 13% 的性能下降。
• 报告大小: 生成的报告大小约为 20MB/GPU/秒。

使用示例:
以下代码展示了如何使用 PyTorch Profiler。通过 with profile(...) 上下文管理器包装训练代码，并使用 record_function 添加自定义的代码区域标注。

代码中配置了 CPU 和 CUDA 活动的追踪，并使用 tensorboard_trace_handler 来处理输出。右侧的截图是分析器生成的时间线视图，可以直观地看到各个操作的执行情况。

PyTorch Profiler 提供了详细的性能分析视图。通过启用不同的选项，可以获取丰富的信息：
- with_stack：记录并显示导致操作的Python和C++调用栈。
- record_shapes：记录操作输入张量的维度。
- profile_memory：记录内存分配/释放事件。

如下图所示，分析器可以展示从顶层Python代码（如torch/nn/modules/module.py中的_call_impl）到底层CUDA核函数（如void sl::native::scaled_elementwise_kernel<...>...）的完整调用链。右侧面板则详细列出了算子的输入类型、维度、步长信息，以及内存分析器的统计数据，如总预留内存（Total Reserved）、总分配内存（Total Allocated）、设备ID（Device ID）等。

7. NVIDIA Resiliency Extension Straggler Detection (NVIDIA 弹性扩展：落后者检测)

7.1. 落后者检测 (Straggler Detection)

识别滞后的单个 GPU

• 识别性能缓慢的 rank 或节点：确保没有单个表现不佳的节点显著拖慢整个分布式进程。
• 在所有 rank 间保持一致的性能：通过检测并可能终止性能不佳的节点，确保整个系统高效运行。

关键优势
- 可扩展性 (Scalability)：随着 rank 数量的增加，落后者检测变得愈发关键，以维持整体的可扩展性。
- 可靠性 (Reliability)：通过确保进程不会因性能不佳的节点而中止，分布式系统的可靠性得到增强。

吞吐量 vs. 步数 (Throughput vs. Step Number)
下图展示了在GCP上使用8个NVIDIA L4s GPU训练Llama 8B模型时的吞吐量变化。可以看到，在某个训练步骤后，吞吐量急剧下降，这表明可能出现了落后者节点。

7.2. 与 PyTorch 集成

• 检测流程:

  • 计算分数 (Compute Scores): 计算相对和个体性能分数。
  • 识别落后者 (Identify Stragglers): 检查分数是否低于指定的阈值。

• 性能分数的两种类型:

  • 相对分数 (Relative): 使用最快的 rank 作为参考基准。
  • 个体分数 (Individual): 使用该 rank 自身历史最佳性能作为参考。

代码集成示例
下图展示了如何在PyTorch训练循环中集成落后者检测器。通过 straggler.Detector.section 上下文管理器包裹训练的关键部分（如前向传播、后向传播、优化器步骤），并在每个epoch或特定间隔后调用 report.identify_stragglers 来生成报告并识别落后者。

8. 内存优化工具 (Tools for Memory Optimization)

8.1. PyTorch 内存分析器：案例研究 (PyTorch Memory Profiler: Case Study)

可视化随时间变化的所有内存分配

• PyTorch API:

  • 开始记录: torch.cuda.memory.record_memory_history(max_entries=100000)
  • 保存快照: torch.cuda.memory.dump_snapshot(file_name)
  • 停止记录: torch.cuda.memory.record_memory_history(enabled=None)

• 与 Megatron 集成:

  • 使用命令行参数 --record-memory-history 可以在最后一个 rank 中记录内存历史。

使用内存分析器定位 OOM (Out-of-Memory) Bug

• 问题背景: 社区报告了一个问题（https://github.com/NVIDIA/Megatron-LM/issues/1566），指出在使用 cuda-graph 时，即使有大量可用内存，也会导致 OOM 错误。
• 调试步骤:
  1. 减小模型尺寸，使得只有在使用 cuda-graph 时才会出现 OOM。
  2. 使用 --record-memory-history 捕获第一次迭代的内存使用情况，并与不使用 cuda-graph 的情况进行比较。

代码示例:

# Add memory recording if specified
if [ "$RECORD_MEMORY" = "true" ]; then
    TRAINING_ARGS+=(--record-memory-history)
    TRAINING_ARGS+=(--memory-snapshot-path ${ROOT_DIR}/snapshot_cuda-graph-${USE_CUDA_GRAPH}_gbs-${GLOBAL_BATCH_SIZE}.pickle)
fi

通过对比内存使用情况，可以观察到：

• 不使用 CUDA Graph (w/o CUDA GRAPH)：内存使用量较为平稳。
• 使用 CUDA Graph (CUDA GRAPH, GBS=4)：内存使用图中出现了明显的“锯齿状”模式（sawtooth patterns）。起初怀疑这是否与全局批量大小（GBS）有关。

进一步实验，将全局批量大小（GBS）从4增加到8，并进行比较：
- GBS=4 vs GBS=8: 当 GBS 增加到 8 时，“锯齿状”模式变得更加明显和频繁。
- 问题根源: 深入分析发现，cuda_graph_capture 会为每一个微批次（micro batch）实例化一个新的输入/输出缓冲区，这并非预期行为，并导致了内存的持续增长，最终引发 OOM。下图中的调用栈信息指出了问题的具体代码位置。

9. 总结 (Summary)

9.1. 各工具的适用场景

下表总结了不同分析工具的最佳使用场景及其优缺点。

工具对比详情:

• NVIDIA Resiliency Extension (NVIDIA 弹性扩展)

  • 最佳用途: 在进行详细性能分析之前，检查是否存在落后者问题。
  • 优点: 易于使用（通过在代码中添加注解或启用Megatron选项），对训练和推理生命周期没有性能影响。
  • 缺点: 不提供直接的性能分析功能。

• Nsight Systems Multi Report Analysis (Nsight Systems 多报告分析)

  • 最佳用途: 比较/聚合多节点性能分析数据。
  • 优点: 适用于多节点训练/推理的回归跟踪。
  • 缺点: 缺乏 rank 间的交互连接分析。

• Nsight Systems

  • 最佳用途: 系统级的 GPU/CPU 时间线分析。
  • 优点: 可通过 NVTX 显示模型信息。
  • 缺点: CPU-GPU 交互和内核级分析，不提供 Python 代码行级别的信息。

• PyTorch Profiler

  • 最佳用途: 模型瓶颈分析、代码追溯和内存分析。
  • 优点: 显示代码位置信息，显示形状和模型信息（需要添加装饰器）。
  • 缺点: 难以合并多节点报告，内存分析文件可能因 OOM 而损坏。

• CUPTI

  • 最佳用途: 为无需修改代码的场景定制指标。
  • 优点: 提供低级别的 CUDA 事件追踪。
  • 缺点: 需要额外的开发工作，缺乏模型上下文信息。

10. 未来工作讨论 (Future Work Discussion)

10.1. LLM 友好的工具

为了更好地支持大规模语言模型（LLM）的开发和调试，未来的工具将关注以下特性：

Feature (特性)	Description (描述)
LLM Component-Level Statistics (LLM 组件级统计)	通过 NVTX 提供 Attr, MLP, DP/TP/EP 通信指标和图表，关注模型模块而非仅仅是低级内核。
Model Information on Timeline (时间线上的模型信息)	在时间线上显示 LLM 特定信息，如张量形状、层号等。
Multi-Rank Communication Correlation (多 Rank 通信关联)	链接跨 rank 的通信/计算，点击一个 rank 中的模块可以连接到所有相关的 rank。
Robust Memory Profiling (鲁棒的内存分析)	防止文件损坏，保证在 OOM 发生前完成 pickle 文件的保存。
Official Support & Maintenance (官方支持与维护)	由 NVIDIA 提供和维护，以进行广泛推广和长期维护。