第 9 章：CUDA OOM 调试完全指南

总参数内存 = 视觉编码器 + 语言模型 + 连接层

视觉编码器 (CLIP-ViT-L/14): 304M × 2 bytes = 608 MB
语言模型 (Vicuna-7B): 7B × 2 bytes = 14 GB
MLP 连接层: 20M × 2 bytes = 40 MB
总计: 约 14.6 GB

可训练参数梯度 = 总参数 - 冻结参数

优化器状态 = 14 GB × 2 = 28 GB
总内存需求 = 14.6 (参数) + 14.6 (梯度) + 28 (优化器) + 激活值
           > 57.2 GB + 激活值

import torch

def diagnose_memory():
    # 1. 检查当前内存使用
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3

    print(f"已分配: {allocated:.2f} GB")
    print(f"已预留: {reserved:.2f} GB")

    # 2. 打印详细内存快照
    print(torch.cuda.memory_summary())

    # 3. 定位大张量
    for obj in gc.get_objects():
        if torch.is_tensor(obj) and obj.is_cuda:
            print(f"{obj.shape}, {obj.dtype}, {obj.element_size() * obj.nelement() / 1024**2:.2f} MB")

# 开启 gradient checkpointing
model.gradient_checkpointing_enable()

# 对于 VLM，可以选择性开启
vision_encoder.gradient_checkpointing_enable()  # 视觉编码器
language_model.gradient_checkpointing_enable()   # 语言模型

def find_optimal_batch_size(model, initial_bs=32):
    batch_size = initial_bs

    while batch_size > 0:
        try:
            # 尝试前向传播
            dummy_batch = create_dummy_batch(batch_size)
            loss = model(dummy_batch)
            loss.backward()

            print(f"最佳 batch size: {batch_size}")
            return batch_size

        except RuntimeError as e:
            if "out of memory" in str(e):
                # 清理缓存
                torch.cuda.empty_cache()
                # 减半重试
                batch_size = batch_size // 2
            else:
                raise e

    return 1  # 最小 batch size

# 目标：等效 batch size = 32
actual_batch_size = 4  # 受限于显存
accumulation_steps = 32 // 4  # 累积 8 步

for step, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps
    loss.backward()

    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast(dtype=torch.float16):
    outputs = model(inputs)
    loss = criterion(outputs, targets)

# 缩放梯度防止下溢
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# DeepSpeed ZeRO-Offload 配置
ds_config = {
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        }
    }
}

# Pipeline 并行示例
from torch.distributed.pipeline.sync import Pipe

# 将模型分割为两部分
model = nn.Sequential(
    vision_encoder,    # GPU 0
    language_model     # GPU 1
)

# 创建 pipeline
model = Pipe(model, balance=[1, 1], devices=[0, 1])

def analyze_memory_detailed():
    # 获取完整内存报告
    summary = torch.cuda.memory_summary(device=0, abbreviated=False)
    print(summary)

    # 关键指标解读
    stats = torch.cuda.memory_stats()

    print("\n=== 内存使用细分 ===")
    print(f"当前分配: {stats['allocated_bytes.all.current'] / 1024**3:.2f} GB")
    print(f"峰值分配: {stats['allocated_bytes.all.peak'] / 1024**3:.2f} GB")
    print(f"预留内存: {stats['reserved_bytes.all.current'] / 1024**3:.2f} GB")
    print(f"活跃内存块: {stats['active_bytes.all.current'] / 1024**3:.2f} GB")

    # 内存碎片分析
    fragmentation = 1 - (stats['allocated_bytes.all.current'] / 
                        stats['reserved_bytes.all.current'])
    print(f"内存碎片率: {fragmentation * 100:.1f}%")

    # OOM 次数
    print(f"OOM 重试次数: {stats['num_ooms']}")

# 1. 持续监控（每 0.1 秒刷新）
nvidia-smi -l 0.1

# 2. 只显示内存信息
nvidia-smi --query-gpu=memory.used,memory.free,memory.total \
           --format=csv,noheader,nounits -l 1

