ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

显存优化后GPU利用率仍低?LLM训练性能瓶颈排查与调优指南

显存优化后GPU利用率仍低?LLM训练性能瓶颈排查与调优指南

1. 显存省下来了,GPU 为什么还是吃不满

显存优化做到位了,nvidia-smi一看显存占用从 22GB 降到 9GB,心里美滋滋,结果watch -n 0.5 nvidia-smi盯了十分钟,GPU 利用率死活卡在 35% 到 50% 之间来回跳,功耗也上不去,风扇都不怎么转。这个场景我见过太多次了,不管是自己搭的 LLM Training Lab,还是帮别人看训练脚本,十有八九都会撞上这个坎。

先把结论摆出来:显存和算力是两条独立的流水线,显存省下来只代表你能塞进更大的 batch 或者更长的序列,不代表 SM(流式多处理器)就有活干。GPU 吃不满,本质上是数据供给、计算调度、通信同步这三条链路里至少有一条堵住了。显存优化解决的是"装得下"的问题,而 GPU 利用率解决的是"喂得饱"的问题,两者根本不是一回事。

这篇内容适合谁看?如果你正在用 PyTorch 跑 LLM 训练,已经折腾过 gradient checkpointing、ZeRO、LoRA、量化这些显存优化手段,但发现 GPU 利用率上不去、训练速度没提升、step time 反而变长了,那这篇就是写给你的。我会从数据管线、算子调度、通信开销、CPU-GPU 协同这几个角度,把"显存省了但 GPU 吃不满"这件事拆开讲清楚,每个环节都给出可复现的排查方法和调优参数。

需要提前说明的是,下面涉及的具体参数和配置,一部分来自我自己的实验记录,一部分是基于 PyTorch 官方文档和常见工程实践做的合理推演,你在自己环境里复现时,数值可能需要根据实际硬件做微调。

2. 先搞清楚 GPU 利用率到底在测什么

2.1 nvidia-smi 的利用率是个"谎言"

很多人把nvidia-smi里的GPU-Util当成 GPU 的真实负载,这个理解有偏差。这个百分比的含义是:在采样周期内,至少有一个 kernel 在 SM 上处于活跃状态的时间占比。注意关键词——"至少一个"。哪怕你的 GPU 有 132 个 SM,只有 1 个 SM 在跑一个很小的 kernel,利用率照样显示 100%。

这就解释了为什么有时候你看到利用率 95%,但训练速度还是很慢——因为大部分 SM 在摸鱼,只有少数几个在干活。反过来,利用率 40% 也不一定代表 GPU 闲,可能是 kernel 之间有间隙,间隙里 SM 确实没活干。

真正能反映 GPU 是否"吃饱"的指标,是这几个:

  • SM Occupancy(SM 占用率):每个 SM 上活跃 warp 数量与理论最大值的比值,用 Nsight Compute 或者torch.cuda的 profiler 能看
  • Tensor Core 利用率:如果你的模型用了混合精度,这个指标比 GPU-Util 有意义得多
  • Memory Throughput:显存带宽利用率,LLM 训练里这个经常是真正的瓶颈
  • Power Draw:功耗上不去,基本可以断定 GPU 没吃满

我一般会同时开三个窗口:一个nvidia-smi -l 1看粗粒度,一个nvidia-smi dmon -s pucvmet看细粒度,再挂一个 PyTorch profiler 抓 timeline。三个对不上,问题基本就定位了。

2.2 显存优化和算力利用是两码事

这里要澄清一个常见的认知误区。显存优化手段大致分几类:

优化手段省的是什么对 GPU 利用率的影响
Gradient Checkpointing激活值显存增加计算量,可能降低利用率
ZeRO / FSDP优化器状态、梯度分片增加通信,可能降低利用率
LoRA / QLoRA可训练参数量减少计算量,利用率可能下降
量化(INT8/INT4)权重和激活显存需要反量化,可能增加开销
Flash AttentionAttention 中间矩阵通常提升利用率
梯度累积等效 batch 显存不改变单步利用率

