PyTorch性能调优实战：从Profiling到torch.compile与分布式扩展

干我们这行的都清楚，模型“能跑”和“能打”完全是两码事。同样的训练脚本，A 同学一卡一天跑完，B 同学三卡三天还在等，差距往往不在算法而在工程细节。PyTorch 性能调优这件事，很多同学是等项目卡住了才想起来做，结果一头扎进 Profiling 日志里翻来覆去，越看越迷茫。这篇文章我会把 PyTorch 性能调优的完整思路走一遍，围绕 Profiling、torch.compile 和分布式扩展三条主线来讲，配合实际踩过的坑和一些可以直接抄走的参数组合，希望能帮还在“调优门外”的同学少走几段弯路。

这篇学习笔记比较适合两类读者：一类是模型已经跑通但训练速度不满意，想系统学习如何定位瓶颈的算法工程师；另一类是准备把单卡脚本改成多卡分布式训练的工程同学。我会尽量把“为什么这么做”也讲透，而不只是丢给你一堆配置项。

1. 性能调优的整体思路：先测量，再动手

1.1 别凭感觉猜瓶颈，先用数据说话

我做性能调优的头一条原则就是：不 Profiling 不优化。这句话听起来像废话，但实际操作中几乎所有人都会犯“凭经验猜测”的毛病。

比如训练变慢了，第一反应往往是“是不是并行度不够？加卡吧！”结果一测，GPU 利用率只有 30%，瓶颈根本不在算力，而在数据加载或者 CPU 端的预处理逻辑。再比如有人一遇到显存不够就想到梯度累积，结果 Profiling 一看，爆炸显存的其实是某个中间激活值，改一下 checkpointing 策略就解决了，压根不用动训练策略。

所以在优化之前，先花十几分钟跑一次完整的 Profiling，把 CPU 时间、GPU 时间、显存占用、算子耗时分布、数据加载时间这些数据摆到桌面上。有了这些数据，你才知道自己的时间该花在哪个环节，而不是像无头苍蝇一样到处瞎调。

为什么很多人不愿意先做测量？因为 Profiling 看起来有点门槛，输出结果密密麻麻不知道看哪里。但现在的 PyTorch Profiler 做得已经相当友好了，能自动出表格、火焰图，甚至直接输出到 TensorBoard。花十几分钟学会看这些数据，后面省下来的时间是以“天”为单位的。

1.2 性能问题的三个层面

我习惯把 PyTorch 性能问题分成三个层面来排查，顺序不能乱：

第一层是存储和数据流动。包括 DataLoader 的加载速度、预处理是否占用了太多 CPU 时间、数据在 CPU 和 GPU 之间的拷贝是否频繁。很多训练任务慢，根本不是算力不够，而是“饭”端不上来，GPU 在空等数据。

第二层是计算效率。包括算子本身的实现效率、是否有大量小算子导致的 kernel launch 开销、是否存在不必要的同步点。这一层的优化通常靠算子融合、把动态图变成静态图、减少 GPU 和 CPU 之间的同步。

第三层是扩展效率。单卡已经调得很好了，但上多卡之后发现加速比远低于理论值，往往是通信开销、负载不均衡、梯度同步策略有问题。这一层需要在分布式层面做手脚。

每一层的问题要用不同的工具和手段去解决，这也是为什么我不建议一上来就套用网上现成的“性能优化十连”配置。你得先弄清楚自己的问题在哪一层，再对症下药。

2. Profiling：定位瓶颈的第一把手术刀

2.1 PyTorch Profiler 的基本用法与关键参数

PyTorch 自带的torch.profiler已经是我日常工作里最常用的性能分析工具了。它在 PyTorch 1.8 之后成为官方推荐方案，替代了老旧的torch.autograd.profiler，用起来非常简单：

import torch from torch.profiler import profile, ProfilerActivity, record_function def run_training_step(model, batch, optimizer, criterion): optimizer.zero_grad() outputs = model(batch) loss = criterion(outputs, batch["labels"]) loss.backward() optimizer.step() return loss.item() # 用 profiler 包住训练循环 with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), on_trace_ready=torch.profiler.tensorboard_trace_handler("./profile_logs"), record_shapes=True, profile_memory=True, ) as prof: for step in range(10): loss = run_training_step(model, batch, optimizer, criterion) prof.step() # 通知 profiler 一个 step 结束 # 打印统计表 print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

