GPU 显存管理与模型量化：从 OOM 到上线的工程路径

模型能跑了不代表能上线。

你在 Colab 上用 FP16 跑通了一个 7B 模型的 demo，兴高采烈地准备部署。然后发现生产环境只有 24GB 的 4090，你要跑 70B 的模型。24GB 装不下 140GB 的权重，物理现实不讲道理。

那就量化。INT8 把显存砍一半，INT4 再砍一半。听起来完美。然后你发现量化后模型开始输出乱码，或者在特定任务上准确率暴降 20%。

这篇文章讲的就是这个：从理解显存的每一个字节去了哪里，到选择正确的量化方法让模型在有限显存里跑出可接受的精度。不是论文综述，是工程路径。

一、GPU 显存的组成

要解决 OOM，首先得知道显存里装了什么。LLM 推理阶段的显存消耗分为四大块：

模型权重

这是最大的固定开销。一个参数在不同精度下占用的字节数：

计算公式很直接：

显存（字节） = 参数量 × 每参数字节数

FP32 下一个 70B 模型需要 280GB——这已经超过了单张 H100 的 80GB。所以生产环境几乎没有人用 FP32 做推理，FP16 是起点，量化是必选项。

优化器状态

训练时，Adam 优化器为每个参数维护两个状态（一阶矩和二阶矩），在 FP32 下额外增加 参数量 × 8 字节。但推理阶段不需要优化器，这一块为零。提到它是因为很多 OOM 问题出现在 fine-tuning 场景——同样的卡，训练比推理多吃 2-3 倍显存，优化器是主要原因。

激活值

前向计算的中间结果。推理时只需要当前层的激活值（计算完就丢），所以显存占用很小——通常只有几百 MB。训练时为了反向传播要保留所有层的激活值，显存占用与模型深度成正比。

KV-Cache

推理阶段的显存大头（除权重外）。Transformer 的自注意力机制需要缓存之前所有 token 的 Key 和 Value 向量，避免重复计算。

KV-Cache 的显存公式：

KV-Cache（字节） = 2 × num_layers × num_kv_heads × head_dim × seq_len × batch_size × bytes_per_param

以 LLaMA-2-70B 为例（80 层，8 个 KV head，128 维，FP16）：

# LLaMA-2-70B KV-Cache 计算
num_layers = 80
num_kv_heads = 8      # GQA: 8 个 KV head
head_dim = 128
seq_len = 4096
batch_size = 16
bytes_per_param = 2   # FP16

kv_cache_bytes = 2 * num_layers * num_kv_heads * head_dim * seq_len * batch_size * bytes_per_param
kv_cache_gb = kv_cache_bytes / (1024 ** 3)
print(f"KV-Cache: {kv_cache_gb:.1f} GB")  # KV-Cache: 20.0 GB

20GB 的 KV-Cache。如果序列长度翻倍到 8192，KV-Cache 也翻倍到 40GB。这就是为什么长上下文场景下，KV-Cache 比模型权重还吃显存。

总览

推理阶段的显存预算：

总显存 ≈ 模型权重 + KV-Cache + 激活值 + 框架开销（~500MB-1GB）

以 70B 模型、4096 序列、batch=16 为例：

INT4 量化把总显存从 162GB 砍到 57GB——从需要 3 张 H100 变成 1 张就够。这就是量化的价值。

二、显存碎片和 OOM

知道了显存的组成，下一个问题是：明明还有剩余显存，为什么还是 OOM？

PyTorch 的 CUDA 内存分配器

PyTorch 不直接调用 cudaMalloc/cudaFree——这些 CUDA API 太慢了（毫秒级）。它用了一个 caching allocator：

1. 第一次分配时向 CUDA 申请一大块内存（通常是 2MB 的倍数）
2. 后续的小分配从这个缓存池里切割
3. 释放时不归还给 CUDA，而是标记为”可重用”
4. 只有调用 torch.cuda.empty_cache() 才把空闲块归还给 CUDA

这个设计在正常场景下很高效，但有一个致命问题：碎片化。

碎片化：小块散落导致大分配失败

假设你先后分配了 [1GB, 256MB, 1GB, 256MB, 1GB]，然后释放所有 1GB 的块。池子里有 3GB 空闲，但它们被 256MB 的块隔开了——不连续。这时如果你需要一个 2GB 的连续分配，OOM。

