如何定位PyTorch训练过程中耗时最长的模块？以OpenPCDet分布式场景为例

问题原因 #

你之前使用torch.autograd.profiler.profile仅得到CPU耗时数据，是因为该接口默认关闭CUDA耗时统计功能，没有添加use_cuda=True参数就不会采集GPU侧的执行数据。

方案1：使用新版PyTorch Profiler（推荐，PyTorch 1.8及以上版本支持） #

新版Profiler功能更完善，支持直观的可视化、调用栈追溯，适配单/多GPU场景，使用方式如下：

import torch

# 先执行3-5轮warmup，排除CUDA初始化、显存分配、算子编译的干扰
for _ in range(3):
    _ = model(input)
torch.cuda.synchronize()

# 开启Profiler采集
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True, # 记录张量形状，辅助定位对应模块
    with_stack=True # 记录调用栈，可直接关联到具体代码行
) as prof:
    # 采集5轮完整训练步的耗时（包含前向、反向、参数更新）
    for _ in range(5):
        output = model(input)
        loss = output.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        prof.step() # 标记每轮迭代边界

采集完成后可以通过两种方式分析结果：

1. 直接打印耗时排名，按GPU总耗时降序排列，快速定位高耗时模块：

print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20
))

2. 导出为Chrome Trace文件，可视化查看全链路执行过程：

prof.export_chrome_trace("train_trace.json")

导出后在Chrome浏览器地址栏输入chrome://tracing，加载生成的json文件即可看到每个算子、每个模块的执行时间、执行顺序，还能筛选CUDA侧的执行队列排查阻塞。

方案2：旧版autograd.profiler正确用法 #

如果使用的PyTorch版本低于1.8，可以用旧版接口，开启CUDA采集参数即可：

torch.cuda.synchronize()
with torch.autograd.profiler.profile(enabled=True, use_cuda=True) as prof:
    output = model(input)
    # 可补充反向传播、参数更新逻辑统计完整训练步
torch.cuda.synchronize()

# 按CUDA总耗时排序输出结果
print(prof.key_averages().table(sort_by="cuda_time_total"))

注意事项 #

• 统计前必须做warmup，否则初始迭代的额外开销会导致结果偏差
• 如果只需要统计前向传播耗时，可以去掉反向传播、优化器更新的相关代码
• 多GPU场景下可以每个rank单独生成trace文件，分别分析对应GPU的耗时情况

内容的提问来源于stack exchange，提问作者zheyuanWang