vLLM内核探秘-第7章 模型加载与权重管理

杨艺韬2026-04-16 9:38

《vLLM 内核探秘》完整目录

第7章 模型加载与权重管理

"Moving data is the bottleneck, not computing on it." -- Jeff Dean

:::tip 本章要点

  • 理解模型加载的三阶段流水线：发现 → 加载 → 分配
  • 深入 BaseModelLoader 的策略模式：8 种加载器适配不同场景
  • 掌握 _prepare_weights 的格式检测逻辑：safetensors → PyTorch → 回退链
  • 深入张量并行下的权重切分策略（列切分 vs 行切分）
  • 理解量化模型的特殊加载流程和注册机制
  • 认识模型注册机制：vLLM 如何用统一接口支持数百种模型架构 :::

源码版本：本章基于 vLLM v0.8.5 源码分析。获取方式：

bash 复制代码 
git clone --branch v0.8.5 https://github.com/vllm-project/vllm.git

7.1 问题的规模：为什么模型加载是个工程问题

在讨论实现细节之前，先感受一下问题的规模。

一个 Llama-3-70B 模型，FP16 格式下：

makefile 复制代码 
参数量:     70 × 10⁹
每参数字节: 2 (FP16)
总权重:     140 GB
分片文件:   ~30 个 safetensors 文件，每个 ~4.7 GB

把 140 GB 从磁盘搬到 GPU 显存，面临的工程挑战：

graph TD D["磁盘\n140 GB safetensors"] -->|"磁盘IO\n~2 GB/s SSD"| CPU["CPU 内存\n需要足够大"] CPU -->|"PCIe 传输\n~25 GB/s"| GPU["GPU 显存\n80 GB A100"] Challenge1["挑战1: 内存\n140GB > 单卡80GB\n需要多卡切分"] -.-> CPU Challenge2["挑战2: 格式\nsafetensors/PT/GGUF/量化\n几十种格式"] -.-> D Challenge3["挑战3: 并行\n张量并行需要精确切分\n每层切法不同"] -.-> GPU

如果采用朴素方法（torch.load → tensor.cuda()），加载 70B 模型需要：

  • 280 GB CPU 内存（PyTorch load 会创建副本）
  • ~70 秒磁盘读取 + ~6 秒 PCIe 传输
  • 无法在单卡上运行（80 GB < 140 GB）

vLLM 的模型加载子系统就是为了解决这些问题而设计的。

7.2 加载器的策略模式

vLLM 的加载逻辑不是一个巨大的 if-else，而是一个优雅的策略模式 。核心是 BaseModelLoader 抽象基类，定义在 vllm/model_executor/model_loader/loader.py:193：

python 复制代码 
# vllm/model_executor/model_loader/loader.py:193-208
class BaseModelLoader(ABC):
    """Base class for model loaders."""

    def __init__(self, load_config: LoadConfig):
        self.load_config = load_config

    @abstractmethod
    def download_model(self, model_config: ModelConfig) -> None:
        """Download a model so that it can be immediately loaded."""
        raise NotImplementedError

    @abstractmethod
    def load_model(self, *, vllm_config: VllmConfig) -> nn.Module:
        """Load a model with the given configurations."""
        raise NotImplementedError

接口极简------只有两个抽象方法：download_model（下载权重）和 load_model（加载到 GPU）。每种加载格式实现一个具体的 Loader 子类：

classDiagram class BaseModelLoader { <> +load_config: LoadConfig +download_model()* +load_model()* } class DefaultModelLoader { +_prepare_weights() +_get_weights_iterator() +load_model() } class DummyModelLoader { +load_model() } class TensorizerLoader { +load_model() } class ShardedStateLoader { +load_model() } class BitsAndBytesModelLoader { +load_model() } class GGUFModelLoader { +load_model() } class RunaiModelStreamerLoader { +load_model() } BaseModelLoader <|-- DefaultModelLoader BaseModelLoader <|-- DummyModelLoader BaseModelLoader <|-- TensorizerLoader BaseModelLoader <|-- ShardedStateLoader BaseModelLoader <|-- BitsAndBytesModelLoader BaseModelLoader <|-- GGUFModelLoader BaseModelLoader <|-- RunaiModelStreamerLoader

