Remove GraniteMoeLinearExperts

src/llmcompressor/modeling/moe/linear_experts.py

@@ -139,7 +139,7 @@ class LinearExperts2D(torch.nn.ModuleList):
 
     # 1. try for mappings (efficient load)
     # 2. try for standardized moe, convert after load
-    # 3. Explicit replacement (GraniteMoeLinearExperts)
+    # 3. Explicit replacement (Llama4LinearExperts)
 
     """
 
@@ -159,7 +159,6 @@ class LinearExperts2D(torch.nn.ModuleList):
     def get_registration(
         cls, key: type[torch.nn.Module], default: Any = None
     ) -> type["LinearExperts2D"]:
-        from .granitemoe import GraniteMoeLinearExperts  # noqa: F401
         from .llama4 import Llama4LinearExperts  # noqa: F401
 
         return cls._registry.get(key, default)

src/llmcompressor/modifiers/pruning/reap/utils.py

@@ -19,7 +19,6 @@
 from torch import distributed as dist
 
 from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag
-from llmcompressor.modeling.moe.granitemoe import GraniteMoeLinearExperts
 from llmcompressor.modeling.moe.linear_experts import ExpertMLP, LinearExperts2D
 from llmcompressor.modeling.moe.llama4 import Llama4LinearExperts
 
@@ -155,19 +154,13 @@ def get_moe_attrs(model: nn.Module, ignore: list[str]) -> MoeModelAttrs | None:
                 continue
             experts = getattr(module, experts_attr)
             # REAP currently only supports LinearExperts2D experts, as they receive the
-            # top_k indices and weights from the router in their forward pass.
-            # Granite and Llama4 experts diverge from this behavior, so they are
-            # unsupported for now.
+            # top_k indices and weights from the router in their forward pass. Llama4
+            # experts diverge from this behavior, so they are unsupported for now.
             if not isinstance(experts, LinearExperts2D):
                 logger.warning(
                     f"Skipping layer {name}: experts module is not LinearExperts2D"
                 )
                 continue
-            if isinstance(experts, GraniteMoeLinearExperts):
-                logger.warning(
-                    f"Skipping unsupported GraniteMoeLinearExperts layer: {name}"
-                )
-                continue
             if isinstance(experts, Llama4LinearExperts):
                 logger.warning(
                     f"Skipping unsupported Llama4LinearExperts layer: {name}"
@@ -179,7 +172,7 @@ def get_moe_attrs(model: nn.Module, ignore: list[str]) -> MoeModelAttrs | None:
         raise ValueError(
             "Could not find any supported MoE layers with experts in "
             "LinearExperts2D format. Make sure the model has MoE layers "
-            "(excluding GraniteMoeLinearExperts and Llama4LinearExperts), "
+            "(excluding Llama4LinearExperts), "
             "and that the name of its experts module is in EXPERTS_ATTRS "
             "and it the name of its router module is in ROUTER_ATTRS in "
             "reap/utils.py"

tests/llmcompressor/modeling/test_linearize.py

@@ -153,9 +153,7 @@ def forward(self, *args, **kwargs):
 
 @torch.no_grad()
 @requires_gpu
-@pytest.mark.parametrize(
-    "model_type", list(ARCH_TO_IMPORT_PATHS.keys() - {"llama4", "granitemoe"})
-)
+@pytest.mark.parametrize("model_type", list(ARCH_TO_IMPORT_PATHS.keys() - {"llama4"}))
 def test_linearize_moe(model_type):
     config_path, experts_path = ARCH_TO_IMPORT_PATHS[model_type]
     config_cls = import_or_none(config_path)
@@ -200,43 +198,6 @@ def test_linearize_moe(model_type):
         assert torch.nn.functional.mse_loss(calib_outputs, true_outputs) < MODULE_MSE
 
 
-def test_linearize_moe_granite():
-    try:
-        from transformers.models.granitemoe.configuration_granitemoe import (
-            GraniteMoeConfig,
-        )
-        from transformers.models.granitemoe.modeling_granitemoe import (
-            GraniteMoeParallelExperts,
-        )
-    except ImportError:
-        pytest.skip("GraniteMoeParallelExperts has been removed")
-
-    config = GraniteMoeConfig(hidden_size=512, intermediate_size=1024)
-    experts = GraniteMoeParallelExperts(
-        config.num_local_experts, config.hidden_size, config.intermediate_size
-    )
-    init.normal_(experts.weight, mean=0.0, std=config.initializer_range)
-
-    mock_model = DummyModel(experts, config)
-    linearize_moe(mock_model)
-    assert mock_model.module is not experts
-
-    hidden_states = torch.randn(NUM_TEST_TOKENS, config.hidden_size, dtype=config.dtype)
-    expert_size = [
-        (NUM_TEST_TOKENS // config.num_local_experts)
-        for _ in range(config.num_local_experts)
-    ]
-    expert_size[-1] += NUM_TEST_TOKENS % config.num_local_experts
-    true_outputs = experts(hidden_states, expert_size)
-    outputs = mock_model(hidden_states, expert_size)
-    with moe_calibration_context():
-        calib_outputs = mock_model(hidden_states, expert_size)
-
-    assert torch.any(true_outputs != 0), "Bad test setup, output is all zeros"
-    assert torch.nn.functional.mse_loss(outputs, true_outputs) < MODULE_MSE
-    assert torch.nn.functional.mse_loss(calib_outputs, true_outputs) < MODULE_MSE
-
-
 def test_linearize_moe_llama4():
     from transformers.models.llama4.configuration_llama4 import (
         Llama4Config,