vLLM 推理内核深度解析

第7章 模型加载与权重管理:从磁盘到 GPU 的 140 GB 旅程

作者 杨艺韬 · 12,025 字 · 发布于 · 更新于

140 GB 权重、30 个分片文件、25 GB/s 的 PCIe——模型加载不是算法问题,是一道纯粹的数据管道工程题。

本章要点

  • 用 Llama-3-70B 真实数据量化:140 GB 权重 × 30 个分片文件 × PCIe 25 GB/s × TP 切分 × 量化变体 = 一个纯粹的"数据管道工程"
  • 掌握 BaseModelLoader 策略模式下的 7 种加载器:各自适配什么场景、什么时候该用哪个
  • 读懂 safetensors 为什么是默认优先路径:按 shard、按 tensor 迭代,避免把完整 checkpoint 一次性反序列化进内存
  • 理解 _prepare_weights 的四步流水线:远程解析 → 文件下载 → 分片去重 → 格式确认
  • 看懂 stacked_params_mapping 的妙用:HuggingFace 的 q_proj/k_proj/v_proj 如何被 vLLM 合并成单一 qkv_proj
  • 掌握 Megatron 式权重切分:Column vs Row 在 load 时就完成,不是 load 完再切
  • 理解量化权重加载的三重差异:bit packing、scale/zero_point 处理、kernel 匹配
  • 学会 Tensorizer / Run:ai streamer / ShardedState 三种加速首次加载的手段
  • 知道如何为 vLLM 添加一个新模型:三个方法接口 + 一行注册表
  • 拿到四个场景的加载优化配方:单机快速启动 / 多副本冷启动 / K8s 滚动重启 / 极速开发迭代

7.1 问题规模:为什么"加载模型"是一个工程问题

打开 HuggingFace 上的 meta-llama/Meta-Llama-3-70B-Instruct,你会看到:

config.json                         1 KB
tokenizer.json                      9 MB
model-00001-of-00030.safetensors    4.6 GB
model-00002-of-00030.safetensors    4.7 GB
...
model-00030-of-00030.safetensors    2.1 GB
model.safetensors.index.json        60 KB

30 个分片文件,前 29 个每个 4.6–5.0 GB、最后一个约 2.1 GB,总计约 141 GB。这就是一次启动要搬运的原始数据量。

对比 Web 服务的"启动":Web 服务启动 = 几秒钟(拉几 MB 代码到内存);LLM 服务启动 = 几分钟(搬几十上百 GB 权重到 GPU)。启动是一类完全不同的工程问题。

7.1.1 搬运路径的三级带宽

权重从磁盘走到 GPU 显存要穿越三级带宽:

graph LR
    Disk["💾 磁盘<br/>HDD 0.2 GB/s<br/>SATA SSD 0.5 GB/s<br/>NVMe SSD 3-7 GB/s"]
    Disk -->|"磁盘 I/O"| CPU["🧠 CPU 内存<br/>DDR5 ~60 GB/s"]
    CPU -->|"PCIe"| GPU["🔥 GPU HBM<br/>A100 2 TB/s<br/>H100 3.35 TB/s"]
    Note[带宽瓶颈在最慢的一环:<br/>NVMe → PCIe → GPU 的链路<br/>实际速度由 NVMe 和 PCIe 决定]
    Disk -.-> Note
    style Disk fill:#94a3b8,color:#333,stroke:none
    style CPU fill:#2563eb,color:#fff,stroke:none
    style GPU fill:#10b981,color:#333,stroke:none

实际数字:

  • 磁盘 → CPU(NVMe SSD 顺序读):3-7 GB/s
  • CPU → GPU(PCIe 4.0 x16 单向):32 GB/s 理论,实际 ~25 GB/s
  • GPU → GPU(H100 的 NVLink 4,单卡双向总带宽):900 GB/s
  • 机间(InfiniBand NDR):50 GB/s

Llama-70B 140 GB 在理论峰值下:

  • 磁盘读:140 / 5 = 28 秒
  • PCIe 传输:140 / 25 = 5.6 秒
  • 两段串行相加:约 34 秒(若读盘与 PCIe 传输完全重叠,下限就是较慢的那段,约 28 秒)

实际:首次加载 70B 级模型的耗时通常远大于这个下限,具体值取决于磁盘/网络、checkpoint 格式、TP 与量化配置,本机无 GPU 也无法给出可复现的实测数,所以这里不写死秒数。中间的时间去哪了?这就是本章要解释的东西——加载不是一次顺序传输,中间有解析、反序列化、切分、合并、量化重排、weight mapping 等一系列工序。

7.1.2 朴素方法会怎么崩

如果用最直观的方式:

# 朴素实现 —— 会崩(示意;假设 checkpoint 是 pickle 格式的 .bin 分片)
import torch
model = build_model_on_cpu(config)                         # 先在 CPU 上建出 70B 模型:约 140 GB
for shard in shards:                                       # pytorch_model-000xx-of-000yy.bin
    state = torch.load(shard, map_location="cpu")          # 每次把整个 shard 反序列化到 CPU
    model.load_state_dict(state, strict=False)
model = model.to("cuda")                                   # 70B 一张 80GB 卡放不下

崩点:

  1. torch.load 把整个 shard 读进 CPU 内存——模型本身已在 CPU 上占约 140 GB,再叠加正在反序列化的 shard;如果像不少脚本那样先把所有 shard 合并成一个完整 state_dict 再加载,峰值会接近权重体积的 2 倍(约 280 GB)
  2. 单卡放不下——Llama-70B FP16 = 140 GB,超过 A100/H100 的 80 GB 显存
  3. weight name 不对齐——HuggingFace 权重名(如 model.layers.0.self_attn.q_proj.weight)和 vLLM 内部参数名不同
  4. TP 切分问题——多卡场景下每张卡只需要一份切片,不是完整权重

vLLM 的模型加载子系统就是为了解决这四个问题。

7.2 BaseModelLoader 策略模式:7 种加载器

vLLM 的加载逻辑组织成一个策略模式,所有具体加载器都继承 BaseModelLoader:

# vllm/model_executor/model_loader/loader.py:193
class BaseModelLoader(ABC):
    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

七个具体实现:

graph TB
    Base[BaseModelLoader]
    Base --> DM["DefaultModelLoader<br/>标准 safetensors/PT 加载"]
    Base --> Dum["DummyModelLoader<br/>随机权重 (benchmark 用)"]
    Base --> Ten["TensorizerLoader<br/>CoreWeave 序列化格式"]
    Base --> Shard["ShardedStateLoader<br/>预分片 state_dict"]
    Base --> BnB["BitsAndBytesModelLoader<br/>运行时量化为 NF4 / 加载预量化 bnb 权重"]
    Base --> GGUF["GGUFModelLoader<br/>llama.cpp 兼容"]
    Base --> Runai["RunaiModelStreamerLoader<br/>Run:ai S3 并行流式加载"]
    style DM fill:#10b981,color:#333,stroke:none
    style Shard fill:#2563eb,color:#fff,stroke:none
    style Ten fill:#2563eb,color:#fff,stroke:none
    style Runai fill:#2563eb,color:#fff,stroke:none
    style BnB fill:#f59e0b,color:#333,stroke:none
    style GGUF fill:#f59e0b,color:#333,stroke:none
    style Dum fill:#94a3b8,color:#333,stroke:none

适用场景速查表:

