PyTorch在Ubuntu上的内存管理

PyTorch在Ubuntu上的内存管理

一 核心概念与快速诊断

• 理解两个关键指标：Allocated（当前已分配显存）与Reserved（缓存分配器为加速分配预留的显存）。当出现reserved 远大于 allocated且伴随偶发 OOM，常见原因是显存碎片化；此时可通过调整分配器参数缓解。监控方式建议同时观察系统级与框架级指标：
  • 系统级：使用nvidia-smi实时查看显存占用与进程列表。
  • 框架级：使用**torch.cuda.memory_allocated() / torch.cuda.memory_reserved() / torch.cuda.memory_summary()**定位分配与缓存状态。
  • 典型现象与对策：若 reserved 持续高于 allocated 且波动大，优先考虑降低max_split_size_mb以减少大块分裂带来的碎片。

二 环境变量与缓存分配器调优

• 设置环境变量PYTORCH_CUDA_ALLOC_CONF可在进程启动前调整 CUDA 缓存分配器行为（仅对新进程生效）：
  • 临时生效（当前终端）：
    • 示例：export PYTORCH_CUDA_ALLOC_CONF=“max_split_size_mb:128”
  • 永久生效（登录即加载）：
    • 写入配置：echo ‘export PYTORCH_CUDA_ALLOC_CONF=“max_split_size_mb:128”’ > > ~/.bashrc
    • 使配置立即生效：source ~/.bashrc
  • 常用键值与作用：
    • max_split_size_mb：限制单次分配的最大块大小，减小该值可降低大块分裂与碎片，典型范围可从32–512按场景试探；当 reserved > > allocated 时可优先尝试减小该值。
    • garbage_collection_threshold：触发缓存回收的阈值，示例：garbage_collection_threshold:0.8。
  • 验证是否生效：echo $PYTORCH_CUDA_ALLOC_CONF；在 Python 中可用 torch.cuda.memory_stats() 辅助观察分配细节。

三 训练与推理阶段的显存释放与优化

• 计算图与中间变量管理：
  • 推理/验证使用with torch.no_grad()；循环或生成任务中对不需要梯度的中间状态使用**detach()**截断计算图，避免无谓累积。
  • 删除不再使用的张量/模型引用后，按需调用**torch.cuda.empty_cache()**清理未使用缓存块（注意：empty_cache 不会释放被引用张量，也不会降低 allocated，主要用于释放缓存）。
• 优化器状态与梯度：
  • 使用**optimizer.zero_grad(set_to_none=True)**可更彻底释放梯度缓存占用（相较默认置零）。
• 混合精度训练（AMP）：
  • 通过torch.cuda.amp.autocast() / GradScaler在保持精度的同时降低显存占用并提升吞吐，适用于大规模训练与推理。
• 典型优化效果参考：
  • AMP在 NVIDIA A100 上常见带来约1.5–2倍显存效率提升；Gradient Checkpointing常用以用约**20%–30%**计算时间换取显著显存节省。

四 多进程与分布式训练的内存要点

• 启动方式：优先使用spawn而非 fork，避免子进程继承父进程显存状态导致的不确定性与泄漏风险。
• 进程同步：在关键阶段（如模型切换、评估前后）使用**torch.distributed.barrier()**同步，确保所有进程完成缓存/显存清理后再继续。
• 并行策略：在单机多卡场景优先采用DistributedDataParallel（DDP），相较 DataParallel 具备更好的显存与吞吐表现；注意根据模型结构合理设置find_unused_parameters以避免不必要的显存开销。

五 常见场景与操作清单

• 场景A：出现“CUDA out of memory”，且观察到reserved > > allocated
  • 对策组合：减小max_split_size_mb（如 128/256/512 逐步试探）、在关键点调用torch.cuda.empty_cache()、降低batch size、开启AMP、必要时使用Gradient Checkpointing。
• 场景B：模型/任务切换后显存未回落
  • 对策组合：在切换前del旧模型与中间变量，调用gc.collect()与torch.cuda.empty_cache()，确保无引用残留后再加载新模型。
• 场景C：推理服务长时间运行后显存缓慢增长
  • 对策组合：全程使用torch.no_grad()，避免在推理循环中累积张量或计算图；对循环状态使用detach()；定期记录memory_allocated / memory_reserved定位异常增长点。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！