4. 权重加载器：掌握自定义权重加载器（WeightLoader）的编写方法

权重加载器，说白了就是告诉 vLLM：「嘿，我的模型权重放在哪，该怎么读进来」。

我刚开始接触 vLLM 时，以为它只支持 HuggingFace 的标准格式。后来有个项目需要加载自己训练的稀疏化模型，权重文件是 .bin 格式但结构完全不同。嗯，那时候我才意识到，自定义 WeightLoader 是个绕不开的技能。

4.1 为什么需要自定义 WeightLoader？

vLLM 默认支持 HuggingFace 的 transformers 格式。但现实世界很残酷：

• 你的模型可能来自内部训练框架，权重格式是自定义的
• 你可能做了量化、剪枝、蒸馏，权重结构变了
• 你也许想从本地路径加载，而不是每次都去 HF 下载

我遇到过最典型的情况：客户给了一个 70B 的模型，权重文件是分片的 .safetensors，但命名规则和 HF 标准完全不一样。当时我花了半天写了个 WeightLoader，问题就解决了。

4.2 WeightLoader 的核心接口

vLLM 的 WeightLoader 其实就是一个类，核心方法就两个：

class BaseWeightLoader(ABC):
    @abstractmethod
    def load_weights(
        self,
        model: nn.Module,
        weights: Dict[str, torch.Tensor]
    ) -> None:
        """将权重加载到模型中"""
        pass

    @abstractmethod
    def get_weights(
        self,
        model_name_or_path: str
    ) -> Dict[str, torch.Tensor]:
        """从源获取权重字典"""
        pass

你想想看，这个设计其实很优雅。你只需要实现两个方法：get_weights 负责「从哪里读」，load_weights 负责「怎么塞进去」。

4.3 实战：从 HuggingFace 加载

vLLM 内置了 HuggingFaceWeightLoader，但我们可以自己写一个来理解原理：

class MyHFWeightLoader(BaseWeightLoader):
    def __init__(self, revision: str = "main"):
        self.revision = revision

    def get_weights(self, model_name_or_path: str) -> Dict[str, torch.Tensor]:
        from transformers import AutoModelForCausalLM

        # 我个人习惯用 from_pretrained 时加上 low_cpu_mem_usage
        hf_model = AutoModelForCausalLM.from_pretrained(
            model_name_or_path,
            revision=self.revision,
            low_cpu_mem_usage=True,
            torch_dtype="auto"
        )
        return dict(hf_model.state_dict())

    def load_weights(self, model: nn.Module, weights: Dict[str, torch.Tensor]):
        # 这里要注意：vLLM 的模型可能对 key 做了重命名
        # 我建议先打印 model.state_dict().keys() 对比一下
        missing, unexpected = model.load_state_dict(weights, strict=False)
        if missing:
            print(f"警告：以下权重未加载: {missing}")

4.4 实战：从本地路径加载

本地加载其实更常见。我经常把模型下载到本地，然后直接指定路径：

class LocalWeightLoader(BaseWeightLoader):
    def __init__(self, file_pattern: str = "*.safetensors"):
        self.file_pattern = file_pattern

    def get_weights(self, model_name_or_path: str) -> Dict[str, torch.Tensor]:
        import glob
        import os
        from safetensors import safe_open

        weights = {}
        # 找到所有匹配的文件
        files = glob.glob(os.path.join(model_name_or_path, self.file_pattern))

        if not files:
            raise FileNotFoundError(f"在 {model_name_or_path} 下没找到权重文件")

        for file_path in files:
            with safe_open(file_path, framework="pt") as f:
                for key in f.keys():
                    weights[key] = f.get_tensor(key)

        return weights

    def load_weights(self, model: nn.Module, weights: Dict[str, torch.Tensor]):
        # 这里可以加一些校验逻辑
        # 我曾经遇到过权重 key 和模型 key 不匹配的情况
        model_keys = set(model.state_dict().keys())
        weight_keys = set(weights.keys())

        if model_keys != weight_keys:
            print(f"模型有 {len(model_keys)} 个参数，权重有 {len(weight_keys)} 个")
            # 可以尝试自动匹配
            self._auto_match_and_load(model, weights)
        else:
            model.load_state_dict(weights)

