Linux上PyTorch内存如何管理

Linux上PyTorch内存管理实战指南

一 内存构成与监控要点

1. GPU显存由三部分构成：allocated（已分配）、reserved（缓存池预留）、active（当前活跃）。释放张量后显存常进入缓存池，不会立即归还系统，因此需用torch.cuda.memory_summary()、torch.cuda.memory_allocated()/memory_reserved()区分“真实使用”和“缓存占用”。出现“CUDA out of memory”但 allocated 远小于总显存，多为显存碎片化或缓存泄漏。系统层面用nvidia-smi -l 1实时监控，Python层可结合psutil观察主机内存。必要时用**torch.cuda.empty_cache()**清理缓存，仅在确认不再使用相关张量后调用，避免频繁清理带来性能抖动。

二 GPU显存优化策略

1. 混合精度训练：使用torch.cuda.amp.autocast()与GradScaler，在保持精度的同时通常可减少约30–50%显存占用，并提升吞吐。若硬件支持（如Ampere及以上），优先BF16以获得更大动态范围与稳定性。
2. 梯度检查点：以计算换显存，显著削减中间激活占用，典型可降至O(√n)量级，代价是训练时间增加约~20%。
3. 梯度累积：用小批次多次前向/反向累积梯度，模拟大批次效果，显存友好但会延长训练时间。
4. 分布式与分片：大模型采用FSDP（Fully Sharded Data Parallel）将参数/梯度/优化器状态跨设备分片，单卡显存压力显著下降，常与其他技术叠加使用。
5. 优化器与激活管理：选择更省显存的优化器（如SGD替代Adam，牺牲部分收敛速度可用学习率调度补偿）；必要时将部分激活/参数卸载到CPU（如配合DeepSpeed等）。
6. 数据加载：启用pin_memory=True与合适的num_workers，加速主机到设备的传输并减少CPU端内存拷贝瓶颈。

三 系统层面与主机内存管理

1. 快速定位与临时缓解：用free -h / htop查看系统内存，关闭不必要进程或重启以缓解内存碎片化；训练脚本中及时del无用张量并调用**gc.collect()与torch.cuda.empty_cache()**释放资源。
2. 动态批处理：从batch_size=1起步，逐步倍增，遇到OOM则折半，自动寻找当前环境下的最大稳定批次。
3. 交换分区与OOM兜底：内存不足时创建Swap（示例：fallocate创建8GB交换文件并启用），作为临时兜底方案，注意会显著增加I/O时间。
4. 内核与内存策略：在CentOS/RHEL等发行版可结合工作负载调节vm.swappiness、vm.min_free_kbytes、vm.overcommit_memory，降低换页倾向或调整内存过量使用策略，以平衡稳定性与性能。

四 实战代码模板

1. 混合精度 + 梯度累积 + 清理

import torch, gcfrom torch.cuda.amp import autocast, GradScalermodel = model.cuda()optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)scaler = GradScaler()accum_steps = 4for epoch in range(epochs):for i, (inputs, targets) in enumerate(train_loader):inputs, targets = inputs.cuda(), targets.cuda()with autocast():outputs = model(inputs)loss = criterion(outputs, targets) / accum_stepsscaler.scale(loss).backward()if (i + 1) % accum_steps == 0:scaler.step(optimizer)scaler.update()optimizer.zero_grad()# 释放本轮临时张量del outputs, loss, inputs, targetsif (i + 1) % 50 == 0:# 适度清理，避免频繁gc.collect()torch.cuda.empty_cache()

1. 自适应批大小

def find_max_batch(model, input_shape, max_mem=8*1024**3, start=1):bsz = startwhile True:try:inp = torch.randn(*input_shape, device='cuda')with torch.cuda.amp.autocast():_ = model(inp[:bsz])used = torch.cuda.max_memory_allocated()if used > 0.9 * max_mem:return max(1, bsz - 1)bsz *= 2except RuntimeError:return max(1, bsz // 2)

1. 显存与系统监控

# GPUprint(f"Alloc: {torch.cuda.memory_allocated()/1024**2:.1f}MB"f"Reserved: {torch.cuda.memory_reserved()/1024**2:.1f}MB")print(torch.cuda.memory_summary())# 系统# watch -n 1 'free -h'

1. 交换分区示例（仅在必要时启用）

sudo fallocate -l 8G /swapfilesudo chmod 600 /swapfilesudo mkswap /swapfilesudo swapon /swapfileecho '/swapfile none swap sw 0 0' | sudo tee -a /etc/fstab

以上模板可与梯度检查点、FSDP、更轻量优化器等按需组合，以在有限硬件上获得更稳健的训练流程。