用torch.profiler来监控大模型推理这样可以得到GPU/CPU 使用情况、时间消耗、内存占用比简单的psutil更精确。下面完整示例importtorchfromtransformersimportAutoModelForCausalLM,AutoTokenizerimporttime# 1️⃣ 模型和 tokenizermodel_namedistilgpt2tokenizerAutoTokenizer.from_pretrained(model_name)modelAutoModelForCausalLM.from_pretrained(model_name)devicecudaiftorch.cuda.is_available()elsecpumodel.to(device)# 输入 promptpromptHello, I am a small language model. I caninputstokenizer(prompt,return_tensorspt).to(device)# 2️⃣ 使用 torch.profiler 监控推理withtorch.profiler.profile(scheduletorch.profiler.schedule(wait1,warmup1,active2,repeat1),on_trace_readytorch.profiler.tensorboard_trace_handler(./log),record_shapesTrue,profile_memoryTrue,with_stackTrue)asprof:# 模拟多步推理方便 profiler 捕捉forstepinrange(4):outputsmodel.generate(**inputs,max_length50,do_sampleTrue,temperature0.7,top_k50,top_p0.95,no_repeat_ngram_size2)prof.step()# 标记 step方便分析# 3️⃣ 打印统计信息print(prof.key_averages().table(sort_byself_cpu_time_total,row_limit20))# 4️⃣ TensorBoard 查看print(Profiler traces saved to ./log, run:)print(tensorboard --logdir./log)✅ 功能说明Profiler 配置record_shapesTrue→ 记录每个操作的 tensor 形状profile_memoryTrue→ 记录显存/内存占用with_stackTrue→ 打印调用堆栈定位耗时操作Schedulescheduletorch.profiler.schedule(wait1,warmup1,active2,repeat1)wait等待 step 数不记录warmup热身 step不计入统计active记录 step 数repeat重复几次 schedule用于多步生成或批量生成时统计更准确Profiler 输出prof.key_averages()→ CPU/GPU 每个操作耗时、显存消耗可以排序self_cpu_time_total/cuda_time_total找到瓶颈TensorBoard 可视化tensorboard--logdir./log可以看到每步操作的 GPU/CPU 时间、内存曲线直观分析模型推理性能优化建议如果模型更大7B可以结合torch.autocast(cuda)做混合精度节省显存并提升速度withtorch.autocast(device_typecuda,dtypetorch.float16):outputsmodel.generate(...)可以在 profiler 的on_trace_ready回调中写自定义分析比如打印每层显存占用。访问https://ui.perfetto.dev/#!/viewer?local_cache_key0-json