Extend ARCH_TO_2D_MAPPINGS to Qwen3.5-MoE

src/llmcompressor/modeling/moe/conversion_mappings.py

@@ -147,6 +147,12 @@
         "transformers.models.qwen3_5_moe.configuration_qwen3_5_moe.Qwen3_5MoeTextConfig",
         "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe.Qwen3_5MoeExperts",
     ),
+    # Text-only checkpoints (`Qwen3_5MoeForCausalLM`) report `qwen3_5_moe_text`
+    # as their top-level model type, and share the multimodal expert classes.
+    "qwen3_5_moe_text": (
+        "transformers.models.qwen3_5_moe.configuration_qwen3_5_moe.Qwen3_5MoeTextConfig",
+        "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe.Qwen3_5MoeExperts",
+    ),
     "qwen3_moe": (
         "transformers.models.qwen3_moe.configuration_qwen3_moe.Qwen3MoeConfig",
         "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeExperts",
@@ -216,9 +222,45 @@
 }
 
 
+# Qwen3.5-MoE stores 2D per-expert tensors and fuses them with the same rules
+# as Qwen2-MoE, so it reuses that 2D body. Both released spellings are
+# registered: `qwen3_5_moe` for the multimodal wrapper, and `qwen3_5_text` for
+# the text-only checkpoint, which `_MODEL_TO_CONVERSION_PATTERN` remaps
+# `qwen3_5_moe_text` to.
+ARCH_TO_2D_MAPPINGS["qwen3_5_moe"] = ARCH_TO_2D_MAPPINGS["qwen2_moe"]
+ARCH_TO_2D_MAPPINGS["qwen3_5_text"] = ARCH_TO_2D_MAPPINGS["qwen2_moe"]
+
+
+def _resolve_checkpoint_conversion_mapping(
+    model_type: str,
+) -> list[WeightTransform] | None:
+    """
+    Resolve the transformers checkpoint conversion mapping for a model type.
+
+    `ARCH_TO_IMPORT_PATHS` is keyed on the top-level model type, but transformers
+    may register the conversion rules on the text tower instead: a multimodal
+    Qwen3.5-MoE config reports `qwen3_5_moe`, while the rules live under
+    `qwen3_5_moe_text`. Fall back to the text spelling so those architectures
+    resolve instead of returning `None`.
+
+    :param model_type: top-level model type reported by the config
+    :return: conversion mapping, or None if neither spelling is registered
+    """
+    mapping = get_checkpoint_conversion_mapping(model_type)
+    if mapping is None and not model_type.endswith("_text"):
+        mapping = get_checkpoint_conversion_mapping(f"{model_type}_text")
+    return mapping
+
+
 def has_linearize_load_mappings(model_type: str) -> bool:
     remapped_type = _MODEL_TO_CONVERSION_PATTERN.get(model_type, model_type)