vLLM v0.8.5 中有 7 个具体的加载器 （loader.py 中 grep class.*BaseModelLoader 可以看到）：

加载器类 源文件行号 适用场景
DefaultModelLoader loader.py:210 标准 HuggingFace 模型（safetensors/PT）
DummyModelLoader loader.py:476 性能测试，用随机数填充权重
TensorizerLoader loader.py:503 CoreWeave Tensorizer 格式（序列化加速）
ShardedStateLoader loader.py:603 预分片的 state dict（跳过运行时切分）
BitsAndBytesModelLoader loader.py:792 bitsandbytes 量化（NF4/FP4）
GGUFModelLoader loader.py:1317 GGUF 格式（llama.cpp 兼容）
RunaiModelStreamerLoader loader.py:1415 Run:ai 流式加载

这种策略模式的核心价值：添加新的加载格式只需要实现一个新的 Loader 子类，不需要修改任何已有代码。

7.3 DefaultModelLoader：深入最常用的加载路径

大多数用户使用的是 DefaultModelLoader。它的加载流程可以拆成四步：

flowchart TD A["1. _prepare_weights\n定位权重文件\n确定格式(safetensors/PT)"] --> B["2. _get_weights_iterator\n创建权重迭代器\n按需逐张量读取"] B --> C["3. model.load_weights()\n模型类的自定义加载\n处理名称映射"] C --> D["4. device_loading_context\nCPU → GPU 传输\n张量并行切分"] style A fill:#3b82f6,color:#fff,stroke:none style B fill:#6366f1,color:#fff,stroke:none style C fill:#8b5cf6,color:#fff,stroke:none style D fill:#10b981,color:#fff,stroke:none

7.3.1 权重发现：_prepare_weights

_prepare_weights（loader.py:270）是加载的入口。它做三件事：

  1. 判断是本地路径还是远程模型 ID，如果是远程就下载到本地
  2. 确定加载格式------safetensors 优先，PyTorch 兜底
  3. 过滤权重文件------去除重复分片、非推理所需的文件

关键的格式检测逻辑（loader.py:288-303）：

python 复制代码 
# vllm/model_executor/model_loader/loader.py:288-303
if load_format == LoadFormat.AUTO:
    allow_patterns = ["*.safetensors", "*.bin"]        # 自动检测
elif load_format == LoadFormat.SAFETENSORS:
    use_safetensors = True
    allow_patterns = ["*.safetensors"]                  # 强制 safetensors
elif load_format == LoadFormat.MISTRAL:
    use_safetensors = True
    allow_patterns = ["consolidated*.safetensors"]      # Mistral 专用格式
    index_file = "consolidated.safetensors.index.json"
elif load_format == LoadFormat.PT:
    allow_patterns = ["*.pt"]                            # 强制 PyTorch
elif load_format == LoadFormat.NPCACHE:
    allow_patterns = ["*.bin"]                            # NumPy 缓存

注意 LoadFormat.AUTO 模式下的回退链 ：先尝试 *.safetensors，找不到才用 *.bin。这是因为 safetensors 支持内存映射（mmap），加载时不需要将整个文件读入 CPU 内存，对于 140 GB 的模型这个差异是致命的。

7.3.2 分片文件去重

大模型通常被分片为多个文件。但有时一个模型仓库中同时存在分片文件和合并文件，直接加载会导致权重重复。filter_duplicate_safetensors_files（weight_utils.py）通过读取 model.safetensors.index.json 索引文件来去重：

python 复制代码 
# weight_utils.py 中的去重逻辑（简化）
def filter_duplicate_safetensors_files(files, folder, index_file):
    index_path = os.path.join(folder, index_file)
    if os.path.isfile(index_path):
        with open(index_path) as f:
            index = json.load(f)
        # 只保留索引中引用的文件
        needed = set(index["weight_map"].values())
        return [f for f in files if os.path.basename(f) in needed]
    return files

7.3.3 权重迭代器：惰性加载的关键

_get_weights_iterator 不是一次性加载所有权重，而是返回一个惰性迭代器 ------每次 yield 一个 (name, tensor) 对。这意味着：

  • 对于 safetensors：直接 mmap 读取单个张量，内存占用约等于单个张量大小
  • 对于 PyTorch：需要加载整个文件，但通过 np_cache_weights_iterator 可以将反序列化后的权重缓存为 NumPy 数组，避免重复加载