import torch

# 模拟碎片化
tensors = []
for i in range(20):
    big = torch.randn(256, 1024, 1024, device='cuda')     # ~1GB
    small = torch.randn(64, 1024, 1024, device='cuda')     # ~256MB
    tensors.append((big, small))

# 只释放大 tensor
for big, small in tensors:
    del big

# 此时有大量空闲显存，但都是碎片
torch.cuda.empty_cache()

# 尝试分配一个大的连续块 — 可能 OOM
try:
    huge = torch.randn(512, 1024, 1024, device='cuda')  # ~2GB
except torch.cuda.OutOfMemoryError as e:
    print(f"OOM! 但空闲显存有 {torch.cuda.mem_get_info()[0]/1e9:.1f}GB")

torch.cuda.empty_cache() 的局限

很多人 OOM 了就加一行 empty_cache()，然后发现没用。原因：

• empty_cache() 只释放 PyTorch 标记为空闲的块——如果 tensor 还被引用着，它不是空闲的
• 它不能合并被占用块之间的碎片
• 它把内存归还给 CUDA，但下次分配时又要重新 cudaMalloc，有性能开销

真正有用的是找到谁在持有不该持有的显存。

显存泄漏的常见原因

# 1. 计算图未释放 — 忘记 torch.no_grad() 或 .detach()
output = model(input)        # output 持有整个计算图
loss_value = output.mean()   # loss_value 也持有计算图
# 正确做法：
with torch.no_grad():
    output = model(input)

# 2. 列表/字典累积 tensor
history = []
for batch in dataloader:
    output = model(batch)
    history.append(output)  # 每个 output 都留在 GPU 上！
# 正确做法：
    history.append(output.detach().cpu())

# 3. 全局变量引用
global_ref = model(some_input)  # 永远不会被 GC 回收

# 4. DataLoader 的 pin_memory 泄漏
# pin_memory=True 会分配 page-locked 内存，某些情况下不会自动释放

nvidia-smi 和 PyTorch 内存统计的差异

你经常会发现 nvidia-smi 显示用了 20GB，但 PyTorch 说只分配了 15GB。差在哪？

# nvidia-smi 看的是进程级 GPU 内存
nvidia-smi --query-gpu=memory.used --format=csv

# PyTorch 看的是自己分配器管理的内存
python3 -c "
import torch
print(f'已分配: {torch.cuda.memory_allocated()/1e9:.2f} GB')
print(f'已缓存: {torch.cuda.memory_reserved()/1e9:.2f} GB')
print(f'最大分配: {torch.cuda.max_memory_allocated()/1e9:.2f} GB')
"

差异来源：

调试技巧：用 torch.cuda.memory_summary() 打印完整的内存分配报告，包括碎片信息：

print(torch.cuda.memory_summary(device='cuda:0', abbreviated=False))

输出会显示每个 block size 的分配/释放次数、当前活跃块数、碎片率。这比 nvidia-smi 有用得多。

OOM 快速诊断流程图

OOM 了别慌，按这个决策树走：

Step 1: nvidia-smi 看全局
  └─ GPU 显存占用 > 95%？
      ├─ 是 → 进入 Step 2
      └─ 否 → 大概率是碎片化问题，跳到 Step 3

Step 2: torch.cuda.memory_summary() 看 PyTorch 内部
  └─ reserved 和 allocated 差距大（>2GB）？
      ├─ 是 → 碎片化严重。尝试 torch.cuda.empty_cache()
      │         然后减小 batch size 或用 gradient checkpointing
      └─ 否 → 确实是模型/数据太大，进入 Step 3

Step 3: 找到最大的显存消费者
  └─ 用下面的诊断脚本找 top-10 tensor
      ├─ 发现意外的大 tensor → 检查是否忘记 .detach() 或 torch.no_grad()
      └─ 所有 tensor 都合理 → 进入 Step 4

Step 4: 检查梯度累积泄漏
  └─ 训练循环里是否有 loss.backward() 后没有 optimizer.zero_grad()？
      ├─ 是 → 梯度不断累积，每轮多吃一份显存
      └─ 否 → 进入 Step 5