# 3. 监控特定进程
nvidia-smi pmon -i 0

# 4. 导出详细日志用于分析
nvidia-smi --query-gpu=timestamp,name,memory.used,memory.free,utilization.gpu \
           --format=csv -l 1 > gpu_log.csv

import subprocess
import pandas as pd

def monitor_gpu_memory():
    result = subprocess.run([
        'nvidia-smi', 
        '--query-gpu=memory.used,memory.free,memory.total',
        '--format=csv,noheader,nounits'
    ], capture_output=True, text=True)

    lines = result.stdout.strip().split('\n')
    for i, line in enumerate(lines):
        used, free, total = map(int, line.split(', '))
        usage_percent = (used / total) * 100
        print(f"GPU {i}: {used}/{total} MB ({usage_percent:.1f}%)")

        if usage_percent > 90:
            print(f"⚠️  GPU {i} 内存使用超过 90%！")

from torch.profiler import profile, ProfilerActivity, record_function

def profile_memory_usage(model, dataloader):
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        profile_memory=True,
        record_shapes=True,
        with_stack=True
    ) as prof:

        for i, batch in enumerate(dataloader):
            if i >= 3:  # 只分析前 3 个 batch
                break

            with record_function("forward"):
                outputs = model(batch)

            with record_function("loss"):
                loss = compute_loss(outputs, batch['labels'])

            with record_function("backward"):
                loss.backward()

            with record_function("optimizer"):
                optimizer.step()
                optimizer.zero_grad()

    # 输出分析结果
    print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))

    # 生成 Chrome 追踪文件
    prof.export_chrome_trace("memory_trace.json")

    # 找出内存热点
    for evt in prof.key_averages():
        if evt.cuda_memory_usage > 100 * 1024 * 1024:  # > 100MB
            print(f"内存热点: {evt.key}, 使用: {evt.cuda_memory_usage / 1024**2:.1f} MB")

import threading
import time
import matplotlib.pyplot as plt
from collections import deque

class MemoryMonitor:
    def __init__(self, interval=1.0, max_history=100):
        self.interval = interval
        self.max_history = max_history
        self.memory_history = deque(maxlen=max_history)
        self.time_history = deque(maxlen=max_history)
        self.running = False
        self.thread = None

    def start(self):
        self.running = True
        self.thread = threading.Thread(target=self._monitor_loop)
        self.thread.start()

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join()

    def _monitor_loop(self):
        start_time = time.time()
        while self.running:
            # 记录内存使用
            allocated = torch.cuda.memory_allocated() / 1024**3
            self.memory_history.append(allocated)
            self.time_history.append(time.time() - start_time)

            # 检测异常
            if allocated > 0.9 * torch.cuda.get_device_properties(0).total_memory / 1024**3:
                print(f"⚠️  内存告警: {allocated:.2f} GB")
                self._dump_tensors()

            time.sleep(self.interval)

    def _dump_tensors(self):
        """输出占用内存最大的张量"""
        tensors = []
        for obj in gc.get_objects():
            if torch.is_tensor(obj) and obj.is_cuda:
                tensors.append((
                    obj.numel() * obj.element_size(),
                    str(obj.shape),
                    str(obj.dtype)
                ))

        tensors.sort(reverse=True)
        print("\n=== Top 5 内存占用张量 ===")
        for size, shape, dtype in tensors[:5]:
            print(f"{size / 1024**2:.1f} MB: {shape} ({dtype})")

    def plot(self):
        plt.figure(figsize=(10, 4))
        plt.plot(self.time_history, self.memory_history)
        plt.xlabel('时间 (秒)')
        plt.ylabel('显存使用 (GB)')
        plt.title('训练过程显存监控')
        plt.grid(True)
        plt.show()

# 使用示例
monitor = MemoryMonitor(interval=0.5)
monitor.start()

# 训练代码
train_model()

monitor.stop()
monitor.plot()

