问题域/PD-450

模型权重加载与分片

Model Weight Loading & Sharding

解决 HuggingFace 模型权重到自定义并行架构的映射问题,通过声明式映射实现自动权重融合与分片

子问题

1.QKV 权重融合映射

2.Gate/Up 权重合并

3.张量并行分片加载

4.自定义 weight_loader 回调

5.词表并行分片加载(Embedding/LMHead)

6.for-else 模式的映射命中与 fallback 路径分离

各项目的解法1 solutions

Signals

横向对比

维度nano-vllm
权重格式safetensors 逐文件遍历,CPU 模式 safe_open 避免 GPU OOM
映射机制packed_modules_mapping 类属性声明式映射,for-else 查表
分片策略参数级 weight_loader 回调,narrow+chunk 零拷贝分片
融合支持QKV 三合一 + gate/up 二合一,支持 GQA 不等头数
并行层体系5 层继承体系:Base→Column→Merged/QKV + Row + Replicated
词表处理VocabParallelEmbedding 按 rank 均分词表行

最佳实践

1.用 packed_modules_mapping 声明式定义权重融合规则

2.每个参数附带 weight_loader 回调处理分片逻辑

3.LinearBase 基类统一绑定 weight_loader,子类只需覆写分片逻辑

4.safe_open 以 CPU 模式打开避免 GPU 内存峰值

5.narrow() 零拷贝视图定位偏移,chunk() 取 TP rank 分片