Step 5: 检查 DataLoader num_workers 共享内存
  └─ /dev/shm 满了吗？（docker 默认只有 64MB）
      ├─ 是 → docker run --shm-size=8g 或减少 num_workers
      └─ 否 → 更复杂的问题，上 torch.cuda.memory_snapshot() 做完整分析

配合使用的诊断脚本——打印当前 GPU 上最大的 10 个 tensor：

import torch
import gc

def print_gpu_top10():
    """打印 GPU 上最大的 10 个 tensor，帮助定位 OOM"""
    tensors = []
    for obj in gc.get_objects():
        try:
            if torch.is_tensor(obj) and obj.is_cuda:
                tensors.append((obj.element_size() * obj.nelement(), obj.shape, obj.dtype))
            elif hasattr(obj, 'data') and torch.is_tensor(obj.data) and obj.data.is_cuda:
                tensors.append((obj.data.element_size() * obj.data.nelement(), obj.data.shape, obj.data.dtype))
        except Exception:
            pass

    tensors.sort(key=lambda x: x[0], reverse=True)
    print(f"{'排名':>4} | {'大小':>12} | {'形状':<30} | {'类型'}")
    print("-" * 70)
    for i, (size, shape, dtype) in enumerate(tensors[:10]):
        print(f"{i+1:>4} | {size/1e6:>9.1f} MB | {str(list(shape)):<30} | {dtype}")

# 在 OOM 前调用：
# print_gpu_top10()

这个脚本遍历 Python GC 里所有活跃对象，所以不要在热路径上用。OOM 排查时手动调用一次就够了。

三、MIG 和 MPS：多租户 GPU 隔离

一张 A100 80GB 跑一个 7B 的 INT8 模型只用 7GB 显存，剩下 73GB 空着。浪费。你想在同一张卡上同时跑多个模型或多个用户的任务。

这就引出了 GPU 的多租户隔离问题。和 CPU 侧的资源隔离思路一样——Cgroups 通过内核机制限制进程的 CPU、内存、IO 用量——GPU 也需要类似的隔离手段。

MIG（Multi-Instance GPU）：硬件分区

MIG 是 NVIDIA 在 A100/H100 上引入的硬件级分区技术。它把一张 GPU 切分成最多 7 个独立实例，每个实例有：

• 独立的显存（物理隔离，不是软件限制）
• 独立的 SM（Streaming Multiprocessor）计算单元
• 独立的 L2 cache 切片
• 独立的显存控制器

# 启用 MIG 模式
sudo nvidia-smi -i 0 -mig 1

# 查看可用的 GPU Instance profile
nvidia-smi mig -lgip

# 创建 GPU Instance（以 A100 80GB 为例）
# 3g.40gb = 3 个 SM slice + 40GB 显存
sudo nvidia-smi mig -cgi 9,9 -i 0  # 创建 2 个 3g.40gb 实例

# 创建 Compute Instance
sudo nvidia-smi mig -cci -i 0

# 查看 MIG 设备
nvidia-smi mig -lgi

A100 80GB 的分区粒度：

关键限制：MIG 分区是静态的。创建/销毁实例需要先终止所有 GPU 进程，不能在运行时动态调整。这意味着你不能根据负载动态扩缩分区。

MPS（Multi-Process Service）：时分复用

MPS 是更早的多任务方案。它不做硬件分区，而是让多个进程共享同一个 CUDA context，通过时分复用方式交替执行 kernel：

# 启动 MPS 服务
export CUDA_MPS_PIPE_DIRECTORY=/run/nvidia-mps
export CUDA_MPS_LOG_DIRECTORY=/var/log/nvidia-mps
nvidia-cuda-mps-control -d

# 设置线程百分比（限制算力）
echo "set_default_active_thread_percentage 50" | nvidia-cuda-mps-control

# 停止 MPS
echo quit | nvidia-cuda-mps-control

MPS 的优点是零开销共享，不需要重启 GPU。缺点是没有显存隔离——一个进程可以吃掉所有显存，导致其他进程 OOM。另外，一个进程的 GPU fault 会影响所有共享进程。

与 CPU 的 cgroup 资源隔离类比

CPU 的 cgroup 隔离可以在不重启进程的情况下动态调整资源限制——echo 2G > memory.max 就生效了。MIG 做不到这一点，这是它最大的局限。

Kubernetes 中的 GPU 调度