Loader 适用 加载速度 典型场景
DefaultModelLoader 标准 HF 模型 标准 绝大多数场景,生产默认
TensorizerLoader 预先 tensorize 过的模型 取决于序列化介质和网络 经常重启的生产环境
ShardedStateLoader 预先 TP-shard 过的 state_dict 避免每次重切 TP shard K8s 滚动部署
RunaiModelStreamerLoader S3 云存储 并行下载 无本地 SSD 的环境
BitsAndBytesModelLoader FP16 checkpoint 运行时量化(NF4),或预量化 bnb checkpoint 运行时量化时较慢(量化耗时) 想快速试量化效果
GGUFModelLoader GGUF 格式(llama.cpp) 标准 复用 llama.cpp 生态
DummyModelLoader 不加载真实权重 秒级 Benchmark / debug

添加新加载器只需实现两个方法——典型的可插拔设计。LoadConfig.load_format 甚至可以直接传一个 loader 类,get_model_loader() 遇到类型对象会直接实例化它(vllm/model_executor/model_loader/loader.py:1518-1519)。整个子系统的入口 get_model() 也只有两行:先 get_model_loader(vllm_config.load_config),再 loader.load_model(vllm_config=vllm_config)(vllm/model_executor/model_loader/__init__.py:12-14)。

这里要把两个数字分清。LoadFormat 在 vllm/config.py:1475-1488 里列了 13 个枚举值,包括 auto、safetensors、fastsafetensors、mistral、npcache 这类格式选择,也包括 tensorizer、sharded_state、bitsandbytes、gguf、runai_streamer 这类专门路径。真正落到 get_model_loader() 时,vllm/model_executor/model_loader/loader.py:1516-1542 会把这些配置路由到 7 个 concrete loader;没有命中特殊路径的配置统一回到 DefaultModelLoader。所以工程上不是"一个格式一个类",而是"少数加载策略覆盖多个格式分支"。

这个设计的好处是职责边界很清楚:LoadConfig 负责表达用户意图,get_model_loader() 负责把意图映射到策略类,具体 loader 再决定是否下载、如何迭代权重、何时执行后处理。新增一个特殊存储格式时,如果它仍然能产出 (name, tensor) 迭代器,就应该尽量挂在 DefaultModelLoader 的分支里;只有当它改变了模型初始化、权重来源或 TP shard 的语义,才值得新增 loader。

7.3 DefaultModelLoader:生产主路径深度拆解

绝大多数部署用的都是 DefaultModelLoader。它的加载流程分四步:

flowchart TB
    Start[load_model 开始]
    Step1["1. _prepare_weights<br/>远程解析 → 下载 → 分片去重 → 格式选定"]
    Step2["2. _get_weights_iterator<br/>lazy 返回 (name, tensor) 迭代器"]
    Step3["3. model.load_weights()<br/>weight name 映射 + QKV/gate_up 合并 + TP 切分"]
    Step4["4. _process_weights_after_loading<br/>量化/注意力层后处理"]
    Ready[模型就绪]
    Start --> Step1 --> Step2 --> Step3 --> Step4 --> Ready
    style Step1 fill:#2563eb,color:#fff,stroke:none
    style Step2 fill:#7c3aed,color:#fff,stroke:none
    style Step3 fill:#db2777,color:#fff,stroke:none
    style Step4 fill:#10b981,color:#333,stroke:none

7.3.1 _prepare_weights:定位 + 下载 + 格式确认

# vllm/model_executor/model_loader/loader.py:270(节选,省略 ModelScope、mistral/pt/npcache 分支与索引文件下载)
def _prepare_weights(self, model_name_or_path, revision, fall_back_to_pt,
                     allow_patterns_overrides) -> Tuple[str, List[str], bool]:
    is_local = os.path.isdir(model_name_or_path)
    load_format = self.load_config.load_format
    use_safetensors = False
    index_file = SAFE_WEIGHTS_INDEX_NAME
    # (1) 按 load_format 定候选 pattern(auto:safetensors 在前,bin 在后)
    if load_format == LoadFormat.AUTO:
        allow_patterns = ["*.safetensors", "*.bin"]
    # ... 其余格式分支 ...
    if fall_back_to_pt:
        allow_patterns += ["*.pt"]

    # (2) 远程模型:下载第一个在 Hub 上有匹配的格式
    if not is_local:
        hf_folder = download_weights_from_hf(
            model_name_or_path, self.load_config.download_dir,
            allow_patterns, revision,
            ignore_patterns=self.load_config.ignore_patterns)
    else:
        hf_folder = model_name_or_path

    # (3) 按 pattern 顺序 glob,命中第一个就停
    hf_weights_files: List[str] = []
    for pattern in allow_patterns:
        hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
        if len(hf_weights_files) > 0:
            if pattern == "*.safetensors":
                use_safetensors = True
            break

    # (4) 去重 / 过滤
    if use_safetensors:
        hf_weights_files = filter_duplicate_safetensors_files(
            hf_weights_files, hf_folder, index_file)
    else:
        hf_weights_files = filter_files_not_needed_for_inference(
            hf_weights_files)

    if len(hf_weights_files) == 0:
        raise RuntimeError(
            f"Cannot find any model weights with `{model_name_or_path}`")
    return hf_folder, hf_weights_files, use_safetensors

几个关键细节:

细节 1:safetensors 优先,但不是硬编码盲选。在 auto 模式下,vllm/model_executor/model_loader/loader.py:288-289 同时允许 *.safetensors 和 *.bin;真正下载时,download_weights_from_hf() 会先向 Hub 列文件,再按 allow_patterns 的顺序选择第一个有匹配结果的格式(vllm/model_executor/model_loader/weight_utils.py:252-263)。本地目录也走同样的 pattern 顺序扫描(vllm/model_executor/model_loader/loader.py:322-328)。这比简单检查目录里有没有某个扩展名稳,因为远程仓库、本地缓存和离线模式都能复用同一套规则。

细节 2:去重逻辑。一些模型仓库既提供 consolidated safetensors,又提供 sharded safetensors;两套文件同时交给 loader 会重复加载同名权重。filter_duplicate_safetensors_files() 会读 model.safetensors.index.json 的 weight_map,只保留索引里引用的文件(vllm/model_executor/model_loader/weight_utils.py:325-346)。如果索引文件不存在,它不会猜测文件关系,而是原样返回匹配结果。

细节 3:allow_patterns 是下载边界,不是清理脚本。远程下载只拉当前格式需要的文件;本地加载只扫描匹配 pattern 的文件。ignore_patterns 默认会排除 original/**/*(vllm/config.py:1557-1563),避免某些 Llama 仓库里的原始 checkpoint 被重复纳入。对非 safetensors 路径,filter_files_not_needed_for_inference() 还会排除 optimizer.pt、scheduler.pt、training_args.bin 这类训练产物(vllm/model_executor/model_loader/weight_utils.py:349-367)。

细节 4:失败要早,不能半加载。如果扫描完没有任何权重文件,_prepare_weights() 直接抛 RuntimeError(vllm/model_executor/model_loader/loader.py:349-351)。加载完成后,DefaultModelLoader.load_model() 会收集模型参数名,调用 model.load_weights(),并在非量化模型上检查有没有参数没从 checkpoint 初始化(vllm/model_executor/model_loader/loader.py:454-469)。这一步很重要:模型加载最危险的错误不是进程崩掉,而是某个参数静默保持随机初始化,然后服务还能启动。

7.3.2 safetensors 为什么是默认优先格式

一个 70B 模型如果用朴素 torch.load 一次性读完整 checkpoint,CPU 内存峰值会非常高;如果按 safetensors shard 迭代,每次只把当前 tensor 暴露给上层加载逻辑,内存曲线会更接近"当前 tensor + page cache"。这里不需要夸大成固定倍数,源码里真正可验证的是两条路径的形态差异。

torch.load 的行为:

1. open file
2. read all bytes into CPU memory     ← 整个文件全塞进内存
3. unpickle                              ← 反序列化再造一份 tensor
4. return state_dict