python 复制代码 
# 根据格式选择迭代器（loader.py:381-395 简化）
if use_safetensors:
    if load_format == LoadFormat.FASTSAFETENSORS:
        weights_iterator = fastsafetensors_weights_iterator(...)
    else:
        weights_iterator = safetensors_weights_iterator(...)
elif load_format == LoadFormat.NPCACHE:
    weights_iterator = np_cache_weights_iterator(...)
else:
    weights_iterator = pt_weights_iterator(...)

7.4 权重名映射与模型类的 load_weights

HuggingFace 模型的权重名（如 model.layers.0.self_attn.q_proj.weight）和 vLLM 内部模型的参数名通常不同。每个模型实现类通过 load_weights 方法处理这种映射。

以 Llama 为例（vllm/model_executor/models/llama.py）：

python 复制代码 
# 简化的 load_weights 逻辑
class LlamaForCausalLM(nn.Module):
    def load_weights(self, weights: Iterable[Tuple[str, Tensor]]):
        # 定义名称映射和合并规则
        stacked_params_mapping = [
            # (vllm参数名, HF权重前缀, 分片索引)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

        for name, loaded_weight in weights:
            # 处理 QKV 合并：HF 的 q_proj/k_proj/v_proj
            # → vLLM 的单一 qkv_proj（减少内核调用）
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name in name:
                    name = name.replace(weight_name, param_name)
                    param = self.state_dict()[name]
                    weight_loader = param.weight_loader
                    weight_loader(param, loaded_weight, shard_id)
                    break
            else:
                # 非合并权重，直接加载
                param = self.state_dict()[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)

这里有一个精妙的优化：QKV 合并 。HuggingFace 的 Llama 将 Q、K、V 存为三个独立的 Linear 层，而 vLLM 将它们合并为一个 qkv_proj。合并的好处是：

  1. 一次矩阵乘法代替三次 → 减少 GPU 内核启动开销
  2. 一次内存读取代替三次 → 更好的 GPU 内存带宽利用
  3. 在张量并行时，三个投影可以一起按列切分

7.5 张量并行下的权重切分

当模型太大无法放入单卡时，需要张量并行（Tensor Parallelism）------将权重切分到多张 GPU。切分策略取决于层的类型，核心思想来自 Megatron-LM 论文。

graph TB subgraph "列切分 (Column Parallel)" W1["权重 [H, 4H]\n一整个矩阵"] W1 --> |"GPU 0 取前半列"| G0["[H, 2H]"] W1 --> |"GPU 1 取后半列"| G1["[H, 2H]"] end subgraph "行切分 (Row Parallel)" W2["权重 [4H, H]\n一整个矩阵"] W2 --> |"GPU 0 取前半行"| G2["[2H, H]"] W2 --> |"GPU 1 取后半行"| G3["[2H, H]"] end

vLLM 中的并行层定义在 vllm/model_executor/layers/linear.py：

层类型 vLLM 类 切分维度 原理
Q/K/V 投影 QKVParallelLinear 按列（输出维度） 每卡处理部分注意力头
Attention 输出投影 RowParallelLinear 按行（输入维度） 配合 QKV 列切分后的 AllReduce
MLP gate/up 投影 MergedColumnParallelLinear 按列 分割中间维度
MLP down 投影 RowParallelLinear 按行 配合 gate/up 列切分后的 AllReduce
词嵌入层 按行切分（词表维度） 按行 每卡只存部分词表的 embedding

切分在加载时完成------每张卡只加载自己负责的那部分权重，无需先加载完整权重再切分：

python 复制代码 
# 每张卡根据自己的 rank 选取对应的权重切片
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()

# 列切分示例：输出维度按 TP 切分
shard_size = param.shape[0] // tp_size
start = tp_rank * shard_size
end = start + shard_size
loaded_weight = loaded_weight[start:end, :]

param.data.copy_(loaded_weight)

7.6 量化模型的特殊加载流程

量化模型的权重格式与原始浮点模型完全不同。以 GPTQ 4-bit 为例：

yaml 复制代码 
原始 FP16:  每参数 16 bit → [4096, 4096] 占 32 MB
GPTQ 4-bit: 每参数 4 bit + scale + zero_point → ~8 MB (4x 压缩)

vLLM 通过量化配置注册机制 支持多种量化格式。核心在 vllm/model_executor/layers/quantization/ 目录下：

python 复制代码 
# 量化方法注册（简化）
@register_quantization_config("gptq")
class GPTQConfig(QuantizationConfig):
    def get_quant_method(self, layer, prefix):
        return GPTQMarlinLinearMethod(self)

@register_quantization_config("fp8")
class Fp8Config(QuantizationConfig):
    def get_quant_method(self, layer, prefix):
        return Fp8LinearMethod(self)

@register_quantization_config("awq")
class AWQConfig(QuantizationConfig):
    def get_quant_method(self, layer, prefix):
        return AWQMarlinLinearMethod(self)

量化加载的特殊之处：

  1. 权重解包：4-bit 整数可能被 pack 成 int32（8 个 4-bit 值 packed 在一个 int32 中），加载时需要按特定布局解包
  2. 量化参数加载：除了权重本身，还需加载 scale（缩放因子）和 zero_point（零点），它们通常存在同一个 safetensors 文件中
  3. 内核匹配：不同的量化格式需要不同的 CUDA 内核执行矩阵乘法。例如 GPTQ 模型优先使用 Marlin 内核（高性能 4-bit 矩阵乘），不可用时回退到 ExLlama 内核

BitsAndBytesModelLoader（loader.py:792）是一个特殊的加载器，它在加载时实时量化：读取 FP16 权重，然后在 GPU 上执行 NF4/FP4 量化。这意味着你可以对任何 FP16 模型使用 bitsandbytes 量化，无需预先量化。

7.7 模型注册与可插拔架构

vLLM 支持数百种模型架构。核心机制是 vllm/model_executor/models/registry.py 中的模型注册表。

python 复制代码 
# 注册表将 HuggingFace 架构名映射到 vLLM 实现
# 每条记录: "HF架构名" -> ("模块路径", "类名")
_TEXT_GENERATION_MODELS = {
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),  # 架构相同，复用
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
    "DeepseekV3ForCausalLM": ("deepseek_v3", "DeepseekV3ForCausalLM"),
    # ... 数百种模型
}

