fuse_for_vllm.py 源码精读:gpt-oss-20b-BF16-w8a8-llmcompressor 如何重写 MoE 专家权重布局以无缝适配 vLLM 加载器 fuse_for_vllm.py 源码精读gpt-oss-20b-BF16-w8a8-llmcompressor 如何重写 MoE 专家权重布局以无缝适配 vLLM 加载器【免费下载链接】gpt-oss-20b-BF16-w8a8-llmcompressor项目地址: https://ai.gitcode.com/hf_mirrors/amd/gpt-oss-20b-BF16-w8a8-llmcompressorgpt-oss-20b-BF16-w8a8-llmcompressor 是 AMD 基于 LLM Compressor 产出的 W8A8 量化版本磁盘体积从 39.0 GiB 压到 20.6 GiB。它的 MoE 专家权重布局无法被 vLLM 直接加载仓库中的 fuse_for_vllm.py 就是来解决这件事的把逐专家存储的专家张量重写拼接、转置、重新分片成 vLLM 期望的融合布局无需修改任何模型文件。模型速览为什么专家权重是大头项目数值基座模型gpt-oss-20bBF1639.0 GiB量化方案W8A8INT8 权重 per-token 动态 INT8 激活对称 per-channel 缩放量化后体积约 20.6 GiB-47%MoE 结构24 层 × 32 专家每 token 选 top-4 专家experts_per_token: 4推理目标AMD EPYC CPUvLLM 0.26.0从 config.json 可以看到num_hidden_layers: 24、num_local_experts: 32。量化配置见 config.json 与 recipe.yaml只量化Linear层MoE 路由器、lm_head、embed_tokens保持 BF16。问题出在布局上。gpt-oss 原生的专家是一个融合的 3D 张量量化时必须先用 shim 模块拆成逐专家的独立nn.Linear见 README.md 的说明量化后检查点里每个专家存成gate_proj/up_proj/down_proj三组独立键。但 vLLM 的 gpt-oss 加载路径要求gate 与 up 合并为一个投影内部叫 w13且读取后要做一次固定的截取 转置——所以存储侧必须提前把转置后的结果存好两边组合起来才恰好是加载器期望的方向源码注释 fuse_for_vllm.py。 仓库里同时保留了两组 safetensorsmodel-*-of-00005是量化产出的原始布局model-*-of-00006是融合后的产物可以直接对照键名差异。转换流程总览一条命令四步走运行方式见文件头注释 fuse_for_vllm.pypython fuse_for_vllm.py SRC_DIR DST_DIR步骤动作代码位置1️⃣ 键分类读 model.safetensors.index.json把 5019 个张量分成专家键与普通键fuse_for_vllm.py2️⃣ 逐专家融合gateup 按行拼接后转置每专家产出 6 个融合张量fuse_for_vllm.py3️⃣ 增量分片缓冲到约 4 GiB 即落盘最后统一改名并重建索引fuse_for_vllm.py4️⃣ 文件搬运非权重文件config、tokenizer 等原样复制到目标目录fuse_for_vllm.py第 1 步用一条正则识别 MoE 专家权重键整个脚本的路由靠一条正则完成fuse_for_vllm.py_EXPERT_KEY re.compile( r^(?Pprefix.*\.mlp\.experts)\.(?Pexpert\d)\. r(?Pprojgate_proj|up_proj|down_proj)\. r(?Pkindweight|weight_scale|bias)$)四个命名组把键拆解为层前缀prefix、专家编号expert、投影名proj、张量类型kind。匹配不上的键注意力投影、LayerNorm、24 个 attention sinks 等 411 个张量走直通通道原样复制。若某个键含mlp.experts.却又匹配不上脚本会直接报错避免静默丢权重fuse_for_vllm.py。第 2 步gate/up 融合与转置——每个专家的 6 个键以本仓库的 W8A8 格式为例每个专家从 9 个键变成 6 个原始键逐专家融合后键变换gate_proj.weightup_proj.weightgate_up_proj.weight沿第 0 维拼接 → 转置gate_proj.weight_scaleup_proj.weight_scalegate_up_proj.weight_scale拼接 → 转置down_proj.weightdown_proj.weight转置down_proj.weight_scaledown_proj.weight_scale转置gate_proj.biasup_proj.biasgate_up_proj.bias拼接不转置down_proj.biasdown_proj.bias原样核心代码就一行fuse_for_vllm.pytorch.cat([gate_w, up_w], 0).t().contiguous()为什么要转置vLLM 加载器读到的张量会被再转置一次预存转置结果才能让最终方向正确——这是无缝适配的关键一步。⚠️W4A16 例外两种量化方案的参数轴向顺序相反——INT8 的 w13 逻辑上是[2I, K]packed INT4 的 w13 是[K/8, 2I]。因此 W4A16 分支直接按拼接原样存储不转置fuse_for_vllm.py、fuse_for_vllm.py。此外 W8A8 的gate_up_proj.weight靠 vLLM 的hf_to_vllm_mapper映射到w13_weight而 W4A16 因无对应映射规则直接以w13_weight_packed命名未匹配的键会原样穿透映射器fuse_for_vllm.py。第 3 步4 GiB 增量分片与索引重建懒加载读取_Source 按需打开 safetensors 分片句柄并缓存不会把整个检查点读进内存。缓冲写出_ShardWriter 累计到约 4 GiB可用--shard-size-gib调整就先落盘为part-00001.safetensorsfinalize时统一改名为model-00001-of-00006.safetensors并生成weight_map。一个贴心的细节save_file新建文件权限是 0600放在 NFS 上其他用户会读不了所以脚本显式chmod 0o644fuse_for_vllm.py。重建索引重写 model.safetensors.index.json写入total_size本仓库约 20.6 GiB与排序后的weight_mapfuse_for_vllm.py。使用前的三个限制先知道再动手仅支持 TP1更高的张量并行度下down 投影的 per-channel scale 无法正确传到非零 rankfuse_for_vllm.py。融合布局不属于 compressed-tensors 惯例所以源检查点仍是对外发布的正式产物融合版只是 vLLM 专用的派生形态fuse_for_vllm.py。W4A16 还需 vLLM 侧配合必须把w13_weight_packed/w2_weight_packed加入 vLLMgpt_oss.py的_GPT_OSS_STREAMED_EXPERT_SUFFIX_TO_SHARD精确匹配检查点侧无绕法W8A8 则不需要fuse_for_vllm.py。总结与参考文件一句话总结fuse_for_vllm.py用正则分类 → 拼接转置 → 增量分片 → 重建索引四步把 W8A8/W4A16 量化检查点的 4608 个 MoE 专家张量重排成 vLLM 可直读的融合布局全程流式处理、不改模型文件。核心脚本fuse_for_vllm.py融合后检查点索引model.safetensors.index.json模型架构与量化配置config.json量化配方recipe.yaml模型卡与量化流程说明README.md许可证LICENSE如需完整运行先获取包含全部权重文件的仓库副本git clone https://gitcode.com/hf_mirrors/amd/gpt-oss-20b-BF16-w8a8-llmcompressor【免费下载链接】gpt-oss-20b-BF16-w8a8-llmcompressor项目地址: https://ai.gitcode.com/hf_mirrors/amd/gpt-oss-20b-BF16-w8a8-llmcompressor创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考