Extend ARCH_TO_2D_MAPPINGS to Qwen3.5-MoE - #3050
Conversation
Qwen3.5-MoE checkpoints already store 2D per-expert tensors, but
`has_linearize_load_mappings` returned False for them, so `load_quantizable_moe`
fell back to loading normally and calling `linearize_moe`. That performs a
2D -> 3D -> 2D round trip and leaves an extra copy of every expert resident,
which is fatal during load for a large MoE on a memory-constrained host.
Register both released spellings - `qwen3_5_moe` for the multimodal wrapper and
`qwen3_5_text` for the text-only checkpoint - reusing the Qwen2-MoE 2D body,
whose fusing rules are identical.
Resolving the transformers mapping needs care: it is registered on the text
tower (`qwen3_5_moe_text`), so `get_checkpoint_conversion_mapping("qwen3_5_moe")`
returns None while llm-compressor keys on the top-level type. Add
`_resolve_checkpoint_conversion_mapping`, which falls back to the text spelling,
and require a resolvable mapping in `has_linearize_load_mappings` so a future 2D
entry cannot pass the predicate and then hand None to the mapping consumer.
Taking the mapping from the text tower also preserves its
`model.language_model.*` prefix rule, which copying the Qwen2-MoE rules verbatim
would have dropped.
Fixes vllm-project#3037
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for Qwen3.5-MoE models (both multimodal and text-only checkpoints) by mapping them to the Qwen2-MoE 2D per-expert tensor structure and resolving their checkpoint conversion mappings with fallback logic for text-only suffixes. Unit tests are added to verify these mappings. The review feedback identifies a potential type mismatch and runtime error in get_linearize_load_mappings where _resolve_checkpoint_conversion_mapping could return None, and suggests adding a defensive check to raise a ValueError if the mapping is not found.
| experts_cls = import_or_none(expert_paths) | ||
|
|
||
| mapping: list[WeightTransform] = get_checkpoint_conversion_mapping(model_type) | ||
| mapping: list[WeightTransform] = _resolve_checkpoint_conversion_mapping(model_type) |
There was a problem hiding this comment.
The function _resolve_checkpoint_conversion_mapping can return None. Assigning its result directly to mapping (annotated as list[WeightTransform]) without a None check will cause a static type checker error (type mismatch) and can lead to a TypeError: 'NoneType' object is not iterable at runtime when mapping is iterated over in the list comprehension below if get_linearize_load_mappings is called directly with an unsupported model type.
Adding a defensive None check and raising a clear ValueError resolves both the type mismatch and prevents runtime crashes.
| mapping: list[WeightTransform] = _resolve_checkpoint_conversion_mapping(model_type) | |
| mapping = _resolve_checkpoint_conversion_mapping(model_type) | |
| if mapping is None: | |
| raise ValueError(f"No checkpoint conversion mapping found for model type: {model_type}") |
`_resolve_checkpoint_conversion_mapping` is Optional, so assigning it straight to a `list[WeightTransform]` was a type mismatch, and a caller reaching `get_linearize_load_mappings` for a type with no mapping would fail on a downstream lookup rather than on the actual cause. Raise a ValueError naming the model type and pointing at `has_linearize_load_mappings`, which is what gates this call.
|
Fixed in 0df550c — the Optional return being assigned straight to a Small correction to the failure mode, for the record: with the current tables the crash is a Either way the error pointed away from the cause, so the guard raises a |
|
See comment: #3037 (comment) |
2 participants
SUMMARY:
Fixes #3037. Qwen3.5-MoE checkpoints already store 2D per-expert tensors, but
has_linearize_load_mappingsreturned False for them, soload_quantizable_moefell back to loading normally and callinglinearize_moe— a 2D → 3D → 2D round trip that leaves an extra copy of every expert resident for the process lifetime. On a memory-constrained host that is fatal during load, before calibration starts.Registers both released spellings, reusing the Qwen2-MoE 2D body (its fusing rules are identical):
qwen3_5_moe— multimodal wrapper (Qwen3_5MoeForConditionalGeneration)qwen3_5_text— what_MODEL_TO_CONVERSION_PATTERNremaps the text-only checkpoint'sqwen3_5_moe_textto, plus itsARCH_TO_IMPORT_PATHSentryThe issue notes that adding a 2D entry alone is not sufficient, and that is confirmed: transformers registers the rules on the text tower, so
get_checkpoint_conversion_mapping("qwen3_5_moe")returnsNonewhile llm-compressor keys on the top-level type. Of the three options raised in the issue this takes the "resolve the text spelling" one, as the least invasive —_resolve_checkpoint_conversion_mappingfalls back tof"{model_type}_text", andhas_linearize_load_mappingsnow additionally requires a resolvable mapping, so a future 2D entry cannot satisfy the predicate and then handNonetoget_linearize_load_mappings. Happy to switch to registering both spellings or normalising the suffix if you prefer a different convention.One thing worth flagging: taking the mapping from the text tower also preserves its
model.language_model.*prefix rule, which the Qwen2-MoE rules do not carry — copying that body verbatim would have silently dropped it. There is a test pinning it.TEST PLAN:
New CPU-only unit tests in
tests/llmcompressor/modeling/test_moe_conversion_mappings.py(15 cases, transformers 5.15.0): the predicate for both Qwen3.5-MoE spellings, noWeightConvertersurviving in the load/save mappings (a survivor means fusion on load, i.e. the round trip), the per-expert renames, thelanguage_modelprefix rule, the shared 2D body, and guards thatqwen2_moe/qwen3_moe/qwen3_next/deepseek_v4/hy_v3still resolve whileqwen3_vl_moekeeps using post-load conversion (its rules are identity, so it has no round trip to avoid).main, and the 6 that pass either way are the regression guards.has_linearize_load_mappingsover all 32ARCH_TO_IMPORT_PATHSkeys before and after: exactlyqwen3_5_moeandqwen3_5_moe_textflip to True, no architecture regresses.ruff checkandruff format --checkclean on both files.