CentOS中PyTorch内存管理技巧有哪些

CentOS上PyTorch内存管理技巧

一 训练期显存优化

• 使用自动混合精度 AMP：以FP16/BF16进行大部分计算，关键处保留FP32，显著降低显存占用并通常提速。示例：
  • with torch.cuda.amp.autocast(): outputs = model(inputs); loss = criterion(outputs, targets)
  • scaler = torch.cuda.amp.GradScaler(); scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()
• 降低Batch Size与梯度累积：在保持“虚拟批量”的同时控制峰值显存，累积步数=目标批量//实际批量。
• 启用梯度检查点：用计算换显存，典型可将激活内存降低约40%–50%；对Transformer优先对FFN等显存密集层打点。
• 选择更省显存的优化器与原地操作：如使用8-bit Adam等低占用优化器；在不影响正确性的前提下使用inplace（如add_、relu_）。
• 使用FSDP 全分片数据并行：将参数/梯度/优化器状态分片至多GPU，单卡显存可大幅下降，必要时结合CPU Offload进一步缓解。
• 数据与传输优化：DataLoader开启pin_memory=True与合适的num_workers，加速Host→Device传输并减少CPU端内存压力。

二 推理期显存优化

• 使用静态图编译：如以torch.compile(…, mode=“inductor”)或torch._dynamo.optimize(“inductor”)编译模型，推理阶段一次性申请workspace，减少显存分配/释放与碎片，常见场景峰值显存可下降20%–40%。
• 结合AMP/BF16推理：在不牺牲精度的前提下进一步压缩显存占用。
• 控制中间张量生命周期：推理循环中及时del不再使用的中间结果，并在关键点调用**torch.cuda.empty_cache()**清理未使用缓存。

三 系统与CUDA层面优化

• 合理设置vm.swappiness、vm.min_free_kbytes、vm.overcommit_memory：降低对交换分区的依赖、保障关键内存余量，避免训练过程因内存回收抖动导致性能劣化。
• 监控与诊断：用nvidia-smi观察显存与进程；用free/top/ps查看系统内存；在PyTorch内用**torch.cuda.memory_summary()与torch.cuda.memory._snapshot()**定位常驻大块与峰值。
• 谨慎清理系统缓存：必要时执行sync && echo 3 > /proc/sys/vm/drop_caches（仅在非生产、可容忍短暂停顿时），避免频繁清理导致I/O抖动。
• 环境与健康：保持PyTorch/CUDA/cuDNN版本匹配与更新；使用conda隔离环境；安装包时加**–no-cache-dir减少磁盘占用；必要时增加物理内存或配置Swap**作为兜底。

四 快速排查与落地清单

• 基准与定位：固定随机种子，记录当前batch_size、精度、是否AMP/Checkpoint/FSDP下的显存峰值（nvidia-smi与memory_summary）；用**memory._snapshot()**导出JSON做可视化分析，找出“常驻大块”。
• 分层优化顺序：优先尝试AMP/BF16 → 降低Batch → 梯度累积 → 梯度检查点；仍不足时上FSDP(+CPU Offload)；推理阶段再叠加torch.compile。
• 数据管道：确保pin_memory=True、num_workers合理；避免在主进程堆积过多预取数据。
• 资源与系统：关闭无关前台/后台进程，确保驱动/CUDA/cuDNN与框架版本匹配；必要时扩容内存或配置Swap，作为稳定性兜底。