📅 创建时间：2026-06-03 🏷️ 标签：#内存规划 #Buffer #Arena #MemoryPool #显存管理 #内存复用 #ActivationReuse 📚 前置知识：[[/04-ai/01-llm-engineering/07-llm-evolution]]（LLM 发展脉络） 📚 前置知识：[[07-graph-optimization-passes]]（图优化 Pass） 📚 前置知识：[[08-operator-fusion]]（算子融合） 📚 相关知识：[[03-memory-optimization]]（../training-infra/03-memory-optimization） 📚 相关知识：[[10-layout-optimization]]（Layout 优化）
┌──────────────────────────────────────────────────────────────────────────────┐ │ 🔍 场景：CUDA OOM 但显存只用了 60%？ │ ├──────────────────────────────────────────────────────────────────────────────┤ │ 你的推理服务在 GPU 上跑，每次请求 tensor shape 不同： │ │ - Request A: batch=1, seq_len=128 │ │ - Request B: batch=4, seq_len=512 │ │ - Request C: batch=8, seq_len=256 │ │ │ │ 你发现 GPU 显存碎片化严重，明明 nvidia-smi 显示只用了 60GB/80GB， │ │ PyTorch 却报 CUDA OOM： │ │ "Tried to allocate 2.0 GiB (GPU 0; 80.00 GiB total capacity; │ │ 59.87 GiB already allocated; 20.13 GiB free; │ │ but largest free contiguous block is only 0.5 GiB)" │ └──────────────────────────────────────────────────────────────────────────────┘
内存规划——Buffer 分配与生命周期管理 / Memory Planning, Buffer Allocation, and Lifetime Management
第1节：显存分配的问题——PyTorch 的 CUDA 缓存分配器行为
1.1 CUDA 缓存分配器的工作原理
PyTorch 使用 cudnnAllocCapped 和自定义缓存分配器来管理显存：
python
# PyTorch 显存分配器行为模拟
class PyTorchCUDAMemoryAllocator:
"""
PyTorch CUDA 显存分配器行为
关键特点：
1. 缓存已释放的显存块供后续使用
2. 不立即释放显存给 CUDA（避免重复 cudaFree/cudaMalloc 开销）
3. 按大小桶（bucket）管理缓存
4. 当缓存不足时，向 CUDA 请求新显存
"""
def __init__(self):
# 缓存块，按大小分组
self.cached_blocks = {} # size -> [block1, block2, ...]
self.total_allocated = 0
self.total_cached = 0
def allocate(self, size):
"""
分配显存
策略：
1. 先在缓存中查找大小合适的块
2. 如果找到，直接复用
3. 如果没找到，向 CUDA 申请新块
"""
# 对齐到 16 字节
size = self._align_size(size)
# 检查缓存
if size in self.cached_blocks and self.cached_blocks[size]:
block = self.cached_blocks[size].pop()
self.total_cached -= block.size
return block
# 向 CUDA 申请新显存
block = self._cuda_malloc(size)
self.total_allocated += size
return block
def deallocate(self, block):
"""
释放显存
注意：不会立即归还给 CUDA，而是放入缓存
"""
# 放入缓存
if block.size not in self.cached_blocks:
self.cached_blocks[block.size] = []
self.cached_blocks[block.size].append(block)
self.total_cached += block.size
def _align_size(self, size):
"""对齐到最小分配单元（通常是 16 或 512 字节）"""
alignment = 512
return ((size + alignment - 1) // alignment) * alignment1.2 碎片化问题
python
# 碎片化产生的原因
# 时间线：
# t=1: 分配 100MB, 90MB, 80MB → [100MB][90MB][80MB][remaining]
# t=2: 释放 90MB 块 → [100MB][ 90MB ][80MB][remaining]
# t=3: 需要分配 95MB → 找不到连续的 95MB！
# → 虽然剩余: 90+80+remaining > 95MB
# → 但最大的连续块只有 90MB < 95MB
# → OOM！
# 碎片化类型
# 1. 外部碎片化：可用显存分散在不连续的区域
# 2. 内部碎片化：分配块比实际需求大（对齐导致）
# 解决方案：
# 1. 内存池（Memory Pool）：预先分配大块，切分使用
# 2. 内存碎片整理（Defragmentation）：定期整理碎片
# 3. 分页显存管理（Paged Memory）：类似操作系统虚拟内存1.3 查看 PyTorch 显存状态
python
import torch
# 查看显存摘要
print(torch.cuda.memory_summary())
# 输出示例：
"""
------------------------------------------------------------------------
__PyTorchCUDAAllocator Report
------------------------------------------------------------------------
Device: 0
Backend: CUDA
Allocated memory: 53747 MB
Reserved memory: 58923 MB
Active memory: 53747 MB
Inactive memory: 5176 MB (major fragmentation)
Free memory: 23077 MB
------------------------------------------------------------------------
Process: python (PID: 12345)
------------------------------------------------------------------------
Allocation requests: 2847295
In-use allocator requests: 2340
Cached allocator requests: 1120
Failed allocations: 3
------------------------------------------------------------------------
Segment counts:
Small segments: 123
Large segments: 45
------------------------------------------------------------------------
Memory usage by size class:
1MB: 5 blocks, 5000 MB total
16MB: 3 blocks, 3000 MB total
64MB: 2 blocks, 2000 MB total
...
"""python
# 详细的内存统计
def print_memory_stats():
"""打印详细的显存统计"""
print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"Reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
print(f"Max allocated: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
print(f"Max reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB")
# 查看每个 tensor 的显存占用
for obj in torch.cuda.memory_summary()['allocation_history'].values():
if obj['size'] > 1024 * 1024: # > 1MB
print(f" {obj['name']}: {obj['size'] / 1024**2:.2f} MB")
# 重置峰值统计
torch.cuda.reset_peak_memory_stats()
torch.cuda.reset_accumulated_memory_stats()第2节：Buffer 生命周期分析
2.1 张量生命周期
python
# 张量生命周期分析
class BufferLifecycleAnalyzer:
"""
分析计算图中每个 Buffer 的生命周期
目标：
1. 确定每个 tensor 何时创建
2. 确定每个 tensor 最后被使用的时间
3. 确定每个 tensor 何时可以释放
"""
def analyze(self, execution_trace):
"""
分析执行轨迹
execution_trace: [(timestamp, op_name, input_ids, output_ids), ...]
"""
# 记录每个 tensor 的创建时间和最后使用时间
created_at = {} # tensor_id -> creation_time
last_used = {} # tensor_id -> last_use_time
for timestamp, op_name, inputs, outputs in execution_trace:
# 记录输出 tensor 的创建时间
for out_id in outputs:
created_at[out_id] = timestamp
# 更新输入 tensor 的最后使用时间
for inp_id in inputs:
last_used[inp_id] = timestamp
# 计算每个 tensor 的存活时间
lifetime = {}
for tensor_id in set(created_at.keys()) | set(last_used.keys()):
start = created_at.get(tensor_id, 0)
end = last_used.get(tensor_id, float('inf'))
lifetime[tensor_id] = (start, end)
return lifetime
def identify_reusable_buffers(self, lifetime):
"""
识别可以复用的 buffer
如果两个 buffer 的生命周期不重叠，可以复用同一块显存
"""
reusable = []
tensors = list(lifetime.items())
tensors.sort(key=lambda x: x[1][0]) # 按创建时间排序
# 贪心算法
free_intervals = [(0, float('inf'))]
for tensor_id, (start, end) in tensors:
# 查找可以复用的区间
for i, (free_start, free_end) in enumerate(free_intervals):
if free_end > start:
# 可以复用
reusable.append((tensor_id, free_start, start))
# 更新空闲区间
if free_start < start:
free_intervals[i] = (free_start, start)
else:
del free_intervals[i]
break
return reusable2.2 静态分析 vs 动态分析
python
# 静态生命周期分析
class StaticLifetimeAnalysis:
"""
编译期分析 buffer 生命周期
适用于静态 shape 的情况
"""
def analyze_static_shape(self, graph):
"""
静态 shape 下的生命周期分析
假设：
1. 已知所有 tensor 的 shape
2. 图结构已知
3. 可以进行精确的生命周期分析
"""
# 后向扫描确定每个节点的最后使用位置
last_use = {}
# 从输出反向遍历
output_node = graph.output_node
self._backward_pass(output_node, last_use)
# 前向遍历确定每个节点的最早释放位置
first_free = {}
self._forward_pass(graph.input_node, first_free)
return last_use, first_free
# 动态生命周期分析
class DynamicLifetimeTracking:
"""
运行时跟踪 buffer 生命周期
适用于动态 shape 或复杂控制流
"""
def __init__(self):
self.allocated = {} # tensor_id -> {'ptr': ptr, 'size': size, 'created_at': time}
self.reference_counts = {}
def track_allocation(self, tensor_id, ptr, size):
self.allocated[tensor_id] = {
'ptr': ptr,
'size': size,
'created_at': time.time()
}
self.reference_counts[tensor_id] = 1
def track_reference(self, tensor_id):
"""引用计数 +1"""
self.reference_counts[tensor_id] += 1
def track_dereference(self, tensor_id):
"""引用计数 -1，当引用为 0 时可以释放"""
self.reference_counts[tensor_id] -= 1
if self.reference_counts[tensor_id] == 0:
self._free_buffer(tensor_id)第3节：Arena 分配器——预分配大块显存
3.1 Arena 分配器原理
python
class ArenaAllocator:
"""
Arena（竞技场）分配器
原理：
1. 预先分配一大块显存（Arena）
2. 在 Arena 内部按顺序分配小 buffer
3. Arena 销毁时一次性释放所有显存
4. 避免频繁的 cudaMalloc/cudaFree 调用
优点：
- 减少 CUDA API 调用开销
- 避免碎片化
- 更好的局部性
"""
def __init__(self, device, size_gb):
self.device = device
self.total_size = size_gb * 1024 * 1024 * 1024
# 预分配 Arena
self.arena_ptr = torch.cuda.caching_allocator_alloc(
self.total_size, device=device
)
self.arena_used = 0
self.arena_end = self.total_size
# 分配记录
self.allocations = {} # ptr -> {'offset', 'size', 'name'}
def allocate(self, size, name=''):
"""
从 Arena 分配
对齐到 256 字节
"""
# 对齐
size = ((size + 255) // 256) * 256
if self.arena_used + size > self.total_size:
raise OutOfMemoryError(f"Arena exhausted: {self.arena_used}/{self.total_size}")
offset = self.arena_used
ptr = self.arena_ptr + offset
self.allocations[ptr] = {
'offset': offset,
'size': size,
'name': name
}
self.arena_used += size
return ptr
def release(self):
"""
释放整个 Arena
"""
torch.cuda.caching_allocator_raw_delete(self.arena_ptr)
self.arena_ptr = None
self.arena_used = 03.2 分层 Arena
python
class LayeredArenaAllocator:
"""
分层 Arena 分配器
不同类型的 buffer 使用不同的 Arena：
1. 权重 Arena：存储模型权重，只分配一次
2. Activation Arena：存储中间激活值，按需分配
3. Scratch Arena：临时工作空间
"""
def __init__(self, device):
self.weight_arena = ArenaAllocator(device, size_gb=2) # 2GB 权重
self.activation_arena = ArenaAllocator(device, size_gb=40) # 40GB 激活
self.scratch_arena = ArenaAllocator(device, size_gb=1) # 1GB scratch
def allocate_weight(self, size, name):
"""分配权重 buffer"""
return self.weight_arena.allocate(size, f"weight_{name}")
def allocate_activation(self, size, name):
"""分配激活值 buffer"""
return self.activation_arena.allocate(size, f"act_{name}")
def allocate_scratch(self, size, name):
"""分配临时 buffer"""
return self.scratch_arena.allocate(size, f"scratch_{name}")
def release_all(self):
"""释放所有 Arena"""
self.weight_arena.release()
self.activation_arena.release()
self.scratch_arena.release()第4节：内存复用策略
4.1 静态分配
python
class StaticMemoryPlanner:
"""
静态内存规划
适用于：
1. 编译期确定所有 shape
2. 模型的计算图固定
3. 没有动态控制流
"""
def plan(self, model, example_input):
"""
静态规划内存分配
步骤：
1. 追踪模型执行，记录所有 tensor
2. 分析生命周期
3. 规划 buffer 复用
"""
# Step 1: 执行并记录
execution_trace = []
tensor_map = {} # node_name -> tensor_info
def forward_hook(module, input, output):
node_name = module.name
tensor_map[node_name] = {
'shape': output.shape,
'dtype': output.dtype,
'size': output.nelement() * output.element_size()
}
execution_trace.append({
'name': node_name,
'size': tensor_map[node_name]['size']
})
# 注册 hook
hooks = []
for name, module in model.named_modules():
h = module.register_forward_hook(forward_hook)
hooks.append(h)
# 执行一次
model(example_input)
# 移除 hook
for h in hooks:
h.remove()
# Step 2: 生命周期分析
lifetime = self._analyze_lifetime(execution_trace)
# Step 3: 内存规划
allocation_plan = self._plan_allocation(tensor_map, lifetime)
return allocation_plan
def _plan_allocation(self, tensor_map, lifetime):
"""
规划分配
使用 First-Fit 策略
"""
# 按创建时间排序
sorted_nodes = sorted(lifetime.items(), key=lambda x: x[1][0])
allocation_plan = {}
current_offset = 0
for node_name, (start, end) in sorted_nodes:
size = tensor_map[node_name]['size']
# 对齐
aligned_size = ((size + 255) // 256) * 256
# 分配
allocation_plan[node_name] = {
'offset': current_offset,
'size': aligned_size
}
current_offset += aligned_size
return allocation_plan4.2 动态分配
python
class DynamicMemoryAllocator:
"""
动态内存分配
适用于：
1. 运行时 shape 未知
2. 动态控制流
3. 输入长度变化
"""
def __init__(self, max_memory_gb):
self.max_memory = max_memory_gb * 1024**3
self.current_offset = 0
self.allocated = [] # [(offset, size), ...]
def allocate(self, size):
"""
动态分配
简单策略：顺序分配，定期 GC
"""
# 对齐
size = ((size + 255) // 256) * 256
# 查找空闲位置
offset = self._find_free_slot(size)
if offset is None:
# GC：整理碎片
self._garbage_collect()
offset = self._find_free_slot(size)
if offset is None:
raise OOMError("Cannot allocate")
self.allocated.append((offset, size))
return offset
def _find_free_slot(self, size):
"""First-Fit 查找"""
sorted_allocs = sorted(self.allocated)
# 检查从 0 开始的空间
if not sorted_allocs or sorted_allocs[0][0] >= size:
if sorted_allocs and sorted_allocs[0][0] >= size:
return 0
elif not sorted_allocs:
return 0
# 检查块间间隙
for i in range(len(sorted_allocs) - 1):
gap_start = sorted_allocs[i][0] + sorted_allocs[i][1]
gap_end = sorted_allocs[i+1][0]
if gap_end - gap_start >= size:
return gap_start
# 检查最后一块之后
if sorted_allocs:
last_end = sorted_allocs[-1][0] + sorted_allocs[-1][1]
if last_end + size <= self.max_memory:
return last_end
return None
def _garbage_collect(self):
"""碎片整理"""
# 将所有分配移动到连续区域
sorted_allocs = sorted(self.allocated)
new_offset = 0
compacted = []
for offset, size in sorted_allocs:
compacted.append((new_offset, size))
new_offset += size
self.allocated = compacted
self.current_offset = new_offset4.3 内存池化
python
class MemoryPool:
"""
内存池
核心思想：
1. 预分配多个固定大小的 buffer
2. 不同请求复用相同的 buffer
3. 减少动态分配
"""
def __init__(self):
# 按 size 分桶
self.buckets = {} # size -> [buffer1, buffer2, ...]
self.in_use = {} # buffer -> user_info
def acquire(self, size):
"""
获取 buffer
"""
# 找最小满足条件的 bucket
for bucket_size in sorted(self.buckets.keys()):
if bucket_size >= size and self.buckets[bucket_size]:
buffer = self.buckets[bucket_size].pop()
self.in_use[buffer] = {'size': bucket_size}
return buffer
# 没有可用 buffer，创建新的
new_buffer = self._create_buffer(size)
self.in_use[new_buffer] = {'size': size}
return new_buffer
def release(self, buffer):
"""
归还 buffer
"""
info = self.in_use.pop(buffer, None)
if info:
size = info['size']
if size not in self.buckets:
self.buckets[size] = []
self.buckets[size].append(buffer)
def _create_buffer(self, size):
"""创建新 buffer"""
# 分配 GPU 显存
num_elements = size // 4 # 假设 float32
return torch.empty(num_elements, dtype=torch.float32, device='cuda')第5节：Activation 重用
5.1 Forward Activation 重用
python
class ActivationReusePlanner:
"""
Activation 重用规划
原理：
1. 在 forward pass 中识别可重用的 activation
2. 复用同一块显存存储多个 activation
3. 需要精确的生命周期分析
"""
def identify_reusable_activations(self, graph):
"""
识别可重用的 activation
"""
# 反向扫描确定每个节点的 lifetime
last_use = {}
self._backward_analysis(graph, last_use)
# 找出生命周期不重叠的节点对
reusable_pairs = []
nodes = sorted(last_use.items(), key=lambda x: x[1][0])
for i, (node_a, (_, end_a)) in enumerate(nodes):
for node_b, (start_b, _) in nodes[i+1:]:
if start_b >= end_a:
# 生命周期不重叠，可以复用
size_a = graph[node_a].output_size
size_b = graph[node_b].output_size
if size_a >= size_b:
reusable_pairs.append((node_a, node_b, size_b))
return reusable_pairs
def plan_reuse(self, reusable_pairs):
"""
规划复用
为每个可复用的节点分配 buffer slot
"""
plan = {}
slots = [] # (offset, size, current_owner)
for node_a, node_b, size in reusable_pairs:
# 分配 slot
slot = self._allocate_slot(slots, size)
plan[node_b] = {'slot': slot, 'reuse_from': node_a}
return plan
def _allocate_slot(self, slots, size):
"""分配 slot"""
for i, (offset, slot_size, owner) in enumerate(slots):
if slot_size == size and owner is None:
slots[i] = (offset, slot_size, 'allocated')
return offset
# 创建新 slot
offset = sum(s[0] + s[1] for s in slots)
slots.append((offset, size, 'allocated'))
return offset5.2 Checkpointing（激活重计算）
python
class GradientCheckpointing:
"""
梯度检查点（Activation Checkpointing）
策略：
1. 不保存所有中间 activation
2. 在 backward 时重新计算某些 activation
3. 用计算换显存
"""
def forward_with_checkpoint(self, x, layers):
"""
梯度检查点实现
不保存所有层的输出，只保存部分 checkpoint
"""
# 选择性保存 checkpoint
# 例如：每隔 3 层保存一个
checkpoints = []
current = x
for i, layer in enumerate(layers):
if i % 3 == 0:
# 保存 checkpoint
checkpoints.append(current)
current = layer(current)
# 保存最后一个
checkpoints.append(current)
return checkpoints
def backward_with_recompute(self, grad_output, checkpoints, layers):
"""
包含重计算的 backward
"""
# 从后向前
current = grad_output
for i in reversed(range(len(layers))):
layer = layers[i]
if i % 3 == 0 and i < len(layers) - 1:
# 需要重计算
x = checkpoints[i]
# 前向重计算
for j in range(i, min(i+3, len(layers))):
x = layers[j](x)
else:
x = checkpoints[i]
# 计算梯度
grad_input = layer.backward(grad_output, x)
grad_output = grad_input
return grad_input第6节：PagedAttention（vLLM 的显存管理）
6.1 PagedAttention 原理
python
class PagedAttention:
"""
PagedAttention：分页式 KV Cache 管理
问题：
1. 传统方式：为每个序列预分配最大长度的 KV Cache
2. 浪费严重：短序列浪费大量显存
解决方案：
1. 将 KV Cache 分页管理（类似操作系统虚拟内存）
2. 按需分配，不预分配
3. 支持不连续的物理块
"""
def __init__(self, num_blocks, block_size):
"""
初始化
num_blocks: 物理块数量
block_size: 每个块的 token 数（通常是 16 或 32）
"""
self.num_blocks = num_blocks
self.block_size = block_size
# 物理块存储
self.blocks = {} # block_id -> (k_cache, v_cache)
# 虚拟到物理的映射
self.block_tables = {} # sequence_id -> [physical_block_ids]
# 空闲块管理
self.free_blocks = set(range(num_blocks))
def allocate_sequence(self, sequence_id, max_length):
"""
为序列分配 KV Cache
按需分配物理块
"""
num_blocks_needed = (max_length + self.block_size - 1) // self.block_size
allocated_blocks = []
for _ in range(num_blocks_needed):
if not self.free_blocks:
# 需要 evict
self._evict_lru()
block_id = self.free_blocks.pop()
allocated_blocks.append(block_id)
# 分配物理块
self.blocks[block_id] = self._allocate_physical_block()
self.block_tables[sequence_id] = allocated_blocks
def free_sequence(self, sequence_id):
"""
释放序列的 KV Cache
"""
if sequence_id in self.block_tables:
for block_id in self.block_tables[sequence_id]:
self.free_blocks.add(block_id)
del self.blocks[block_id]
del self.block_tables[sequence_id]
def append_token(self, sequence_id, token_id, k_cache, v_cache):
"""
追加 token 的 K/V
"""
block_table = self.block_tables[sequence_id]
num_blocks = len(block_table)
# 检查是否需要新块
current_length = num_blocks * self.block_size
new_pos = current_length + 1 # 下一个位置
if new_pos > current_length:
# 需要分配新块
new_block_id = self.free_blocks.pop()
self.blocks[new_block_id] = self._allocate_physical_block()
block_table.append(new_block_id)
# 写入最后一页
last_block_id = block_table[-1]
offset_in_block = (new_pos - 1) % self.block_size
self.blocks[last_block_id]['k_cache'][offset_in_block] = k_cache
self.blocks[last_block_id]['v_cache'][offset_in_block] = v_cache6.2 碎片避免
python
class MemoryFragmentationAvoidance:
"""
避免显存碎片化
"""
def __init__(self, total_memory):
self.total_memory = total_memory
self.free_chunks = [(0, total_memory)] # [(start, size), ...]
self.allocated = {} # ptr -> (start, size)
def allocate_best_fit(self, size):
"""
Best-Fit 分配策略
找到最小满足条件的空闲块，减少碎片
"""
best_chunk = None
best_size = float('inf')
for start, chunk_size in self.free_chunks:
if chunk_size >= size and chunk_size < best_size:
best_chunk = (start, chunk_size)
best_size = chunk_size
if best_chunk is None:
raise OOMError("No suitable chunk")
start, chunk_size = best_chunk
self.free_chunks.remove(best_chunk)
# 如果有剩余，分割块
if chunk_size > size:
self.free_chunks.append((start + size, chunk_size - size))
self.allocated[start] = (start, size)
return start
def allocate_worst_fit(self, size):
"""
Worst-Fit 分配策略
分配最大的空闲块，减少碎片
"""
worst_chunk = max(self.free_chunks, key=lambda x: x[1])
if worst_chunk[1] < size:
raise OOMError("No suitable chunk")
self.free_chunks.remove(worst_chunk)
start, chunk_size = worst_chunk
if chunk_size > size:
self.free_chunks.append((start + size, chunk_size - size))
self.allocated[start] = (start, size)
return start第7节：调试 OOM——如何分析显存分配
7.1 OOM 诊断流程
python
def diagnose_oom():
"""
OOM 诊断流程
"""
print("=== OOM 诊断 ===\n")
# 1. 查看当前显存状态
print("1. 显存状态:")
print(f" Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f" Reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
print(f" Max allocated: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
# 2. 查看最大的 tensor
print("\n2. 最大的 Tensors:")
memory_by_tensor = torch.cuda.memory_summary()['allocation_history']
largest = sorted(memory_by_tensor.items(),
key=lambda x: x[1]['size'], reverse=True)[:10]
for tid, info in largest:
print(f" {info.get('name', tid)}: {info['size'] / 1024**2:.2f} MB")
# 3. 查看是否有显存泄漏
print("\n3. 显存泄漏检测:")
leaks = detect_memory_leaks()
if leaks:
print(f" 发现 {len(leaks)} 个潜在泄漏")
for leak in leaks:
print(f" - {leak}")
else:
print(" 未发现明显泄漏")
# 4. 碎片化分析
print("\n4. 碎片化分析:")
fragmentation = analyze_fragmentation()
print(f" 碎片化程度: {fragmentation:.1%}")
if fragmentation > 0.3:
print(" ⚠️ 碎片化严重，建议使用 torch.cuda.empty_cache()")
def detect_memory_leaks():
"""检测显存泄漏"""
# 方法1：比较 allocated 和 reserved
allocated = torch.cuda.memory_allocated()
reserved = torch.cuda.memory_reserved()
leak = reserved - allocated
leaks = []
if leak > 1024**3: # > 1GB
leaks.append(f"缓存中有 {leak / 1024**3:.2f} GB 未使用")
# 方法2：检查是否有 tensor 没有被释放
summary = torch.cuda.memory_summary()
return leaks
def analyze_fragmentation():
"""分析碎片化程度"""
summary = torch.cuda.memory_summary()
# 简单估算：inactive / (inactive + free)
inactive = summary.get('inactive_split_bytes', 0)
free = summary.get('active_bytes', 0) - summary.get('allocated_bytes', 0)
if inactive + free == 0:
return 0.0
return inactive / (inactive + free)7.2 常见 OOM 原因和解决方案
| 问题 | 原因 | 解决方案 |
|---|---|---|
| 碎片化 OOM | 显存可用但碎片化 | 使用 torch.cuda.empty_cache()，或重启进程 |
| 真实 OOM | 显存真的不够 | 减少 batch size，使用 gradient checkpointing |
| 缓存泄漏 | PyTorch 缓存未释放 | 调用 torch.cuda.empty_cache() |
| 预分配过大 | 初始化时分配太多 | 检查模型的 register_buffer |
| 中间 activation 过多 | 融合不够 | 使用融合算子，减少中间结果 |
升华
┌─────────────────────────────────────────────────────────────────────────────┐
│ 内存规划核心原则 │
├─────────────────────────────────────────────────────────────────────────────┤
│ 1. 碎片化是隐形的 OOM │
│ → nvidia-smi 显示还有显存，但 largest contiguous block 不足 │
│ 2. Arena 分配是性能关键 │
│ → 预分配大块，按需切分，避免 cudaMalloc/cudaFree 开销 │
│ 3. 生命周期分析是复用的基础 │
│ → 知道何时创建、何时释放，才能复用 buffer │
│ 4. PagedAttention 是 KV Cache 的最佳实践 │
│ → 分页管理避免碎片化，支持动态长度 │
│ 5. Profiler 驱动优化 │
│ → 使用 memory_summary 和 profiler 找到泄漏和碎片 │
└─────────────────────────────────────────────────────────────────────────────┘"AI 可查 vs 必须理解"清单
必须理解（不理解就等于不会）：
- 🔴 碎片化 OOM 的本质：明明有足够显存但无法分配（largest contiguous block 不足）
- 🔴 PyTorch 缓存分配器行为：释放的显存不会立即还给 CUDA，放入缓存复用
- 🔴 Arena 分配器的原理：预分配大块，按顺序切分，减少碎片
- 🔴 Buffer 生命周期：知道每个 tensor 的创建时间和最后使用时间才能复用
- 🔴 PagedAttention 的核心思想：将 KV Cache 分页管理，避免预分配浪费
AI 可查（知道去哪查就行）：
- ✅ 特定 GPU 的显存大小和特性（如 A100 80GB vs H100 80GB）
- ✅ PyTorch 缓存分配器的具体配置参数
- ✅ vLLM 或 TGI 的显存管理实现细节
- ✅ 特定场景的 batch size 选择建议
- ✅ 其他框架的内存优化技巧
学习状态：🟡 开始学习