看出来了吧?大部分显存优化手段,要么增加计算量,要么增加通信量,要么减少单步计算量。它们让你能跑更大的模型,但不保证跑得更快。Gradient checkpointing 是最典型的例子——它用 30% 左右的额外计算换显存,如果你的瓶颈本来就在计算,那开了它只会更慢。

所以"显存省下来了,GPU 还是吃不满"这个现象,很多时候不是 bug,而是这些优化手段的固有代价。真正要做的,是找到当前配置下的真实瓶颈,然后针对性解决。

3. 数据管线:最常见的"饿死"GPU 的元凶

3.1 DataLoader 的 num_workers 和 pin_memory

我敢说,LLM 训练里 GPU 吃不满,一半以上的原因是 DataLoader 喂数据太慢。尤其是当你做指令微调、多轮对话训练的时候,数据预处理(tokenize、padding、mask 构造)的耗时经常被低估。

先看几个关键参数:

from torch.utils.data import DataLoader dataloader = DataLoader( dataset, batch_size=8, num_workers=8, # 关键参数 pin_memory=True, # 关键参数 prefetch_factor=4, # 关键参数 persistent_workers=True,# 关键参数 shuffle=True, drop_last=True, )

num_workers设多少合适?经验公式是CPU 物理核心数 / GPU 数量,但实际要看数据预处理复杂度。如果你的__getitem__里要做复杂的 tokenize 和 mask 构造,可以设到核心数的 1.5 倍。设太多反而会因为进程切换开销导致变慢。

pin_memory=True的作用是把数据放到锁页内存(pinned memory),这样 CPU 到 GPU 的拷贝可以用 DMA 直接传输,速度能快 2 到 3 倍。这个参数很多人忘了开,开了之后 H2D 拷贝时间能明显下降。

prefetch_factor控制每个 worker 预取多少个 batch,默认是 2。在数据预处理慢的场景下,调到 4 到 8 能有效平滑数据供给。

persistent_workers=True让 worker 进程在 epoch 之间不销毁重建,省掉每次 epoch 开始的进程启动开销。数据量小、epoch 多的时候效果明显。

3.2 用 profiler 确认数据管线是不是瓶颈

光看参数没用,得实测。PyTorch profiler 能直接告诉你时间花在哪:

import torch from torch.profiler import profile, ProfilerActivity with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2), on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'), record_shapes=True, profile_memory=True, ) as prof: for step, batch in enumerate(dataloader): if step >= 10: break outputs = model(**batch) loss = outputs.loss loss.backward() optimizer.step() optimizer.zero_grad() prof.step()

跑完之后看 trace,重点看DataLoader相关的 CPU 操作和cudaMemcpyAsync之间的间隙。如果 timeline 上 GPU 有明显的空白段,而空白段对应的 CPU 侧在跑__getitem__或者 tokenize,那基本可以确诊是数据管线瓶颈。

我踩过的一个坑:数据集用了 HuggingFace 的datasets库,map的时候没设num_proc,结果 tokenize 是单进程跑的,训练时每个 batch 都要现场 tokenize,GPU 等得花儿都谢了。后来改成预处理阶段用num_proc=16一次性 tokenize 完,训练时直接读 token id,利用率从 45% 拉到 85%。

3.3 数据格式和 IO 的隐藏开销

数据存成什么格式,对读取速度影响巨大。我做过对比测试,同样 100 万条指令数据:

存储格式单次读取耗时备注
JSON Lines120ms需要逐行解析
Parquet35ms列式存储,压缩好
Arrow / HuggingFace datasets18ms内存映射,零拷贝
预 tokenize 的 numpy memmap8ms最快,但灵活性差
WebDataset (tar shards)12ms适合大规模分布式

如果你的数据还在用 JSON Lines 存,每次__getitem__都要json.loads,那 GPU 不饿才怪。建议至少换成 Parquet 或者 Arrow 格式,预处理阶段把 tokenize 做完,训练时只做 tensor 拼接。

