PyTorch显存管理基础：从分配到释放的全流程

PyTorch的显存管理由CUDA内存分配器（如cudaMalloc和cudaFree）驱动，其核心逻辑分为动态分配与碎片整理两个阶段。当执行张量操作时，PyTorch会通过torch.cuda模块的底层接口向GPU申请连续内存块，例如：

import torch
x = torch.randn(10000, 10000).cuda()  # 动态分配约400MB显存

此过程依赖默认的cudaMalloc策略，但频繁的小规模分配易导致显存碎片化。例如，连续创建100个10MB的张量后，可能因碎片无法分配200MB的连续空间，即使总空闲显存充足。

显存预留机制：torch.cuda.memory_reserved()与empty_cache()

1. 预留显存的底层原理

PyTorch通过cudaMalloc的预分配机制实现显存预留。调用torch.cuda.memory_reserved()可查看当前预留的显存总量（单位：字节），该值通常大于实际使用的显存，因为分配器会保留部分空闲块以加速后续分配。例如：

print(torch.cuda.memory_reserved())  # 输出类似1073741824（1GB）

预留策略由PyTorch的CACHING_ALLOCATOR控制，其通过维护空闲内存池减少系统调用开销。当池中无合适块时，才会触发真正的cudaMalloc。

2. 手动清理预留显存：empty_cache()的适用场景

torch.cuda.empty_cache()函数用于释放分配器缓存中的空闲块，但不会减少物理显存占用。其典型应用场景包括：

• 模型切换时：从大模型切换到小模型前清理碎片

  # 训练大模型后切换任务
  model_large.eval()
  torch.cuda.empty_cache()  # 清理残留碎片
  model_small.train()

• 调试显存泄漏：配合torch.cuda.memory_summary()定位异常分配
• 多任务共享GPU：在任务间隙释放未使用的预留块

需注意，频繁调用empty_cache()可能导致性能下降，因其需遍历并合并空闲块。

高级显存控制：set_per_process_memory_fraction()与自定义分配器

1. 限制进程显存配额

通过torch.cuda.set_per_process_memory_fraction()可设置进程最大显存使用比例（相对于总显存），例如限制为50%：

torch.cuda.set_per_process_memory_fraction(0.5, device=0)

此功能在多进程训练中尤为关键，可避免单个进程独占资源。

2. 自定义分配器实现

对于特定场景（如稀疏张量存储），可通过继承torch.cuda.memory.Allocator实现自定义分配逻辑。例如，实现一个优先使用小内存块的分配器：

class SmallBlockAllocator(torch.cuda.memory.Allocator):
    def allocate(self, size):
        # 优先分配小于1MB的块
        if size < 1e6:
            return super().allocate(size)
        else:
            # 大块分配延迟处理
            pass

需谨慎使用自定义分配器，因其可能破坏PyTorch的内存池优化。

实战技巧：显存优化五步法

1. 监控显存使用

使用torch.cuda.memory_stats()获取详细分配信息：

stats = torch.cuda.memory_stats()
print(f"Active bytes: {stats['active.all.bytes']}")
print(f"Reserved bytes: {stats['reserved.all.bytes']}")

重点关注active.all.bytes（实际使用）与reserved.all.bytes（预留总量）的比值，理想情况下应低于70%。

2. 梯度累积替代大batch

当显存不足时，可通过梯度累积模拟大batch训练：

accumulation_steps = 4
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

此方法将batch分4次处理，梯度累积后更新，显存需求降低至1/4。

3. 混合精度训练

使用torch.cuda.amp自动管理半精度训练：

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

FP16运算可减少50%显存占用，同时通过动态缩放保持数值稳定性。

4. 模型并行与张量并行

对于超大规模模型（如GPT-3），采用模型并行将不同层分配到不同GPU：

# 简单示例：将模型前半部分放在GPU0，后半部分放在GPU1
model_part1 = nn.Sequential(*list(model.children())[:3]).cuda(0)
model_part2 = nn.Sequential(*list(model.children())[3:]).cuda(1)

更复杂的实现可参考PyTorch的DistributedDataParallel与Megatron-LM框架。

5. 显存碎片整理

当遇到CUDA out of memory错误且memory_reserved()显示大量碎片时，可尝试：

1. 重启内核释放所有显存
2. 使用torch.cuda.ipc_collect()（需PyTorch 1.10+）强制整理碎片
3. 降低torch.backends.cudnn.benchmark为False减少动态算法选择带来的分配波动

常见问题与解决方案

Q1: 为什么empty_cache()后显存占用未下降？

empty_cache()仅清理分配器缓存，物理显存仍由GPU驱动管理。若需释放物理显存，需删除所有张量并调用torch.cuda.empty_cache()：

del x  # 删除张量
torch.cuda.empty_cache()  # 此时可能释放物理显存

Q2: 如何检测显存泄漏？

通过周期性记录torch.cuda.memory_allocated()和torch.cuda.memory_reserved()，若allocated持续增长而reserved稳定，则可能存在泄漏。使用objgraph或torch.autograd.profiler进一步定位。

Q3: 多GPU训练时如何均衡显存？

使用DistributedDataParallel的bucket_cap_mb参数控制梯度桶大小，或通过torch.cuda.set_device()显式指定设备，避免默认分配策略导致的不均衡。

总结与最佳实践

PyTorch的显存管理需结合动态分配、预留机制与手动控制。关键建议包括：

1. 监控memory_allocated()与memory_reserved()的比值，保持健康状态
2. 在模型切换或任务间隙调用empty_cache()
3. 对大模型采用梯度累积或混合精度训练
4. 使用set_per_process_memory_fraction()限制多进程显存
5. 复杂场景下考虑模型并行或自定义分配器

通过合理利用这些工具，开发者可在有限GPU资源下实现高效深度学习训练，避免因显存问题导致的训练中断或性能下降。