-    return model_type in ARCH_TO_IMPORT_PATHS and remapped_type in ARCH_TO_2D_MAPPINGS
+    return (
+        model_type in ARCH_TO_IMPORT_PATHS
+        and remapped_type in ARCH_TO_2D_MAPPINGS
+        # a 2D entry alone is not enough: without conversion rules to strip the fused
+        # targets from, `get_linearize_load_mappings` has nothing to build on
+        and _resolve_checkpoint_conversion_mapping(model_type) is not None
+    )
 
 
 def get_linearize_load_mappings(
@@ -228,7 +270,13 @@ def get_linearize_load_mappings(
     _config_paths, expert_paths = ARCH_TO_IMPORT_PATHS[model_type]
     experts_cls = import_or_none(expert_paths)
 
-    mapping: list[WeightTransform] = get_checkpoint_conversion_mapping(model_type)
+    mapping = _resolve_checkpoint_conversion_mapping(model_type)
+    if mapping is None:
+        raise ValueError(
+            "No checkpoint conversion mapping is registered for model type "
+            f"`{model_type}`, so linearized load mappings cannot be derived. "
+            "Gate this call on `has_linearize_load_mappings`."
+        )
     model_type = _MODEL_TO_CONVERSION_PATTERN.get(model_type, model_type)
     remove_targets, new_mappings = ARCH_TO_2D_MAPPINGS[model_type]
 

tests/llmcompressor/modeling/test_moe_conversion_mappings.py

@@ -0,0 +1,96 @@
+import pytest
+from transformers.core_model_loading import WeightConverter
+
+from llmcompressor.modeling.moe.conversion_mappings import (
+    ARCH_TO_2D_MAPPINGS,
+    get_linearize_load_mappings,
+    has_linearize_load_mappings,
+)
+
+# Qwen3.5-MoE checkpoints store 2D per-expert tensors, so they must take the direct-load
+# pathway. Both released spellings are covered: `qwen3_5_moe` is reported by
+# the multimodal wrapper config, `qwen3_5_moe_text` by the text-only checkpoint.
+QWEN3_5_MOE_TYPES = ["qwen3_5_moe", "qwen3_5_moe_text"]
+
+
+@pytest.mark.parametrize("model_type", QWEN3_5_MOE_TYPES)
+def test_qwen3_5_moe_has_linearize_load_mappings(model_type):
+    assert has_linearize_load_mappings(model_type)
+
+
+@pytest.mark.parametrize("model_type", QWEN3_5_MOE_TYPES)
+def test_qwen3_5_moe_load_mappings_avoid_conversion(model_type):
+    """
+    A remaining `WeightConverter` means weights are fused on load, which is the
+    2D -> 3D -> 2D round trip this pathway exists to avoid.
+    """
+    experts_cls, load_mappings, save_mappings = get_linearize_load_mappings(model_type)
+
+    assert experts_cls is not None
+    assert not any(isinstance(mapping, WeightConverter) for mapping in load_mappings)
+    assert not any(isinstance(mapping, WeightConverter) for mapping in save_mappings)
+
+
+@pytest.mark.parametrize("model_type", QWEN3_5_MOE_TYPES)
+def test_qwen3_5_moe_load_mappings_keep_expert_renames(model_type):
+    """The 2D body must contribute a per-expert rename for each projection."""
+    _experts_cls, load_mappings, _save_mappings = get_linearize_load_mappings(
+        model_type
+    )
+    patterns = [
+        pattern for mapping in load_mappings for pattern in mapping.source_patterns
+    ]
+
+    for projection in ("gate_proj", "up_proj", "down_proj"):
+        assert any(projection in pattern for pattern in patterns), projection
+
+
+@pytest.mark.parametrize("model_type", QWEN3_5_MOE_TYPES)
+def test_qwen3_5_moe_keeps_language_model_prefix_rule(model_type):
+    """
+    Transformers registers Qwen3.5-MoE's rules on the text tower, and they include a
+    `model.language_model.*` prefix rule that the Qwen2-MoE rules do not have.
+    Resolving the mapping through the wrong spelling would silently drop it.
+    """
+    _experts_cls, load_mappings, _save_mappings = get_linearize_load_mappings(
+        model_type
+    )
+    patterns = [
+        pattern for mapping in load_mappings for pattern in mapping.source_patterns
+    ]
+
+    assert any("language_model" in pattern for pattern in patterns)
+
+
+def test_qwen3_5_moe_reuses_qwen2_moe_2d_body():
+    for model_type in ("qwen3_5_moe", "qwen3_5_text"):
+        assert ARCH_TO_2D_MAPPINGS[model_type] == ARCH_TO_2D_MAPPINGS["qwen2_moe"]
+
+
+@pytest.mark.parametrize(
+    "model_type", ["qwen2_moe", "qwen3_moe", "qwen3_next", "deepseek_v4", "hy_v3"]
+)
+def test_existing_architectures_still_resolve(model_type):
+    assert has_linearize_load_mappings(model_type)
+
+
+def test_model_type_without_conversion_mapping_raises_clearly():
+    """
+    `gpt_oss` is in `ARCH_TO_IMPORT_PATHS` but transformers registers no conversion
+    mapping for it, so `has_linearize_load_mappings` is False and the direct-load
+    pathway is never taken. Called directly it should say so, rather than failing on
+    a downstream lookup or iterating `None`.
+    """
+    assert not has_linearize_load_mappings("gpt_oss")
+
+    with pytest.raises(ValueError, match="No checkpoint conversion mapping"):
+        get_linearize_load_mappings("gpt_oss")
+
+
+def test_qwen3_vl_moe_still_uses_post_load_conversion():
+    """
+    Qwen3-VL-MoE's conversion rules are identity, so its checkpoints are already
+    3D and it must keep falling back to `linearize_moe` rather than claiming a
+    direct-load pathway.
+    """
+    assert not has_linearize_load_mappings("qwen3_vl_moe")