还有一个容易忽略的点:数据在机械硬盘上还是 SSD 上。LLM 训练数据动辄几十上百 GB,如果放在机械盘上,随机读取的 IOPS 根本跟不上。我见过一个案例,训练脚本本身没问题,就是数据放在了一块老机械盘上,GPU 利用率死活上不去,换到 NVMe SSD 之后直接翻倍。

4. 算子调度与 kernel 间隙:GPU 在等 CPU 发号施令

4.1 Python 解释器和 GIL 的锅

PyTorch 的训练循环是 Python 写的,每个 step 里 CPU 要做的事情包括:从 DataLoader 取 batch、把 tensor 搬到 GPU、构造计算图、启动各个 kernel、更新优化器状态。这些操作里,kernel 启动是异步的,但 Python 侧的调度是同步的。

当模型比较小、单个 kernel 执行时间很短的时候,CPU 启动 kernel 的速度可能跟不上 GPU 执行的速度,GPU 就会在 kernel 之间出现间隙。这就是所谓的kernel launch bound。

判断方法:用 Nsight Systems 抓 timeline,看 GPU 上 kernel 之间是否有明显空隙。如果有,且空隙对应的 CPU 侧在跑 Python 代码,那就是 launch bound。

缓解手段有几个:

  • 增大 batch size:每个 kernel 执行时间变长,相对间隙占比下降
  • 使用 CUDA Graph:把整个 step 的 kernel 序列捕获成一个图,一次启动,消除逐个 kernel 的启动开销
  • torch.compile:PyTorch 2.0 之后的编译器,能把 Python 侧的调度开销大幅降低
  • 减少 CPU 侧同步操作:比如.item()、.cpu()、print(loss)这些都会强制同步

CUDA Graph 在 LLM 训练里效果特别明显。我实测过一个 7B 模型的训练,开 CUDA Graph 之前 step time 是 1.8s,开了之后降到 1.3s,GPU 利用率从 55% 拉到 78%。用法也不复杂:

# 简单示意,实际使用需要处理动态 shape 和随机性 g = torch.cuda.CUDAGraph() static_input = torch.zeros_like(batch) with torch.cuda.graph(g): static_output = model(static_input) loss = static_output.loss loss.backward() optimizer.step() # 训练循环里 static_input.copy_(batch) g.replay()

注意 CUDA Graph 对动态 shape 不友好,LLM 训练里如果序列长度变化大,需要做 padding 到固定长度,或者用分段捕获。

4.2 小算子过多导致的调度开销

LLM 里有很多逐元素操作(element-wise),比如 LayerNorm、激活函数、残差连接、dropout。这些操作单独看都很快,但数量多了之后,kernel 启动开销和显存访问开销会累积。

一个典型的 Transformer block 里,逐元素操作的 kernel 数量可能占到 60% 以上,但计算量占比不到 10%。这就是为什么算子融合(operator fusion)在 LLM 训练里这么重要。

torch.compile能自动做一部分融合,但效果取决于后端。更彻底的做法是用融合 kernel,比如:

  • Flash Attention:把 attention 的多个步骤融合成一个 kernel
  • Fused LayerNorm:Apex 或者 xformers 提供的融合实现
  • Fused Adam:把优化器更新融合成少量 kernel

我一般会先上torch.compile,看能提升多少,如果还不够再手动替换关键算子。torch.compile的用法:

model = torch.compile(model, mode="max-autotune")

max-autotune模式会花更多编译时间,但生成的 kernel 通常更快。第一次跑会慢,因为要编译,之后就好了。注意编译缓存要保留,不然每次重启都要重新编译。

4.3 同步点:那些悄悄拖慢训练的调用

训练循环里有些操作会强制 CPU 和 GPU 同步,每次同步都会让 GPU 停下来等 CPU。常见的同步点包括:

  • loss.item()或者loss.cpu():为了打印 loss
  • print(tensor):会触发 tensor 的 repr,需要同步
  • if loss > threshold:条件判断需要把 GPU 上的值取回 CPU
  • torch.cuda.synchronize():显式同步
  • 某些自定义的 logging 逻辑

