1. 原地拷贝#

1
def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor):
2
    param.data.copy_(loaded_weight)

先看第一个函数default_weight_loader，注意这里用的是param.data.copy_()，在 PyTorch 中，如果直接使用 = 赋值会改变对象的引用，有可能破坏原有的计算图，或者导致 Parameter 对象丢失 requires_grad 等属性。通过 .data.copy_() 原地拷贝，就可以直接将底层的 Tensor 内存数据进行覆盖。

2. 流式加载#

1
def safetensors_weights_iterator(hf_folder: str) -> Iterable[tuple[str, torch.Tensor]]:
2
    files = glob.glob(os.path.join(hf_folder, "*.safetensors"))
3
    if not files:
4
        raise FileNotFoundError(f"No safetensors files found in {hf_folder}")
5
    for st_file in tqdm(files, desc="Loading weights"):
6
        with safe_open(st_file, framework="pt") as f:
7
            for name in f.keys():
8
                yield name, f.get_tensor(name)

再来看流式加载safetensors_weights_iterator。

这样做是防止把几百 GB 的权重一次性塞进内存。

以 Qwen3-0.6B 为例，权重大约 1.2 GB；但如果是 72B，光权重文件就有 140 GB 以上，而且通常被切成多个 .safetensors 分片。如果用 return 把所有 tensor 收集到一个 list 里再返回，内存峰值 ≈ 全部权重之和。还没开始推理，内存就先爆了。

用 yield 之后，整个加载链路变成了流水线：

1
磁盘 (.safetensors)
2
  │  mmap 映射，不立刻读入
3
  ▼
4
safe_open → f.get_tensor(name)    ← 按需读取单个 tensor
5
  │  yield (name, tensor)          ← 吐出一个，暂停
6
  ▼
7
load_weights 的 for 循环           ← 立刻 copy_ 到模型参数
8
  │  default_weight_loader / QKV 合并
9
  ▼
10
模型 param.data 就位               ← 这个 tensor 可以被 GC 回收

这意味着在任意时刻，内存里最多只躺着当前正在处理的那一个权重 tensor，处理完就释放，峰值内存从 O(全部权重) 降到了 O(单个最大 tensor)。

还有一个工程上的好处是解耦。safetensors_weights_iterator 只负责”从磁盘一个一个吐权重”，load_weights 只负责”接到一个、写入一个”。以后如果要支持从网络流式下载、或者从量化后的不同格式加载，只需要换一个新的 iterator，消费端的 load_weights 完全不用改。这正是 vLLM 真实代码里的设计思路。

3.把权重填进模型#

iterator 生成的流式输出最终要放在我们(即将)搭建好的模型骨架当中。

1
def load_weights(self, weights):
2
    params = dict(self.named_parameters())
3
    loaded = set()
4

5
    for name, w in weights: # ← 消费 yield 出来的每一个权重
6
        # HuggingFace 存了 q_proj / k_proj / v_proj 三个矩阵
7
        # 我们推理时合并成了一个 qkv_proj，需要按偏移量拼回去
8
        for hf, our, offset, size in [
9
            ("q_proj", "qkv_proj", 0, _q_size),
10
            ("k_proj", "qkv_proj", _q_size, _kv_size),
11
            ("v_proj", "qkv_proj", _q_size + _kv_size, _kv_size),
12
        ]:
13
            if hf in name:
14
                p = params[name.replace(hf, our)]
15
                p.data[offset:offset + size].copy_(w)  # 切片写入，不是整段覆盖
16
                break
17
        else:
18
            if name in params:
19
                default_weight_loader(params[name], w)  # 名字对得上，直接拷贝

这里有两个值得注意的细节：

1. QKV 合并加载。HuggingFace 的 checkpoint 里，q_proj、k_proj、v_proj 是三个独立的权重文件条目；但我们在 Qwen3Attention 里为了推理效率，把它们合并成了一个 qkv_proj。所以加载时不能整段 copy_，而是按偏移量把 Q、K、V 分别塞进 qkv_proj 的前、中、后三段。

2. 名字映射与跳过。HF 的权重名和我们 named_parameters() 的名字不一定完全一致（比如 model.layers.0.self_attn.q_proj.weight vs 我们内部的命名）。对不上的权重会被记录到 skipped 列表里打印出来，方便排查。

把 safetensors_weights_iterator 和 load_weights 串起来，完整的加载流程就是：

1
def load_model(model: nn.Module, hf_folder: str) -> nn.Module:
2
    model.load_weights(safetensors_weights_iterator(hf_folder))
3
    return model.eval()

1
load_model(model, hf_folder)
2
  └─ model.load_weights(safetensors_weights_iterator(hf_folder))
3
       └─ for (name, tensor) in iterator:   # yield 逐个吐出
4
            └─ copy_ 到对应 param             # 消费一个、释放一个