feat: support afmoe and bailing_hybrid expert defusion

README.md

@@ -66,7 +66,7 @@ Defuser currently supports the following `transformers>=5.3.0` `model_type` valu
 
 | Pattern | Supported model types                                                                                                                                                                                                               | Defused op performed ⚙️ |
 | --- |-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| --- |
-| Standard routed expert tensors 🧱 | `deepseek_v2`, `diffusion_gemma`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `laguna`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open`, `solar_open2` | Splits fused expert tensors or registered expert buffers into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
+| Standard routed expert tensors 🧱 | `afmoe`, `bailing_hybrid`, `deepseek_v2`, `diffusion_gemma`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `laguna`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open`, `solar_open2` | Splits fused expert tensors or registered expert buffers into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
 | Mixed sparse and shared experts | `deepseek_v3`, `deepseek_v32`, `deepseek_v4`, `glm5_next`, `glm_moe_dsa`, `qwen3_5_moe`, `qwen3_5_moe_text`, `qwen4_exp`, `qwen4_exp_text`                                                                                                              | Runtime expert tensor defusion for routed experts while preserving the model's shared-expert path. |
 | Transposed or packed expert tensors | `gpt_oss`, `phimoe`                                                                                                                                                                                                                 | Splits transposed fused expert `gate_up_proj` tensors into per-expert `gate_proj` + `up_proj`, preserves expert bias when present, and converts expert tensors into numbered expert `nn.Linear` modules. |
 | Flattened expert layout | `dbrx`                                                                                                                                                                                                                              | Rebuilds the flattened DBRX expert FFN weights into numbered expert `gate_proj`, `up_proj`, and `down_proj` `nn.Linear` modules. |

defuser/model_registry.py