这里有几个参数值得解释一下：

• schedule参数里的wait、warmup、active和repeat是用来控制采样节奏的。wait=1表示前 1 个 step 不记录，这是为了让模型处于稳定状态；warmup=1表示接下来 1 个 step 用于预热，机制是为了让缓存分配、cuDNN autotune 这些一次性开销先被触发掉；active=3才是真正采集数据的 3 个 step。这套组合拳打下来，拿到的是稳定状态下的性能数据，而不是包含各种初始化开销的数据。

• record_shapes=True可以记录每个算子的输入张量形状。这个参数在分析内存分配和算子融合效果时非常关键，但也有一点代价：会让 Profiling 变慢一点。我通常是先不开它跑一遍看整体，再单独开它看某个可疑算子的输入形状。

• profile_memory=True会记录内存分配情况，定位显存瓶颈时很好用，但也会增加开销。实际使用中我一般只在排查显存问题时才打开。

• on_trace_ready可以指定 trace 处理函数。tensorboard_trace_handler会把结果输出到指定目录，之后通过 TensorBoard 的 PyTorch Profiler 插件查看火焰图、算子时间线。这个可视化能力非常强大，能看到每个 kernel 在 GPU 上的执行顺序、间隙、以及 CPU 与 GPU 之间的依赖关系。

注意：Profiling 本身是有开销的，尤其是开启record_shapes和profile_memory之后，采样期间的速度会比正常训练慢不少。所以不要用 Profiling 阶段的数据去衡量真实训练速度，它只用于定位瓶颈。真实速度的评估，还是用简单的time.time()包住整个训练循环更靠谱。

2.2 三张表看懂输出：先看全局，再追细节

prof.key_averages().table()输出是一张很大的表，刚接触的同学往往盯着前几行就懵了。我一般只看三组关键指标：

第一组是Self CPU time total（算子自身的 CPU 耗时，不含子算子）和Self CUDA time total（算子自身的 GPU 耗时）。这两个指标能帮你识别最耗时的算子，是关注的重点。

第二组是CUDA time total和CPU time total的比值。如果所有算子的 GPU 耗时加在一起远小于 CPU 耗时，说明 GPU 在等 CPU 喂数据或者 Launch Kernel，这种情况大概率有 CPU 端的瓶颈。

第三组是Call Stack（调用栈）和Input Shapes。当定位到某个耗时的算子时，点进调用栈看它是从哪个模块发起的，这样才能把问题映射到代码层面。

看表的顺序也很重要：我不建议从耗时最高的算子开始一个一个优化，而是先看整体占比结构。比如某个算子占了 60% 的 GPU 时间，这通常不是什么问题，除非它的理论耗时不该这么高。真正值得警惕的，是那种单个耗时占比不大、但数量极多的小算子——它们加起来的总开销非常可观，而且往往意味着可以做算子融合来消除。

下面这张表我整理了几个关键指标的解读思路，方便你对着自己的 Profiling 结果做对照判断：

2.3 实际案例：一个 Transformer 训练脚本的 Profiling 现象

举一个我实际处理过的例子。有个朋友跑来问我，他的 BERT 微调脚本在 A100 上 GPU 利用率只有 40% 左右，训练 10 万步预计要跑两周，急得不行。

我让他先跑一次 Profiling，结果非常典型：CPU 侧耗时几乎等于 GPU 侧耗时，而且DataLoader的迭代时间占了 CPU 总耗时的 37%。进一步看时间线，GPU 在一段 kernel 执行完之后会有一段很长的空白，然后才等来下一批数据。

问题立刻清楚了：数据加载和预处理跟不上 GPU 的消费速度。再往下细查，他的数据管道里有个自定义的 tokenizer 做了非常重的文本清理操作，但 DataLoader 的num_workers还停留在默认值 0。也就是说，所有预处理都在主进程里跑，GPU 只能在旁边干等着。

解决方案也很直接：