def detect_memory_leak(model, dataloader, num_iterations=50):
    """检测训练过程中的内存泄漏"""

    memory_usage = []

    for i, batch in enumerate(dataloader):
        if i >= num_iterations:
            break

        # 训练步骤
        outputs = model(batch)
        loss = compute_loss(outputs, batch['labels'])
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # 记录内存
        torch.cuda.synchronize()
        memory_usage.append(torch.cuda.memory_allocated())

        # 每 10 步检查一次
        if i > 0 and i % 10 == 0:
            # 计算内存增长率
            recent_memory = memory_usage[-10:]
            growth_rate = (recent_memory[-1] - recent_memory[0]) / recent_memory[0]

            if growth_rate > 0.05:  # 增长超过 5%
                print(f"⚠️  可能存在内存泄漏！步骤 {i}, 增长率: {growth_rate:.2%}")

                # 尝试定位泄漏源
                for name, param in model.named_parameters():
                    if param.grad is not None and param.grad.data_ptr() != 0:
                        # 检查梯度是否异常累积
                        if hasattr(param, '_grad_accumulation_count'):
                            if param._grad_accumulation_count > 1:
                                print(f"  梯度累积异常: {name}")

    return memory_usage

# 常见泄漏原因及解决方案
"""

1. 保存了计算图：使用 loss.item() 而不是 loss
2. 列表累积张量：定期清理或使用 .detach()
3. 自定义 autograd 函数：确保正确实现 backward
4. hook 未释放：训练结束后调用 handle.remove()
"""

# 问题代码示例
def process_images(images, vision_encoder):
    # 危险！所有图像同时编码
    features = []
    for img in images:  # images: [B, N, C, H, W]
        feat = vision_encoder(img)  # 每次都保留在显存中
        features.append(feat)
    return torch.stack(features)

单张图像 tokens = (H/patch_size) × (W/patch_size)
ViT-L/14: 1024×1024 图像 → 5184 tokens！
内存 = B × N_images × tokens × hidden_dim × 4 bytes
     = 1 × 4 × 5184 × 1024 × 4 = 84.9 MB（仅激活值）

# 方案 1：批处理优化
def process_images_optimized(images, vision_encoder, max_batch=2):
    B, N, C, H, W = images.shape
    features = []

    # 分批处理
    for i in range(0, N, max_batch):
        batch_images = images[:, i:i+max_batch]
        with torch.cuda.amp.autocast():  # 使用混合精度
            feat = vision_encoder(batch_images)
        features.append(feat)

        # 及时清理
        if i + max_batch < N:
            torch.cuda.empty_cache()

    return torch.cat(features, dim=1)

# 方案 2：动态分辨率策略
def adaptive_resolution(image, base_resolution=336):
    """根据显存动态调整分辨率"""
    available_memory = torch.cuda.mem_get_info()[0] / 1024**3  # GB

    if available_memory < 4:
        return F.interpolate(image, size=(base_resolution, base_resolution))
    elif available_memory < 8:
        return F.interpolate(image, size=(base_resolution*2, base_resolution*2))
    else:
        return image  # 原始分辨率

# 注意力矩阵大小计算
def attention_memory(seq_len, num_heads, batch_size):
    # Q @ K^T 的大小
    memory_bytes = batch_size * num_heads * seq_len * seq_len * 4
    return memory_bytes / 1024**3  # GB

# 示例：2048 tokens, 32 heads, batch_size=1
print(f"注意力矩阵: {attention_memory(2048, 32, 1):.2f} GB")
# 输出: 0.50 GB（单层！）

# 方案 1：Flash Attention
from flash_attn import flash_attn_func

class FlashAttentionVLM(nn.Module):
    def forward(self, q, k, v):
        # Flash Attention：内存从 O(L^2) 降至 O(L)
        return flash_attn_func(q, k, v, causal=False)

