背景：Checkpoint 模型在 ComfyUI 里的“老大难”
第一次把 SDXL 的 6.5 GB checkpoint 拖进 ComfyUI 时，我差点被 30 s 的加载时间劝退。更尴尬的是，一张 24 GB 显存的 A10 居然在跑 2048×2048 图时直接 OOM。
痛点总结下来就三句话：
- 模型文件越来越大，磁盘 IO 成为首屏瓶颈
- 完整加载后常驻显存，导致高分辨率批次推理寸步难行
- 多 GPU 环境缺乏原生亲和性策略，单机多卡利用率低
于是我把“让 checkpoint 跑得动、跑得快、还能热更新”当成迭代目标，踩了两个月坑，最终把端到端延迟从 28 s 压到 7 s，显存占用峰值下降 42%。下面把趟过的路写成可抄作业的代码。
技术对比：三种加载策略谁更香
| 维度 | 完整加载 | 分片加载 | 动态加载 |
|---|---|---|---|
| 首屏延迟 | 高（一次性读大文件） | 中（按需读分片） | 低（只拉必要层） |
| 显存峰值 | 高（全权重常驻） | 中（片内常驻） | 低（用后即焚） |
| 代码复杂度 | 低 | 中 | 高 |
| 推理并发 | 差 | 好 | 极好 |
| 热更新 | 需重启进程 | 可局部替换 | 单 Layer 替换 |
一句话结论：
- 离线批跑追求吞吐 → 分片加载
- 在线服务追求低延迟 → 动态加载
- 教学 demo 一键跑通 → 完整加载也无妨
实现方案：从torch.load到分布式 Pipeline
1. 安全加载：设备映射 + 校验
import hashlib, torch, contextlib, os, json from pathlib import Path CKPT_PATH = Path("/data/models/sd_xl_base_1.0.ckpt") SHA256_ETALON = "7c819b6e..." # 官方给出的哈希 def _check_sha256(path: Path, etalon: str): sha = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(1 << 20), b""): sha.update(chunk) assert sha.hexdigest() == etalon, "checksum fail" @contextlib.contextmanager def load_ckpt_safe(ckpt_path: Path, device="cpu"): _check_sha256(ckpt_path, SHA256_ETALON) ckpt = torch.load(ckpt_path, map_location=device, weights_only=True) yield ckpt del ckpt torch.cuda.empty_cache()weights_only=True屏蔽恶意 pickle- 上下文管理器保证显存及时释放
2. 分片加载：把大模型切成 2 GB 一块
思路：
- 提前用脚本把 checkpoint 按
state_dictkey 做“层”级分片，每片 ≤ 2 GB - 推理时只加载本次采样所需的层片
class ShardLoader: def __init__(self, index_file: Path, device="cuda:0"): with open(index_file) as f: self.index = json.load(f) # {"unet": "unet_00.pth", ...} self.device = device self.cache = {} # 简易 LRU 可自己加 def load_layer(self, name: str): if name in self.cache: return self.cache[name] path = Path(self.index[name]) state = torch.load(path, map_location=self.device, weights_only=True) self.cache[name] = state return state def flush(self): self.cache.clear() torch.cuda.empty_cache()内存监控：
import psutil, threading, time def monitor_ram(interval=1): def _run(): while True: print("[RAM]", psutil.virtual_memory()._asdict()) time.sleep(interval) threading.Thread(target=_run, daemon=True).start()跑推理前monitor_ram()，可实时观察系统内存，防止把宿主机 OOM。
3. 分布式 Pipeline：多 GPU 流水线
ComfyUI 原生只认单卡，我们借torch.distributed做“图内并行”：
- UNet 放 cuda:0，VAE 放 cuda:1，CLIP 留在 CPU
- 用
torch.cuda.Stream事件同步，避免空等
import torch.multiprocessing as mp def worker(rank, world_size, queue_in, queue_out): torch.cuda.set_device(rank) # 初始化子模型 if rank == 0: unet = load_layer("unet").half().cuda(rank) elif rank == 1: vae = load_layer("vae").half().cuda(rank) while True: data = queue_in.get() if data is None: break latents = data["latents"] if rank == 0: latents = unet(latents) # 伪代码 elif rank == 1: images = vae.decode(latents) queue_out.put({"latents" if rank == 0 else "images": locals()[["latents", "images"][rank]}) def spawn_pipeline(): mp.set_start_method("spawn", force=True) q1, q2 = mp.Queue(), mp.Queue() procs = [mp.Process(target=worker, args=(r, 2, q1, q2)) for r in range(2)] for p in procs: p.start() return procs, q1, q2- 生产环境可换成
torchrun+ RPC，更优雅 - 注意
half()降低带宽，但需验证 NAN/INF
性能数据：A100 vs V100 实测
| 策略 | 硬件 | 首 token 延迟 | 2048×2048 吞吐 | 峰值显存 |
|---|---|---|---|---|
| 完整加载 | V100 32 GB | 28 s | 0.12 img/s | 30.1 GB |
| 分片加载 | V100 32 GB | 9 s | 0.35 img/s | 17.4 GB |
| 动态加载 | A100 40 GB | 7 s | 0.51 img/s | 11.2 GB |
测试条件：
- batch=1，采样步 20，Euler a
- 分片 2 GB/片，动态加载仅拉 9 层 UNET
- 分布式版本额外节省 1.8 s 的 VAE decode
避坑指南：三个隐形炸弹
文件校验
下载完 checkpoint 一定先做 SHA256，血泪教训：一次 NFS 异常导致文件尾部 4 KB 全是 0，结果推理图全是噪点。多 GPU 亲和性
别轻信CUDA_VISIBLE_DEVICES，在 Docker 里可能和nvidia-smi顺序不一致。推荐torch.cuda.get_device_name()打印确认。热更新
直接覆盖文件会被 mmap 报错“text file busy”。正确姿势：- 写新文件 → 原子 mv → 发 USR1 信号给进程 → 内重新
torch.load - 或者上
fuser -k简单粗暴，但会断当前请求
- 写新文件 → 原子 mv → 发 USR1 信号给进程 → 内重新
可直接复用的完整示例
把下面脚本保存为shard_inference.py，改路径就能跑：
#!/usr/bin/env python import torch, json, time, contextlib from pathlib import Path from shard_loader import ShardLoader def main(prompt: str): loader = ShardLoader(Path("/data/sdxl_shards/index.json")) with contextlib.ExitStack() as stack: # 按需加载 text_encoder = stack.enter_context(loader.load_layer("clip")) unet = stack.enter_context(loader.load_layer("unet")) vae = stack.enter_context(loader.load_layer("vae")) # 伪推理 c = text_encoder(prompt) z = torch.randn(1, 4, 128, 128).half().cuda() for _ in range(20): z = unet(z, c) pixels = vae.decode(z) print("done", pixels.shape) if __name__ == "__main__": main("a cute robot")跑前export CUDA_VISIBLE_DEVICES=0，显存稳稳地停在 10 GB 左右。
延伸思考
- 分片粒度到底多细才合适？片越大 IO 少但显存高，片越小 IO 多却调度碎，如何自动权衡？
- 动态加载已把“用后即焚”做到极致，但频繁
torch.load会触发 Python GIL，未来有无可能把权重池放到共享内存或 GPU Direct Storage，进一步削掉 IO 延迟？
把上面的代码全部跑通后，我的 ComfyUI 服务终于可以在 8 卡 A100 上同时给 50 个设计师出图而不掉链子。虽然脚本里还有不少 hardcode，比如 LRU 大小、分片键值规则，但至少证明了 checkpoint 不是非得“全量进显存”才能玩得转。下一步打算把动态加载做成 ComfyUI 的自定义节点，让社区里更多非 Python 出身的玩家也能一键提速。