每个模型实现类遵循统一的三方法接口：

classDiagram class ModelInterface { <> +__init__(vllm_config) +forward(input_ids, positions, kv_caches) logits +load_weights(weights_iterator) } class LlamaForCausalLM { +__init__() 构建计算图 +forward() 前向传播 +load_weights() 权重映射+加载 } class Qwen2ForCausalLM { +__init__() +forward() +load_weights() } ModelInterface <|.. LlamaForCausalLM ModelInterface <|.. Qwen2ForCausalLM

这种设计实现了模型层和引擎层的完全解耦 。EngineCore、调度器、KV Cache 管理器都不需要知道模型的具体架构------它们只关心统一的 forward(input_ids, positions, kv_caches) → logits 签名。

添加新模型的步骤：

  1. 在 vllm/model_executor/models/ 下创建新文件
  2. 实现 __init__、forward、load_weights 三个方法
  3. 在注册表中添加一行映射

7.8 性能优化：加载速度的工程细节

7.8.1 safetensors 的 mmap 优势

safetensors 格式支持 mmap（内存映射文件），这意味着读取单个张量时：

  • 不需要将整个文件加载到内存
  • 操作系统按需从磁盘读取页面
  • 多个进程可以共享同一份 mmap

对于 140 GB 的模型，这将 CPU 内存需求从 280 GB（PyTorch load 的两倍开销）降到几乎为零。

7.8.2 device_loading_context：智能的设备迁移

device_loading_context（loader.py:68）是一个 context manager，用于在加载权重时管理 CPU → GPU 的迁移：

python 复制代码 
# loader.py:68-91
@contextmanager
def device_loading_context(module: torch.nn.Module,
                           target_device: torch.device):
    if target_device.type == "cpu":
        yield module
        return

    original_device_states: Dict[str, torch.device] = {}

    # 先把所有参数移到 CPU（节省 GPU 显存）
    for name, p in module.named_parameters():
        if p.device.type == "cpu":
            original_device_states[name] = p.device

    yield module

    # 加载完毕后，把参数移到 GPU
    for name, p in module.named_parameters():
        if name in original_device_states:
            p.data = p.data.to(target_device)

