第1章模型加载优化

模型加载优化：从文件到显存的最后一公里

模型加载，听起来很简单对吧？不就是把文件读进来嘛。但如果你真的部署过百亿参数的大模型，就会明白——加载这一步，往往是整个推理链路的第一个瓶颈。我个人习惯把模型加载比作「搬砖」：砖头（权重）就堆在那，怎么搬得快、搬得稳，是门学问。

核心观点：模型加载优化的本质，是让 I/O 和计算尽可能重叠，减少 GPU 的空转等待时间。

1.1 模型文件格式：PyTorch vs SafeTensors

先说文件格式。你训练完模型，PyTorch 默认存成 .pt 或 .pth 文件。这玩意儿本质是 Python 的 pickle 序列化。嗯，这里要注意——pickle 有安全隐患。我在项目中遇到过，加载一个从网上下载的 .pt 文件，结果执行了恶意代码，整个环境被污染了。从那以后，我对 pickle 格式就多留了个心眼。

SafeTensors 就是为解决这个问题诞生的。它只存张量的原始字节，不存任何 Python 对象。说白了，就是「纯数据」文件。

特性	PyTorch (.pt/.pth)	SafeTensors
序列化方式	Pickle（可执行代码）	纯数据（无代码）
加载速度	较慢（需反序列化）	快（直接内存映射）
安全性	低（可植入恶意代码）	高（只读数据）
零拷贝加载	不支持	支持

你想想看，一个 70B 的模型，光权重就 140GB 左右。如果用 PyTorch 的 torch.load()，它会先把整个文件读到内存，再反序列化成 Python 对象，最后转成张量。这中间的内存峰值可能是模型大小的 2-3 倍。我建议生产环境优先用 SafeTensors 格式。

# PyTorch 加载方式（不推荐用于大模型）
import torch
model = torch.load('model.pth')  # 内存峰值高

# SafeTensors 加载方式（推荐）
from safetensors.torch import load_file
weights = load_file('model.safetensors')  # 零拷贝，内存友好

1.2 内存映射加载：让操作系统帮你搬砖

内存映射（mmap）是个好东西。它的原理很简单：把磁盘文件直接映射到进程的虚拟地址空间。你访问某个地址时，操作系统才真正去磁盘读数据。这就实现了「按需加载」。

我记得第一次用 mmap 加载一个 30GB 的模型文件时，mmap() 调用瞬间返回，我吓了一跳——以为没加载成功。后来才明白，mmap 只是建立了映射关系，真正的 I/O 发生在你访问张量数据的那一刻。

小技巧：用 mmap 加载模型时，配合 madvise(MADV_WILLNEED) 可以提前预取数据，减少首次访问的延迟。

import mmap
import os

# 内存映射加载
with open('model.safetensors', 'rb') as f:
    # 映射整个文件
    mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    # 告诉内核我们即将访问这些数据
    os.madvise(mm.fileno(), 0, mm.size(), os.MADV_WILLNEED)
    # 现在可以像访问数组一样访问数据
    data = mm[offset:offset+size]

为什么会这样？因为 mmap 利用了操作系统的页缓存机制。你第一次读数据时，会有一次缺页中断，但后续访问就直接走内存了。说白了，mmap 把 I/O 调度交给了内核，而内核在这方面比我们任何用户态代码都做得好。

1.3 分布式加载策略：多机多卡怎么分

单机显存不够怎么办？分布式加载。这里有两个主流策略：

张量并行（Tensor Parallelism）：把单个张量切分到多张卡上。比如一个 4096x4096 的权重矩阵，切成 4 份，每张卡拿 1024x4096。
流水线并行（Pipeline Parallelism）：按层切分。比如 40 层 Transformer，前 10 层放 GPU0，中间 20 层放 GPU1，后 10 层放 GPU2。

我在项目中遇到过一个问题：用张量并行时，如果切分不均匀，会导致某些卡的计算量远大于其他卡，形成「木桶效应」。我建议用 动态负载均衡 策略，根据每张卡的实际算力动态调整切分比例。

# 伪代码：分布式加载策略
def distributed_load(model_path, world_size, rank):
    # 1. 读取模型元信息
    metadata = read_metadata(model_path)
    
    # 2. 根据 rank 计算本卡需要加载的权重范围
    if strategy == 'tensor_parallel':
        # 按维度切分
        shard_size = metadata['total_size'] // world_size
        start = rank * shard_size
        end = (rank + 1) * shard_size if rank != world_size - 1 else metadata['total_size']
    elif strategy == 'pipeline_parallel':
        # 按层切分
        layers_per_gpu = metadata['num_layers'] // world_size
        start_layer = rank * layers_per_gpu
        end_layer = (rank + 1) * layers_per_gpu
    
    # 3. 只加载本卡需要的部分
    weights = load_shard(model_path, start, end)
    return weights

1.4 模型权重预热：别让第一次推理拖后腿

权重预热，说白了就是让模型「跑起来」一次。为什么需要预热？因为 PyTorch 的 JIT 编译、CUDA 的 kernel 加载、显存分配等，都有惰性初始化的特点。第一次推理时，这些开销会叠加，导致延迟飙升。

我曾经部署一个 GPT 模型，第一次推理花了 12 秒，第二次只要 800 毫秒。嗯，这就是没做预热的后果。从那以后，我每次部署都会加一个预热步骤。

注意：预热时要用真实数据形状，不要用随机形状。因为 PyTorch 的 JIT 会缓存编译结果，如果预热数据和实际数据形状不同，缓存就失效了。

def warmup_model(model, input_shape, device, num_warmup=10):
    """
    模型预热函数
    """
    # 创建与真实数据形状一致的 dummy input
    dummy_input = torch.randn(input_shape, device=device)
    
    # 先跑一次，触发所有 CUDA kernel 加载
    with torch.no_grad():
        _ = model(dummy_input)
    
    # 再跑几次，让 JIT 完成编译优化
    for _ in range(num_warmup):
        _ = model(dummy_input)
    
    # 同步 CUDA，确保所有操作完成
    torch.cuda.synchronize()
    
    print(f"模型预热完成，共执行 {num_warmup + 1} 次推理")

知识体系总览

下面这张图，是我梳理的模型加载优化全链路。你可以把它当作一张「地图」，后续深入每个环节时，随时回来对照。

我的建议：如果你刚开始做模型部署，先从 SafeTensors + mmap 入手。这两个组合能解决 80% 的加载问题。等业务规模上去了，再考虑分布式加载和预热策略。

模型加载优化，说白了就是让数据更快地到达 GPU。文件格式决定了数据怎么存，mmap 决定了怎么读，分布式决定了怎么分，预热决定了第一次跑多快。这四个环节环环相扣，任何一个短板都会拖慢整体。

我曾经在一个项目中，只优化了文件格式和 mmap，就把模型加载时间从 45 秒降到了 8 秒。你想想看，如果每次部署都要等将近一分钟，那体验得多糟糕。优化之后，整个团队的迭代效率都上来了。

📚 课程目录

模型加载优化：从文件到显存的最后一公里

1.1 模型文件格式：PyTorch vs SafeTensors

1.2 内存映射加载：让操作系统帮你搬砖

1.3 分布式加载策略：多机多卡怎么分

1.4 模型权重预热：别让第一次推理拖后腿

知识体系总览