如何用PyTorch Profiler分析GPT-J等模型的各层耗时？

问题描述 #

我正在学习使用PyTorch Profiler分析模型，官方示例运行正常，但换成GPT-J这类模型后，输出结果只显示辅助函数，看不到q_proj、k_proj、v_proj、softmax等标准PyTorch层的耗时信息。

查看模型代码后确认这些层都是标准的PyTorch组件（比如nn.Linear、nn.functional.softmax），但用Profiler分析时得到的输出如下：

---

Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  

---

forward        90.29%     558.000us        94.34%     583.000us     583.000us             1  
       aten::zeros         5.02%      31.000us         5.66%      35.000us      35.000us             1  
      aten::unbind         1.62%      10.000us         2.43%      15.000us      15.000us             1  
      aten::detach         0.49%       3.000us         1.29%       8.000us       8.000us             1  
      aten::select         0.65%       4.000us         0.81%       5.000us       5.000us             1  
            detach         0.81%       5.000us         0.81%       5.000us       5.000us             1  
       aten::empty         0.65%       4.000us         0.65%       4.000us       2.000us             2  
       aten::zero_         0.16%       1.000us         0.16%       1.000us       1.000us             1  
  aten::as_strided         0.16%       1.000us         0.16%       1.000us       1.000us             1  
          aten::to         0.16%       1.000us         0.16%       1.000us       1.000us             1  
aten::resolve_conj         0.00%       0.000us         0.00%       0.000us       0.000us             1  
 aten::resolve_neg         0.00%       0.000us         0.00%       0.000us       0.000us             1  

---

Self CPU time total: 618.000us

我使用的代码如下：

import torch
import torchvision.models as models
from torch.profiler import profile, record_function, ProfilerActivity

from transformers import AutoModelForCausalLM, AutoTokenizer

model     = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

prompt = ("In a shocking finding, scientists discovered a herd of unicorns living in a remote, "
           "previously unexplored valley, in the Andes Mountains. Even more surprising to the "
           "researchers was the fact that the unicorns spoke perfect English."
        )

input_ids = tokenizer(prompt, return_tensors="pt").input_ids
gen_tokens = model.generate(input_ids,
                            do_sample=True,
                            temperature=0.9,
                            max_length=100)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    with record_function("forward"):
       gen_text = tokenizer.batch_decode(gen_tokens)[0]


print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=100))

print ("----- Group by input shape")
print(prof.key_averages(group_by_input_shape=True).table(sort_by="cpu_time_total", row_limit=10))

prof.export_chrome_trace("trace.json")

请问针对这类模型，正确使用PyTorch Profiler的方式是什么？是否需要修改模型？

解决方法 #

1. 修正Profiler的作用范围 #

核心问题是Profiler没有包裹模型的推理过程——当前代码里，model.generate在Profiler上下文之外执行，Profiler只捕获了tokenizer.batch_decode的耗时，这就是输出只有辅助函数的原因。必须把模型的generate或forward调用放在Profiler的with块内。

2. 启用详细追踪选项 #

为了捕获到q_proj、k_proj等层的耗时，需要给Profiler添加以下参数：

• with_stack=True：记录调用栈，方便定位层的来源
• profile_memory=True：可选，同时分析内存使用
• 对于生成式模型的循环推理，使用schedule参数控制Profiler的采样时机，避免重复采样

3. 不需要修改模型 #

这类模型使用的是标准PyTorch层，无需修改模型代码，只要Profiler配置正确就能捕获到层的耗时。

修正后的代码示例 #

import torch
from torch.profiler import profile, record_function, ProfilerActivity, schedule

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和tokenizer
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

prompt = ("In a shocking finding, scientists discovered a herd of unicorns living in a remote, "
          "previously unexplored valley, in the Andes Mountains. Even more surprising to the "
          "researchers was the fact that the unicorns spoke perfect English."
         )
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

# 定义Profiler调度：跳过前2次迭代（预热），记录接下来的3次
def trace_handler(prof):
    print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=20))
    prof.export_chrome_trace(f"trace_{prof.step_num}.json")

with profile(
    activities=[ProfilerActivity.CPU],
    record_shapes=True,
    with_stack=True,
    schedule=schedule(wait=2, warmup=1, active=3, repeat=1)
) as prof:
    # 将generate放在Profiler上下文内
    gen_tokens = model.generate(
        input_ids,
        do_sample=True,
        temperature=0.9,
        max_length=100,
        # 每生成一个token后调用profiler.step()，配合schedule使用
        callback=lambda step: prof.step()
    )

# 打印所有关键指标
print("=== 完整耗时统计 ===")
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=50))

# 过滤查看指定层的耗时
print("\n=== q_proj/k_proj/v_proj耗时 ===")
for evt in prof.key_averages():
    if "q_proj" in evt.key or "k_proj" in evt.key or "v_proj" in evt.key:
        print(f"{evt.key}: CPU total = {evt.cpu_time_total / 1e6:.2f} ms, 调用次数 = {evt.count}")

print("\n=== Softmax耗时 ===")
for evt in prof.key_averages():
    if "softmax" in evt.key:
        print(f"{evt.key}: CPU total = {evt.cpu_time_total / 1e6:.2f} ms, 调用次数 = {evt.count}")

4. 查看结果的技巧 #

• 使用filter_name方法快速筛选目标层：

  q_proj_stats = prof.key_averages().filter_name("q_proj")
  print(q_proj_stats.table())
  

• 导出Chrome Trace后，在Chrome浏览器中打开chrome://tracing/，导入trace文件，搜索层名即可查看详细的时间线

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