1. 自动混合精度训练(AMP)
结合FP16与FP32浮点格式的优势,在保持模型精度的同时减少内存占用。PyTorch通过torch.cuda.amp.autocast()
自动选择计算精度,配合GradScaler
处理梯度缩放,避免数值溢出或下溢。例如,在训练循环中包裹autocast()
上下文,可显著降低激活值和梯度的存储需求。
2. 梯度检查点(Gradient Checkpointing)
前向传播时仅存储部分中间激活值,反向传播时重新计算缺失的激活值,以“计算换内存”。通过torch.utils.checkpoint.checkpoint
包装模型层,可将中间激活值的内存占用减少40%-50%,尤其适用于深层模型。
3. 梯度累积(Gradient Accumulation)
通过多次迭代累积小批量的梯度,再更新模型参数,实现“虚拟大批次”训练。例如,设置accumulation_steps=4
,则每处理4个小批量才更新一次参数,相当于将实际批量大小扩大4倍,但不增加单次内存消耗。这种方法适用于内存有限但需要较大批量大小的场景。
4. 低精度训练(Lower-Precision Training)
使用FP16或BF16(Brain Floating Point)格式替代FP32,减少内存占用。FP16适合大多数深度学习任务,而BF16提供更大动态范围,更适合NVIDIA Ampere及更新架构的GPU。通过torch.cuda.is_bf16_supported()
检查GPU支持情况,可有效降低内存使用。
5. 清理与释放内存
torch.cuda.empty_cache()
清空PyTorch缓存的无用内存(如未使用的张量),释放GPU内存;del x
),并调用gc.collect()
触发Python垃圾回收,避免内存泄漏。6. 优化数据加载
通过torch.utils.data.DataLoader
的参数优化数据加载效率:
num_workers>0
(如num_workers=4
),启用多进程加载数据,避免数据预处理阻塞主进程;pin_memory=True
,将数据固定在主机内存中,加速数据从CPU到GPU的传输;batch_size
,避免一次性加载过多数据导致内存溢出。7. 使用原地操作(In-place Operations)
原地操作直接修改现有张量,而非创建新张量,减少临时内存分配。例如,使用x.add_(y)
代替x = x + y
,x.mul_(0.5)
代替x = x * 0.5
。原地操作可降低内存碎片和总体内存占用,尤其在迭代训练循环中效果显著。
8. 分布式训练与张量分片
torch.nn.parallel.DistributedDataParallel
(DDP)将模型分布到多个GPU上,并行处理数据和计算,降低单GPU内存负担;fsdp_model = FSDP(model)
即可启用分片训练。9. 监控内存使用
通过torch.cuda.memory_summary(device=None, abbreviated=False)
打印详细的GPU内存使用报告(包括已用内存、缓存内存、碎片情况),或使用nvidia-smi
命令实时监控GPU内存占用,及时发现内存瓶颈并调整策略。