[Fix] Gate HF export key remapping behind explicit --vllm flag#63
[Fix] Gate HF export key remapping behind explicit --vllm flag#63yubofredwang merged 2 commits intomainfrom
Conversation
Preserve native TorchSpec checkpoint keys by default so sglang exports keep midlayer-prefixed weights. Add an explicit --vllm option for the layers.0 remap path needed by vLLM-compatible exports.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 19d966a39a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| logger.info("Exporting native checkpoint keys") | ||
| return dict(tensors) |
There was a problem hiding this comment.
When export_for_vllm is false, this new branch emits native midlayer.* weights, but the saved config is still rewritten by _fixup_export_config in both save paths, and that helper unconditionally increments eagle_aux_hidden_state_layer_ids for model_type == "kimi_k2" (a vLLM-specific adjustment per its comment). This leaves default/native exports with off-by-one aux-layer IDs in config.json, which can select wrong hidden-state layers for non-vLLM consumers despite using native keys.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Pull request overview
This PR updates the HuggingFace export tool to preserve TorchSpec-native checkpoint key names by default, and only apply the midlayer.* -> layers.0.* key remapping when explicitly exporting for vLLM.
Changes:
- Added an explicit
--vllmCLI flag to opt into vLLM-compatible key remapping during export. - Centralized tensor preparation behind
_prepare_export_tensors()and threaded the export mode through both vocab-pruned and non-pruned save paths.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Only apply the Kimi K2 layer-id rewrite when exporting vLLM-compatible checkpoints so native sglang exports preserve the original config semantics alongside the original weight keys.
Summary
midlayer.*--vllmflag to opt into themidlayer.* -> layers.0.*remap needed for vLLM-compatible exportsTest plan
python3 tools/convert_to_hf.py --helppython3 -m py_compile tools/convert_to_hf.pyiter_0014701_hfwithout--vllmand verified exported safetensors containedmidlayer.*keys and nolayers.0.*keys