在 K8s 集群中，NVIDIA 的 device plugin 可以直接暴露 MIG 实例为独立资源：

# Pod 请求 MIG 实例
apiVersion: v1
kind: Pod
metadata:
  name: llm-inference
spec:
  containers:
  - name: model-server
    image: vllm/vllm-openai:latest
    resources:
      limits:
        nvidia.com/mig-3g.40gb: 1   # 请求一个 3g.40gb MIG 实例

# 查看集群中的 MIG 资源
kubectl describe node gpu-node-01 | grep nvidia.com/mig
#  nvidia.com/mig-1g.10gb:  7
#  nvidia.com/mig-3g.40gb:  2

MIG vs MPS vs 独占：选型对比

四、模型量化基础

解决了显存管理的问题，接下来看如何从根源上减少显存占用——量化。

量化的数学本质

量化就是把浮点数映射到低精度整数。核心公式：

x_quant = round(x / scale) + zero_point
x_dequant = (x_quant - zero_point) × scale

其中 scale 和 zero_point 决定了映射关系。以 INT8 为例，把 [-1.0, 1.0] 范围的浮点数映射到 [-128, 127] 的整数：

import torch

def quantize_tensor(x: torch.Tensor, num_bits: int = 8):
    """对称量化示例"""
    qmin = -(2 ** (num_bits - 1))
    qmax = 2 ** (num_bits - 1) - 1

    x_max = x.abs().max()
    scale = x_max / qmax

    x_quant = torch.clamp(torch.round(x / scale), qmin, qmax).to(torch.int8)
    return x_quant, scale

def dequantize_tensor(x_quant: torch.Tensor, scale: float):
    """反量化"""
    return x_quant.float() * scale

# 演示
weight = torch.randn(4, 4)
w_quant, scale = quantize_tensor(weight)
w_dequant = dequantize_tensor(w_quant, scale)

print(f"原始:   {weight[0]}")
print(f"量化后: {w_quant[0]}")
print(f"反量化: {w_dequant[0]}")
print(f"误差:   {(weight - w_dequant).abs().mean():.6f}")

对称量化 vs 非对称量化

对称量化：zero_point = 0，浮点零点精确映射到整数零点。公式简单，计算快，但如果权重分布不对称（比如 ReLU 后全是正数），会浪费一半的表示范围。

非对称量化：zero_point ≠ 0，可以覆盖不对称的分布。多一次减法运算，但对激活值的量化通常更友好。

def asymmetric_quantize(x: torch.Tensor, num_bits: int = 8):
    """非对称量化"""
    qmin, qmax = 0, 2 ** num_bits - 1
    x_min, x_max = x.min(), x.max()

    scale = (x_max - x_min) / (qmax - qmin)
    zero_point = round((-x_min / scale).item())
    zero_point = max(qmin, min(qmax, zero_point))

    x_quant = torch.clamp(torch.round(x / scale + zero_point), qmin, qmax).to(torch.uint8)
    return x_quant, scale, zero_point

Per-tensor vs Per-channel vs Per-group

量化粒度决定了精度-开销的权衡：

• Per-tensor：整个 tensor 共享一个 scale。最快，但如果不同 channel 的数值范围差异大（比如一个 channel 范围 [-0.01, 0.01]，另一个 [-10, 10]），小数值 channel 的精度几乎为零
• Per-channel：每个输出 channel 一个 scale。LLM 的线性层通常用这种——每列权重有自己的缩放因子，精度好很多
• Per-group：把每个 channel 再分成组（通常 128 个元素一组），每组一个 scale。GPTQ 和 AWQ 默认用 group_size=128。精度最好，但 scale 参数的额外显存开销也最大

动态量化 vs 静态量化 vs 量化感知训练

LLM 场景的现实：QAT 几乎不可行（70B 模型训练一次要几百张卡跑几周），所以主流方法都是 PTQ（Post-Training Quantization），也就是拿训练好的模型直接量化。

五、主流量化方法实战

理论讲完了，看实际用什么工具。

GPTQ：用 Hessian 补偿量化误差

GPTQ（2022）的核心思路：逐层量化权重时，用 Hessian 矩阵的信息来补偿量化引入的误差。量化一列权重后，调整剩余未量化列来最小化整层的输出误差。