我见过最离谱的一个案例:训练循环里每个 step 都print(loss.item()),而且 print 还带了一堆格式化字符串。GPU 每个 step 都要停下来等 Python 打印完,利用率直接腰斩。改成每 100 步记录一次,且用loss.detach()避免保留计算图,利用率立刻上去了。

正确的做法是:把需要同步的操作攒起来,批量处理。比如每 100 步记录一次 loss,用torch.cuda.Event做计时,日志写到内存里最后统一落盘。

5. 通信开销:分布式训练里的隐形杀手

5.1 数据并行下的 all-reduce

如果你用 DDP(DistributedDataParallel)做数据并行,每个 step 结束都要做一次梯度 all-reduce。这个通信量和模型参数量成正比,7B 模型用 FP16 就是 14GB 的通信量。如果网络带宽不够,GPU 就会在等通信。

判断方法:看 timeline 上 all-reduce kernel 的耗时占比。如果超过 20%,说明通信是瓶颈。

缓解手段:

  • 梯度累积:多个 step 才做一次 all-reduce,通信频率降低
  • 通信与计算重叠:DDP 默认会尝试 overlap,但效果取决于 bucket 大小
  • 调整 bucket_cap_mb:DDP 的bucket_cap_mb参数控制梯度分桶大小,默认 25MB,调大能减少通信次数但增加显存
  • 用 NCCL 的调优参数:比如NCCL_ALGO、NCCL_PROTO
model = DDP( model, device_ids=[local_rank], bucket_cap_mb=100, # 调大 bucket gradient_as_bucket_view=True, # 省显存 static_graph=True, # 如果模型结构固定,开启能优化 )

static_graph=True这个参数很多人不知道,开启后 DDP 能做一些图级别的优化,通信效率能提升 10% 到 20%。前提是你的模型在训练过程中计算图不变,LLM 训练一般满足。

5.2 ZeRO / FSDP 的通信代价

ZeRO 和 FSDP 通过分片优化器状态、梯度、参数来省显存,但代价是更多的通信。ZeRO-3 每个 step 都要 all-gather 参数,通信量比 DDP 大好几倍。

如果你的显存优化用的是 ZeRO-3 或者 FSDP 的 full shard,GPU 吃不满很可能是通信瓶颈。这时候可以考虑:

  • 降级到 ZeRO-2:只分片优化器状态和梯度,通信量小很多
  • 混合分片:FSDP 支持sharding_strategy配置,可以选SHARD_GRAD_OP而不是FULL_SHARD
  • 增大计算通信比:增大 batch size 或者序列长度,让每次通信对应的计算量更大

我个人的经验是:单机多卡场景下,ZeRO-2 通常比 ZeRO-3 更划算,除非模型大到 ZeRO-2 装不下。多机场景下通信更贵,要更谨慎地选择分片策略。

5.3 通信和计算的 overlap

理想情况下,通信应该和计算重叠,GPU 在等通信的时候应该还有别的活干。但实际中 overlap 效果往往不理想,原因包括:

  • 通信和计算依赖同一个 stream,无法并行
  • bucket 划分不合理,通信粒度太细
  • 网络带宽本身就不够

PyTorch 2.x 之后,FSDP 的 overlap 做得比之前好,但还是要看具体配置。可以用torch.profiler看通信 kernel 和计算 kernel 是否真的重叠了。如果 timeline 上通信和计算是串行的,那 overlap 没生效,需要检查配置。

6. 实操排查流程:从现象到根因

6.1 第一步:建立基线

在开始调优之前,先建立一个干净的基线。用一个最简单的配置跑一遍,记录以下指标:

# 终端 1:粗粒度监控 nvidia-smi -l 1 # 终端 2:细粒度监控 nvidia-smi dmon -s pucvmet -d 1 # 终端 3:训练脚本 python train.py 2>&1 | tee baseline.log

基线配置建议:

  • batch size 设为能跑起来的最小值
  • 关闭所有显存优化(gradient checkpointing、ZeRO 等)
  • 关闭混合精度
  • 数据管线用最简单的配置

记录 step time、GPU 利用率、显存占用、功耗。这个基线是后续所有对比的参照。

6.2 第二步:逐项开启优化,观察变化

每次只开一个优化,观察指标变化。我一般按这个顺序:

  1. 开混合精度(AMP),看 step time 和利用率
  2. 开 Flash Attention,看 attention 部分耗时
  3. 开 gradient checkpointing,看显存和 step time 的 trade-off
  4. 开 torch.compile,看整体提升
  5. 开分布式(如果有),看通信开销

每开一项,记录 step time、GPU 利用率、显存占用。如果某一项开了之后利用率反而下降,说明这项优化在当前配置下不划算,要么调整参数,要么放弃。

6.3 第三步:用 profiler 定位瓶颈

基线建立之后,用 profiler 抓一次详细的 timeline:

import torch from torch.profiler import profile, ProfilerActivity, tensorboard_trace_handler def train_step(model, batch, optimizer): outputs = model(**batch) loss = outputs.loss loss.backward() optimizer.step() optimizer.zero_grad() return loss with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=torch.profiler.schedule(wait=2, warmup=2, active=6, repeat=1), on_trace_ready=tensorboard_trace_handler('./profiler_log'), record_shapes=True, profile_memory=True, with_stack=True, ) as prof: for step, batch in enumerate(dataloader): if step >= 20: break loss = train_step(model, batch, optimizer) prof.step()