# 方案 2：滑动窗口注意力
def sliding_window_attention(q, k, v, window_size=512):
    """只计算局部窗口内的注意力"""
    B, H, L, D = q.shape
    attention_scores = []

    for i in range(0, L, window_size // 2):  # 50% 重叠
        start = max(0, i - window_size // 2)
        end = min(L, i + window_size)

        q_window = q[:, :, start:end]
        k_window = k[:, :, start:end]
        v_window = v[:, :, start:end]

        scores = torch.matmul(q_window, k_window.transpose(-2, -1))
        scores = F.softmax(scores / math.sqrt(D), dim=-1)
        out = torch.matmul(scores, v_window)
        attention_scores.append(out)

    return combine_windows(attention_scores)

# 方案 3：稀疏注意力
class SparseAttentionVLM(nn.Module):
    def __init__(self, sparsity_ratio=0.9):
        super().__init__()
        self.sparsity_ratio = sparsity_ratio

    def forward(self, q, k, v):
        # 只保留 top-k 注意力权重
        scores = torch.matmul(q, k.transpose(-2, -1))

        # 保留 top 10% 的值
        k_val = int((1 - self.sparsity_ratio) * scores.shape[-1])
        topk_scores, topk_indices = torch.topk(scores, k_val, dim=-1)

        # 创建稀疏矩阵
        sparse_scores = torch.zeros_like(scores)
        sparse_scores.scatter_(-1, topk_indices, topk_scores)

        attn_weights = F.softmax(sparse_scores, dim=-1)
        return torch.matmul(attn_weights, v)

# 问题代码
def batch_images_naive(image_list):
    # 所有图像 pad 到最大尺寸 → 内存浪费！
    max_h = max(img.shape[-2] for img in image_list)
    max_w = max(img.shape[-1] for img in image_list)

    padded_images = []
    for img in image_list:
        pad_h = max_h - img.shape[-2]
        pad_w = max_w - img.shape[-1]
        padded = F.pad(img, (0, pad_w, 0, pad_h))
        padded_images.append(padded)

    return torch.stack(padded_images)

# 方案 1：分组批处理
def group_by_resolution(images, num_groups=3):
    """按分辨率分组，减少 padding 浪费"""
    # 计算每张图像的像素数
    resolutions = [img.shape[-2] * img.shape[-1] for img in images]

    # K-means 聚类
    groups = defaultdict(list)
    sorted_indices = np.argsort(resolutions)

    for i, idx in enumerate(sorted_indices):
        group_id = i * num_groups // len(sorted_indices)
        groups[group_id].append(images[idx])

    # 每组单独处理
    processed_groups = []
    for group_images in groups.values():
        batch = batch_images_naive(group_images)  # 组内 padding
        processed_groups.append(batch)

    return processed_groups

# 方案 2：动态分块处理
class DynamicPatchProcessor:
    def __init__(self, base_size=224, max_patches=16):
        self.base_size = base_size
        self.max_patches = max_patches

    def process(self, image):
        H, W = image.shape[-2:]

        # 计算需要的 patch 数量
        n_h = math.ceil(H / self.base_size)
        n_w = math.ceil(W / self.base_size)

        if n_h * n_w > self.max_patches:
            # 降采样以满足内存限制
            scale = math.sqrt(self.max_patches / (n_h * n_w))
            new_h = int(H * scale)
            new_w = int(W * scale)
            image = F.interpolate(image, size=(new_h, new_w))
            n_h = math.ceil(new_h / self.base_size)
            n_w = math.ceil(new_w / self.base_size)

        # 分块处理
        patches = []
        for i in range(n_h):
            for j in range(n_w):
                patch = image[..., 
                             i*self.base_size:(i+1)*self.base_size,
                             j*self.base_size:(j+1)*self.base_size]
                patches.append(patch)

        return patches, (n_h, n_w)

# 交叉注意力内存计算
def cross_attention_memory(text_len, image_tokens, num_layers, hidden_dim):
    # 每层都需要存储 K, V
    kv_memory = 2 * image_tokens * hidden_dim * 4  # bytes

    # 注意力矩阵
    attn_memory = text_len * image_tokens * 4  # bytes

    total = num_layers * (kv_memory + attn_memory)
    return total / 1024**3  # GB

# 示例：1024 text tokens, 576 image tokens, 24 layers
memory = cross_attention_memory(1024, 576, 24, 4096)
print(f"交叉注意力内存: {memory:.2f} GB")

# 方案 1：共享 KV cache
class SharedCrossAttention(nn.Module):
    def __init__(self, num_layers, hidden_dim):
        super().__init__()
        # 只在第一层计算 image KV，后续层复用
        self.image_proj_k = nn.Linear(hidden_dim, hidden_dim)
        self.image_proj_v = nn.Linear(hidden_dim, hidden_dim)
        self.layers = nn.ModuleList([
            CrossAttentionLayer(hidden_dim) for _ in range(num_layers)
        ])

    def forward(self, text_hidden, image_features):
        # 一次性计算所有层的 KV
        image_k = self.image_proj_k(image_features)
        image_v = self.image_proj_v(image_features)

        for layer in self.layers:
            text_hidden = layer(text_hidden, image_k, image_v)

        return text_hidden

# 方案 2：门控交叉注意力
class GatedCrossAttention(nn.Module):
    """只在必要时进行交叉注意力"""
    def __init__(self, hidden_dim, threshold=0.5):
        super().__init__()
        self.gate = nn.Linear(hidden_dim, 1)
        self.threshold = threshold
        self.cross_attn = CrossAttentionLayer(hidden_dim)

    def forward(self, text_hidden, image_features):
        # 计算门控值
        gate_scores = torch.sigmoid(self.gate(text_hidden.mean(dim=1)))

        if gate_scores.mean() > self.threshold:
            # 执行交叉注意力
            return self.cross_attn(text_hidden, image_features)
        else:
            # 跳过，节省内存
            return text_hidden

# 方案 3：低秩分解
class LowRankCrossAttention(nn.Module):
    """使用低秩分解减少参数和内存"""
    def __init__(self, hidden_dim, rank=64):
        super().__init__()
        self.rank = rank

        # 分解 W_q, W_k, W_v
        self.q_down = nn.Linear(hidden_dim, rank, bias=False)
        self.q_up = nn.Linear(rank, hidden_dim, bias=False)

        self.kv_down = nn.Linear(hidden_dim, rank * 2, bias=False)
        self.kv_up = nn.Linear(rank * 2, hidden_dim * 2, bias=False)

    def forward(self, text_hidden, image_features):
        # 低秩投影
        q = self.q_up(self.q_down(text_hidden))
        kv = self.kv_up(self.kv_down(image_features))
        k, v = kv.chunk(2, dim=-1)

        # 标准注意力计算
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
        attn = F.softmax(scores, dim=-1)
        return torch.matmul(attn, v)

class MemoryOptimizedVLM:
    """集成所有内存优化技术的 VLM"""

    def __init__(self, config):
        self.config = config
        self.setup_memory_optimization()

    def setup_memory_optimization(self):
        # 1. 启用梯度检查点
        if self.config.gradient_checkpointing:
            self.model.gradient_checkpointing_enable()

        # 2. 使用 Flash Attention
        if self.config.use_flash_attention:
            replace_attention_with_flash_attention(self.model)

        # 3. 混合精度训练
        if self.config.mixed_precision:
            self.scaler = torch.cuda.amp.GradScaler()

        # 4. 内存监控
        self.memory_monitor = MemoryMonitor(interval=10)

    def train_step(self, batch):
        # 动态调整 batch size
        if self.should_reduce_batch_size():
            batch = self.split_batch(batch)

        # 分组处理多分辨率图像
        image_groups = self.group_images_by_resolution(batch['images'])

        total_loss = 0
        for images in image_groups:
            with torch.cuda.amp.autocast(enabled=self.config.mixed_precision):
                # 前向传播
                outputs = self.model(images, batch['text'])
                loss = self.criterion(outputs, batch['labels'])

            # 反向传播
            if self.config.mixed_precision:
                self.scaler.scale(loss).backward()
            else:
                loss.backward()

            total_loss += loss.item()

            # 及时清理
            if len(image_groups) > 1:
                torch.cuda.empty_cache()

        return total_loss / len(image_groups)

    def should_reduce_batch_size(self):
        """动态检测是否需要减小 batch size"""
        memory_usage = torch.cuda.memory_allocated() / torch.cuda.max_memory_allocated()
        return memory_usage > 0.9

RuntimeError: CUDA out of memory. Tried to allocate 2.50 GiB 
(GPU 0; 23.69 GiB total capacity; 21.45 GiB already allocated; 
1.89 GiB free; 21.50 GiB reserved in total by PyTorch)

def forward(self, images, text_ids):
    # 视觉编码
    B, N, C, H, W = images.shape
    all_features = []
    for i in range(B):
        img_features = []
        for j in range(N):
            feat = self.vision_encoder(images[i, j])
            img_features.append(feat)
        all_features.append(torch.stack(img_features))
    vision_features = torch.stack(all_features)

    # 文本嵌入
    text_embeds = self.text_embedder(text_ids)

    # 交叉注意力
    for layer in self.cross_attention_layers:
        text_embeds = layer(text_embeds, vision_features)

    return text_embeds

Step 100: Loss=2.34, Memory=18.2GB
Step 200: Loss=2.11, Memory=18.5GB  
Step 300: Loss=1.98, Memory=18.9GB
Step 400: Loss=1.87, Memory=19.4GB
Step 500: Loss=1.76, Memory=20.1GB
Step 600: Loss=1.65, Memory=20.9GB
Step 700: Loss=1.54, Memory=21.8GB
Step 800: Loss=1.43, Memory=22.9GB
Step 900: Loss=1.32, Memory=24.2GB
Step 1000: CUDA OOM

# 错误理解
print(f"已用内存: {torch.cuda.memory_reserved()}")  # 错！这是预留的

# 正确
print(f"实际使用: {torch.cuda.memory_allocated()}")
print(f"PyTorch 预留: {torch.cuda.memory_reserved()}")
print(f"可用于分配: {torch.cuda.memory_reserved() - torch.cuda.memory_allocated()}")

# 危险：注意力内存是 O(L^2)
seq_len = 4096
memory_gb = (seq_len ** 2 * 4) / 1024**3  # 64 MB 仅单个头！

# 安全：使用 Flash Attention 或分块

# 问题：batch 中一张大图导致整体 OOM
images = [img1_224x224, img2_224x224, img3_1024x1024]  # 第三张导致 OOM

# 解决：预先排序和分组
images.sort(key=lambda x: x.shape[-2] * x.shape[-1])

# 错误：对小模型使用反而更慢
tiny_model.gradient_checkpointing_enable()  # 2 层模型，收益为负

# 正确：只对深层模型使用
if model.num_layers >= 12:
    model.gradient_checkpointing_enable()

# 容易忽视：Adam 需要 2 倍参数内存
# 7B 模型 + Adam = 14GB (参数) + 14GB (梯度) + 28GB (优化器) = 56GB！

# 考虑使用 8-bit Adam 或 Adafactor

# 慢：频繁的小批量传输
for img in images:
    img = img.cuda()  # 每次传输开销大

# 快：批量传输
images = torch.stack(images).cuda()  # 一次传输

# 导致碎片化：频繁分配不同大小
for size in [100, 1000, 10, 10000, 1]:
    tensor = torch.randn(size).cuda()

# 监控碎片化
fragmentation = 1 - (torch.cuda.memory_allocated() / torch.cuda.memory_reserved())
if fragmentation > 0.3:
    torch.cuda.empty_cache()  # 整理内存

# 错误：每个进程都加载完整模型
model = load_model()  # 每个 GPU 都有完整副本

# 正确：使用 DDP 或 FSDP
model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)