Remove TraceableMistralForCausalLM

src/llmcompressor/transformers/tracing/__init__.py

@@ -1,7 +1,6 @@
 from .llava import (
     LlavaForConditionalGeneration as TraceableLlavaForConditionalGeneration,
 )
-from .mistral import MistralForCausalLM as TraceableMistralForCausalLM
 from .mllama import (
     MllamaForConditionalGeneration as TraceableMllamaForConditionalGeneration,
 )
@@ -12,6 +11,5 @@
 __all__ = [
     "TraceableLlavaForConditionalGeneration",
     "TraceableMllamaForConditionalGeneration",
-    "TraceableMistralForCausalLM",
     "TraceableQwen2VLForConditionalGeneration",
 ]

src/llmcompressor/transformers/tracing/llava.py

@@ -19,20 +19,14 @@
 from typing import List, Optional, Tuple, Union
 
 import torch
-from transformers import AutoModel, AutoModelForCausalLM, LlavaForConditionalGeneration
+from transformers import LlavaForConditionalGeneration
 from transformers.models.llava.configuration_llava import LlavaConfig
 from transformers.models.llava.modeling_llava import (
     LlavaCausalLMOutputWithPast,
-    LlavaMultiModalProjector,
-    LlavaPreTrainedModel,
     logger,
 )
-from transformers.models.mistral.configuration_mistral import MistralConfig
 from transformers.utils.fx import HFProxy
 
-# TRACING: Reuse traceable subclass
-from .mistral import MistralForCausalLM as TraceableMistralForCausalLM
-
 
 # TRACING: The shape of image_features is known and documented by
 # LlavaForConditionalGeneration.get_image_features
@@ -75,22 +69,6 @@ def maybe_install_metadata_inputs_embeds_masked(
 
 # TRACING: override `__init__` and `forward`
 class LlavaForConditionalGeneration(LlavaForConditionalGeneration):
-    def __init__(self, config: LlavaConfig):
-        super(LlavaPreTrainedModel, self).__init__(config)
-        self.vision_tower = AutoModel.from_config(config.vision_config)
-
-        self.multi_modal_projector = LlavaMultiModalProjector(config)
-        self.vocab_size = config.text_config.vocab_size
-
-        # TRACING: Must use TraceableMistralForCausalLM which wraps an untraceable function
-        if isinstance(config.text_config, MistralConfig):
-            self.language_model = TraceableMistralForCausalLM(config.text_config)
-        else:
-            self.language_model = AutoModelForCausalLM.from_config(config.text_config)
-
-        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
-        self.post_init()
-
     def forward(
         self,
         input_ids: torch.LongTensor = None,