# 使用 AutoGPTQ 量化模型
from transformers import AutoModelForCausalLM, AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

model_name = "meta-llama/Llama-2-7b-hf"
quant_config = BaseQuantizeConfig(
    bits=4,
    group_size=128,
    desc_act=False,       # 激活值顺序无关（更快）
    damp_percent=0.01,    # Hessian 正则化
)

# 加载模型和校准数据
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoGPTQForCausalLM.from_pretrained(model_name, quant_config)

# 准备校准数据（通常 128 条样本就够）
calibration_data = [
    tokenizer(text, return_tensors="pt", max_length=2048, truncation=True)
    for text in calibration_texts[:128]
]

# 量化（7B 模型 ~15-30 分钟）
model.quantize(calibration_data)
model.save_quantized("llama-2-7b-gptq-4bit")

优点：INT4 下精度损失小，生态成熟（Hugging Face 直接支持）。 缺点：量化过程慢（需要逐层计算 Hessian），校准数据的选择影响结果。

AWQ：保护重要权重

AWQ（Activation-Aware Weight Quantization，2023）观察到一个现象：只有 ~1% 的权重对模型输出影响极大（“salient weights”）。这些权重对应的激活值通道 magnitude 特别大。

AWQ 的做法：不均匀量化，对重要权重保留更高精度（通过 per-channel scaling）。

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_name = "meta-llama/Llama-2-7b-hf"
quant_path = "llama-2-7b-awq-4bit"

# 加载
model = AutoAWQForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# 量化配置
quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM",  # GEMM kernel，比 GEMV 快
}

# 量化（比 GPTQ 快 ~3-5x）
model.quantize(tokenizer, quant_config=quant_config)
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

优点：量化速度快（不需要 Hessian），INT4 精度通常优于 GPTQ，和 vLLM 配合好。 缺点：依赖激活值分布的统计——如果校准集和实际分布差异大，效果会下降。

SmoothQuant：平滑激活值

SmoothQuant（2022）解决的问题是：权重分布通常比较均匀（容易量化），但激活值经常有 outlier（某些 channel 的值比其他 channel 大 100 倍）。直接量化激活值会严重丢精度。

核心思路：在量化前，把激活值的”尖峰”通过数学等价变换（per-channel scaling）转移到权重上。权重变得不那么均匀了，但比激活值容易处理。

Y = X · W = (X · diag(s)^{-1}) · (diag(s) · W)

其中 s 是 per-channel 的平滑因子，让 X' 和 W' 的量化难度更均衡。

GGUF / llama.cpp：CPU 友好的量化

不是所有人都有 GPU。GGUF（GPT-Generated Unified Format）是 llama.cpp 使用的量化格式，主要面向 CPU 推理和消费级硬件：

# 安装 llama.cpp
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp && make -j

# 转换 HF 模型到 GGUF
python3 convert_hf_to_gguf.py ../Llama-2-7b-hf/ --outtype f16

# 量化（多种精度可选）
./llama-quantize ../Llama-2-7b-hf/model-f16.gguf \
    ../Llama-2-7b-hf/model-Q4_K_M.gguf Q4_K_M

# 运行推理
./llama-cli -m ../Llama-2-7b-hf/model-Q4_K_M.gguf \
    -p "Explain GPU memory management" -n 256

GGUF 的量化类型命名规则：Q{bits}_{group_type}_{size}。常用的：

• Q4_K_M：4-bit，K-quant，Medium size（推荐平衡点）
• Q5_K_M：5-bit，精度更好，稍大
• Q8_0：8-bit，精度接近 FP16，大小翻倍
• Q2_K：2-bit，精度很差，只适合测试

各方法精度-速度权衡

实际部署代码：vLLM + AWQ

# 使用 vLLM 部署 AWQ 量化模型
from vllm import LLM, SamplingParams

# vLLM 自动检测 AWQ 格式并使用优化 kernel
llm = LLM(
    model="TheBloke/Llama-2-70B-Chat-AWQ",
    quantization="awq",
    tensor_parallel_size=2,       # 2 卡并行
    gpu_memory_utilization=0.90,  # 显存利用率上限
    max_model_len=4096,
    enforce_eager=False,          # 使用 CUDA Graph 加速
)

