A PyTorch training profiler with support for CUDA, NVTX, NCCL and COS upload
Project description
tiprof - PyTorch 分布式训练性能分析工具
tiprof 是一个专为 PyTorch 分布式训练设计的性能分析工具,支持 CUDA、NVTX 和 NCCL 性能追踪。无需修改训练代码,即可对 PyTorch 训练任务进行全面的性能分析。
目录
核心特性
- ✅ 无代码侵入 - 无需修改训练代码,通过命令行包装器即可启用性能分析
- ✅ 多分析器支持 - 支持 PyTorch Profiler 和 Nsight Systems 两种分析工具
- ✅ 分布式训练支持 - 完美支持多 GPU、多节点的分布式训练场景
- ✅ 灵活的分析范围 - 支持按训练步骤(steps)或时间范围(durations)进行分析
- ✅ NCCL 性能追踪 - 支持 TCCL 增强版 NCCL,提供详细的通信性能数据
- ✅ 云存储集成 - 支持将分析结果直接上传到腾讯云 COS
- ✅ 多框架兼容 - 兼容 PyTorch、DeepSpeed、Horovod、ColossalAI 等主流训练框架
支持的框架
- PyTorch
安装
系统依赖
- Python 3.8 或更高版本
- CUDA 11.0 或更高版本(用于 GPU 训练)
- PyTorch 2.0 或更高版本
使用 pip 安装
pip install tiprofiler
验证安装
tiprof --version
快速开始
分析特定训练步骤
分析第 10-15 步和第 20-25 步:
tiprof --steps 5-10 --nsys python train.py
按时间范围分析
分析训练开始后 30-40,以第一个step作为开始时间
tiprof --durations 30-40 python train.py
查看分析结果
-
PyTorch Profiler 结果:使用 TensorBoard 查看
pip install torch_tb_profiler tensorboard --logdir=<your trace result directory>
-
PyTorch Profiler 结果:使用 perfetto 查看
打开perfetto网站,点击左上角导航栏Open trace file打开生成的profiler文件即可查看timeline
- Nsight Systems 结果:使用 Nsight Systems GUI 打开
.nsys-rep文件
把profile文件下载到本地,使用客户端打开即可
命令行参数
| 参数 | 类型 | 默认值 | 描述 |
|---|---|---|---|
--steps |
str | None | 指定要分析的训练步骤范围,格式:"start-end,start-end",例如 "10-15,20-25" |
--durations |
str | None | 指定要分析的时间范围(秒),格式:"start-end,start-end",例如 "30-60,90-120" |
--ranks |
str | "0" | 指定要分析的 rank,多个 rank 用逗号分隔,例如 "0,1,2" 或 "all" 分析所有 rank |
--nsys |
flag | False | 启用 Nsight Systems |
--output |
str | "file:///mnt/data/profiler" | 输出位置,支持本地路径 file://path 或腾讯云 COS cos://bucket/path |
--tccl |
flag | False | 启用 TCCL 增强版 NCCL,提供更详细的通信性能数据 |
--profiler-args |
str | None | 传递给 PyTorch Profiler或Nsight System 的额外参数,格式:"key1=value1;key2=value2" |
注意:--steps 和 --durations 参数互斥,只能使用其中一个。
参数详细说明
--steps 参数
指定要分析的训练步骤范围。支持多个范围,用逗号分隔:
# 分析第 5-10 步
tiprof --steps 5-10 python train.py
# 分析第 5-10 步和第 20-30 步
tiprof --steps 5-10,20-30 python train.py
--durations 参数
指定要分析的时间范围(从训练开始计时,单位:秒):
# 分析训练开始后 30-40 秒
tiprof --durations 30-40 python train.py
# 分析多个时间段
tiprof --durations 30-40,100-110 python train.py
--ranks 参数
在分布式训练中,指定要分析的 rank:
# 只分析 rank 0
tiprof --ranks 0 python train.py
# 分析 rank 0, 1, 2
tiprof --ranks 0,1,2 python train.py
--nsys 参数
启用 Nsight Systems 进行系统级性能分析。Nsight Systems 提供更底层的 CUDA kernel、内存传输等信息:
tiprof --nsys true python train.py
PyTorch Profiler vs Nsight Systems:
- PyTorch Profiler:专注于 PyTorch 算子级别的性能分析,适合分析模型结构和算子性能
- Nsight Systems:提供系统级性能分析,包括 CUDA kernel、内存传输、CPU 活动等,适合深入分析底层性能
--output 参数
指定分析结果的输出位置:
# 输出到本地目录
tiprof --output file:///path/to/output python train.py
# 输出到腾讯云 COS
tiprof --output cos://my-bucket/profiling-results python train.py
--profiler-args 参数
PyTorch Profiler支持的参数
默认使用的参数
tiprof --profiler-args "record_shapes=True;with_stack=True" python train.py
更多参数参考torch profiler官方文档:
-
activities (iterable) – list of activity groups (CPU, CUDA) to use in profiling, supported values: torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, torch.profiler.ProfilerActivity.XPU. Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA or (when available) ProfilerActivity.XPU.
-
record_shapes (bool) – save information about operator’s input shapes.
-
profile_memory (bool) – track tensor memory allocation/deallocation (see export_memory_timeline for more details).
-
with_stack (bool) – record source information (file and line number) for the ops.
-
with_flops (bool) – use formula to estimate the FLOPS of specific operators (matrix multiplication and 2D convolution).
-
with_modules (bool) – record module hierarchy (including function names) corresponding to the callstack of the op. e.g. If module A’s forward call’s module B’s forward which contains an aten::add op, then aten::add’s module hierarchy is A.B Note that this support exist, at the moment, only for TorchScript models and not eager mode models.
Nsight system支持的参数
默认nsight system使用的参数
# torch 2.4及以上版本
tiprof --profiler-args "--trace=cuda,nvtx,osrt;--pytorch=autograd-nvtx,functions-trace;--trace-fork-before-exec=true;--stop-on-exit=true" python train.py
# torch 2.3以下版本
tiprof --profiler-args "--trace=cuda,nvtx,osrt;--pytorch=autograd-nvtx;--trace-fork-before-exec=true;--stop-on-exit=true" python train.py
添加额外参数:
tiprof --profiler-args "--python-sampling=true;--python-backtrace=cuda;--cudabacktrace=all" python train.py
完整参数参考官方文档
输出结果
tiprof 会生成多种类型的性能分析文件,具体取决于使用的分析器和配置。
输出文件类型
-
PyTorch Profiler 结果(默认)
- 文件格式:TensorBoard 格式(
.pt.trace.json) - 位置:
<output_dir>/<task_id>/profile_<node_name>.<timestamp>.<rank>.<timestamp>.pt.trace.json - 查看方法:使用 TensorBoard
- 文件格式:TensorBoard 格式(
-
Nsight Systems 结果(使用
--nsys时)- 文件格式:
.nsys-rep - 位置:
<output_dir>/<task_id>/profile_<node_name>.<timestamp>.nsys-rep - 查看方法:使用 Nsight Systems GUI
- 文件格式:
- task_id 为TI_INSTANCE_ID环境变量的值,由TI平台自动注入
- node_name 为NODE_NAME环境变量的值
- rank为RANK环境变量的值
- timestamp为当前时间戳
查看 PyTorch Profiler 结果
使用 TensorBoard 查看分析结果:
# 启动 TensorBoard
tensorboard --logdir=./
# 在浏览器中打开 http://localhost:6006
TensorBoard 提供的视图:
- Overview:整体性能概览
- Operator View:算子级别的性能统计
- Kernel View:CUDA kernel 性能统计
- Trace View:时间线视图,显示算子执行顺序
查看 Nsight Systems 结果
- 下载并安装 Nsight Systems
- 打开 Nsight Systems GUI
- 选择
File->Open-> 选择.nsys-rep文件 - 分析时间线、CUDA kernel、内存传输等信息
常见问题
Q1: --steps 和 --durations 可以同时使用吗?
A: 不可以。这两个参数是互斥的,只能选择其中一个。
- 使用
--steps当你知道要分析的具体训练步骤 - 使用
--durations当你想按时间范围进行分析
Q2: 分析结果文件未生成
A: 可能的原因和解决方案:
-
训练步骤不足:确保训练运行的步骤数大于指定的分析范围
# 如果指定 --steps 10-20,确保训练至少运行 20 步 -
权限问题:检查输出目录的写权限
ls -la ./ chmod 755 ./output_dir
-
磁盘空间不足:检查磁盘空间
df -h -
查看日志: 检查日志中是否有start profiler和stop profiler输出
Q3: 性能分析会影响训练性能吗?
A: 会有一定影响,但可以控制:
- PyTorch Profiler:通常会增加 10-30% 的开销
- Nsight Systems:开销较小,通常在 5-15%
- 建议:
- 只分析少量步骤(如 5-10 步)
- 只分析部分 rank(如
--ranks 0) - 在验证集或测试集上进行分析,而不是整个训练过程
Q4: 如何在多机训练中使用 tiprof?
A: 在每个节点上都使用 tiprof 包装训练命令:
# 节点 0
tiprof --ranks 0 torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
--master_addr="192.168.1.1" --master_port=29500 train.py
# 节点 1
tiprof --ranks 0 torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 \
--master_addr="192.168.1.1" --master_port=29500 train.py
建议:
- 每个节点只分析一个 rank(如 rank 0)以减少开销
- 使用共享存储或 COS 收集所有节点的分析结果
Q5: 如何分析内存使用情况?
A: 使用 profile_memory 参数:
tiprof --profiler-args "profile_memory=True" --steps 10-15 python train.py
然后在 TensorBoard 中查看 Memory View。
Q6: 分析结果文件太大怎么办?
A: 减小分析范围:
-
减少分析的步骤数:
tiprof --steps 10-12 python train.py # 只分析 3 步
-
减少分析的 rank 数:
tiprof --ranks 0 python train.py # 只分析 rank 0
-
不记录额外信息:
# 不使用 record_shapes, with_stack 等参数 tiprof --steps 10-15 python train.py
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file tiprofiler-1.0.1-py3-none-any.whl.
File metadata
- Download URL: tiprofiler-1.0.1-py3-none-any.whl
- Upload date:
- Size: 15.4 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.10.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
7a3bb510403af7d9673c5bd6d7d62373a6b611ad341237a71dd16437bae9eef0
|
|
| MD5 |
9da6c947cf19c74572b99dfd2d376b68
|
|
| BLAKE2b-256 |
57c532a7cf1d3f2554e0619c147e86a2c92fa0214848914d04158e4db75579a8
|