CPU 内存通常会承受完整 state dict 的压力,而且 pickle 反序列化无法像 safetensors 一样先读 header 再按 tensor 拉取。对几十到上百 GB 的模型,这种路径很容易把启动瓶颈从 GPU 显存转移到主机内存和 page cache 抖动上。

safetensors 按 tensor 读取的行为(vLLM 默认路径):

1. safe_open(file)
2. 读 header(描述所有 tensor 的 offset + shape + dtype)
3. 对每个 tensor:
     f.get_tensor(name)              ← 只读这个 tensor 的字节区间,得到一个 CPU tensor
     weight_loader(param, tensor)    ← 按 TP 切片后 copy_ 进 GPU 上的参数

safetensors_weights_iterator() 的真实代码也非常短:遍历 shard 文件,safe_open(st_file, framework="pt"),再对 f.keys() 逐个 get_tensor() 并 yield(vllm/model_executor/model_loader/weight_utils.py:430-444)。这说明 vLLM 没有先构造一个巨大的 state_dict;它把 checkpoint 变成流,后面的 load_weights() 边读边映射、边切分、边拷贝。

因此,safetensors 是默认优先路径的原因不是一个单点性能神话,而是三个工程属性同时成立:格式有 header,适合按 tensor 定位;文件不可执行,避免 pickle 反序列化的代码执行风险;迭代器天然能和 TP 切分、QKV 合并、量化参数注册衔接。

7.3.3 权重迭代器:lazy 读取

_get_weights_iterator 返回一个 Python generator,每次 yield (name, tensor) 对:

def safetensors_weights_iterator(hf_weights_files):
    for shard_file in hf_weights_files:
        with safe_open(shard_file, framework="pt") as f:
            for name in f.keys():
                yield name, f.get_tensor(name)  # 按需读出单个 tensor

上层用:

for name, weight in weights_iterator:
    # 处理这一个 tensor
    model_param = model.state_dict()[mapped_name]
    load_into(model_param, weight)
    # tensor 出作用域自动 GC

迭代器额外占用的 CPU 内存由"最大的单个 tensor"决定,而不是整个 shard。以 Llama-3-70B 为例,最大的是词表嵌入和 lm_head,各 128256 × 8192 × 2 字节 ≈ 2.1 GB;单个 MLP 投影约 0.47 GB。

7.3.4 加载流程的状态机视角

把 DefaultModelLoader 展开成状态机,会比"下载然后 load"更准确:

状态 进入条件 退出条件 失败边界
解析来源 model_config.model 是本地路径、HF ID 或 ModelScope ID 得到本地 hf_folder 远程不可达、revision 不存在
选择格式 LoadFormat.AUTO 或用户显式指定格式 得到 allow_patterns 与 use_safetensors 指定了 vLLM 不认识的格式
下载/扫描 远程走 snapshot_download,本地走 glob 得到候选文件列表 没有任何匹配权重
去重/过滤 safetensors 读 index,bin/pt 过滤训练产物 得到最终文件列表 index 指向文件缺失
构造迭代器 根据格式选 safetensors、fastsafetensors、pt、npcache 产生 (name, tensor) 流 单个 shard 损坏或 tensor shape 不匹配
模型接收 model.load_weights() 消费迭代器 返回已加载参数集合 参数名映射失败、shape 检查失败、非量化模型参数缺失

这个状态机解释了为什么模型加载代码里有大量小的 guard:加载不是一个原子动作,而是一串跨网络、文件系统、格式解析、模型结构和分布式切分的边界。如果任何一步选择"容错继续",后面的错误会变得更隐蔽;vLLM 在关键位置选择 fail fast,是生产推理系统更可靠的做法。

这也是排障时的阅读顺序:先确认格式选择,再确认文件列表,最后才看模型类映射。顺序反了,容易把文件缺失误判成模型实现问题。

7.4 load_weights:模型类的名称映射 + 权重合并魔法

HuggingFace 模型的参数名和 vLLM 内部的参数名不一致。为什么?因为 vLLM 做了一系列工程优化——最典型的是权重合并。

7.4.1 HuggingFace 的 Llama:三个独立投影

# HuggingFace Transformers 的 Llama 模型里
class LlamaAttention:
    self.q_proj = nn.Linear(hidden, num_heads * head_dim, bias=False)     # Q 投影
    self.k_proj = nn.Linear(hidden, num_kv_heads * head_dim, bias=False)  # K 投影(GQA 下头数更少)
    self.v_proj = nn.Linear(hidden, num_kv_heads * head_dim, bias=False)  # V 投影

权重名(shape 为 [out, in]):

model.layers.0.self_attn.q_proj.weight    [num_heads * head_dim, hidden]
model.layers.0.self_attn.k_proj.weight    [num_kv_heads * head_dim, hidden]
model.layers.0.self_attn.v_proj.weight    [num_kv_heads * head_dim, hidden]

forward 时:

q = self.q_proj(x)  # kernel launch 1
k = self.k_proj(x)  # kernel launch 2
v = self.v_proj(x)  # kernel launch 3

三次 matmul kernel launch。在 decode 阶段每 step 要跑,吞吐瓶颈之一。

7.4.2 vLLM 的合并:一个大投影

# vLLM 内部
class LlamaAttention:
    self.qkv_proj = QKVParallelLinear(
        hidden_size=hidden_size,
        head_size=self.head_dim,
        total_num_heads=self.total_num_heads,
        total_num_kv_heads=self.total_num_kv_heads,
        bias=bias,
        quant_config=quant_config,
        prefix=f"{prefix}.qkv_proj",
    )
    # 本 rank 的权重 shape: [q_size + 2 * kv_size, hidden](q_size/kv_size 已按 TP 切过)

权重名:

model.layers.0.self_attn.qkv_proj.weight    [q+k+v 拼接的大矩阵]

forward 时:

qkv = self.qkv_proj(x)               # 一次 kernel launch
q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)

合并成一次 matmul,减少投影层数量,也让 TP 切分更集中:一次列并行布局就能覆盖 Q/K/V 的目标位置。具体收益会随 batch、head 数、硬件和 kernel 实现变化,章节里不需要给一个脱离实验环境的固定微秒数;真正稳定的事实是 vLLM 的内部参数布局与 HF checkpoint 布局不同,加载阶段必须完成这个映射。

7.4.3 stacked_params_mapping:weight name 自动转换

但 HuggingFace 的权重文件里存的是 q_proj / k_proj / v_proj 三份。怎么加载到 qkv_proj 里?

# vllm/model_executor/models/llama.py:381(LlamaModel.load_weights,节选)
# LlamaForCausalLM.load_weights(:569)用 AutoWeightsLoader 按前缀把 "model.*" 权重转交到这里,
# 因此这里看到的名字已去掉 "model." 前缀。省略了 rotary 缓存跳过、KV cache scale、
# GPTQ 多余 bias、流水线并行缺层等分支。
class LlamaModel(nn.Module):
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".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),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        for name, loaded_weight in weights:
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                # 比如 "layers.0.self_attn.q_proj.weight"
                #   → "layers.0.self_attn.qkv_proj.weight"
                name = name.replace(weight_name, param_name)
                param = params_dict[name]
                # weight_loader 知道怎么把 Q/K/V 填到 qkv_proj 的对应切片里
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # 非合并权重,直接加载
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

注意取参数用的是 named_parameters() 而不是 state_dict():weight_loader 是挂在 Parameter 对象上的属性,返回的 loaded_params 集合则供 §7.3.1 细节 4 的"参数是否都被初始化"检查使用。

weight_loader 是每个 parameter 带的一个小函数(通常在 __init__ 时由 QKVParallelLinear 等层注入),它知道:

  • 目标权重的 shard 布局(Q 占哪几行、K 占哪几行、V 占哪几行)
  • 当前 rank 在 TP 中的位置(只拷对应切片)
  • 量化权重的 pack 方式(如果有)