跑完之后用 TensorBoard 打开./profiler_log,看 timeline。重点看:

  • GPU 上是否有空白段
  • 空白段对应的 CPU 侧在做什么
  • 哪些 kernel 耗时最长
  • 通信 kernel 占比多少
  • 显存分配和释放的模式

6.4 常见瓶颈速查表

现象可能原因排查方法解决方向
GPU 利用率低,timeline 有空白数据管线慢看 DataLoader 耗时增 num_workers、预处理数据
GPU 利用率低,kernel 间隙大kernel launch bound看 kernel 间空隙CUDA Graph、torch.compile
GPU 利用率低,通信 kernel 占比高通信瓶颈看 all-reduce 耗时梯度累积、调 bucket、换分片策略
GPU 利用率高但速度慢小算子过多看 kernel 数量和耗时算子融合、torch.compile
显存够但 batch 上不去激活值显存看显存分配gradient checkpointing
step time 波动大数据或 IO 不稳定看 step time 分布换存储、预处理
功耗上不去GPU 没吃满看 power draw综合排查上述各项

7. 几个容易被忽略的细节

7.1 混合精度里的 dtype 转换开销

AMP(自动混合精度)能省显存、加速计算,但 dtype 转换本身有开销。如果模型里频繁在 FP16 和 FP32 之间转换,开销会累积。

torch.cuda.amp.autocast会自动管理转换,但有些操作(比如 LayerNorm、softmax)在 FP16 下数值不稳定,需要保持 FP32。这些"例外"操作会导致额外的转换。

我一般会用torch.compile配合 AMP,让编译器做更激进的融合和转换优化。另外,bfloat16在数值稳定性上比float16好,如果硬件支持(Ampere 及以上),优先用bfloat16。

from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() with autocast(dtype=torch.bfloat16): outputs = model(**batch) loss = outputs.loss scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

注意bfloat16不需要 GradScaler,因为它的动态范围和 FP32 一样。用float16才需要。

7.2 优化器状态的内存布局

Adam 优化器每个参数要存两个状态(一阶矩和二阶矩),如果用 FP32 存,显存占用是参数量的 8 倍。7B 模型就是 56GB,这也是为什么大家要用 ZeRO 或者 8-bit Adam。

但优化器状态的内存布局也会影响性能。如果状态张量在显存里是碎片化的,访问效率会下降。PyTorch 的foreach实现会把多个参数的状态合并操作,减少 kernel 数量:

