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

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

本章要点

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

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

git clone --branch v0.8.5 https://github.com/vllm-project/vllm.git

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

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

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

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

把 140 GB 从磁盘搬到 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：

# 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 子类：

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。它的加载流程可以拆成四步：

7.3.1 权重发现：_prepare_weights

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

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

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

# 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 索引文件来去重：

# 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 数组，避免重复加载

# 根据格式选择迭代器（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）：

# 简化的 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 论文。

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

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

# 每张卡根据自己的 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 为例：

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

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

# 量化方法注册（简化）
@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 中的模型注册表。

# 注册表将 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"),
    # ... 数百种模型
}

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

这种设计实现了模型层和引擎层的完全解耦。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 的迁移：

# 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）支持加载预先按张量并行切分的权重。这跳过了运行时切分步骤——每张卡直接加载自己那份权重文件。适合需要频繁重启的生产环境：

# 首次加载后保存分片（一次性开销）
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/