sampling = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=512)
outputs = llm.generate(["解释 GPU 显存管理的关键概念"], sampling)
print(outputs[0].outputs[0].text)

# 或者用 vLLM 的 OpenAI 兼容服务器
python3 -m vllm.entrypoints.openai.api_server \
    --model TheBloke/Llama-2-70B-Chat-AWQ \
    --quantization awq \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.90 \
    --max-model-len 4096 \
    --port 8000

# 测试
curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{"model": "TheBloke/Llama-2-70B-Chat-AWQ", "prompt": "Hello", "max_tokens": 50}'

量化精度 vs 性能权衡

上面的表格是按量化方法分的。换个视角——如果你关心的是具体 benchmark 上掉了多少分，这张表更直接：

⚠️ 以上数据基于 Llama-2-70B 的公开评测，不同模型差异很大。7B 模型对量化更敏感（冗余少），数学/代码任务比对话任务掉分更明显。永远要在你自己的任务上验证。

几个反直觉的结论：

1. AWQ 通常比 GPTQ 精度更好——尽管两者都是 4-bit。AWQ 保护了对精度影响最大的权重通道。
2. FP8 是 H100 时代的甜点——精度损失几乎可以忽略，但速度接近 INT8。唯一限制是需要 Ada/Hopper 架构。
3. INT4 在 GSM8K 上掉分最明显——数学推理对精度损失最敏感，如果你的场景涉及计算，三思后再用 4-bit。

六、量化选型和最佳实践

工具都有了，怎么选？

INT8 vs INT4：什么场景用什么精度

用 INT8 的场景： - 对精度要求高的任务（代码生成、数学推理、长文档总结） - 显存预算允许（比如 2×A100 跑 70B） - 需要 SmoothQuant 处理激活值 outlier

用 INT4 的场景： - 显存受限（单卡 24GB 要跑 70B） - 对话/聊天/文本生成等对轻微精度损失不敏感的任务 - 需要最大化并发 batch size

经验法则：先试 INT8，如果显存不够再降到 INT4。不要直接跳到 INT4——你可能白丢了精度。

FP8：H100 时代的新选择

H100 的 Transformer Engine 原生支持 FP8 运算，这是 INT8 和 FP16 之间一个重要的新精度选项。FP8 有两种格式：

为什么 FP8 值得关注：

• 精度优于 INT8 PTQ：浮点格式天然适配神经网络权重分布，不需要校准数据集也能获得不错的精度
• 速度接近 INT8：H100 的 FP8 Tensor Core 吞吐量与 INT8 相当，远超 FP16
• 硬件原生支持：不是模拟的，H100/H200/B100 Transformer Engine 直接执行 FP8 矩阵乘法

实际使用：

# vLLM 中使用 FP8 量化
python3 -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-70b-chat-hf \
    --dtype fp8 \
    --tensor-parallel-size 2

# TensorRT-LLM 中使用 FP8
trtllm-build --checkpoint_dir ./llama-70b \
    --use_fp8 \
    --output_dir ./llama-70b-fp8-engine

FP8 在整体精度-速度光谱中的位置：

经验法则：如果你有 H100，先试 FP8 再考虑 INT8——同等显存节省，更好的精度，更少的量化工程。

量化后的评估

量化不是一锤子买卖。你需要验证量化模型的质量：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

def compute_perplexity(model, tokenizer, dataset, max_length=2048):
    """计算困惑度（perplexity）— 越低越好"""
    model.eval()
    total_loss = 0
    total_tokens = 0

    for sample in dataset:
        inputs = tokenizer(
            sample["text"], return_tensors="pt",
            max_length=max_length, truncation=True
        ).to(model.device)

        with torch.no_grad():
            outputs = model(**inputs, labels=inputs["input_ids"])

        total_loss += outputs.loss.item() * inputs["input_ids"].shape[1]
        total_tokens += inputs["input_ids"].shape[1]

    ppl = torch.exp(torch.tensor(total_loss / total_tokens))
    return ppl.item()

# 对比 FP16 和 INT4 的困惑度
wikitext = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")

fp16_ppl = compute_perplexity(fp16_model, tokenizer, wikitext)
int4_ppl = compute_perplexity(int4_model, tokenizer, wikitext)