所有这些逻辑对上层调用者透明。load_weights 的代码只负责 name → name 的字符串映射。

7.4.4 MLP 的 gate_up_proj 合并

类似的合并也发生在 MLP:

HuggingFace: gate_proj + up_proj + down_proj(3 个 Linear)
vLLM: gate_up_proj + down_proj(2 个 Linear,合并 gate+up)

forward:
    gate_up = gate_up_proj(x)         # 一次 matmul
    gate, up = gate_up.chunk(2, dim=-1)
    return down_proj(silu(gate) * up)

省一次 kernel launch。对 decode 场景累积收益可观。

7.5 Megatron 式权重切分:在 load 时就完成

多卡 TP 场景下,每张卡只需要完整权重的 1/tp_size。朴素做法是"先全量 load 再切分"——但这意味着每张卡临时需要完整 140 GB 的 CPU 内存,不可接受。

vLLM 的做法:权重在进入 GPU 之前就按 TP rank 切好。

7.5.1 列并行权重的加载

# QKVParallelLinear 的 weight_loader(概念性简化)
def qkv_weight_loader(param, loaded_weight, shard_id):
    tp_rank = get_tensor_model_parallel_rank()
    tp_size = get_tensor_model_parallel_world_size()

    # 按 shard_id 决定目标在 qkv_proj 里的偏移
    if shard_id == "q":
        target_offset = 0
        target_size = num_heads_per_rank * head_dim
    elif shard_id == "k":
        target_offset = q_size
        target_size = num_kv_heads_per_rank * head_dim
    else:  # v
        target_offset = q_size + kv_size
        target_size = num_kv_heads_per_rank * head_dim

    # TP 切分:每张卡只拿对应的 heads
    # (简化:真实代码中 K/V 的片号是 tp_rank // num_kv_head_replicas,
    #   tp_size 大于 KV 头数时多个 rank 复制同一组 KV 头)
    source_size = loaded_weight.shape[0]
    shard_offset = tp_rank * (source_size // tp_size)
    shard_end = shard_offset + (source_size // tp_size)
    loaded_weight_sharded = loaded_weight[shard_offset:shard_end]

    # 拷到目标位置
    param[target_offset:target_offset + target_size].copy_(loaded_weight_sharded)

关键点:

  • loaded_weight[shard_offset:shard_end] —— 真实代码用的是 loaded_weight.narrow(output_dim, start_idx, shard_size),在已经读进 CPU 的 tensor 上取视图,不额外复制
  • 目标参数 param 已经是切好尺寸的([q_size + 2 * kv_size, hidden],按本 rank 的头数计)
  • 从 CPU 拷到 GPU 的只有本 rank 那一片,不是完整权重

对 TP=8 的 70B 模型,每张 rank 最终只在 GPU 上保留约 140/8 = 17.5 GB 的权重。但在 v0.8.5 的默认路径里,读盘量并不随之缩小:safetensors_weights_iterator() 对每个 tensor 调 get_tensor() 取出整块 CPU tensor,切片发生在之后的 weight_loader 里,所以每个 rank 仍要把整份 checkpoint 读一遍。vllm/model_executor/model_loader/weight_utils.py:558-570 的 convert_pyslice_to_tensor() 注释提到 safetensors 的 PySafeSlice 可以先索引再读、减少读入内存的数据量,但这个函数在 v0.8.5 全树没有调用方,各迭代器也不用 get_slice()。想让每个 rank 只读自己那份,要走 §7.7.2 的 ShardedStateLoader。

7.5.2 行并行权重的加载

类似地,RowParallelLinear 按输入维度切分:

def row_parallel_weight_loader(param, loaded_weight):
    tp_rank = get_tensor_model_parallel_rank()
    tp_size = get_tensor_model_parallel_world_size()

    # 按行切(按输入维度)
    input_size = loaded_weight.shape[1]
    shard_size = input_size // tp_size
    start = tp_rank * shard_size
    end = start + shard_size
    loaded_weight_sharded = loaded_weight[:, start:end]

    param.copy_(loaded_weight_sharded)

Column + Row 的切分方式跟第 14 章要讲的 Megatron TP 数学完全对应。

7.6 量化模型的特殊加载路径

量化权重和 FP16 权重存储方式完全不同。以 GPTQ 4-bit 为例:

7.6.1 Bit packing 的处理

GPTQ 把 4-bit 权重 pack 进 int32——8 个 4-bit 值塞进一个 int32 里:

INT32 存储:  [b31..b28|b27..b24|...|b3..b0]
对应 8 个 4-bit 权重

读的时候需要解包或重排。vLLM 的量化路径通常会让 checkpoint 里的 packed 权重以量化参数的形式注册到 layer 上,再由后处理或 kernel 消费;这比"先在 CPU 还原成 FP16,再传 GPU"更符合推理场景。

加载时权重直接按 packed 格式搬到 GPU:

# vllm/model_executor/layers/quantization/gptq_marlin.py:194
class GPTQMarlinLinearMethod(LinearMethodBase):
    def apply(self, layer, x, bias=None):
        return self.kernel.apply_weights(layer, x, bias)

# 选中 MarlinLinearKernel 时,最终落到 vllm/model_executor/layers/quantization/utils/marlin_utils.py
def apply_gptq_marlin_linear(input, weight, weight_scale, weight_zp,
                             g_idx, g_idx_sort_indices, workspace,
                             wtype, output_size_per_partition,
                             input_size_per_partition, has_zp, is_k_full,
                             bias=None, use_fp32_reduce=USE_FP32_REDUCE_DEFAULT):
    ...

注意 v0.8.5 的 GPTQ Marlin 已经把具体 kernel 收进 self.kernel(MPLinearKernel 体系),GPTQMarlinLinearMethod.apply() 和 process_weights_after_loading() 都只是转调 kernel 对象(vllm/model_executor/layers/quantization/gptq_marlin.py:335-345);具体用哪个 kernel 由 create_weights() 里的 choose_mp_linear_kernel() 按平台与量化参数选定,Marlin 只是候选之一。

7.6.2 scale / zero_point 的加载

量化权重不是单一张量——还有配套的 scale 和 zero_point。load_weights 时要把这些张量都处理好:

model.layers.0.self_attn.q_proj.qweight      # int32 packed
model.layers.0.self_attn.q_proj.qzeros       # int32 packed
model.layers.0.self_attn.q_proj.scales       # FP16
model.layers.0.self_attn.q_proj.g_idx        # int32,act-order 的分组索引

vLLM 的 GPTQMarlinLinearMethod.create_weights 会预先注册 qweight、g_idx、scales、qzeros 四个 parameter(vllm/model_executor/layers/quantization/gptq_marlin.py:324-327),加载时按 name 自动 route。

7.6.3 kernel 自动选择

不同量化格式需要不同的 kernel。这个选择分两层,都不在 load_weights() 里:

第一层:配置级改写。GPTQMarlinConfig.override_quantization_method() 先用 is_gptq_marlin_compatible() 检查 checkpoint 的 quant_method / bits / group_size / sym / desc_act 以及当前平台能否跑 Marlin,可以就把 quantization 改写成 gptq_marlin(vllm/model_executor/layers/quantization/gptq_marlin.py:132-151、:169-191)。

第二层:层级路由。get_quant_method() 再按 layer 类型分派(vllm/model_executor/layers/quantization/gptq_marlin.py:153-166):

    def get_quant_method(self, layer: torch.nn.Module,
                         prefix: str) -> Optional["QuantizeMethodBase"]:
        if isinstance(layer, FusedMoE):
            from vllm.model_executor.layers.quantization.moe_wna16 import (
                MoeWNA16Config)
            if not check_moe_marlin_supports_layer(layer, self.group_size):
                logger.warning_one(
                    f"Layer '{prefix}' is not supported by GPTQMoeMarlin. "
                    "Falling back to Moe WNA16 kernels.")
                return MoeWNA16Config.from_config(
                    self.full_config).get_quant_method(layer, prefix)
            return GPTQMarlinMoEMethod(self)
        return get_linear_quant_method(self, layer, prefix,
                                       GPTQMarlinLinearMethod)

加载过程对量化的三个处理(pack、scale/zero_point、kernel)完全封装在 QuantizationConfig 子类里,上层 load_weights 不需要知道。

7.6.4 量化加载不能只看 dtype

量化 checkpoint 容易被误解成"把 FP16 换成 int4/int8 文件"。从 loader 视角看,这个说法太粗。真正复杂的是:每种量化方法都会改变 layer 上注册的参数集合、参数 shape、后处理步骤和 kernel 入口。

以 AWQ Marlin 为例,create_weights() 会注册 qweight、qzeros、scales 三类参数(vllm/model_executor/layers/quantization/awq_marlin.py:185-248),MoE 路径还会注册 w13_qweight、w2_qweight、w13_scales、w2_scales 等专家维度参数(vllm/model_executor/layers/quantization/awq_marlin.py:324-394)。这些名字不是普通 FP16 模型里的 weight;所以 load_weights() 的名字映射必须能把 checkpoint tensor 路由到这些量化参数上。

再看后处理:AWQ Marlin 的 process_weights_after_loading() 会把 checkpoint 中的布局 repack 成 Marlin kernel 需要的布局,并替换 layer 上的参数对象(vllm/model_executor/layers/quantization/awq_marlin.py:258-293)。这就是为什么 §7.8 的 device_loading_context() 不能被解释成普通加载阶段的显存优化;它真正服务的是这种"参数已经加载完,但还要按 kernel 需求变形"的阶段。

维度 FP16 普通权重 量化权重
参数名 通常是 weight / bias 可能是 qweight、qzeros、scales 或专家参数
shape 语义 直接对应矩阵维度 常包含 pack、group size、专家数或 scale 粒度
加载后处理 多数是直接可用 可能要 repack、permute scales、替换 parameter
失败形态 shape mismatch 比较直观 可能在后处理或 kernel 调用阶段才暴露

因此,判断一个量化模型是否"支持 vLLM"不能只看它的文件能不能被下载,也不能只看 quantization 字段是否被识别。更可靠的检查顺序是:量化配置类是否存在;目标 Linear / MoE layer 是否注册了对应参数;模型类的 load_weights() 是否处理这些名字;process_weights_after_loading() 是否能在当前设备和并行拓扑下完成;最后才是 kernel 是否支持当前 GPU 架构。

7.7 三种加速首次加载的手段

大模型冷启动慢是生产部署的痛。vLLM 提供三类经常会被拿来优化首次加载的路径。它们的共同点不是承诺固定加速倍数,而是把"每次启动都做的工作"迁移到更便宜的位置:预序列化、预切分,或者并行读取。

有一点要先说清:在 v0.8.5 里,--load-format tensorizer 和 --load-format sharded_state 都被 _is_v1_supported_oracle() 判为 V1 不支持(vllm/engine/arg_utils.py:1298-1303)——未设 VLLM_USE_V1 时自动回退 V0 引擎,显式设 VLLM_USE_V1=1 则直接报错。runai_streamer 不在这个名单里。

7.7.1 Tensorizer:序列化加速

CoreWeave 的 Tensorizer 把模型序列化成一种连续字节流格式。加载时可以做:

  • 并行流式读(多线程读多个字节段)
  • 直接 S3 → GPU(skip CPU,如果是云存储)
  • 压缩(可选 zstd)

TensorizerLoader 在 vllm/model_executor/model_loader/loader.py:503-590 中有两条路径:如果输入已经是 vLLM tensorized 模型,就直接 load_with_tensorizer();否则仍然要初始化模型,再通过 tensorizer 的权重迭代器调用 model.load_weights()。前者才是"把昂贵转换前置"的路径,后者只是换了读取介质。

# 先一次性序列化(v0.8.5 走的是仓库里的示例脚本,不是内置 CLI 子命令;
# 需在仓库根目录执行,TP>1 时按 rank 存成 model-rank-%03d.tensors)
python -m examples.other.tensorize_vllm_model \
    --model meta-llama/Meta-Llama-3-70B \
    --tensor-parallel-size 4 \
    serialize \
    --serialized-directory /local/nvme \
    --suffix v1

# 后续加载:--model 仍给原模型(读 config 与 tokenizer),
# 序列化文件经 --model-loader-extra-config 的 tensorizer_uri 指定
vllm serve meta-llama/Meta-Llama-3-70B \
    --tensor-parallel-size 4 \
    --load-format tensorizer \
    --model-loader-extra-config \
    '{"tensorizer_uri": "/local/nvme/vllm/meta-llama/Meta-Llama-3-70B/v1/model-rank-%03d.tensors"}'

7.7.2 ShardedStateLoader:预切分

对 TP 部署,每次启动都要重新"按 TP rank 切分全量权重"。如果能先做好切分、保存 sharded state dict,后续启动就省去切分时间:

# 一次性:先按 TP=4 切好保存(v0.8.5 用的是示例脚本,没有 vllm save-sharded-state 子命令;
# 脚本要求 --model 是本地目录,给 HF ID 会直接报错)
python examples/offline_inference/save_sharded_state.py \
    --model /models/Meta-Llama-3-70B \
    --tensor-parallel-size 4 \
    --output /fast-ssd/llama-70b-tp4/

# 后续启动:每张卡直接加载自己那份
vllm serve /fast-ssd/llama-70b-tp4/ \
    --tensor-parallel-size 4 \
    --load-format sharded_state

ShardedStateLoader 的源码注释说得很直白:它直接加载每个 worker 自己的 state dict,让大 TP 模型的每个 worker 只读自己的 shard,而不是完整 checkpoint(vllm/model_executor/model_loader/loader.py:603-609)。加载时会用当前 tensor parallel rank 拼出 model-rank-{rank}-part-{part}.safetensors 这样的文件名(vllm/model_executor/model_loader/loader.py:612、vllm/model_executor/model_loader/loader.py:698-710)。这条路径适合 TP 拓扑稳定的集群;如果 TP size、PP 划分或模型结构经常变,它的预处理产物反而容易变成运维负担。

7.7.3 Run:ai Model Streamer:S3 并行流式

对于模型存在 S3 的场景(K8s 部署常见),Run:ai 的 Model Streamer 提供多连接并行流式加载:

vllm serve s3://my-models/llama-70b/ \
    --load-format runai_streamer

Run:ai 路径在源码里分两层:普通 runai_streamer 走 RunaiModelStreamerLoader,runai_streamer_sharded 则复用 ShardedStateLoader(load_config, runai_model_streamer=True)(vllm/model_executor/model_loader/loader.py:1536-1540)。权重迭代器最终落到 runai_safetensors_weights_iterator(),它用 SafetensorsStreamer 逐文件 stream,再 yield streamer 解析出的 tensors(vllm/model_executor/model_loader/weight_utils.py:447-460)。这条路径的价值在于远端对象存储的流式读取,不在于本地 NVMe 场景。

三种加速手段对比:

方案 优化点 前置成本 适用场景
Tensorizer 减少普通 checkpoint 解析和转换 一次性序列化 本地或对象存储上的重复冷启动
ShardedState 每个 rank 直接读取自己的 TP shard 按固定并行拓扑保存 固定 TP 配置的 K8s
Run:ai Streamer 对象存储流式读取 safetensors 依赖 streamer 环境 S3 / 云存储部署

经验上的搭配:Tensorizer + 本地 NVMe 适合单机重复冷启动;ShardedState 适合 TP 固定的 K8s 滚动重启;Run:ai Streamer 适合没本地 SSD 的云原生环境。前两者在 v0.8.5 上都会让引擎回退到 V0,选型时要一并考虑。

7.8 device_loading_context:后处理阶段的设备边界

这一节必须按源码重新理解。device_loading_context() 不是"加载前把 GPU 参数挪到 CPU,再最后搬回 GPU"。在 vLLM v0.8.5 的 vllm/model_executor/model_loader/loader.py:68-108 里,它做的是相反方向的临时迁移:如果某个模块的参数原本在 CPU,而量化后处理需要目标设备,就把这些 CPU 参数临时搬到 target_device;后处理结束后,再复制回原来的 CPU 设备,必要时使用 pinned memory。

# 简化:省略 target_device 为 CPU 时直接 yield 返回的分支;
# copy_back_to_cpu_or_original_device 代表源码里"原设备是 CPU 就拷进
# (可用时为 pinned 的)CPU tensor,否则 .to(原设备)"的那段逻辑,并非真实函数名
@contextmanager
def device_loading_context(module, target_device):
    original_device_states = {}
    for name, p in module.named_parameters():
        if p.device.type == "cpu":
            original_device_states[name] = p.device
            p.data = p.data.to(target_device)
    try:
        yield module
    finally:
        for name, p in module.named_parameters():
            if name in original_device_states:
                p.data = copy_back_to_cpu_or_original_device(p.data)

它的调用点在 _process_weights_after_loading():遍历 model.named_modules(),如果模块的 quant_method 是 QuantizeMethodBase,就进入 device_loading_context(),调用 quant_method.process_weights_after_loading(module)(vllm/model_executor/model_loader/loader.py:164-180)。换句话说,它服务的是加载之后的量化重排、repack、scale 处理,而不是普通 FP16 safetensors 拷贝。

这个边界和 CPU offload 关系很大。部分部署会把一部分权重留在 CPU,或者先在 CPU 上保存某些量化参数;但 AWQ Marlin、FP8 等后处理往往需要调用 GPU kernel 或至少要求 tensor 位于目标设备。device_loading_context() 让后处理函数看到的是目标设备上的参数,同时不永久改变模块原来的放置策略。它保护的是"后处理语义和 offload 策略的兼容性",不是一个通用的 2 倍显存节省开关。

这也解释了为什么 _process_weights_after_loading() 放在 model.load_weights() 后面。load_weights() 负责名字映射和数据填充;后处理负责把已经填好的量化权重改造成 kernel 想要的布局。例如某些 Marlin 路径会在 process_weights_after_loading() 中 repack qweight、permute scales、替换参数对象。把这两步拆开,loader 不需要知道每种量化 kernel 的内部格式,量化配置类也不需要接管下载和 checkpoint 格式选择。

7.9 模型注册表:146 个唯一架构如何被支持

本章对照 vLLM v0.8.5 源码。registry.py 的分类条目不是简单相加后的最终支持数,因为 _EMBEDDING_MODELS 会通过 dict 展开复用一部分 Llama 家族条目,后续 _VLLM_MODELS 合并时相同 key 会覆盖。按源码计算,分类条目合计 163 个,合并后 _VLLM_MODELS 有 146 个唯一 architecture key,对应 125 个唯一实现类、114 个被引用的模型实现文件。「125」按 (模块, 类名) 二元组算——按类名去重则是 124,差的那一个是 MixtralForCausalLM:它同时挂在 mixtral 和 mixtral_quant 两个模块下,是同名不同实现,所以计入两条。

_TEXT_GENERATION_MODELS = {
    # HF 架构名 → (模块名, 类名)
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),  # 复用 Llama
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
    "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),  # 与 V2 同文件
    # ...(Qwen2-VL 等多模态模型在 _MULTIMODAL_MODELS 里)
}

加载时并不是直接拿 _TEXT_GENERATION_MODELS[hf_arch]。真实路径是:

architectures = model_config.hf_config.architectures
model_cls, arch = ModelRegistry.resolve_model_cls(architectures)

get_model_architecture() 从 config.json 的 architectures 字段开始(vllm/model_executor/model_loader/utils.py:87-109)。如果用户显式要求 Transformers 后端,或者当前 architecture 没有 vLLM 原生实现但 Transformers backend 兼容,就先把 architecture 改写成 TransformersForCausalLM。最终再由 ModelRegistry.resolve_model_cls() 走 lazy import。

7.9.1 新增一个模型的三步

假设你想给 vLLM 加一个新模型 MyLLMForCausalLM:

Step 1 实现模型类(在 vllm/model_executor/models/ 目录下新建 my_llm.py):

class MyLLMForCausalLM(nn.Module):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        # 定义 Transformer 结构,复用 vllm.model_executor.layers 里的组件
        ...

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # 前向传播
        ...

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # weight name mapping + stacked_params_mapping;返回已加载的参数名集合
        ...

Step 2 在 registry.py 的 _TEXT_GENERATION_MODELS 字典里加一行:

    "MyLLMForCausalLM": ("my_llm", "MyLLMForCausalLM"),

Step 3 提 PR 到 vllm-project/vllm。不想改 vLLM 源码的话,也可以在自己的包里做 out-of-tree 注册:

from vllm import ModelRegistry

# my_pkg 是你自己的包,模型类放在 my_pkg/my_llm.py 里
ModelRegistry.register_model("MyLLMForCausalLM", "my_pkg.my_llm:MyLLMForCausalLM")

register_model() 的第二个参数既可以是类本身,也可以是 <module>:<class> 字符串——后者会被包成 _LazyRegisteredModel,用到时才 import,避免在主进程里过早初始化 CUDA(vllm/model_executor/models/registry.py:368-409)。

更准确地说,新增一个模型至少需要三类接口:构造函数接受 vllm_config 和 prefix,forward 能对接 vLLM 的输入张量约定,load_weights() 能把 HF checkpoint 名字映射到内部参数布局。_initialize_model() 会优先检查新式构造函数签名(vllm/model_executor/model_loader/loader.py:128-140);旧式类还能被兼容,但会走 warning 路径。注册表只解决"找到类",不解决"类是否符合执行层约定"。

7.9.2 对齐现有模型家族

很多模型架构相似(比如 Mistral 和 Llama 几乎一样),vLLM 的策略是直接复用:

"MistralForCausalLM": ("llama", "LlamaForCausalLM"),

Mistral、CodeLlama、Llama-2、Llama-3 等全系列都复用 LlamaForCausalLM。差异通过 config.json 里的参数表达(num_key_value_heads、head_dim 等),不需要单独的模型类。

这让 vLLM 的模型目录保持"一对多复用"。在这份源码里,146 个唯一 architecture key 最终映射到 114 个模型实现文件;同一个实现文件经常覆盖多个 HF architecture 名。例如 AquilaModel、InternLMForCausalLM、MistralForCausalLM、XverseForCausalLM 都路由到 llama.py 的 LlamaForCausalLM。差异通过 config 字段表达,而不是为每个名称复制一套模型类。

7.9.3 六个独立注册表:不是一个大 dict

上面伪代码把注册表描述成单个 _TEXT_GENERATION_MODELS。真实 registry.py 按模型用途分成 6 个独立 dict(vllm/model_executor/models/registry.py:33-226):

_TEXT_GENERATION_MODELS       # 主流 LLM:Llama/Qwen/Phi/Gemma/DeepSeek/Mistral 等
_EMBEDDING_MODELS             # 嵌入模型:E5/BGE/Jina 等
_CROSS_ENCODER_MODELS         # 重排模型(检索场景)
_MULTIMODAL_MODELS            # VLM:Qwen2-VL/InternVL/LLaVA 等(第 15 章展开)
_SPECULATIVE_DECODING_MODELS  # 投机解码专用(EAGLE/Medusa/MTP 等,第 12 章展开)
_TRANSFORMERS_MODELS          # HuggingFace transformers 万能回退

_VLLM_MODELS = {
    **_TEXT_GENERATION_MODELS,
    **_EMBEDDING_MODELS,
    **_CROSS_ENCODER_MODELS,
    **_MULTIMODAL_MODELS,
    **_SPECULATIVE_DECODING_MODELS,
    **_TRANSFORMERS_MODELS,
}

拆成六个 dict 而非一个大 dict 的好处(源码没有写明动机,以下是从结构上能看出的作用):

  • 文档化模型能力分类——从注册表就能看出 "哪些是 LLM、哪些是 embedding、哪些是投机解码 helper"
  • 改动隔离——加一个新 VLM 只改 _MULTIMODAL_MODELS,不会手滑改到 _TEXT_GENERATION_MODELS 里
  • 迁移与 deprecation 单独处理——比如未来移除某一整类(某种投机解码方案 EOL),只需删一个 dict

当前版本的分类计数如下:

分类 源码位置 条目数说明
_TEXT_GENERATION_MODELS vllm/model_executor/models/registry.py:33-120 80 个显式文本生成条目
_EMBEDDING_MODELS vllm/model_executor/models/registry.py:122-159 25 个显式条目,再通过 dict 展开复用 8 个 Llama 实现条目
_CROSS_ENCODER_MODELS vllm/model_executor/models/registry.py:161-169 4 个重排/交叉编码条目
_MULTIMODAL_MODELS vllm/model_executor/models/registry.py:171-213 39 个多模态条目
_SPECULATIVE_DECODING_MODELS vllm/model_executor/models/registry.py:215-222 6 个投机解码 helper 条目
_TRANSFORMERS_MODELS vllm/model_executor/models/registry.py:224-226 1 个 Transformers backend 回退条目

这些分类合起来是 163 个条目,但有 17 个 architecture key 在多个分类中重复出现,所以 _VLLM_MODELS 的最终唯一 key 是 146 个。这个数字比"把分类行数相加"更适合写进正文,因为运行时查找的就是合并后的 dict。

_SPECULATIVE_DECODING_MODELS 展开(line 215-222):

# vllm/model_executor/models/registry.py:215
_SPECULATIVE_DECODING_MODELS = {
    "EAGLEModel":                   ("eagle",       "EAGLE"),
    "EagleLlamaForCausalLM":        ("llama_eagle", "EagleLlamaForCausalLM"),
    "Eagle3LlamaForCausalLM":       ("llama_eagle3","Eagle3LlamaForCausalLM"),
    "DeepSeekMTPModel":             ("deepseek_mtp","DeepSeekMTP"),
    "MedusaModel":                  ("medusa",      "Medusa"),
    "MLPSpeculatorPreTrainedModel": ("mlp_speculator","MLPSpeculator"),
}

这 6 种对应第 12 章要讲的不同投机解码方案的 helper 模型——每种方案的草稿模型/预测头有自己的 class。它们由投机解码配置加载,不是给用户当主模型用的;单列一个 dict 让这层区别在注册表里一目了然。不过合并进 _VLLM_MODELS 后它们和其他条目一样可以被解析,单列本身并不构成访问限制。

7.9.4 TransformersForCausalLM 万能回退

_TRANSFORMERS_MODELS 只有一条(line 224-226):

_TRANSFORMERS_MODELS = {
    "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
}

这个类的作用是用 HuggingFace transformers 库里的建模代码跑那些没有 vLLM 原生实现的模型。流程:

  1. 用户指定的模型,其 config.json 里的 architecture 不在 vLLM 注册表里
  2. model_impl 为默认的 auto 时,若该模型的 Transformers 实现声明兼容 vLLM 后端,就自动改写成 TransformersForCausalLM;不兼容则直接报错
  3. 这个 wrapper 用 AutoModel.from_config(..., attn_implementation="vllm") 构造模型(vllm/model_executor/models/transformers.py:138-147),注意力走 vLLM 自己的 Attention 层,线性层按模型的 TP plan 替换成 vLLM 的并行 Linear
  4. 因此 PagedAttention 的 KV cache 和连续批处理照样生效;缺的是原生实现里手工做的融合与适配(如 qkv_proj / gate_up_proj 合并),源码 warning 也只说"部分特性可能不支持、性能可能不是最优"

"vLLM 支持哪些模型" 的真实含义要分两层:合并后的 146 个 key 是 vLLM 明确知道如何解析的 architecture;注册表外但 Transformers backend 兼容的模型,可以被改写成 TransformersForCausalLM 路径(vllm/model_executor/model_loader/utils.py:73-83、vllm/model_executor/models/registry.py:449-452)。源码里的 warning 也说得很谨慎:部分特性可能不支持,性能也可能不是最优。生产上评估一个新模型时,不能只问"能不能启动",还要问它走的是原生 vLLM 模型类,还是 Transformers 回退类。

这种"精细优化 + 兜底兼容"的双轨是 vLLM 跟上 HF 生态的关键。新 architecture 不一定一开始就有专用 load_weights()、packed module mapping、量化后处理和多模态输入处理;先走 Transformers backend 可以降低可用性门槛,后续再为热门模型补原生实现。

7.9.5 _ModelInfo 12 字段元数据与 subprocess 隔离

上面的 dict 只记录"架构名 → Python 类路径"映射,每个模型真正的能力标签存在 _ModelInfo 里(line 247-277):

# vllm/model_executor/models/registry.py:247
@dataclass(frozen=True)
class _ModelInfo:
    architecture: str
    is_text_generation_model: bool
    is_pooling_model: bool
    supports_cross_encoding: bool
    supports_multimodal: bool
    supports_pp: bool                 # pipeline parallel
    has_inner_state: bool              # Mamba/RWKV 类带状态的
    is_attention_free: bool            # Mamba 完全无 attention
    is_hybrid: bool                    # Jamba 类 attn+mamba 混合
    has_noops: bool
    supports_transcription: bool       # Whisper 类
    supports_v0_only: bool             # 暂时没迁 V1 的模型

除 architecture 外的 11 个布尔字段标注这个模型类在 vLLM 里具备哪些能力。from_model_cls 静态方法通过一系列 supports_*() 检测器自动填充这些字段——用户不需要手写、靠运行时反射模型类的 mixin 继承关系判断。

_SUBPROCESS_COMMAND 的 subprocess 隔离(line 242-244)是另一个精致细节:

_SUBPROCESS_COMMAND = [
    sys.executable, "-m", "vllm.model_executor.models.registry"
]

对还没在主进程 import 过的模型(_LazyRegisteredModel),vLLM 在 "检查这个模型类具备哪些能力" 时会启动 subprocess 运行 registry.py 自己,在子进程里 import 并生成 _ModelInfo 再传回——源码注释写的理由是"避免初始化 CUDA":import 模型模块可能顺带初始化 CUDA,主进程一旦初始化过 CUDA,之后 fork 出的 worker 就会报 Cannot re-initialize CUDA in forked subprocess(register_model() 的 docstring 点名了这个错误)。隔离到子进程检查后,主进程只在真正需要时才导入模型类。

这就像第 6 章的 multiprocessing.connection.wait(sentinels) 故障监控——用进程边界做隔离是 vLLM 对"Python 全局状态 + CUDA 初始化是有毒的"这一现实的一贯应对。

7.9.6 vllm/model_executor/model_loader/ 6 文件 3193 行的真实拆分

§7.2 给出 7 种 loader,对应 loader.py 单文件里 8 个类(1 base + 7 concrete,line 193 / 210 / 476 / 503 / 603 / 792 / 1317 / 1415)——

文件 行 角色
loader.py 1542 8 个 loader 类全部塞在一文件——BaseModelLoader ABC + Default / Dummy / Tensorizer / ShardedState / BitsAndBytes / GGUF / RunaiModelStreamer 7 种实现
weight_utils.py 749 _prepare_weights(§7.3.1)调用的下载与过滤函数 + safetensors 等权重迭代器 + 文件锁
tensorizer.py 468 Tensorizer(§7.7.1)的独立实现层、和 vllm/model_executor/model_loader/loader.py:503 TensorizerLoader 配合
neuron.py 243 AWS Trainium / Inferentia 的特殊加载路径
utils.py 171 get_model_architecture()、Transformers 回退解析、dtype context 等模型选择工具
__init__.py 20 公共 export

三条值得记住的事实——

  1. loader.py 单文件 1542 行 = 8 个类全部塞一起——这是因为 _initialize_model()、_process_weights_after_loading()、DefaultModelLoader.Source、get_model_loader() 等 helper 被多个 loader 共享,拆成很多小文件未必更清晰
  2. weight_utils.py 749 行——下载、文件锁、safetensors iterator、np cache、TP shard loader 都放在这里,说明"权重来源和文件格式"是跨 loader 的公共基础设施
  3. tensorizer.py 独立 468 行——Tensorizer 需要自己的配置、序列化和 stream 逻辑;loader.py 只保留策略入口,避免把第三方格式的细节扩散到默认加载路径

7.10 四类部署场景的加载优化配方

场景 A:单机 offline 推理(快速启动)

# Tensorizer:一次性序列化,后续用 tensorizer 路径加载
python -m examples.other.tensorize_vllm_model \
    --model meta-llama/Meta-Llama-3-70B --tensor-parallel-size 4 \
    serialize --serialized-directory /local/nvme --suffix v1

python -c "from vllm import LLM; llm = LLM(model='meta-llama/Meta-Llama-3-70B', tensor_parallel_size=4, load_format='tensorizer', model_loader_extra_config={'tensorizer_uri': '/local/nvme/vllm/meta-llama/Meta-Llama-3-70B/v1/model-rank-%03d.tensors'})"

这类部署的核心假设是模型权重和并行配置比较稳定。首次转换可以放到镜像构建、发布流水线或离线准备任务里;在线服务启动时只消费转换后的产物。不要把 Tensorizer 当成所有环境的默认答案:如果模型经常变、checkpoint 来自不同来源、或者排障时需要直接对照 HF 原始权重,普通 safetensors 路径更透明。

场景 B:K8s 滚动部署(多副本冷启动)

# 在 initContainer 里一次性切分到 PVC
initContainers:
- name: prepare-sharded
  command: ["python", "examples/offline_inference/save_sharded_state.py",
            "--tensor-parallel-size=4", "--output=/shared-pvc/llama-70b-tp4/", ...]
  volumeMounts: [mountPath: /shared-pvc]

# 所有副本挂载同一个 PVC
containers:
- name: vllm
  command: ["vllm", "serve", "/shared-pvc/llama-70b-tp4/",
            "--tensor-parallel-size=4",
            "--load-format=sharded_state"]

这类部署最适合固定 TP size 的服务。initContainer 或独立 job 先生成 sharded state,业务容器只读 rank 对应文件。风险点也很明确:TP size、pipeline parallel、模型结构、量化配置任何一个改变,都要重新生成 shard;否则不是变慢,而是直接 shape mismatch 或参数缺失。

场景 C:S3 原生部署

vllm serve s3://ml-models/llama-70b/ \
    --load-format runai_streamer \
    --tensor-parallel-size 4

这类部署的关键不是追求一个固定秒数,而是减少"先完整下载到本地磁盘,再由 vLLM 读取"的中间步骤。对象存储吞吐、跨 AZ 网络、credential 刷新、Pod 调度位置都会影响实际启动时间;因此最好把 streamer 路径纳入压测,而不是只在单机上验证命令能跑。

场景 D:开发迭代(快速试模型)

# FP8 量化 + 本地 SSD,减少搬运数据量(路径为本地 FP8 checkpoint 目录)
vllm serve /models/Meta-Llama-3-70B-Instruct-FP8 \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.9

# 或用 DummyModelLoader 跳过权重加载(仅测性能;参数仍按 BF16 全量分配,70B 需多卡)
vllm serve meta-llama/Meta-Llama-3-70B \
    --tensor-parallel-size 4 \
    --load-format dummy

开发迭代场景下,目标通常不是复现生产性能,而是缩短"改代码 → 启服务 → 看到行为"的回路。dummy 可以跳过真实权重,适合验证调度、API、profile 和内存路径;FP8 或其它量化模型可以降低权重体积,但要接受量化 kernel、后处理和精度行为都会变,不能把它当作 FP16 生产路径的完全替身。

7.11 本章小结

模型加载是 LLM 推理服务的"冷启动"阶段——对用户看不见但工程上极其关键:

  • 问题规模:140 GB 权重 × 30 个分片 × 三级带宽鸿沟(NVMe → PCIe → GPU)
  • 朴素实现崩点:CPU 内存 2× 爆炸、单卡装不下、weight name 不对齐、TP 切分困难
  • 策略模式:7 种 Loader 适配 7 种场景;BaseModelLoader + 两个抽象方法保持核心代码干净
  • safetensors 迭代器:按 shard、按 tensor yield,避免完整 checkpoint 一次性进入 Python state dict
  • lazy 迭代器:按需 yield 单个 tensor,额外内存由最大单 tensor 决定
  • QKV / gate_up 合并:HF 三个 Linear → vLLM 一个大 Linear,省 kernel launch;通过 stacked_params_mapping 自动映射
  • Megatron 式切分:逐 tensor 读入后立即切出本 rank 那片再拷上 GPU,不是先 load 完整模型再切(默认路径下每个 rank 仍读整份 checkpoint)
  • 量化特殊路径:bit packing / scale / kernel 三重差异,封装在 QuantizationConfig 子类
  • 三种加速方案:Tensorizer 预序列化、ShardedState 预切分、Run:ai Streamer 面向对象存储流式读取
  • device_loading_context:服务量化后处理和 CPU offload 的设备边界,不是普通加载的显存魔法
  • 模型支持:本地 v0.8.5 源码中 _VLLM_MODELS 合并后有 146 个唯一 architecture key;注册表数字不能直接等同于性能承诺
  • 四场景配方:单机 Tensorizer / K8s ShardedState / S3 Streamer / FP8 量化

到这里,从请求、调度到 Worker 进程与模型加载,执行前的准备都已就位。下一章我们继续第三篇,进入 ModelRunner——看看加载好的模型如何真正被调用执行。


源码导航

  • Loader 基类与默认实现:vllm/model_executor/model_loader/loader.py
  • 权重工具(mmap、index 解析):vllm/model_executor/model_loader/weight_utils.py
  • 模型注册表:vllm/model_executor/models/registry.py
  • Llama 模型实现(好的阅读起点):vllm/model_executor/models/llama.py
  • 并行 Linear 层(QKV/gate_up 合并在这里定义):vllm/model_executor/layers/linear.py
  • 量化配置注册:vllm/model_executor/layers/quantization/base_config.py
  • Tensorizer 集成:vllm/model_executor/model_loader/tensorizer.py
  • Run:ai Streamer 集成:vllm/model_executor/model_loader/loader.py 与 weight_utils.py

延伸阅读