fix linter

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

Author: mayank31398

lm_engine/hf_models/config/sequence_mixer.py

@@ -16,13 +16,8 @@ class _SoftmaxAttentionArgs(BaseArgs):
     add_bias: bool = False
     attention_multiplier: float | None = None
     sliding_window: int | None = None
-    # needed for Qwen 2 MoE
-    qkv_bias: bool = None
 
     def model_post_init(self, __context: Any) -> None:
-        if self.qkv_bias is None:
-            self.qkv_bias = self.add_bias
-
         assert self.sequence_mixer_type == "softmax_attention"
 
 

lm_engine/hf_models/mixins/dense/main.py

@@ -21,7 +21,6 @@
 
 class CausalLMModelMixin(PreTrainedModelMixin):
     base_model_class = None
-    model_parallel_state_dict_function = None
 
     def __init__(self, config: CommonConfig, **kwargs) -> CausalLMModelMixin:
         super().__init__(config, **kwargs)
@@ -46,9 +45,7 @@ def _init_model(self, config: CommonConfig, **kwargs) -> None:
             self.m_width = config.m_width
 
         self.is_tp_enabled = ProcessGroupManager.is_tensor_parallel_enabled()
-
-        if self.is_tp_enabled:
-            self.tp_mesh = ProcessGroupManager.get_tensor_parallel_mesh()
+        self.tp_mesh = ProcessGroupManager.get_tensor_parallel_mesh() if self.is_tp_enabled else None
 
     def forward(
         self,
@@ -339,19 +336,3 @@ def _get_dummy_intermediate_tensor(
             )
 
         return tensor
-
-    def load_from_safetensors_weights_manager(self, safetensors_weights_manager: SafeTensorsWeightsManager) -> None:
-        with torch.device(torch.cuda.current_device()):
-            position_embedding_type = self.config.position_embedding_type
-
-            if position_embedding_type == "rope":
-                self.transformer.rope.reset_parameters()
-
-        state_dict = self.__class__.model_parallel_state_dict_function(
-            config=self.config,
-            safetensors_weights_manager=safetensors_weights_manager,
-            num_pipeline_stages=self.num_pipeline_stages,
-            pipeline_stage_id=self.pipeline_stage_id,
-        )
-
-        self.load_state_dict(state_dict)

lm_engine/hf_models/mixins/dense_TP/main.py

@@ -32,6 +32,8 @@
 
 
 class CausalLMModelMixin_TP(CausalLMModelMixin):
+    model_parallel_state_dict_function = None
+
     def forward(
         self,
         input_ids: torch.Tensor | list[list[int]] | None = None,
@@ -177,3 +179,19 @@ def from_pretrained(
         model.load_from_safetensors_weights_manager(SafeTensorsWeightsManager(pretrained_model_name_or_path))
 
         return model
+
+    def load_from_safetensors_weights_manager(self, safetensors_weights_manager: SafeTensorsWeightsManager) -> None:
+        with torch.device(torch.cuda.current_device()):
+            position_embedding_type = self.config.position_embedding_type
+
+            if position_embedding_type == "rope":
+                self.transformer.rope.reset_parameters()
+
+        state_dict = self.__class__.model_parallel_state_dict_function(
+            config=self.config,
+            safetensors_weights_manager=safetensors_weights_manager,
+            num_pipeline_stages=self.num_pipeline_stages,
+            pipeline_stage_id=self.pipeline_stage_id,
+        )
+
+        self.load_state_dict(state_dict)

lm_engine/hf_models/model_conversion/__init__.py

@@ -22,12 +22,6 @@
     _import_granitemoeshared_state_dict,
 )
 from .llama import _export_llama_config, _export_llama_state_dict, _import_llama_config, _import_llama_state_dict
-from .qwen2_moe import (
-    _export_qwen2_moe_config,
-    _export_qwen2_moe_state_dict,
-    _import_qwen2_moe_config,
-    _import_qwen2_moe_state_dict,
-)
 
 
 _MODEL_IMPORT_FUNCTIONS = {
@@ -36,7 +30,6 @@
     "granitemoeshared": (_import_granitemoeshared_config, _import_granitemoeshared_state_dict),
     "granitemoehybrid": (_import_granitemoehybrid_config, _import_granitemoehybrid_state_dict),
     "llama": (_import_llama_config, _import_llama_state_dict),
-    "qwen2_moe": (_import_qwen2_moe_config, _import_qwen2_moe_state_dict),
 }
 
 
@@ -77,7 +70,6 @@ def import_from_huggingface(
     "granitemoeshared": (_export_granitemoeshared_config, _export_granitemoeshared_state_dict),
     "granitemoehybrid": (_export_granitemoehybrid_config, _export_granitemoehybrid_state_dict),
     "llama": (_export_llama_config, _export_llama_state_dict),
-    "qwen2_moe": (_export_qwen2_moe_config, _export_qwen2_moe_state_dict),
 }