模型权重加载与分片
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 分片