4、显存复用与张量共享：in-place操作、张量视图(view)机制、显存池化技术
显存这东西，说白了就是钱。你想想看，一张A100 80G要多少钱？更别提集群里几百张卡同时跑。我见过太多团队，模型没多大，显存先爆了，最后只能被迫降低batch size，推理速度直接腰斩。
这一章，我们就来聊聊怎么把显存「榨干」。核心思路就三个：in-place操作、张量视图、显存池化。它们分别解决不同层面的浪费问题，但目标一致——让每一字节显存都发挥价值。
核心观点：显存优化的本质，是减少「临时数据」的生存周期，以及避免「不必要的数据拷贝」。
4.1 in-place操作：原地修改，不占新坑
什么叫in-place？就是直接在原张量上做修改，不创建新的张量。举个例子，a = a + 1 和 a.add_(1) 的区别。前者会创建一个新张量，后者直接在原地址上加1。
我个人习惯，在推理场景中，只要不影响后续计算，能in-place就in-place。尤其是激活函数、归一化层这些中间结果，用完就改，别留着。
# 不推荐：创建新张量
x = torch.relu(x) # 新分配显存
# 推荐：in-place操作
torch.relu_(x) # 原地修改，零额外开销我的经验：有一次优化一个BERT推理管线，发现仅仅把几个激活函数改成in-place版本，显存占用就降了12%。别小看这点，对于长序列推理，12%可能就是能不能跑batch size=32和64的区别。
但要注意，in-place操作有陷阱。PyTorch的自动求导机制会记录计算图，如果你在训练时对需要梯度的张量做in-place修改，会破坏计算图，导致反向传播出错。所以我的原则是：推理场景放心用，训练场景慎用。
避坑指南：我曾经在推理时对一个被多个分支引用的张量做了in-place修改，结果下游分支全乱了。记住：只有确定该张量不再被其他变量引用时，才能安全地in-place。
4.2 张量视图(view)机制：不复制，只改「看法」
张量视图，是PyTorch里最优雅的设计之一。它允许你以不同的「形状」去解读同一块显存数据，而不需要真的复制数据。
你想想看，一个形状为 (batch, seq_len, hidden) 的张量，你想把它变成 (batch * seq_len, hidden)，如果做拷贝，显存直接翻倍。但用 view() 或 reshape()，只是改了张量的元数据，底层数据指针没变。
# 原始张量：占用 1MB 显存
x = torch.randn(1024, 256) # 形状: (1024, 256)
# 视图操作：不拷贝数据
y = x.view(512, 512) # 形状变了，但指向同一块显存
z = x.t() # 转置视图，同样不拷贝
# 验证：修改y会影响x
y[0, 0] = 999
print(x[0, 0]) # 输出 999，证明共享内存这里有个关键点：view() 要求张量在内存中是「连续」的。如果张量经过转置、切片等操作变得不连续，view() 会报错。这时候可以用 reshape()，它会自动处理连续性——但注意，reshape() 在无法返回视图时会默默拷贝数据。
| 操作 | 是否拷贝数据 | 连续性要求 | 推荐场景 |
|---|---|---|---|
| view() | 否 | 必须连续 | 确定张量连续时 |
| reshape() | 可能拷贝 | 无要求 | 不确定连续性时 |
| contiguous() | 可能拷贝 | 强制连续 | 需要连续视图时 |
我的建议：在写推理代码时，尽量用 view() 并确保连续性。如果实在不行，用 reshape() 兜底。但千万别频繁调用 contiguous()，那会触发隐式拷贝，反而增加显存压力。
4.3 显存池化技术：像内存池一样管理显存
显存池化，说白了就是「提前申请，重复使用」。GPU显存的分配和释放是有开销的，频繁的 cudaMalloc 和 cudaFree 不仅慢，还会导致显存碎片化。
PyTorch的CUDA缓存分配器（Caching Allocator）就是干这个的。它会预先申请一大块显存，然后按需分配给张量，用完不立即释放，而是放回池子里复用。这就是为什么你经常看到 nvidia-smi 显示显存占用很高，但实际模型只用了其中一部分——剩下的在缓存池里。
# PyTorch默认启用缓存分配器
# 你可以通过环境变量控制
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
# 查看缓存状态
import torch
print(torch.cuda.memory_summary())这里有个坑：缓存池虽然快，但会「藏」显存。你明明释放了张量，显存占用却没降下来。我见过有人因此误判显存泄漏，其实只是缓存没清空。
避坑指南：我曾经调试一个服务，发现每次推理后显存占用都在涨，以为是内存泄漏。查了两天，最后发现是缓存池的 max_split_size_mb 设置不合理，导致大量碎片无法合并。调小这个值后，问题解决。
除了PyTorch自带的缓存池，你还可以自己实现更精细的池化策略。比如，对于固定大小的中间张量（如注意力矩阵），可以预先分配一个池子，推理时直接从池子里取，用完归还，彻底避免分配开销。
class TensorPool:
"""简单的张量池，用于复用固定大小的中间张量"""
def __init__(self, shape, dtype=torch.float32, device='cuda'):
self.pool = []
self.shape = shape
self.dtype = dtype
self.device = device
def get(self):
if self.pool:
return self.pool.pop()
else:
return torch.empty(self.shape, dtype=self.dtype, device=self.device)
def put(self, tensor):
self.pool.append(tensor)
# 使用示例
pool = TensorPool((1024, 1024))
attn = pool.get() # 从池中取
# ... 计算 ...
pool.put(attn) # 用完归还我的经验：在部署GPT类模型时，我常用这种池化技术管理KV Cache。因为KV Cache的大小在推理过程中是固定的，用池化可以避免每次生成token时都重新分配显存，推理延迟能降5-10%。
4.4 三者的协同：一个完整的优化案例
光讲理论没意思，我们来看一个实际场景。假设你在部署一个Transformer模型做在线推理，输入序列长度可变。怎么把这三招用上？
- 视图机制：把
(batch, seq_len, hidden)的输入通过view()变成(batch * seq_len, hidden)，避免数据拷贝。 - in-place操作：LayerNorm和激活函数用in-place版本，减少中间张量。
- 显存池化：KV Cache和注意力分数矩阵用池化管理，避免频繁分配。
这三步走下来，显存占用通常能降20-30%，而且推理速度还有提升——因为减少了内存分配和拷贝的开销。
总结一下：
- in-place操作：省的是「临时变量」的显存
- 张量视图：省的是「数据拷贝」的显存
- 显存池化：省的是「分配释放」的开销和碎片
三者结合，才能把显存用到极致。
嗯，这一章的内容就到这里。记住，显存优化不是一锤子买卖，而是贯穿整个推理管线的工作。每多省一点，你的模型就能跑更大的batch，或者处理更长的序列——这在生产环境中，就是实打实的成本降低。