@@ -16,9 +16,27 @@ class PATCH(str, Enum):
 
 
 MODEL_CONFIG = {
+    "afmoe": {
+        "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION,
+        PATCH.EXPERTS_DEFUSE: [
+            {
+                "module_class": "transformers.models.afmoe.modeling_afmoe.AfmoeExperts",
+                "forward_impl": "linear_loop",
+            }
+        ],
+    },
     "axk2": {
         "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION,
     },
+    "bailing_hybrid": {
+        "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION,
+        PATCH.EXPERTS_DEFUSE: [
+            {
+                "module_class": "transformers.models.bailing_hybrid.modeling_bailing_hybrid.BailingHybridExperts",
+                "forward_impl": "linear_loop",
+            }
+        ],
+    },
     "cohere2_moe": {
         "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION,
     },

pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "Defuser"
-version = "0.0.28"
+version = "0.0.29"
 description = "Model defuser helper for HF Transformers."
 readme = "README.md"
 requires-python = ">=3.9"

tests/test_convert_model.py

@@ -9,6 +9,13 @@
 from safetensors.torch import save_file
 from torch import nn
 from transformers.core_model_loading import WeightConverter
+from transformers.models.afmoe.configuration_afmoe import AfmoeConfig
+from transformers.models.afmoe.modeling_afmoe import AfmoeExperts, AfmoeForCausalLM
+from transformers.models.bailing_hybrid.configuration_bailing_hybrid import BailingHybridConfig
+from transformers.models.bailing_hybrid.modeling_bailing_hybrid import (
+    BailingHybridExperts,
+    BailingHybridForCausalLM,
+)
 from transformers.models.glm4_moe.modeling_glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM, Glm4MoeMoE
 from transformers.models.glm4_moe_lite.configuration_glm4_moe_lite import Glm4MoeLiteConfig
 from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import (
@@ -85,6 +92,53 @@ def _tiny_moe_config(config_cls, num_hidden_layers: int=1):
     )
 
 
+def _tiny_afmoe_config(num_hidden_layers: int = 1):
+    return AfmoeConfig(
+        num_hidden_layers=num_hidden_layers,
+        # Keep every layer on the MoE path so the routed experts are present.
+        num_dense_layers=0,
+        hidden_size=64,
+        intermediate_size=128,
+        moe_intermediate_size=32,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=16,
+        num_experts=4,
+        num_experts_per_tok=2,
+        num_shared_experts=2,
+        vocab_size=128,
+        max_position_embeddings=64,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+    )
+
+
+def _tiny_bailing_hybrid_config(num_hidden_layers: int = 2):
+    return BailingHybridConfig(
+        num_hidden_layers=num_hidden_layers,
+        # Keep every layer on the MoE path so the routed experts are present.
+        first_k_dense_replace=0,
+        layer_types=(
+            ["linear_attention", "full_attention"] * ((num_hidden_layers + 1) // 2)
+        )[:num_hidden_layers],
+        hidden_size=64,
+        intermediate_size=128,
+        moe_intermediate_size=32,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=16,
+        num_local_experts=4,
+        num_experts_per_tok=2,
+        n_shared_experts=2,
+        vocab_size=128,
+        max_position_embeddings=64,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+    )
+
+
 def _tiny_qwen3_omni_config():
     return Qwen3OmniMoeConfig(
         initializer_range=0.02,
@@ -1251,6 +1305,160 @@ def test_gpt_oss_split_forward_matches_fused_math():
     torch.testing.assert_close(actual, expected)
 
 
+def test_afmoe_registry_and_defusion():
+    from defuser.modeling.moe_experts_interface import linear_loop_experts_forward
+
+    assert MODEL_CONFIG["afmoe"][PATCH.EXPERTS_DEFUSE] == [
+        {
+            "module_class": "transformers.models.afmoe.modeling_afmoe.AfmoeExperts",
+            "forward_impl": "linear_loop",
+        }
+    ]
+
+    model = AfmoeForCausalLM(_tiny_afmoe_config())
+    assert model.config.model_type == "afmoe"
+
+    experts = model.model.layers[0].mlp.experts
+    assert isinstance(experts, AfmoeExperts)
+    num_experts = model.config.num_experts
+    hidden_dim = model.config.hidden_size
+    intermediate_dim = model.config.moe_intermediate_size
+    assert tuple(experts.gate_up_proj.shape) == (num_experts, 2 * intermediate_dim, hidden_dim)
+    assert tuple(experts.down_proj.shape) == (num_experts, hidden_dim, intermediate_dim)
+
+    expected_gate = experts.gate_up_proj[0, :intermediate_dim, :].clone()
+    expected_up = experts.gate_up_proj[0, intermediate_dim:, :].clone()
+    expected_down = experts.down_proj[0, :, :].clone()
+
+    converted = convert_model(model, cleanup_original=False, max_layers=1)
+    assert converted
+
+    experts = model.model.layers[0].mlp.experts
+    assert isinstance(experts, AfmoeExperts)
+    assert not hasattr(experts, "gate_up_proj")
+    assert not hasattr(experts, "down_proj")
+    assert experts.forward.__func__ is linear_loop_experts_forward
+
+    containers = [name for name, _ in experts.named_children() if name.isdigit()]
+    assert sorted(containers, key=int) == [str(idx) for idx in range(num_experts)]
+
+    expert0 = getattr(experts, "0")
+    assert isinstance(expert0.gate_proj, nn.Linear)
+    assert isinstance(expert0.up_proj, nn.Linear)
+    assert isinstance(expert0.down_proj, nn.Linear)
+
+    torch.testing.assert_close(expert0.gate_proj.weight, expected_gate)
+    torch.testing.assert_close(expert0.up_proj.weight, expected_up)
+    torch.testing.assert_close(expert0.down_proj.weight, expected_down)
+
+
+def test_afmoe_split_forward_matches_fused_math():
+    model = AfmoeForCausalLM(_tiny_afmoe_config())
+    fused_experts = model.model.layers[0].mlp.experts
+    assert isinstance(fused_experts, AfmoeExperts)
+
+    input_ids = torch.randint(0, model.config.vocab_size, (2, 7))
+    hidden_states = torch.randn(5, model.config.hidden_size, dtype=torch.float32)
+    top_k_index = torch.randint(0, model.config.num_experts, (hidden_states.size(0), 2))
+    top_k_weights = torch.rand((hidden_states.size(0), 2), dtype=hidden_states.dtype)
+
+    with torch.no_grad():
+        expected_experts = fused_experts(hidden_states, top_k_index, top_k_weights)
+        expected_logits = model(input_ids).logits
+
+    converted = convert_model(model, cleanup_original=False, max_layers=1)
+    assert converted
+
+    split_experts = model.model.layers[0].mlp.experts
+    _assert_unfused_expert_module(split_experts)
+    with torch.no_grad():
+        actual_experts = split_experts(hidden_states, top_k_index, top_k_weights)
+        actual_logits = model(input_ids).logits
+
+    # The defused linear_loop path should reproduce the original fused experts math.
+    torch.testing.assert_close(actual_experts, expected_experts, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_logits, expected_logits, atol=1e-5, rtol=1e-5)
+
+
+def test_bailing_hybrid_registry_and_defusion():
+    from defuser.modeling.moe_experts_interface import linear_loop_experts_forward
+
+    assert MODEL_CONFIG["bailing_hybrid"][PATCH.EXPERTS_DEFUSE] == [
+        {
+            "module_class": (
+                "transformers.models.bailing_hybrid.modeling_bailing_hybrid."
+                "BailingHybridExperts"
+            ),
+            "forward_impl": "linear_loop",
+        }
+    ]
+
+    model = BailingHybridForCausalLM(_tiny_bailing_hybrid_config())
+    assert model.config.model_type == "bailing_hybrid"
+
+    experts = model.model.layers[0].mlp.experts
+    assert isinstance(experts, BailingHybridExperts)
+    num_experts = model.config.num_local_experts
+    hidden_dim = model.config.hidden_size
+    intermediate_dim = model.config.moe_intermediate_size
+    assert tuple(experts.gate_up_proj.shape) == (num_experts, 2 * intermediate_dim, hidden_dim)
+    assert tuple(experts.down_proj.shape) == (num_experts, hidden_dim, intermediate_dim)
+
+    expected_gate = experts.gate_up_proj[0, :intermediate_dim, :].clone()
+    expected_up = experts.gate_up_proj[0, intermediate_dim:, :].clone()
+    expected_down = experts.down_proj[0, :, :].clone()
+
+    converted = convert_model(model, cleanup_original=False, max_layers=1)
+    assert converted
+
+    experts = model.model.layers[0].mlp.experts
+    assert isinstance(experts, BailingHybridExperts)
+    assert not hasattr(experts, "gate_up_proj")
+    assert not hasattr(experts, "down_proj")
+    assert experts.forward.__func__ is linear_loop_experts_forward
+
+    containers = [name for name, _ in experts.named_children() if name.isdigit()]
+    assert sorted(containers, key=int) == [str(idx) for idx in range(num_experts)]
+
+    expert0 = getattr(experts, "0")
+    assert isinstance(expert0.gate_proj, nn.Linear)
+    assert isinstance(expert0.up_proj, nn.Linear)
+    assert isinstance(expert0.down_proj, nn.Linear)
+
+    torch.testing.assert_close(expert0.gate_proj.weight, expected_gate)
+    torch.testing.assert_close(expert0.up_proj.weight, expected_up)
+    torch.testing.assert_close(expert0.down_proj.weight, expected_down)
+
+
+def test_bailing_hybrid_split_forward_matches_fused_math():
+    model = BailingHybridForCausalLM(_tiny_bailing_hybrid_config())
+    fused_experts = model.model.layers[0].mlp.experts
+    assert isinstance(fused_experts, BailingHybridExperts)
+
+    hidden_states = torch.randn(5, model.config.hidden_size, dtype=torch.float32)
+    top_k_index = torch.randint(
+        0, model.config.num_local_experts, (hidden_states.size(0), 2)
+    )
+    top_k_weights = torch.rand((hidden_states.size(0), 2), dtype=hidden_states.dtype)
+
+    with torch.no_grad():
+        expected_experts = fused_experts(hidden_states, top_k_index, top_k_weights)
+
+    converted = convert_model(model, cleanup_original=False, max_layers=1)
+    assert converted
+
+    split_experts = model.model.layers[0].mlp.experts
+    _assert_unfused_expert_module(split_experts)
+    with torch.no_grad():
+        actual_experts = split_experts(hidden_states, top_k_index, top_k_weights)
+
+    # The defused linear_loop path should reproduce the original fused experts
+    # math. The full decoder forward is not exercised here because
+    # ``bailing_hybrid``'s delta-attention path calls ``causal_conv1d``, which
+    # requires a CUDA device.
+    torch.testing.assert_close(actual_experts, expected_experts, atol=1e-5, rtol=1e-5)
+
+
 def test_laguna():
     from transformers.models.laguna.modeling_laguna import LagunaExperts
 

tests/test_meta_model_defusion.py

@@ -197,6 +197,11 @@ def _build_model_config(case: dict):
         config.ffn_config.ffn_act_fn = {"name": "silu"}
     elif model_type == "deepseek_v3":
         config.first_k_dense_replace = 0
+    elif model_type == "bailing_hybrid":
+        # The default config keeps the first layers dense, so the routed
+        # experts would not be constructed. Force every tiny layer onto MoE.
+        config.first_k_dense_replace = 0
+        config.mlp_layer_types = ["sparse"] * config.num_hidden_layers
     elif model_type == "deepseek_v32":
         config.mlp_layer_types = ["sparse"] * config.num_hidden_layers
     elif model_type == "ernie4_5_vl_moe":
@@ -346,6 +351,28 @@ def _validate_defused_module(case: dict, module) -> None:
 
 
 META_MODEL_CASES = [
+    {
+        "model_type": "afmoe",
+        "mode": "convert",
+        "model_module": "transformers.models.afmoe.modeling_afmoe",
+        "model_class": "AfmoeForCausalLM",
+        "config_module": "transformers.models.afmoe.configuration_afmoe",
+        "config_class": "AfmoeConfig",
+        "target_class_paths": ("transformers.models.afmoe.modeling_afmoe.AfmoeExperts",),
+        "validator": "experts",
+    },
+    {
+        "model_type": "bailing_hybrid",
+        "mode": "convert",
+        "model_module": "transformers.models.bailing_hybrid.modeling_bailing_hybrid",
+        "model_class": "BailingHybridForCausalLM",
+        "config_module": "transformers.models.bailing_hybrid.configuration_bailing_hybrid",
+        "config_class": "BailingHybridConfig",
+        "target_class_paths": (
+            "transformers.models.bailing_hybrid.modeling_bailing_hybrid.BailingHybridExperts",
+        ),
+        "validator": "experts",
+    },
     {
         "model_type": "axk2",
         "mode": "convert",