GPU 内存计算和配置¶

本指南介绍了如何计算 GPU 内存需求以及如何为 vLLM-Omni 阶段正确配置 gpu_memory_utilization。

概述¶

gpu_memory_utilization 是一个关键参数，用于控制每个阶段可用的 GPU 内存量。它被指定为一个介于 0.0 和 1.0 之间的分数，其中： - 0.8 表示 GPU 总内存的 80% - 1.0 表示 GPU 总内存的 100%（不推荐，不留缓冲）

内存如何计算¶

内存分配公式¶

对于每个阶段，vLLM-Omni 计算请求的内存为

requested_memory = total_gpu_memory × gpu_memory_utilization

系统会检查

free_memory ≥ requested_memory

如果不满足此条件，该阶段将无法初始化，并显示内存需求的错误消息。

内存组成部分¶

阶段使用的总内存包括

模型权重：加载到 GPU 上的模型参数的大小
KV 缓存：生成过程中存储键值缓存的内存
激活内存：用于中间计算的临时内存
系统开销：CUDA、PyTorch 和其他系统组件使用的内存
非 PyTorch 内存：在 PyTorch 之外分配的内存（例如，CUDA 图）

示例计算¶

对于具有 80GB 总内存的 GPU： - gpu_memory_utilization: 0.8 → 阶段可用 64GB - gpu_memory_utilization: 0.6 → 阶段可用 48GB - gpu_memory_utilization: 0.15 → 阶段可用 12GB

设置 `gpu_memory_utilization`¶

步骤 1: 确定 GPU 内存¶

首先，检查 GPU 的总内存

# Using nvidia-smi
nvidia-smi --query-gpu=memory.total --format=csv

# Or using Python
python -c "import torch; print(f'{torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')"

步骤 2: 估算模型内存需求¶

对于自回归 (AR) 阶段¶

AR 阶段通常需要更多内存，因为： - 大模型权重 - 注意力的 KV 缓存 - 激活缓冲区

对于扩散/生成阶段¶

扩散阶段（如 code2wav）通常需要较少的内存： - 较小的模型组件 - 不同的内存访问模式

典型值： - 对于大多数扩散阶段，为 0.1 - 0.3

步骤 3: 考虑多阶段场景¶

当多个阶段共享同一 GPU 时，必须确保它们的 gpu_memory_utilization 值之和不超过 1.0。

示例：GPU 0 上的两个阶段

stage_args:
  - stage_id: 0
    runtime:
      devices: "0"
    engine_args:
      gpu_memory_utilization: 0.6  # Uses 60% of GPU 0

  - stage_id: 1
    runtime:
      devices: "0"
    engine_args:
      gpu_memory_utilization: 0.3  # Uses 30% of GPU 0
      # Total: 90% of GPU 0 (safe, leaves 10% buffer)

重要提示： 如果阶段运行在不同的 GPU 上，每个 GPU 最多可以独立使用 1.0。

步骤 4: 考虑张量并行¶

当使用 tensor_parallel_size > 1 时，模型被分割到多个 GPU 上，因此每个 GPU 需要更少的内存。

示例：2 路张量并行

stage_args:
  - stage_id: 0
    runtime:
      devices: "0,1"  # Uses both GPUs
    engine_args:
      tensor_parallel_size: 2
      gpu_memory_utilization: 0.6  # 60% per GPU
      # Model is split, so each GPU uses ~30% of model memory

示例¶

Qwen3-Omni-MoE 在 2x H100-80GB 上¶

stage_args:
  - stage_id: 0  # Thinker stage with TP=2
    runtime:
      devices: "0,1"
    engine_args:
      tensor_parallel_size: 2
      gpu_memory_utilization: 0.6  # 48GB per GPU

  - stage_id: 1  # Talker stage
    runtime:
      devices: "1"
    engine_args:
      gpu_memory_utilization: 0.3  # 24GB on GPU 1

  - stage_id: 2  # Code2Wav stage
    runtime:
      devices: "0"
    engine_args:
      gpu_memory_utilization: 0.1  # 8GB on GPU 0

注意： 在此配置中，阶段 0 和阶段 2 共享 GPU 0，但它们在管道中不同的时间运行，因此它们的内存使用不会重叠。

故障排除¶

错误：“可用内存少于期望的 GPU 内存利用率”¶

这意味着在阶段开始时，GPU 没有足够的可用内存。

解决方案： 1. 通过关闭其他进程释放内存 2. 降低此阶段的 gpu_memory_utilization 3. 使用内存更多的 GPU 4. 将阶段移至其他 GPU

错误: 推理期间 OOM¶

该阶段已初始化，但在处理过程中内存不足。

解决方案： 1. 降低 max_num_batched_tokens 2. 降低运行时配置中的 max_batch_size 3. 稍微降低 gpu_memory_utilization 4. 如果支持，启用量化

内存未充分利用¶

如果您看到内存使用率很低，可以： 1. 提高 gpu_memory_utilization 以允许更大的 KV 缓存 2. 提高 max_num_batched_tokens 以实现更好的批处理 3. 检查是否有其他阶段限制了吞吐量

内存计算的有用公式¶

KV 缓存内存¶

KV 缓存的大小取决于： - 批次中的序列数 - 序列长度（提示 + 生成） - 模型隐藏大小 - 注意力头数量 - 层数

近似公式

kv_cache_memory ≈ batch_size × seq_len × hidden_size × num_layers × 2 × dtype_size

2 用于 k 和 v

模型权重内存¶

model_memory ≈ num_parameters × dtype_size

例如： - FP16 中的 7B 参数：约 14GB - FP32 中的 7B 参数：约 28GB - INT8 中的 7B 参数：约 7GB

激活内存¶

激活内存通常较小，但会随以下因素变化： - 批次大小 - 序列长度 - 模型架构

在推理期间，它通常是模型权重内存的 10-30%。