1. 为什么一个深度学习框架值得反复折腾刚接触深度学习那会儿，我对框架的选择其实没什么概念，别人说用什么就用什么。直到有一次跑一个图像分类的小项目，训练到一半显存爆了，报错信息看得我一头雾水，才意识到框架的底层机制不是“能跑就行”这么简单。后来花了不少时间啃PyTorch的文档和源码，踩了无数坑，才慢慢摸清楚它的脾气。这篇文章不是入门教程，市面上入门教程已经够多了，我想聊的是那些文档里不会明说、但实际项目中一定会遇到的东西——动态图的真实代价、显存管理的底层逻辑、数据加载的隐藏瓶颈、混合精度的适用边界，以及怎么把这些东西组合成一套能稳定跑起来的训练流程。PyTorch本质上是一个基于动态计算图的深度学习框架，核心能力是自动求导和GPU加速张量运算。它能做什么？从简单的线性回归到大规模分布式训练，从计算机视觉到自然语言处理，几乎覆盖了深度学习的所有主流场景。适合谁来参考？如果你已经能写基本的训练循环，但在实际项目中经常遇到显存不够、训练速度慢、loss不收敛、多卡跑不起来这些问题，那这篇内容应该能帮你省下不少排查时间。如果你刚入门，也可以看看，至少知道后面会遇到什么。我个人的习惯是，学一个框架不能只学API怎么调，得理解它为什么这么设计。PyTorch选择动态图而不是静态图，这个决策背后有一整套权衡，理解了这些权衡，很多“奇怪”的行为就变得合理了。2. 动态图机制的真实代价与应对策略2.1 动态图到底“动”在哪里PyTorch最核心的设计就是动态计算图。什么意思？每次前向传播的时候，计算图是实时构建的。你写一个for循环，循环次数不同，计算图的结构就不同。这跟静态图框架（先定义图再运行）有本质区别。动态图的好处很直观：调试方便。你可以在forward函数里随便加print，可以用Python的if-else控制流，可以动态改变网络结构。写起来跟普通Python代码没区别，心智负担极低。但代价是什么？每次迭代都要重新构建计算图，这意味着额外的CPU开销。对于小模型，这个开销可能比GPU计算本身还大。我实测过一个简单的全连接网络，batch size设成32的时候，GPU利用率只有30%左右，瓶颈根本不在GPU，而在Python层面的图构建。注意：动态图的构建开销在训练小模型或使用小batch size时尤为明显，这时候GPU往往在等CPU“画图”。2.2 什么时候该考虑静态图PyTorch 2.0之后引入了torch.compile，本质上是在保留动态图编程体验的前提下，通过图捕获和编译来减少Python开销。我试过在一个中等规模的Transformer上开启torch.compile，训练速度提升了大概20%到30%，效果因模型结构而异。但torch.compile不是万能的。它需要图捕获，如果你的forward函数里有大量Python控制流依赖运行时数据，捕获会失败或者频繁重编译。我遇到过一个case，模型里有根据输入长度动态选择不同分支的逻辑，torch.compile直接报了一堆graph break的警告，最后只能放弃。所以我的建议是：先用eager模式跑通，确认模型逻辑没问题，再尝试torch.compile。如果graph break太多，就别硬上了，收益可能还不如优化数据加载。2.3 动态图下的显存管理逻辑动态图还有一个隐性成本：中间激活值的管理。前向传播过程中产生的所有中间结果，只要还需要用于反向传播的梯度计算，就不能释放。这意味着显存占用跟模型深度、batch size、序列长度都强相关。PyTorch的显存分配器用的是缓存机制。当你释放一个张量，显存不会立刻还给GPU，而是留在缓存池里供后续分配复用。这个设计是为了减少频繁的cudaMalloc和cudaFree调用，但也会导致一个现象：nvidia-smi显示的显存占用比你实际需要的多。我踩过的一个坑是，在验证阶段用torch.no_grad()包住了前向传播，但忘了把验证集的batch size调小，结果验证阶段显存反而比训练阶段还高。原因是验证阶段虽然不存梯度，但中间激活值仍然需要存储，batch size太大照样爆。排查显存问题的时候，我一般用torch.cuda.memory_summary()看详细分配情况，比nvidia-smi的信息有用得多。它会告诉你有多少显存被缓存、多少是活跃的、碎片化程度如何。3. 数据加载：最容易被忽视的性能瓶颈3.1 DataLoader的num_workers不是越大越好很多人调DataLoader的时候，习惯性把num_workers设成CPU核心数，觉得越多越快。实际上这个参数需要根据具体情况调。num_workers控制的是数据加载的子进程数量，每个子进程都会复制一份数据集对象。如果数据集本身很大，开太多worker会导致内存爆炸。我一般从4开始试，观察GPU利用率和CPU内存占用。如果GPU利用率上不去，而CPU内存还有余量，就往上加。如果CPU内存吃紧，就往下减。没有一个万能值。还有一个细节：num_workers大于0的时候，主进程和子进程之间的数据传输是有开销的。如果数据预处理逻辑很轻（比如只是读个numpy数组），开多进程反而可能更慢，因为进程间通信的成本超过了并行加载的收益。3.2 pin_memory和non_blocking的正确用法pin_memory=True的作用是把数据加载到锁页内存（pinned memory），这样从CPU传到GPU的时候可以用异步传输。配合non_blocking=True，数据传输和GPU计算可以重叠。但这两个参数要配合使用才有效。只设pin_memory=True而不设non_blocking=True，数据传输仍然是同步的，收益有限。反过来只设non_blocking=True而不设pin_memory=True，异步传输根本不生效，因为普通内存不支持DMA直接访问。我的标准配置是这样的：dataloader = DataLoader( dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True, drop_last=True ) for data, target in dataloader: data = data.cuda(non_blocking=True) target = target.cuda(non_blocking=True) # 训练逻辑drop_last=True在训练时我基本都会开，因为最后一个不完整的batch可能导致BatchNorm层报错（batch size为1时方差为0）。验证时一般设False，保证所有样本都被评估到。3.3 自定义Dataset的常见陷阱写自定义Dataset的时候，__getitem__里尽量不要做太重的计算。如果每次取数据都要做复杂的增强，考虑用缓存或者预计算。我见过有人在__getitem__里做FFT变换，结果数据加载成了整个训练流程的瓶颈。另外，__getitem__返回的张量最好是固定shape的。如果每个样本的shape不一样，DataLoader的collate_fn需要自定义，否则默认的拼接逻辑会报错。变长序列的处理一般用padding，在collate_fn里统一pad到当前batch的最大长度，而不是全局最大长度，这样能省不少显存。4. 混合精度训练：省显存但不省心4.1 AMP的基本原理混合精度训练的核心思想是：前向传播和反向传播用FP16（半精度），参数更新用FP32（单精度）。FP16的显存占用是FP32的一半，计算速度在支持Tensor Core的GPU上也能快不少。PyTorch通过torch.cuda.amp提供了自动混合精度（AMP）的支持。基本用法是：from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for data, target in dataloader: optimizer.zero_grad() with autocast(): output = model(data) loss = criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler的作用是防止FP16下梯度下溢。FP16能表示的最小正数大概是6e-8，很多梯度值比这个还小，直接变成0就训练不动了。GradScaler先把loss放大一个系数，反向传播得到放大后的梯度，更新参数前再缩回去。4.2 什么时候不该用AMPAMP不是所有场景都适用。我遇到过的几个问题：第一，某些自定义的CUDA算子不支持FP16，开了autocast之后直接报错。这时候需要用torch.cuda.amp.custom_fwd和custom_bwd装饰器手动指定精度。第二，loss计算涉及大数相减的时候，FP16的精度不够，会导致数值不稳定。比如计算交叉熵的时候，如果logits的值域很大，FP16下可能会溢出。这种情况一般把loss计算放在autocast外面，用FP32算。第三，小模型上AMP的收益不明显，甚至可能因为类型转换的开销而变慢。我一般只在模型参数量超过1000万或者显存吃紧的时候才开AMP。4.3 GradScaler的调参经验GradScaler有几个参数可以调，但大多数时候默认值就够了。init_scale默认是2的16次方，growth_interval默认是2000。如果训练过程中频繁出现梯度溢出（scaler会跳过更新），可以适当降低init_scale。我一般会监控scaler.get_scale()的值。如果它一直往下掉，说明模型训练不稳定，可能需要检查学习率是不是太大了，或者数据里有没有异常值。5. 多卡训练：从DataParallel到DistributedDataParallel5.1 DataParallel为什么不够用DataParallel（DP）是最简单的多卡方案，一行代码就能把模型放到多张卡上。但它的效率很低，因为每次前向传播都要把模型复制到每张卡上，反向传播再把梯度汇总到主卡。主卡的显存和计算压力都很大，卡越多，主卡瓶颈越明显。我实测过一个模型，单卡训练一个epoch要10分钟，用DP双卡反而要12分钟，因为通信开销超过了并行收益。所以DP基本只适合快速验证，不适合正式训练。5.2 DistributedDataParallel的正确打开方式DistributedDataParallel（DDP）是现在的主流方案。每张卡上有一个独立的进程，各自维护一份模型副本，梯度通过NCCL通信库做all-reduce。没有主卡瓶颈，扩展性也好得多。DDP的启动方式有两种：torchrun和multiprocessing。torchrun是官方推荐的方式，配置简单，支持弹性训练。基本命令是：torchrun --nproc_per_node=4 train.py代码里需要初始化进程组：import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backend='nccl') local_rank = int(os.environ['LOCAL_RANK']) torch.cuda.set_device(local_rank) model = model.to(local_rank) model = DDP(model, device_ids=[local_rank])数据加载需要用DistributedSampler，保证每张卡看到不同的数据子集：sampler = DistributedSampler(dataset) dataloader = DataLoader(dataset, sampler=sampler, batch_size=64)每个epoch开始前要调用sampler.set_epoch(epoch)，否则shuffle的随机种子在每张卡上是一样的，数据划分会重复。5.3 DDP的常见坑第一个坑：BatchNorm在多卡下会出问题。DDP默认每张卡独立计算BN的统计量，如果每张卡的batch size太小，统计量不准。解决方案是用SyncBatchNorm，它在所有卡之间同步均值和方差。但SyncBatchNorm会引入额外的通信开销，batch size小的时候收益不明显。第二个坑：学习率需要按卡数缩放。如果单卡学习率是0.001，4卡训练时理论上应该用0.004。但这不是线性的，实际中我一般用平方根缩放，4卡用0.002，然后根据loss曲线微调。第三个坑：保存checkpoint的时候只需要主进程写文件，其他进程等待。否则多个进程同时写同一个文件会冲突。if local_rank == 0: torch.save(model.module.state_dict(), 'checkpoint.pth') dist.barrier()注意保存的时候要用model.module而不是model，因为DDP包装后的model多了一层。6. 训练调试与问题排查实录6.1 loss不收敛的排查思路loss不收敛是最常见也最头疼的问题。我一般按这个顺序排查：先看数据。把dataloader里的一个batch拿出来，手动检查输入和标签是否对应。我遇到过一次标签和输入错位的问题，原因是自定义Dataset里__getitem__的索引逻辑写错了，导致模型在学随机噪声。再看学习率。学习率太大loss会震荡，太小loss下降极慢。我一般先用一个很小的学习率跑几百步，确认loss能稳定下降，再逐步放大。然后看初始化。PyTorch默认的初始化对大多数情况够用，但某些激活函数（比如ReLU）配合默认初始化可能导致神经元死亡。可以试试Kaiming初始化或者Xavier初始化。最后看梯度。用torch.autograd.gradcheck检查梯度计算是否正确，或者打印每层的梯度范数，看有没有梯度消失或爆炸。6.2 显存溢出的应急处理显存溢出的时候，第一反应不应该是减小batch size，而是先定位哪里占的显存多。torch.cuda.memory_summary()能看到每个张量的分配情况。如果确实是激活值占太多，有几个方案：一是用梯度检查点（gradient checkpointing），用计算换显存，把中间激活值丢掉，反向传播时重新计算；二是用混合精度，激活值用FP16存储；三是减小batch size，但要注意BN层对batch size敏感。梯度检查点的用法很简单：from torch.utils.checkpoint import checkpoint def forward(self, x): x = checkpoint(self.layer1, x) x = checkpoint(self.layer2, x) return x但要注意，checkpoint包裹的层在反向传播时会重新做一次前向，训练时间会增加大概30%。6.3 常见问题速查表问题现象可能原因排查方法解决方案loss变成NaN学习率过大、梯度爆炸打印梯度范数降低学习率、加梯度裁剪GPU利用率低数据加载瓶颈监控CPU和GPU利用率增加num_workers、优化预处理多卡训练速度不升反降通信开销过大检查NCCL配置换用DDP、减少同步操作验证集效果远差于训练集过拟合对比训练和验证loss加正则化、数据增强显存溢出batch size过大memory_summary减小batch、梯度检查点训练中途卡死数据加载死锁检查num_workers设num_workers=0调试7. 一些零散但实用的经验模型保存的时候，我习惯同时保存state_dict和optimizer的状态。只保存state_dict的话，恢复训练时优化器的动量信息丢了，loss可能会有个明显的跳变。学习率调度器用CosineAnnealingLR的时候，要注意T_max的设置。设成总epoch数的话，学习率会在训练结束时降到接近0。如果还要继续训练，需要重新warmup。torch.backends.cudnn.benchmark = True在输入shape固定的情况下能自动选择最优的卷积算法，速度提升明显。但如果输入shape变化频繁，这个设置反而会导致反复benchmark，拖慢训练。日志记录我一般用tensorboard或者wandb，但要注意不要每个step都写，IO开销很大。我一般每100个step写一次，或者每个epoch写一次。最后说一个关于随机种子的经验。PyTorch的随机种子需要同时设置torch.manual_seed、torch.cuda.manual_seed_all和numpy的种子，否则结果不可复现。如果用了DataLoader的num_workers，还需要在worker_init_fn里设置种子，否则每个worker的随机状态不一样。def worker_init_fn(worker_id): np.random.seed(torch.initial_seed() % 2**32 + worker_id) dataloader = DataLoader(..., worker_init_fn=worker_init_fn)这些东西文档里都有，但散落在各个角落，实际项目中遇到问题的时候往往想不起来。我自己的做法是维护一个checklist，每次开新项目的时候过一遍，能避免很多低级错误。
1. 为什么一个深度学习框架值得反复折腾刚接触深度学习那会儿，我对框架的选择其实没什么概念，别人说用什么就用什么。直到有一次跑一个图像分类的小项目，训练到一半显存爆了，报错信息看得我一头雾水，才意识到框架的底层机制不是“能跑就行”这么简单。后来花了不少时间啃PyTorch的文档和源码，踩了无数坑，才慢慢摸清楚它的脾气。这篇文章不是入门教程，市面上入门教程已经够多了，我想聊的是那些文档里不会明说、但实际项目中一定会遇到的东西——动态图的真实代价、显存管理的底层逻辑、数据加载的隐藏瓶颈、混合精度的适用边界，以及怎么把这些东西组合成一套能稳定跑起来的训练流程。PyTorch本质上是一个基于动态计算图的深度学习框架，核心能力是自动求导和GPU加速张量运算。它能做什么？从简单的线性回归到大规模分布式训练，从计算机视觉到自然语言处理，几乎覆盖了深度学习的所有主流场景。适合谁来参考？如果你已经能写基本的训练循环，但在实际项目中经常遇到显存不够、训练速度慢、loss不收敛、多卡跑不起来这些问题，那这篇内容应该能帮你省下不少排查时间。如果你刚入门，也可以看看，至少知道后面会遇到什么。我个人的习惯是，学一个框架不能只学API怎么调，得理解它为什么这么设计。PyTorch选择动态图而不是静态图，这个决策背后有一整套权衡，理解了这些权衡，很多“奇怪”的行为就变得合理了。2. 动态图机制的真实代价与应对策略2.1 动态图到底“动”在哪里PyTorch最核心的设计就是动态计算图。什么意思？每次前向传播的时候，计算图是实时构建的。你写一个for循环，循环次数不同，计算图的结构就不同。这跟静态图框架（先定义图再运行）有本质区别。动态图的好处很直观：调试方便。你可以在forward函数里随便加print，可以用Python的if-else控制流，可以动态改变网络结构。写起来跟普通Python代码没区别，心智负担极低。但代价是什么？每次迭代都要重新构建计算图，这意味着额外的CPU开销。对于小模型，这个开销可能比GPU计算本身还大。我实测过一个简单的全连接网络，batch size设成32的时候，GPU利用率只有30%左右，瓶颈根本不在GPU，而在Python层面的图构建。注意：动态图的构建开销在训练小模型或使用小batch size时尤为明显，这时候GPU往往在等CPU“画图”。2.2 什么时候该考虑静态图PyTorch 2.0之后引入了torch.compile，本质上是在保留动态图编程体验的前提下，通过图捕获和编译来减少Python开销。我试过在一个中等规模的Transformer上开启torch.compile，训练速度提升了大概20%到30%，效果因模型结构而异。但torch.compile不是万能的。它需要图捕获，如果你的forward函数里有大量Python控制流依赖运行时数据，捕获会失败或者频繁重编译。我遇到过一个case，模型里有根据输入长度动态选择不同分支的逻辑，torch.compile直接报了一堆graph break的警告，最后只能放弃。所以我的建议是：先用eager模式跑通，确认模型逻辑没问题，再尝试torch.compile。如果graph break太多，就别硬上了，收益可能还不如优化数据加载。2.3 动态图下的显存管理逻辑动态图还有一个隐性成本：中间激活值的管理。前向传播过程中产生的所有中间结果，只要还需要用于反向传播的梯度计算，就不能释放。这意味着显存占用跟模型深度、batch size、序列长度都强相关。PyTorch的显存分配器用的是缓存机制。当你释放一个张量，显存不会立刻还给GPU，而是留在缓存池里供后续分配复用。这个设计是为了减少频繁的cudaMalloc和cudaFree调用，但也会导致一个现象：nvidia-smi显示的显存占用比你实际需要的多。我踩过的一个坑是，在验证阶段用torch.no_grad()包住了前向传播，但忘了把验证集的batch size调小，结果验证阶段显存反而比训练阶段还高。原因是验证阶段虽然不存梯度，但中间激活值仍然需要存储，batch size太大照样爆。排查显存问题的时候，我一般用torch.cuda.memory_summary()看详细分配情况，比nvidia-smi的信息有用得多。它会告诉你有多少显存被缓存、多少是活跃的、碎片化程度如何。3. 数据加载：最容易被忽视的性能瓶颈3.1 DataLoader的num_workers不是越大越好很多人调DataLoader的时候，习惯性把num_workers设成CPU核心数，觉得越多越快。实际上这个参数需要根据具体情况调。num_workers控制的是数据加载的子进程数量，每个子进程都会复制一份数据集对象。如果数据集本身很大，开太多worker会导致内存爆炸。我一般从4开始试，观察GPU利用率和CPU内存占用。如果GPU利用率上不去，而CPU内存还有余量，就往上加。如果CPU内存吃紧，就往下减。没有一个万能值。还有一个细节：num_workers大于0的时候，主进程和子进程之间的数据传输是有开销的。如果数据预处理逻辑很轻（比如只是读个numpy数组），开多进程反而可能更慢，因为进程间通信的成本超过了并行加载的收益。3.2 pin_memory和non_blocking的正确用法pin_memory=True的作用是把数据加载到锁页内存（pinned memory），这样从CPU传到GPU的时候可以用异步传输。配合non_blocking=True，数据传输和GPU计算可以重叠。但这两个参数要配合使用才有效。只设pin_memory=True而不设non_blocking=True，数据传输仍然是同步的，收益有限。反过来只设non_blocking=True而不设pin_memory=True，异步传输根本不生效，因为普通内存不支持DMA直接访问。我的标准配置是这样的：dataloader = DataLoader( dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True, drop_last=True ) for data, target in dataloader: data = data.cuda(non_blocking=True) target = target.cuda(non_blocking=True) # 训练逻辑drop_last=True在训练时我基本都会开，因为最后一个不完整的batch可能导致BatchNorm层报错（batch size为1时方差为0）。验证时一般设False，保证所有样本都被评估到。3.3 自定义Dataset的常见陷阱写自定义Dataset的时候，__getitem__里尽量不要做太重的计算。如果每次取数据都要做复杂的增强，考虑用缓存或者预计算。我见过有人在__getitem__里做FFT变换，结果数据加载成了整个训练流程的瓶颈。另外，__getitem__返回的张量最好是固定shape的。如果每个样本的shape不一样，DataLoader的collate_fn需要自定义，否则默认的拼接逻辑会报错。变长序列的处理一般用padding，在collate_fn里统一pad到当前batch的最大长度，而不是全局最大长度，这样能省不少显存。4. 混合精度训练：省显存但不省心4.1 AMP的基本原理混合精度训练的核心思想是：前向传播和反向传播用FP16（半精度），参数更新用FP32（单精度）。FP16的显存占用是FP32的一半，计算速度在支持Tensor Core的GPU上也能快不少。PyTorch通过torch.cuda.amp提供了自动混合精度（AMP）的支持。基本用法是：from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for data, target in dataloader: optimizer.zero_grad() with autocast(): output = model(data) loss = criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler的作用是防止FP16下梯度下溢。FP16能表示的最小正数大概是6e-8，很多梯度值比这个还小，直接变成0就训练不动了。GradScaler先把loss放大一个系数，反向传播得到放大后的梯度，更新参数前再缩回去。4.2 什么时候不该用AMPAMP不是所有场景都适用。我遇到过的几个问题：第一，某些自定义的CUDA算子不支持FP16，开了autocast之后直接报错。这时候需要用torch.cuda.amp.custom_fwd和custom_bwd装饰器手动指定精度。第二，loss计算涉及大数相减的时候，FP16的精度不够，会导致数值不稳定。比如计算交叉熵的时候，如果logits的值域很大，FP16下可能会溢出。这种情况一般把loss计算放在autocast外面，用FP32算。第三，小模型上AMP的收益不明显，甚至可能因为类型转换的开销而变慢。我一般只在模型参数量超过1000万或者显存吃紧的时候才开AMP。4.3 GradScaler的调参经验GradScaler有几个参数可以调，但大多数时候默认值就够了。init_scale默认是2的16次方，growth_interval默认是2000。如果训练过程中频繁出现梯度溢出（scaler会跳过更新），可以适当降低init_scale。我一般会监控scaler.get_scale()的值。如果它一直往下掉，说明模型训练不稳定，可能需要检查学习率是不是太大了，或者数据里有没有异常值。5. 多卡训练：从DataParallel到DistributedDataParallel5.1 DataParallel为什么不够用DataParallel（DP）是最简单的多卡方案，一行代码就能把模型放到多张卡上。但它的效率很低，因为每次前向传播都要把模型复制到每张卡上，反向传播再把梯度汇总到主卡。主卡的显存和计算压力都很大，卡越多，主卡瓶颈越明显。我实测过一个模型，单卡训练一个epoch要10分钟，用DP双卡反而要12分钟，因为通信开销超过了并行收益。所以DP基本只适合快速验证，不适合正式训练。5.2 DistributedDataParallel的正确打开方式DistributedDataParallel（DDP）是现在的主流方案。每张卡上有一个独立的进程，各自维护一份模型副本，梯度通过NCCL通信库做all-reduce。没有主卡瓶颈，扩展性也好得多。DDP的启动方式有两种：torchrun和multiprocessing。torchrun是官方推荐的方式，配置简单，支持弹性训练。基本命令是：torchrun --nproc_per_node=4 train.py代码里需要初始化进程组：import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backend='nccl') local_rank = int(os.environ['LOCAL_RANK']) torch.cuda.set_device(local_rank) model = model.to(local_rank) model = DDP(model, device_ids=[local_rank])数据加载需要用DistributedSampler，保证每张卡看到不同的数据子集：sampler = DistributedSampler(dataset) dataloader = DataLoader(dataset, sampler=sampler, batch_size=64)每个epoch开始前要调用sampler.set_epoch(epoch)，否则shuffle的随机种子在每张卡上是一样的，数据划分会重复。5.3 DDP的常见坑第一个坑：BatchNorm在多卡下会出问题。DDP默认每张卡独立计算BN的统计量，如果每张卡的batch size太小，统计量不准。解决方案是用SyncBatchNorm，它在所有卡之间同步均值和方差。但SyncBatchNorm会引入额外的通信开销，batch size小的时候收益不明显。第二个坑：学习率需要按卡数缩放。如果单卡学习率是0.001，4卡训练时理论上应该用0.004。但这不是线性的，实际中我一般用平方根缩放，4卡用0.002，然后根据loss曲线微调。第三个坑：保存checkpoint的时候只需要主进程写文件，其他进程等待。否则多个进程同时写同一个文件会冲突。if local_rank == 0: torch.save(model.module.state_dict(), 'checkpoint.pth') dist.barrier()注意保存的时候要用model.module而不是model，因为DDP包装后的model多了一层。6. 训练调试与问题排查实录6.1 loss不收敛的排查思路loss不收敛是最常见也最头疼的问题。我一般按这个顺序排查：先看数据。把dataloader里的一个batch拿出来，手动检查输入和标签是否对应。我遇到过一次标签和输入错位的问题，原因是自定义Dataset里__getitem__的索引逻辑写错了，导致模型在学随机噪声。再看学习率。学习率太大loss会震荡，太小loss下降极慢。我一般先用一个很小的学习率跑几百步，确认loss能稳定下降，再逐步放大。然后看初始化。PyTorch默认的初始化对大多数情况够用，但某些激活函数（比如ReLU）配合默认初始化可能导致神经元死亡。可以试试Kaiming初始化或者Xavier初始化。最后看梯度。用torch.autograd.gradcheck检查梯度计算是否正确，或者打印每层的梯度范数，看有没有梯度消失或爆炸。6.2 显存溢出的应急处理显存溢出的时候，第一反应不应该是减小batch size，而是先定位哪里占的显存多。torch.cuda.memory_summary()能看到每个张量的分配情况。如果确实是激活值占太多，有几个方案：一是用梯度检查点（gradient checkpointing），用计算换显存，把中间激活值丢掉，反向传播时重新计算；二是用混合精度，激活值用FP16存储；三是减小batch size，但要注意BN层对batch size敏感。梯度检查点的用法很简单：from torch.utils.checkpoint import checkpoint def forward(self, x): x = checkpoint(self.layer1, x) x = checkpoint(self.layer2, x) return x但要注意，checkpoint包裹的层在反向传播时会重新做一次前向，训练时间会增加大概30%。6.3 常见问题速查表问题现象可能原因排查方法解决方案loss变成NaN学习率过大、梯度爆炸打印梯度范数降低学习率、加梯度裁剪GPU利用率低数据加载瓶颈监控CPU和GPU利用率增加num_workers、优化预处理多卡训练速度不升反降通信开销过大检查NCCL配置换用DDP、减少同步操作验证集效果远差于训练集过拟合对比训练和验证loss加正则化、数据增强显存溢出batch size过大memory_summary减小batch、梯度检查点训练中途卡死数据加载死锁检查num_workers设num_workers=0调试7. 一些零散但实用的经验模型保存的时候，我习惯同时保存state_dict和optimizer的状态。只保存state_dict的话，恢复训练时优化器的动量信息丢了，loss可能会有个明显的跳变。学习率调度器用CosineAnnealingLR的时候，要注意T_max的设置。设成总epoch数的话，学习率会在训练结束时降到接近0。如果还要继续训练，需要重新warmup。torch.backends.cudnn.benchmark = True在输入shape固定的情况下能自动选择最优的卷积算法，速度提升明显。但如果输入shape变化频繁，这个设置反而会导致反复benchmark，拖慢训练。日志记录我一般用tensorboard或者wandb，但要注意不要每个step都写，IO开销很大。我一般每100个step写一次，或者每个epoch写一次。最后说一个关于随机种子的经验。PyTorch的随机种子需要同时设置torch.manual_seed、torch.cuda.manual_seed_all和numpy的种子，否则结果不可复现。如果用了DataLoader的num_workers，还需要在worker_init_fn里设置种子，否则每个worker的随机状态不一样。def worker_init_fn(worker_id): np.random.seed(torch.initial_seed() % 2**32 + worker_id) dataloader = DataLoader(..., worker_init_fn=worker_init_fn)这些东西文档里都有，但散落在各个角落，实际项目中遇到问题的时候往往想不起来。我自己的做法是维护一个checklist，每次开新项目的时候过一遍，能避免很多低级错误。