print(f"FP16 PPL: {fp16_ppl:.2f}")
print(f"INT4 PPL: {int4_ppl:.2f}")
print(f"PPL 增长: {(int4_ppl - fp16_ppl) / fp16_ppl * 100:.1f}%")

困惑度增长的经验阈值：

除了困惑度，还要在实际下游任务上评估：MMLU（知识）、HumanEval（代码）、GSM8K（数学）。有些量化方法困惑度很低但在特定任务上精度暴降。

常见问题

问题 1：量化后输出乱码

通常是 group_size 太小或量化方法不匹配。检查清单：

# 1. 确认量化配置
print(model.config.quantization_config)

# 2. 确认 tokenizer 和模型匹配
assert tokenizer.vocab_size == model.config.vocab_size

# 3. 尝试增大 group_size（128 → 256）或提高精度（INT4 → INT8）

问题 2：特定任务精度暴降

模型量化对不同任务的影响不均匀。代码生成和数学推理对量化最敏感（因为一个 token 错就全错），闲聊和摘要最不敏感。如果你的核心场景是代码生成，INT8 比 INT4 安全得多。

问题 3：量化模型加载慢

GPTQ/AWQ 模型首次加载时需要反序列化量化权重，比 FP16 慢。解决方案：

# 用 safetensors 格式替代 bin 格式（加载速度 2-3x）
# 大多数 Hugging Face 量化模型已经提供 safetensors

# vLLM 支持模型预加载到共享内存
python3 -m vllm.entrypoints.openai.api_server \
    --model path/to/model \
    --load-format safetensors

推荐流程

不要一上来就 INT4。遵循这个流程：

FP16 基线  →  评估（PPL + 下游任务）
    ↓
INT8 量化（SmoothQuant 或 BitsAndBytes）  →  评估
    ↓
如果显存够 → 部署 INT8
如果显存不够 ↓
    ↓
INT4 量化（AWQ 或 GPTQ）  →  评估
    ↓
精度可接受 → 部署 INT4
精度不可接受 → 换量化方法 / 换更大的卡 / 用 tensor parallel

每一步都要评估。跳过评估直接部署是生产事故的第一步。

与推理框架的配合

不同的推理框架对量化格式的支持不同：

vLLM + AWQ 是当前 GPU 推理的最优解之一：

# vLLM 的 PagedAttention 会自动管理 KV-Cache 的显存
# 结合 AWQ 量化，70B 模型在 2×A100 上可以：
# - batch_size=32, seq_len=4096
# - 吞吐 ~2000 tokens/s
# - 显存利用率 >90%

TensorRT-LLM 则适合对延迟有极致要求的场景：

# TRT-LLM 构建 INT8 引擎
python3 convert_checkpoint.py \
    --model_dir ./Llama-2-70b-hf \
    --output_dir ./trt_ckpt \
    --dtype float16 \
    --use_smooth_quant \
    --per_token \
    --per_channel

trtllm-build \
    --checkpoint_dir ./trt_ckpt \
    --output_dir ./trt_engine \
    --gemm_plugin float16 \
    --max_batch_size 32 \
    --max_input_len 2048 \
    --max_seq_len 4096

结语

GPU 显存管理不是调一个参数的事。从理解显存的四大组成部分（权重、KV-Cache、激活值、框架开销），到处理碎片化和 OOM，到用 MIG/MPS 做多租户隔离，到选择合适的量化方法——每一步都是工程决策，需要在精度、速度、成本之间取舍。

核心原则：

1. 先量化权重，再优化 KV-Cache——权重是固定开销，KV-Cache 随负载变化
2. INT8 先行，INT4 后备——不要过早丢精度
3. 每次量化后都评估——困惑度 + 下游任务，两个都看
4. 选对推理框架——框架的量化支持比量化算法本身更影响最终性能

从 OOM 到上线不是一条直线，而是一个评估-调整的循环。但只要你理解了显存的每个字节去了哪里，这条路就不再是黑箱。

延伸阅读：

• Cgroups v2：让容器不能吃掉整台机器 — GPU 的 MIG 隔离 vs CPU 的 cgroup 隔离，思路一脉相承
• 让 LLM 帮你写系统代码：哪些能信，哪些会死 — 量化后的模型在代码生成上到底掉多少精度
• AI 生成的”高性能”代码到底有多快 — 用 perf 和 cachegrind 看 AI 代码的真实性能