4.5 实战：自定义格式加载

这是最灵活的场景。假设你的权重文件是二进制格式，每 4 字节一个 float：

class CustomBinaryWeightLoader(BaseWeightLoader):
    def __init__(self, dtype: torch.dtype = torch.float32):
        self.dtype = dtype

    def get_weights(self, model_name_or_path: str) -> Dict[str, torch.Tensor]:
        import struct
        import os

        weights = {}
        # 假设有一个 meta 文件描述结构
        meta_path = os.path.join(model_name_or_path, "meta.txt")
        bin_path = os.path.join(model_name_or_path, "model.bin")

        with open(meta_path, "r") as f:
            for line in f:
                name, shape_str = line.strip().split(":")
                shape = tuple(map(int, shape_str.split(",")))
                weights[name] = shape  # 先存形状信息

        # 读取二进制数据
        with open(bin_path, "rb") as f:
            for name, shape in weights.items():
                num_elements = 1
                for dim in shape:
                    num_elements *= dim
                # 读取对应数量的 float
                data = struct.unpack(f"{num_elements}f", f.read(num_elements * 4))
                weights[name] = torch.tensor(data, dtype=self.dtype).reshape(shape)

        return weights

    def load_weights(self, model: nn.Module, weights: Dict[str, torch.Tensor]):
        # 这里可以按需做转换
        # 比如有些层需要转置
        for name, param in model.named_parameters():
            if name in weights:
                if param.shape != weights[name].shape:
                    # 自动处理形状不匹配
                    weights[name] = weights[name].reshape(param.shape)
                param.data.copy_(weights[name])

4.6 注册你的 WeightLoader

写好了类，怎么让 vLLM 用上？很简单，注册一下：

from vllm.model_executor.weight_loader import WeightLoaderRegistry

# 注册自定义加载器
WeightLoaderRegistry.register("my_custom", MyCustomWeightLoader)

# 使用
from vllm import LLM

llm = LLM(
    model="path/to/model",
    weight_loader="my_custom",  # 指定使用自定义加载器
    # 还可以传参数
    weight_loader_kwargs={
        "dtype": torch.bfloat16,
        "file_pattern": "*.bin"
    }
)

嗯，这里要注意：注册的名字不能和已有的冲突。vLLM 内置了 "hf"、"local" 等名字，你可以用 WeightLoaderRegistry.list_available() 查看。

4.7 避坑指南

我踩过的坑，分享给你：

• Key 不匹配：vLLM 内部可能会对权重 key 做重命名（比如加 model. 前缀）。我建议在 load_weights 里打印一下 key 的映射关系。
• 数据类型：权重文件可能是 float32，但模型期望 bfloat16。记得做类型转换。
• 分片加载：大模型权重通常分片存储。你的 get_weights 要能处理多个文件。
• 内存管理：加载完权重后，及时释放临时变量。我习惯用 del 加上 torch.cuda.empty_cache()。

4.8 知识体系图

下面这张图总结了 WeightLoader 的核心逻辑：

这张图展示了三种加载器都继承自同一个基类，核心就是实现 get_weights 和 load_weights 两个方法。最后通过注册机制让 vLLM 知道你的加载器存在。

4.9 小结

自定义 WeightLoader 其实不复杂。你只需要：

1. 继承 BaseWeightLoader
2. 实现 get_weights 和 load_weights
3. 注册到 WeightLoaderRegistry

就这么三步。我刚开始也觉得这东西很神秘，后来发现其实就是个「读文件 + 塞参数」的流程。你想想看，vLLM 的设计哲学就是「把复杂留给自己，把简单留给用户」。

嗯，下一节我们会聊到模型配置的适配。不过那是另一个话题了，先把 WeightLoader 练熟再说。

---