这避免了在加载过程中 GPU 显存同时存在旧权重和新权重（会导致 OOM）。

7.8.3 预分片加载

ShardedStateLoader（loader.py:603）支持加载预先按张量并行切分的权重。这跳过了运行时切分步骤------每张卡直接加载自己那份权重文件。适合需要频繁重启的生产环境：

bash 复制代码 
# 首次加载后保存分片（一次性开销）
vllm save-sharded-state --model meta-llama/Llama-3-70B --output /models/llama-70b-tp4/

# 后续加载（跳过切分步骤，更快）
vllm serve /models/llama-70b-tp4/ --load-format sharded_state

7.9 本章小结

模型加载是 LLM 推理的"冷启动"阶段，它的质量直接影响服务的启动速度和资源效率：

  1. 策略模式 --- BaseModelLoader + 7 个具体加载器，适配 safetensors/PT/GGUF/Tensorizer/BnB 等格式
  2. 智能格式检测 --- _prepare_weights 实现 safetensors → PyTorch 的自动回退链
  3. 惰性迭代器 --- 按需逐张量读取，避免一次性加载全部权重到 CPU 内存
  4. QKV 合并 --- 将 HuggingFace 的三个独立投影合并为一个，减少 GPU 内核调用
  5. 张量并行切分 --- 遵循 Megatron-LM 方案，列切分 QKV/gate/up，行切分 output/down
  6. 量化注册机制 --- @register_quantization_config 支持可插拔的量化方法（GPTQ/AWQ/FP8/BnB）
  7. 模型注册机制 --- 统一的三方法接口（init/forward/load_weights）实现引擎与模型完全解耦
  8. mmap + 预分片 --- safetensors 内存映射和 ShardedStateLoader 分别优化首次和重复加载

下一章，我们将进入模型前向传播的核心------ModelRunner，看看 CUDA Graph 和持久化批次如何将 GPU 利用率推到极限。

源码导航（vLLM v0.8.5）

  • 加载器基类与策略：vllm/model_executor/model_loader/loader.py
  • 权重工具函数：vllm/model_executor/model_loader/weight_utils.py
  • 模型注册表：vllm/model_executor/models/registry.py
  • 模型实现（Llama）：vllm/model_executor/models/llama.py
  • 并行层定义：vllm/model_executor/layers/linear.py
  • 量化配置注册：vllm/model_executor/layers/quantization/
相关推荐
storyseek3 小时前
拆解 DeerFlowd：一个开源 Super Agent Harness 是怎么做出来的
agent·harness
杨艺韬3 小时前
vLLM内核探秘-第13章 量化引擎：精度与速度的平衡
agent
杨艺韬3 小时前
vLLM内核探秘-第18章 设计模式与架构哲学
agent
杨艺韬3 小时前
vLLM内核探秘-第10章 前缀缓存：零开销的加速
agent
杨艺韬4 小时前
Harness Engineering-第4章 上下文工程：比 Prompt Engineering 更重要的事
agent
杨艺韬4 小时前
vLLM内核探秘-第9章 采样与输出处理
agent
杨艺韬4 小时前
Harness Engineering-前言
agent
杨艺韬4 小时前
Harness Engineering-第2章 Agent 架构模式全景
agent
杨艺韬4 小时前
Harness Engineering-第3章 Agent Loop：心跳与决策循环
agent
热门推荐
012026年4月技术前沿：AI大模型爆发、智能体革命与量子安全新纪元02GitHub 镜像站点03一周AI热点速览（2026.03.31-04.06）：GPT-6曝光、谷歌开源Gemma 4、资本狂飙与模型军备竞赛04GPT-6核心能力解析及与现有主流大模型对比05基于 Docker 部署 Hermes Agent 并接入飞书机器人的完整指南06GPT-6发布日深度解析-Symphony架构200万Token实战07零成本！Ollama本地部署国产大模型全指南（支持Kimi-K2.5/GLM-5/Qwen，新手秒上手）08AI Weekly | 2026年4月第二周 · GitHub热门项目与AI发展趋势深度解析09从限购到畅通：GLM-5.1 Coding Plan接入攻略10CodeBuddy与WorkBuddy深度对比：腾讯两款AI工具差异及实操指南