data_loader = DataLoader( dataset, batch_size=32, shuffle=True, num_workers=8, # 从 0 提到 8，让子进程负责数据加载 prefetch_factor=4, # 每个 worker 预取 4 个 batch pin_memory=True, # 锁页内存，加速 CPU 到 GPU 的拷贝 )

改完再测，GPU 利用率直接从 40% 干到了 85% 以上，训练时间预计压缩到一周以内。

这个案例说明了一个朴素的道理：在你往模型结构上折腾之前，先老老实实跑一次 Profiling，看看瓶颈到底在哪一层。很多时候，问题压根出在数据管道，而不是模型本身。

3. torch.compile：一行代码的“静态化”红利

3.1 torch.compile 到底做了什么

PyTorch 2.0 之后最让人兴奋的改动就是torch.compile。它的原理说起来也不算复杂：通过 TorchDynamo 在 Python 字节码层面捕获模型的前向计算图，然后用 TorchInductor 把这套计算图编译成高性能的 GPU kernel（默认生成 Triton 内核），并在这个过程中做算子融合、显存规划、循环优化等一系列自动优化。

你可以把它理解成把“逐行解释执行”的 Python 代码变成“整体编译执行”的 C 代码。对于 Python 这种解释型语言来说，这个收益是巨大的，尤其是当模型里有很多小算子、很多 Python 层循环的时候。

用起来确实是“一行代码”的事情：

import torch model = MyModel().cuda() model = torch.compile(model) # 这行就够了

但这里有个很重要的预期管理：torch.compile不是所有模型都能白捡加速的。它最擅长的是那种算子小而多、Python 逻辑重的模型，比如基于 Transformer 的各类任务；但对于已经高度优化、单算子就是大头（比如大矩阵乘法、卷积）的模型，收益有限，甚至可能因为编译开销导致整体变慢。

3.2 三种编译模式怎么选

torch.compile提供了几种模式，我实际用下来比较推荐的是：

• default：最均衡的选择。编译时间适中，运行速度大部分场景下比 eager 模式快 20%-50%。适合日常开发和验证。
• reduce-overhead：会在 default 的基础上进一步减少 Python 端到 CUDA 端的启动开销，对小型 batch 场景尤其有效。代价是会增加显存占用，因为它会尽量把中间结果留在 GPU 上。
• max-autotune：编译时会花长时间跑各种 Triton kernel 的 autotune，找到当前硬件上最快的实现。训练速度提升最明显，但编译时间可能长达十几分钟，而且显存占用进一步增加。一般只在最终跑实验、确定不改代码的时候用。

我自己的习惯是：开发期间用default，要出最终训练效果了开一次max-autotune。如果你显存比较紧张，可以优先试试reduce-overhead，它往往能在显存和速度之间找到一个不错的平衡点。

还有dynamic=True这个参数值得单独说一嘴。默认情况下torch.compile会按照第一次运行时的输入形状完成优化，如果后续输入形状变化很大，它会重新编译，反而变慢。如果你的输入序列长度不固定（比如 NLP 任务里的变长 batch），建议显式设置dynamic=True，让它提前对动态形状做优化，减少反复编译带来的额外开销。

3.3 torch.compile 的兼容性“暗坑”

torch.compile虽然好用，但远没到“一行代码走天下”的程度。以下几个坑是我实际踩过的：

第一个是动态控制流的兼容问题。Python里常见的if、for、while依赖运行时张量条件、在循环体内动态改变张量形状，这些情况 TorchDynamo 可能捕获失败。不过好消息是它有一套 fallback 机制，会把无法捕获的部分回退到 eager 模式，并打印警告。出现的警告别直接无视，最好逐条看，有些回退会导致模型核心逻辑没有被编译，性能提升会大打折扣。

第二个是不支持某些第三方算子。比如一些自定义 CUDA extension、部分 HuggingFace 算子库里的复杂算子，都可能触发回退或直接报错。遇到这种情况，可以先看报错信息里是否提示“Graph break”，然后针对性地把不兼容的子模块从编译范围里排除：

model = torch.compile(model) model.layers.unsupported_module = model.layers.unsupported_module.to(memory_format=torch.contiguous_format) # 或者直接对目标子模块“原样保留”，只编译其他部分

第三个是编译时间和启动延迟。大模型首次编译可能要几分钟，这本身不是问题，问题在于每次更换机器、更换 GPU 型号之后缓存失效，又得重新编译。目前 PyTorch 会在磁盘上缓存编译结果，但缓存的有效性依赖于各种环境因素，如果频繁切换环境，建议把关键运行都固定在同一套镜像或环境里。

第四个是显存占用会上升。reduce-overhead和max-autotune模式会倾向于把更多中间结果保存在显存中，这在显存紧张的模型中可能直接 OOM。如果碰上了，要么切回default模式，要么显式开启gradient_checkpointing来换取一部分显存空间。

3.4 torch.compile 叠加 Profiling：1+1 > 2 的效果

有人觉得torch.compile之后就不需要 Profiling 了，这是天大的误解。恰恰相反，编译模式下的 Profiling 价值更高，因为它能直接告诉你哪些算子被融合了、哪些算子还残留着。

我一般会这样操作：先跑一次 eager 模式的 Profiling，记录各个算子的耗时占比；然后开torch.compile，再跑一次同样的 Profiling，对比前后算子的耗时变化。如果某类算子耗时在编译后依然占比很高，说明它没有被很好地融合，你可能需要手动改写这部分逻辑，或者调整编译模式。

我之前做一个小型语音模型训练时，eager模式下LayerNorm占了 GPU 耗时 8%，torch.compile之后直接降到 1% 以下，因为被融合进了前一个算子。但某个自定义的 attention mask 算子怎么都融合不动，一看 Input Shapes，发现它在运行时动态生成了 mask 张量，导致编译器的图优化无法覆盖。把 mask 生成挪到模型外部、变成静态输入之后，编译优化终于追上了这个算子，整体速度又快了一截。

这就是 Profiling 和torch.compile配合着用的意义：编译器能帮你省掉一部分人工优化的功夫，但省到什么程度、还剩什么残留，必须靠数据来判断。

4. 分布式扩展：单卡极限之后的必经之路

4.1 什么时候该上分布式，什么时候不该上

很多同学的直觉是“训练速度慢了，多搞几张卡”。但实际上，分布式训练有启动成本、通信成本、调试成本，在模型规模和单卡算力不匹配的情况下，盲目上多卡反而可能更慢。

我个人的判断标准很简单：如果单卡训练时 GPU 利用率已经超过 90%，且单卡显存还有富余，这时候加卡通常能线性加速，收益明显。但如果单卡利用率不到 50%，优先解决的是单卡效率问题，而不是扩展问题。把单卡从 50% 调到 90%，比多上两张卡的效果都好。

另一个判断标准是模型显存。如果你的模型已经大到单卡放不下，哪怕不用追求极致速度，也要考虑分布式训练，至少可以用数据并行拆分显存占用，或者用模型并行/FSDP 来分配模型参数和梯度。

4.2 DDP 的原理与使用要点

PyTorch 最常用的分布式训练方案是DistributedDataParallel（DDP），它的核心思想是：每个进程持有一份完整的模型副本，各自处理不同的数据 batch；前向传播完成后，所有进程通过 all-reduce 通信同步梯度，保证每个副本都更新到相同的参数。

DDP 的使用看起来很简单：

import torch import torch.distributed as dist import torch.multiprocessing as mp def train_worker(rank, world_size): # 初始化进程组 dist.init_process_group( backend="nccl", init_method="tcp://127.0.0.1:23456", rank=rank, world_size=world_size, ) torch.cuda.set_device(rank) model = MyModel().to(rank) model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) dataset = MyDataset() sampler = torch.utils.data.distributed.DistributedSampler( dataset, num_replicas=world_size, rank=rank ) dataloader = DataLoader(dataset, batch_size=32, sampler=sampler) for epoch in range(num_epochs): sampler.set_epoch(epoch) # 保证每个 epoch 数据打乱方式不同 for batch in dataloader: optimizer.zero_grad() outputs = model(batch) loss = criterion(outputs, batch["labels"]) loss.backward() optimizer.step() dist.destroy_process_group() if __name__ == "__main__": world_size = 4 mp.spawn(train_worker, args=(world_size,), nprocs=world_size)

有几个细节我在实战中反复踩过坑：

• DistributedSampler一定要配合sampler.set_epoch(epoch)使用，否则每个 epoch 的数据顺序不变，模型容易过拟合到固定的 batch 排列上。
• batch_size指的是每个 GPU 上的 batch size，不是全局 batch size。总 batch size 是单卡 batch_size * world_size。这也意味着如果你原本单卡训练用的 batch size 是 64，上 4 卡之后如果每卡还是 64，全局 batch size 就变成 256，学习率通常要跟着调大一些（线性缩放法则），否则收敛效果可能变差。
• DDP 的效率瓶颈往往在梯度同步的通信开销上，尤其是卡间通信带宽有限的时候。如果 4 卡加速比远低于 3.5 倍，先看NCCL_DEBUG=INFO的日志，确认通信是否成了瓶颈。

4.3 FSDP：当模型大到单卡放不下时

DDP 虽然好用，但每个进程都要持有完整的模型副本。当模型参数+优化器状态+梯度超过单卡显存时，DDP 就无能为力了。这时候就要考虑FullyShardedDataParallel（FSDP）。

FSDP 的核心思路是把模型参数、梯度和优化器状态分片到多个 GPU 上，每个 GPU 只保存自己那部分。计算某个算子时需要用到完整的参数，再通过通信把分片聚合起来。这个思想和大模型预训练里常见的 ZeRO 策略一脉相承。

FSDP 的使用方式也很简单：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP model = MyBigModel().to(rank) model = FSDP( model, sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, )

sharding_strategy有三个主要选项：FULL_SHARD（参数、梯度、优化器状态全分片，显存最省但通信最重）、SHARD_GRAD_OP（只分片梯度和优化器状态，参数完整复制，通信稍轻）、NO_SHARD（相当于 DDP，不分片）。

实际选型时我的经验是：能上单卡尽量单卡，单卡放不下且参数在 1B 以下，优先考虑SHARD_GRAD_OP；参数超过 1B，再考虑FULL_SHARD，但这时候通信开销会显著拉低训练效率，最好配合梯度累积来摊薄通信成本。

FSDP 和 DDP 在性能表现上还有个差异：FSDP 的通信是穿插在前向反向计算过程中的，不是像 DDP 那样在反向结束之后统一同步。所以 FSDP 对计算和通信的重叠要求更高，配置不好容易出现“GPU 在等通信”的情况。这就需要借助 Profiling 来看时间线，定位是否存在大段的通信等待间隔。

4.4 分布式性能排查的常见手段

分布式训练出了问题，定位起来比单卡复杂一个量级。我常用的排查手段有三板斧：

第一板斧是看NCCL_DEBUG=INFO日志。NCCL 是英伟达多卡通信库，打开这个环境变量之后，它会输出每个通信算子的耗时、带宽、使用的通信算法等信息。重点看有没有timed out、no peer access、connect to x.x.x.x failed之类的异常。

第二板斧是跑一次torch.distributed.all_reduce的纯通信压力测试，看多卡之间的实际带宽是否正常。如果通信带宽远低于硬件理论值，要么是网络配置问题，要么是通信库版本不匹配，这和模型代码无关。

第三板斧是在多个 worker 上开启 Profiling。分布式场景下，不要只在 rank 0 上跑 profiler，最好让所有 rank 都输出 trace，然后在 TensorBoard 里把时间线对齐看。如果发现某个 rank 的算子执行时间明显偏长，大概率是负载不均衡；如果所有 rank 都在等待同样长度的通信区间，那就是批量太小、通信占比太高。

分布式训练还有一个容易被忽略的“减速带”：checkpoint 保存。如果每个 epoch 都把模型保存到共享文件系统，而模型够大、分布式文件系统带宽不够，那么保存阶段可能比训练本身还慢。我的建议是优先保存 rank 0 的模型即可，并且可以异步保存，不要阻塞训练主循环。

5. 常见问题与调优实战速查

5.1 问题排查速查表

我在项目里积累了一份“踩坑速查表”，每次调优遇到问题先对着查一遍，解决了不少看起来莫名其妙的问题。这里直接分享出来：

5.2 我最常用的一套调优工作流

如果你是新项目，我的建议是不要一上来就把所有调优手段全堆上。这样出了问题你都分不清是谁导致的。推荐按照下面的顺序逐步叠加，每加一步验证一次效果：

1. 先用默认配置跑通一个小数据集，得到 baseline 数据。
2. Profiling 单卡训练，重点解决 GPU 利用率低的问题（数据加载、同步、小算子）。
3. 开启torch.compile(default)，对比 Profiling 前后的算子耗时变化。
4. 如果模型放不下单卡，再上 DDP 或 FSDP，从 2 卡开始验证加速比，确认无异常再加到预期卡数。
5. 训练稳定后，再考虑max-autotune、梯度累积、混合精度等进一步优化。

这套工作流我用了很多次，基本没有翻过车。关键是它保证每一步的变量都足够小，出了问题可以快速二分定位。

5.3 一个“调优工具箱”级的配置参考

最后给大家一套我现在做视觉模型训练时常用的基准配置，可以作为你调参的起点。注意这些不是银弹，不同任务需要微调，但作为参考还是很香的：

# 训练环境相关 torch.backends.cudnn.benchmark = True # 输入 shape 固定时开启，自动选最优卷积算法 torch.backends.cuda.matmul.fp16_allow_bf16_reduced_precision_reduction = True # 混合精度矩阵乘的精度/速度权衡 # DataLoader dataloader = DataLoader( dataset, batch_size=64, num_workers=8, pin_memory=True, prefetch_factor=4, persistent_workers=True, # epoch 间不销毁 worker 进程，减少重建开销 ) # 模型 model = torch.compile(model, mode="reduce-overhead") # 训练 scaler = torch.cuda.amp.GradScaler() # 配合 amp 自动混合精度

混合精度（AMP）这个点值得单独强调一下。开启torch.cuda.amp.autocast配合GradScaler，在大多数模型上能直接带来 1.5-2 倍的提速，显存占用也会明显下降。代价是精度可能掉一点，但对绝大多数任务来说影响很小，属于“性价比”极高的优化。

如果你用的是较新版本的 PyTorch，比如 2.3 之后，AMP 已经默认开启了一些优化，但显式控制GradScaler的使用还是能帮你避开一些混合精度下的梯度下溢问题。

6. 实操中的心得体会

最后分享几个我做性能调优这几年的个人体会，不一定系统，但都很实用。

第一，性能调优的回报曲线是阶梯式的。不是说你每投入一小时调优，都能稳定换到等比例的加速。很多时候，把最明显的那个瓶颈打通之后，速度会跳上一个台阶，然后进入平台期。这时候不要盲目继续抠细节，而是要先问一句：现在的速度够用了吗？实验迭代的时间成本能不能接受？如果已经够用，停手是对的，继续调优的边际收益可能很低。

第二，要养成“调参前先备份”的习惯。尤其是跑分布式训练的时候，一次错误的 batch size 调整或者学习率调整，可能会浪费好几天的实验时间。我一般会给每次实验记录一个配置文件，包含模型结构、batch size、学习率、优化器参数、编译模式、混合精度设置。这样即使效果崩了，也能快速回到之前的稳定点。

第三，善用社区经验，但别盲信。网上很多性能调优的教程和配置看起来很有道理，但硬件环境、模型结构、数据特性不一样，效果可能完全相反。我在实操中就遇到过同样一个torch.compile配置，在自己的 3090 上加速了 40%，换到 A100 上反而慢了 8%。所以任何配置都值得在自己环境里用 Profiling 结果验证一遍。

第四，别忽略版本的隐性问题。PyTorch 的迭代速度极快，同样的代码在不同小版本甚至不同 CUDA 版本下的性能差异可能非常大。如果你遇到奇怪的性能问题，先检查版本配套关系，说不定只是环境不匹配导致的。更新或回退版本之前，务必备份当前可用环境，别切完之后发现新的环境还不如旧的。

性能调优这件事，说难也难，说简单也简单。难在它是“系统性工程”，涉及数据管道、计算框架、硬件资源、算法结构多个层面；说简单，是因为只要按照“测量-分析-优化-再测量”的循环一步步来，绝大多数性能问题都能被定位和解决。希望这篇笔记能帮你在 PyTorch 性能调优的路上少踩几个坑，把宝贵的实验时间花在更有意义的地方。