[GRANITEMOESHARED] drop granitemoeshared model support (#379)

Signed-off-by: Mayank Mishra <mayank31398@gmail.com>

Author: mayank31398

lm_engine/hf_models/model_conversion/__init__.py

@@ -14,18 +14,11 @@
     _import_granitemoehybrid_config,
     _import_granitemoehybrid_state_dict,
 )
-from .granitemoeshared import (
-    _export_granitemoeshared_config,
-    _export_granitemoeshared_state_dict,
-    _import_granitemoeshared_config,
-    _import_granitemoeshared_state_dict,
-)
 from .llama import _export_llama_config, _export_llama_state_dict, _import_llama_config, _import_llama_state_dict
 
 
 _MODEL_IMPORT_FUNCTIONS = {
     "granite": (_import_granite_config, _import_llama_state_dict),
-    "granitemoeshared": (_import_granitemoeshared_config, _import_granitemoeshared_state_dict),
     "granitemoehybrid": (_import_granitemoehybrid_config, _import_granitemoehybrid_state_dict),
     "llama": (_import_llama_config, _import_llama_state_dict),
 }
@@ -64,7 +57,6 @@ def import_from_huggingface(
 
 _MODEL_EXPORT_FUNCTIONS = {
     "granite": (_export_granite_config, _export_llama_state_dict),
-    "granitemoeshared": (_export_granitemoeshared_config, _export_granitemoeshared_state_dict),
     "granitemoehybrid": (_export_granitemoehybrid_config, _export_granitemoehybrid_state_dict),
     "llama": (_export_llama_config, _export_llama_state_dict),
 }

lm_engine/hf_models/model_conversion/granitemoehybrid.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2025, Mayank Mishra
 # **************************************************
 
+import torch
 from transformers import GraniteMoeHybridConfig, GraniteMoeHybridForCausalLM
 
 from ...utils import SafeTensorsWeightsManager, divide_if_divisible
@@ -10,7 +11,12 @@
     split_query_key_value_tensor_for_attention,
 )
 from ..models import GPTBaseConfig
-from .granitemoeshared import _split_and_reorder_for_glu
+
+
+def _split_and_reorder_for_glu(weight: torch.Tensor, dim: int) -> torch.Tensor:
+    x, y = weight.chunk(2, dim=dim)
+    weight = torch.cat([y, x], dim=dim)
+    return weight
 
 
 def _import_granitemoehybrid_config(original_config: GraniteMoeHybridConfig) -> GPTBaseConfig:

lm_engine/hf_models/model_conversion/granitemoeshared.py

@@ -36,26 +36,6 @@ def test_granite_model_conversion(self, device: torch.device, add_bias: bool) ->
             lm_engine_config=lm_engine_config, model_type="granite", device=device, exact_match=False
         )
 
-    @parameterized.expand(TestCommons.get_all_devices())
-    def test_granitemoeshared_model_conversion(self, device: torch.device) -> None:
-        lm_engine_config = self.get_moe_test_config(
-            "rope",
-            add_bias=False,
-            shared_n_inner=64,
-            activation_function="swiglu",
-            normalization_function="rmsnorm",
-            m_emb=2,
-            m_width=2,
-        )
-
-        self.model_conversion_test(
-            lm_engine_config=lm_engine_config,
-            model_type="granitemoeshared",
-            device=device,
-            exact_match=False,
-            compare_loss=False,
-        )
-
     @parameterized.expand(TestCommons.make_args_matrix(TestCommons.get_all_devices(), [True, False]))
     def test_granitemoehybrid_model_conversion(self, device: torch.device, is_moe: bool) -> None:
         if is_moe: