📚 课程目录第4章 AI训练Profiling基础

4. AI训练Profiling基础

1.1 为什么要做Profiling？

说实话，我见过太多团队花几周时间调优模型，结果发现瓶颈根本不在他们想的地方。有人拼命优化数据加载，实际GPU利用率才30%。有人疯狂调学习率，结果发现通信开销占了训练时间的一半。

Profiling说白了就是给你的训练脚本做一次「体检」。它能告诉你：

GPU到底在干活还是在摸鱼？
数据加载是不是拖了后腿？
算子执行时间都花在哪了？
显存有没有被浪费？

我个人习惯是，拿到一个新模型，第一件事就是跑一次profile。这就像医生看病先做CT，总得知道病灶在哪才能对症下药。

1.2 nsys profile 命令入门

NVIDIA Nsight Systems 的命令行工具叫 nsys。它的核心命令就是 nsys profile。用法其实很简单：

nsys profile -o output_profile python train.py

这条命令会启动你的训练脚本，同时采集性能数据。最终生成一个 .nsys-rep 文件。

嗯，这里要注意几个常用参数：

参数	作用	我的建议
`-o`	指定输出文件名	一定要加，不然默认名字很难找
`-t`	指定跟踪的API类型	常用 `cuda,nvtx,osrt`
`-c`	设置采集时长	比如 `-c duration -d 30` 只采30秒
`-w`	设置等待时间	跳过前几步初始化，避免噪声

举个例子，我经常这样用：

nsys profile -o resnet50_profile -t cuda,nvtx -w 5 -c duration -d 60 python train_resnet50.py

这条命令的意思是：等5秒让模型初始化完，然后采集60秒的CUDA和NVTX数据。为什么要等5秒？因为训练刚开始那几步，数据加载和预热会引入很多噪声，你想想看，如果把这些也采进去，分析时还得手动过滤，多麻烦。

1.3 理解 .nsys-rep 文件

生成的 .nsys-rep 文件，本质上是一个压缩包。里面包含了：

时间轴上的所有事件（kernel启动、内存拷贝、API调用等）
GPU硬件计数器（SM利用率、显存带宽等）
CPU线程调度信息
自定义NVTX标记

这个文件可以直接用 Nsight Systems GUI 打开。双击就行，或者用命令：

nsys-ui output_profile.nsys-rep

打开后你会看到一个时间轴视图。上面是CPU活动，下面是GPU活动。我刚开始用的时候，第一反应是「哇，信息真多」，第二反应是「这怎么看啊？」

别急，我教你一个技巧：先看GPU那一行有没有大段的空白。如果有，说明GPU在等数据，瓶颈在数据加载。如果GPU一直很忙，那就看具体哪个kernel占的时间最长。

核心要点：

GPU空闲 → 数据加载或CPU预处理是瓶颈
GPU忙碌但单个kernel时间长 → 算子本身需要优化
GPU忙碌但频繁切换kernel → 可能kernel launch开销大

1.4 实战：对PyTorch训练脚本做Profiling

我们拿一个典型的PyTorch训练脚本举例。假设你有一个 train.py，里面大概长这样：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

model = ResNet50().cuda()
dataloader = DataLoader(dataset, batch_size=64, num_workers=4)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for epoch in range(10):
    for batch in dataloader:
        x, y = batch
        x, y = x.cuda(), y.cuda()
        
        output = model(x)
        loss = criterion(output, y)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

直接跑 nsys profile -o pytorch_profile python train.py 就能采集。但我建议你加点NVTX标记，这样分析时能一眼看出每个阶段：

import nvtx

for epoch in range(10):
    with nvtx.annotate(f"epoch_{epoch}", color="blue"):
        for batch in dataloader:
            with nvtx.annotate("data_load", color="green"):
                x, y = batch
                x, y = x.cuda(), y.cuda()
            
            with nvtx.annotate("forward", color="yellow"):
                output = model(x)
                loss = criterion(output, y)
            
            with nvtx.annotate("backward", color="red"):
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

然后采集时加上 -t cuda,nvtx：

nsys profile -o pytorch_profile -t cuda,nvtx python train.py

打开生成的 .nsys-rep 文件，你就能看到每个epoch、每个阶段的时间分布了。我曾经帮一个团队优化训练，他们一直觉得是模型太大导致慢。结果一profile发现，数据加载阶段占了40%的时间，GPU大部分时间都在空等。后来把 num_workers 从4改成8，训练速度直接翻倍。

小技巧： 如果你只想看前几个step，可以用 --capture-range=cudaProfilerApi 配合 torch.cuda.profiler.start()/stop() 精确控制采集范围。这样文件小，分析也快。

1.5 常见坑与避坑指南

我曾经踩过一个坑：用 nsys profile 采集时，默认会跟踪所有系统调用，导致profile文件巨大无比，打开都卡半天。后来发现加个 -t cuda,nvtx 只跟踪CUDA和NVTX事件，文件大小直接降到原来的十分之一。

还有一次，我帮客户分析一个分布式训练脚本。他们用 nsys profile 采集了10分钟，生成的文件有5个GB。我一看，原来忘了加 -c duration -d 30 限制时长。你想想看，10分钟的数据量，光加载就要等半天。

所以我的建议是：

先跑一个短时间的profile（30秒到1分钟）
确认数据质量没问题
再根据需要调整参数做详细采集

注意： 不要在训练脚本里同时使用 torch.cuda.profiler 和 nsys profile 的自动采集，两者会冲突。要么用 nsys profile 全自动，要么用 torch.cuda.profiler 手动控制，别混着来。

1.6 知识体系总览

下面这张图帮你理清本章的核心逻辑：

这张图展示了整个流程：训练脚本作为输入，经过 nsys profile 采集，生成 .nsys-rep 文件。关键参数控制采集范围和质量，最后用 Nsight Systems GUI 打开分析。

好了，这一章的内容就到这里。记住一句话：没有profile，就没有优化。下次你遇到训练慢的问题，别急着改代码，先跑一次profile看看。