Linux系统下PyTorch内存管理技巧有哪些

Linux系统下PyTorch内存管理技巧

一 显存与计算优化

• 启用自动混合精度 AMP：在大部分算子中使用FP16/BF16计算，关键处保留FP32，通常可减少约30%–50%显存占用并加速训练。示例：使用torch.cuda.amp.autocast与GradScaler完成缩放、反向与更新。若硬件支持，优先选择BF16（如 Ampere 及更高架构的 NVIDIA GPU）。
• 使用梯度检查点：以计算换显存，常见做法可使激活显存降低约40%–50%；在超长序列/大模型（如 BERT-large）中，实际可降低约**70%**激活显存。
• 采用梯度累积：用小批次累积梯度达到“虚拟大批次”，在不增大显存的前提下维持吞吐与收敛特性。
• 选择更省显存的优化器：如 Adafactor（分解二阶矩，显存占用可降约40%）、SGD（无额外状态，参数状态内存约为 Adam 的1/3），在 LLM/大模型场景尤为关键。
• 利用激活/参数卸载：将部分中间激活或参数临时移至CPU 内存（必要时再迁回 GPU），缓解单卡显存瓶颈。

二 数据与数据加载优化

• 启用Pinned Memory与多进程加载：设置 pin_memory=True、合理提升 num_workers，加速主机到设备的传输并减少数据加载阻塞。
• 使用内存映射数据集：通过 np.memmap 按需读取大文件，避免一次性将整个数据集载入 RAM。
• 采用流式/生成式数据管道：用生成器与迭代器逐批读取样本，降低峰值内存。
• 避免不必要的数据复制：优先使用视图与必要的原地操作（谨慎使用，注意数值稳定性与梯度正确性）。

三 运行时与系统级调优

• 理解并利用CUDA 缓存分配器：PyTorch 通过缓存分配器与流并行复用显存，减少频繁分配/释放；可用 torch.cuda.memory_summary() 观测分配与保留情况。
• 谨慎清理缓存：在关键阶段调用 torch.cuda.empty_cache() 释放未使用缓存；结合阈值或阶段性清理，避免频繁清理影响性能。
• 控制CPU 线程数：通过 torch.set_num_threads() 匹配 CPU 核心与 NUMA 拓扑，减少线程争用与内存抖动。
• 启用cuDNN 自动调优：设置 torch.backends.cudnn.benchmark=True 提升卷积等算子的性能（注意可能引入非确定性）。
• 多卡/多机训练：在单卡显存不足时，采用张量并行、流水线并行或FSDP 全分片数据并行；FSDP 可将参数/梯度/优化器状态分片至多卡，显著降低单卡显存占用。
• 关注NUMA 与异构内存：在多插槽服务器上结合 NUMA 绑定与内存亲和性优化；需要时利用 HMM 等异构内存机制提升 CPU-GPU 协同效率。

四 监控与排障

• 实时监控显存：使用 torch.cuda.memory_allocated() / memory_reserved() / memory_summary() 观察分配、保留与碎片情况，定位异常增长。
• 监控主机内存：结合 psutil 与 sys.getsizeof 等工具，跟踪进程 RSS 与对象占用，及早发现内存泄漏或膨胀。
• 排查显存泄漏：核对自定义 Autograd Function 的 backward 是否正确释放中间张量；多进程数据加载注意避免主进程持有 CUDA 张量引用。
• 谨慎使用显存清理：在训练循环中频繁调用 empty_cache() 可能降低性能，建议基于阈值或阶段化策略执行。
• 辅助工具：必要时使用 valgrind 等工具检测内存问题（注意其对 CUDA 的支持与性能开销）。

五 实用配置示例

• 混合精度训练（AMP）
  • 代码示例：
    • scaler = torch.cuda.amp.GradScaler()
    • for data, target in loader:
      • optimizer.zero_grad(set_to_none=True)
      • with torch.cuda.amp.autocast():
        • output = model(data)
        • loss = criterion(output, target)
      • scaler.scale(loss).backward()
      • scaler.step(optimizer)
      • scaler.update()
• 高效数据加载
  • 代码示例：
    • loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
• 梯度检查点
  • 代码示例：
    • from torch.utils.checkpoint import checkpoint
    • def segment(x): return model.layer2(model.layer1(x))
    • y = checkpoint(segment, x)
• 阶段性清理缓存
  • 代码示例：
    • if torch.cuda.memory_reserved() > 0.8 * torch.cuda.get_device_properties().total_memory:
      • torch.cuda.empty_cache()
• 线程与性能调优
  • 代码示例：
    • torch.set_num_threads(8)
    • torch.backends.cudnn.benchmark = True