optimizer = torch.optim.AdamW( model.parameters(), lr=1e-4, foreach=True, # 开启 foreach 实现 fused=True, # 如果可用,用融合 kernel )

foreach=True在参数多的时候效果明显,能把优化器更新的 kernel 数量从几百个降到几个。fused=True需要 CUDA 支持,效果更好但兼容性差一些。

7.3 梯度累积和 batch size 的等效关系

梯度累积是省显存的常用手段,但要注意它和 batch size 的等效关系。假设你想要的等效 batch size 是 64,但显存只够放 8,那就累积 8 次:

accumulation_steps = 8 for step, batch in enumerate(dataloader): outputs = model(**batch) loss = outputs.loss / accumulation_steps loss.backward() if (step + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()

注意 loss 要除以accumulation_steps,否则梯度会放大。另外,BatchNorm 这类依赖 batch 统计的层,在梯度累积下行为和真实大 batch 不一样,LLM 里一般用 LayerNorm 所以没这个问题。

梯度累积的代价是:每个 micro-step 都要做一次前向和反向,但优化器更新频率降低。如果瓶颈在优化器更新,梯度累积能加速;如果瓶颈在前向反向计算,梯度累积不改变单步计算量,利用率不会提升。

7.4 显存碎片和 caching allocator

PyTorch 的 CUDA caching allocator 会缓存显存块,减少cudaMalloc调用。但缓存策略不当会导致碎片,表现为"显存够但分配失败"。

可以用环境变量调整:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

expandable_segments让 allocator 能扩展已分配的段,减少碎片。PyTorch 2.1 之后支持,效果不错。

另外,torch.cuda.memory_summary()能打印详细的显存分配情况,包括碎片率。如果碎片率高,可以考虑定期torch.cuda.empty_cache(),但注意这会释放缓存,下次分配又要重新cudaMalloc,有开销,不要频繁调用。

8. 一个完整的调优案例

8.1 初始状态

假设我们有一个 7B 模型的指令微调任务,单卡 A100 80GB,初始配置:

  • batch size 4,序列长度 2048
  • FP32 训练
  • 无 gradient checkpointing
  • DataLoader num_workers=4
  • 无 torch.compile

实测:step time 3.2s,GPU 利用率 42%,显存占用 68GB,功耗 180W。

8.2 逐项优化

第一步:开 AMP(bfloat16)

step time 降到 2.1s,利用率 58%,显存降到 42GB,功耗 240W。提升明显。

第二步:开 Flash Attention

step time 降到 1.7s,利用率 65%,显存降到 38GB。Attention 部分耗时从 0.6s 降到 0.25s。

第三步:DataLoader 调优

num_workers 从 4 调到 12,pin_memory=True,prefetch_factor=4。step time 降到 1.5s,利用率 72%。timeline 上数据管线的空白段基本消失。

第四步:torch.compile

step time 降到 1.2s,利用率 80%。编译花了 3 分钟,但之后每个 step 都受益。

第五步:CUDA Graph

step time 降到 1.0s,利用率 88%。注意序列长度要 padding 到固定值。

第六步:优化器调优

AdamW 开 foreach 和 fused,step time 降到 0.95s,利用率 90%。

最终:step time 从 3.2s 降到 0.95s,加速 3.4 倍,GPU 利用率从 42% 拉到 90%,功耗 320W。

8.3 关键观察

这个案例里,显存优化(AMP、Flash Attention)和利用率优化(DataLoader、torch.compile、CUDA Graph)是交替进行的。每一步都观察指标,确认有效再进入下一步。如果某一步没效果甚至变慢,就回退。

另外注意,优化是有顺序的。如果先开 CUDA Graph 再调 DataLoader,CUDA Graph 的效果会被数据管线的空白段掩盖,看起来提升不大。正确的顺序是先解决数据供给,再解决 kernel 调度。

9. 一些反直觉的经验

9.1 有时候降低 batch size 反而更快

听起来反直觉,但在某些配置下确实如此。如果 batch size 太大导致显存压力大,caching allocator 频繁触发碎片整理,或者 gradient checkpointing 被强制开启,那降低 batch size 反而能让每个 step 更快。

我遇到过一个大 batch 导致 attention 的中间矩阵超出 SRAM 容量,Flash Attention 退化到慢路径的情况。batch 减半之后,attention 走快路径,整体速度反而提升。

9.2 不是所有优化都要开

gradient checkpointing 省显存但增加计算,如果你的显存本来就够,开了只会变慢。同理,ZeRO-3 在多机场景下通信开销巨大,如果模型能塞进单机,用 DDP 更划算。

优化的原则是:先测量,再优化,只优化瓶颈。不要因为某个技术"听起来很厉害"就无脑上。

9.3 监控本身也有开销

nvidia-smi -l 1每秒采样一次,开销可以忽略。但如果你在训练循环里频繁调用torch.cuda.memory_allocated()或者torch.cuda.utilization(),这些调用会触发同步,拖慢训练。

监控要用外部工具,不要在训练循环里做。日志记录也要攒批,不要每个 step 都写盘。

10. 工具链和版本选择

10.1 PyTorch 版本

PyTorch 2.x 之后,torch.compile和 FSDP 的成熟度大幅提升。如果还在用 1.x,建议升级。但注意 CUDA 版本和驱动的兼容性,升级前先在测试环境验证。

我目前用的组合是 PyTorch 2.3 + CUDA 12.1 + cuDNN 8.9,稳定性不错。torch.compile在这个版本下对 LLM 的支持已经比较完善。

10.2 辅助工具

  • Nsight Systems:抓 timeline,看 kernel 和通信
  • Nsight Compute:分析单个 kernel 的性能
  • PyTorch Profiler:集成在 PyTorch 里,方便
  • Weights & Biases / TensorBoard:记录训练指标
  • nvitop:比 nvidia-smi 更友好的监控工具

nvitop我强烈推荐,界面清晰,能同时看多个 GPU 的利用率、显存、功耗、温度,还能看进程级别的占用。安装pip install nvitop,运行nvitop就行。

10.3 环境变量调优

# NCCL 调优 export NCCL_ALGO=Ring # 或 Tree,取决于拓扑 export NCCL_PROTO=Simple # 或 LL,取决于消息大小 export NCCL_NSOCKS_PERTHREAD=4 export NCCL_SOCKET_NTHREADS=2 # CUDA 调优 export CUDA_DEVICE_MAX_CONNECTIONS=1 # 减少 stream 数量,降低调度开销 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True # 线程调优 export OMP_NUM_THREADS=8 # 根据 CPU 核心数调整

CUDA_DEVICE_MAX_CONNECTIONS=1这个参数在 kernel launch bound 的场景下有效,能减少 stream 之间的切换开销。但注意它可能影响通信和计算的重叠,要实测。

11. 最后再分享几个实操技巧

第一个技巧:用 dummy data 做基准测试。在排查数据管线问题的时候,把 DataLoader 换成一个生成随机 tensor 的 dummy loader,如果 GPU 利用率立刻上去了,那问题就在数据侧。这个方法能快速排除数据管线的嫌疑。

第二个技巧:二分法定位瓶颈。如果 timeline 很复杂,不知道从哪下手,可以把模型截断,只跑前几层,看利用率。如果前几层利用率就低,问题在数据或调度;如果前几层正常,后面低,问题在模型结构或通信。

第三个技巧:关注 step time 的方差,不只看均值。如果 step time 波动大,说明有间歇性的瓶颈,比如数据 IO 抖动、GC 暂停、其他进程干扰。用torch.cuda.Event记录每个 step 的耗时,画个分布图,比看均值有用得多。

第四个技巧:显存优化和利用率优化分开做。先把显存压到能跑,再单独调利用率。两个一起调,变量太多,很难定位问题。

第五个技巧:保留一个能跑的最简配置。调优过程中经常会改坏,保留一个 baseline 脚本,随时能回退对比。我一般会把 baseline 的 commit hash 记下来,出问题就git checkout回去。

这些经验都是踩坑踩出来的,你在自己环境里复现的时候,硬件、驱动、库版本不同,具体数值会有差异,但排查思路是通用的。核心就一句话:先测量,再优化,只优化瓶颈,每次只改一个变量。

返回列表