sglang启动时CUDA OOM如何排查显存占用？

Question

在使用 SGLang 启动大语言模型时，常出现 CUDA Out-of-Memory（OOM）错误。问题多源于显存占用过高，尤其是在加载大规模模型（如7B以上参数量）时。需排查点包括：模型权重是否以 FP16/INT8 量化加载、是否有冗余副本驻留显存、并行策略配置不当导致显存重复分配，以及运行时缓存（如KV Cache）预分配过大。此外，其他进程（如残留训练任务或推理服务）可能占用显卡资源。建议结合 `nvidia-smi` 与 `torch.cuda.memory_allocated()` 实时监控显存，定位峰值占用环节，进而优化模型加载方式或调整批处理大小。

Answer 1

1. 常见现象与初步诊断

在使用 SGLang 启动大语言模型（如 LLaMA-7B、ChatGLM 等）时，CUDA Out-of-Memory (OOM) 错误是高频问题。典型表现为程序启动失败或推理过程中突然崩溃，报错信息如：cuda runtime error (2): out of memory。这类错误通常出现在显存容量有限的 GPU 上（如 24GB 的 A10 或 3090），尤其当加载 FP32 格式的模型权重时，7B 模型即可占用超过 30GB 显存。

初步排查建议从以下三方面入手：

1. 运行 nvidia-smi 查看当前 GPU 显存使用情况，确认是否有其他进程（如残留的训练任务、旧版推理服务）正在占用资源；
2. 检查模型加载脚本是否显式指定了数据类型（如未设置 torch.float16 或量化模式）；
3. 验证批处理大小（batch size）是否过高，特别是在生成长文本时，KV Cache 会随序列长度线性增长。

2. 显存占用构成分析

理解大模型推理过程中的显存分布是解决 OOM 的关键。下表列出了一个典型 7B 参数模型在不同配置下的显存消耗估算：

组件	FP32 (GB)	FP16 (GB)	INT8 (GB)	备注
模型权重	28	14	7	7B × 4 / 2 / 1 bytes
KV Cache	-	4~8	4~8	依赖 batch_size 和 seq_len
激活值（Activations）	-	2~5	2~5	前向传播中间结果
优化器状态	56	0	0	仅训练阶段存在
总估测（推理）	—	~20–25	~13–18	需留出余量给系统开销

可见，通过将模型从 FP32 转为 FP16 或 INT8，可显著降低权重显存占用。此外，KV Cache 预分配策略若未优化（如固定最大长度），极易导致内存溢出。

3. 深层原因排查路径

结合实践经验，OOM 的根本原因往往不止于单一因素，而是多个环节叠加所致。以下是系统化的排查流程图：

graph TD
    A[CUDA OOM Error] --> B{nvidia-smi 是否显示高占用?}
    B -- 是 --> C[检查是否有残留进程并 kill]
    B -- 否 --> D[进入代码级分析]
    D --> E[是否启用 FP16/INT8 加载?]
    E -- 否 --> F[修改 load_policy 为 half 或 quantized]
    E -- 是 --> G[检查并行策略: tensor_parallel_size]
    G --> H{是否存在多卡冗余复制?}
    H -- 是 --> I[调整 parallel_config 避免重复加载]
    H -- 否 --> J[监控 torch.cuda.memory_allocated()]
    J --> K[定位峰值发生在 model.load 还是 generate 阶段]
    K -- load 阶段 --> F
    K -- generate 阶段 --> L[减小 batch_size 或 max_tokens]

4. 关键解决方案与最佳实践

针对上述分析，提出以下可落地的技术方案：

• 量化加载：在 SGLang 中可通过设置 dtype=torch.float16 或启用 AWQ/GPTQ 量化插件实现 INT4/INT8 推理。示例代码如下：

from sglang import Runtime

runtime = Runtime(
    model_path="meta-llama/Llama-2-7b-chat-hf",
    dtype="float16",  # 或 "int8", "awq"
    tensor_parallel_size=2,
    mem_fraction_static=0.8  # 控制 KV Cache 分配上限
)

• 动态批处理与缓存管理：启用 PagedAttention 技术（SGLang 默认支持），避免为每个请求预分配完整 KV Cache。通过 context_length 和 max_num_sequence 限制并发请求数。
• 资源隔离：使用 CUDA_VISIBLE_DEVICES=0 sglang launch ... 明确指定设备，防止跨卡干扰。
• 运行时监控：在关键节点插入显存检测逻辑：

import torch
print(f"Memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

该方法可用于识别 load_model、tokenizer.encode、generate 等阶段的显存跃升点，辅助性能调优。