显存管理基础：理解PyTorch的显存分配机制

PyTorch的显存分配遵循”按需分配+缓存池”机制，CUDA上下文初始化时会预分配基础显存，后续操作通过显存分配器动态申请。开发者可通过torch.cuda.memory_summary()查看详细分配情况。显存占用主要分为三类：模型参数（权重、偏置）、中间激活值（前向传播计算图）、优化器状态（动量、二阶矩）。

典型案例中，ResNet50模型参数约100MB，但批大小为32时中间激活值可达800MB，Adam优化器状态更会翻倍显存需求。这种非线性增长特性要求开发者建立显式显存监控机制：

def print_gpu_memory():
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"Allocated: {allocated:.2f}MB, Reserved: {reserved:.2f}MB")

梯度检查点（Gradient Checkpointing）：以时间换空间的经典方案

该技术通过重构计算图，在反向传播时重新计算前向节点来节省显存。核心原理是将N个操作划分为K个区块，每个区块仅存储输入输出而非中间结果。PyTorch提供torch.utils.checkpoint.checkpoint实现：

from torch.utils.checkpoint import checkpoint
class CheckpointBlock(nn.Module):
    def __init__(self, submodule):
        super().__init__()
        self.submodule = submodule
    def forward(self, x):
        return checkpoint(self.submodule, x)
# 使用示例
model = nn.Sequential(
    nn.Linear(1024, 2048),
    CheckpointBlock(nn.Sequential(
        nn.ReLU(),
        nn.Linear(2048, 4096),
        nn.ReLU()
    )),
    nn.Linear(4096, 1000)
)

实测数据显示，在BERT-base模型上使用梯度检查点可使显存占用从12GB降至4.5GB，但训练时间增加约30%。最佳应用场景为：模型深度大但宽度适中、批处理大小受显存限制时。

混合精度训练：FP16与FP32的完美平衡

NVIDIA A100等现代GPU对FP16运算有显著加速，PyTorch的AMP（Automatic Mixed Precision）通过动态类型转换实现：

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

关键优化点包括：

1. 动态缩放：解决FP16梯度下溢问题，通过GradScaler自动调整损失尺度
2. 主类型选择：对Conv/Linear等计算密集型操作使用FP16，对BatchNorm等数值敏感操作保持FP32
3. 内存节省：FP16参数仅占FP32一半，配合梯度累积可实现更大有效批处理

在Vision Transformer训练中，混合精度可使显存占用减少40%，同时吞吐量提升25%。需注意某些自定义算子可能需要手动指定精度。

模型并行与张量并行：突破单卡显存极限

当模型规模超过单卡容量时，并行策略成为必然选择：

流水线并行（Pipeline Parallelism）

将模型按层分割到不同设备，通过微批（micro-batch）实现流水线执行：

# 简单示例（需配合PyTorch Lightning等框架）
model = nn.Sequential(
    nn.Linear(1024, 4096).to('cuda:0'),
    nn.ReLU(),
    nn.Linear(4096, 2048).to('cuda:1')
)
def forward_pipeline(x, device_map):
    x = x.to('cuda:0')
    x = model[0](x)
    x = x.to('cuda:1')
    return model[2](x)

GPipe算法通过气泡（bubble）优化将设备利用率提升至80%以上，实际实现推荐使用FairScale或Megatron-LM框架。

张量并行（Tensor Parallelism）

对矩阵乘法进行并行分解，典型如Megatron-LM的列并行线性层：

# 简化版张量并行实现
class ColumnParallelLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.world_size = torch.distributed.get_world_size()
        self.rank = torch.distributed.get_rank()
        self.out_features_per_rank = out_features // self.world_size
        self.weight = nn.Parameter(
            torch.randn(self.out_features_per_rank, in_features) / 
            math.sqrt(in_features)
        ).to(f'cuda:{self.rank}')
    def forward(self, x):
        # 跨设备All-Reduce
        x_part = x[:, :, self.rank*self.out_features_per_rank:(self.rank+1)*self.out_features_per_rank]
        output_part = torch.matmul(x_part, self.weight.t())
        # 使用NCCL进行归约
        dist.all_reduce(output_part, op=dist.ReduceOp.SUM)
        return output_part

实测在8卡A100上训练GPT-3 175B参数模型，张量并行可将单步训练时间控制在2分钟内。

显存优化工具链

1. PyTorch Profiler：通过torch.profiler分析显存分配热点

   with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CUDA],
    profile_memory=True
   ) as prof:
    train_step()
   print(prof.key_averages().table(
    sort_by="cuda_memory_usage", row_limit=10))

2. TensorBoard集成：可视化显存随时间变化曲线
3. 自定义分配器：对特定场景可重写torch.cuda.MemoryAllocator

实战建议

1. 渐进式优化：监控→梯度检查点→混合精度→并行化
2. 批处理大小选择：遵循batch_size = floor(total_memory / (model_size + 2*activation_size))
3. 梯度累积：当物理批处理受限时，通过accumulation_steps虚拟增大批处理
4. 激活值压缩：对ReLU输出使用8位量化（需自定义算子）

最新研究显示，结合ZeRO优化器（来自DeepSpeed）和选择性激活检查点，可在不损失精度情况下将千亿参数模型训练显存需求从1.2TB降至480GB。开发者应根据具体硬件配置（如NVLink带宽、HBM容量）选择最优组合策略。