fx symbolic tracing support for mistral
Author: avishka22
Pull request: no
Labels: name: Feature request

name: Good Difficult Issue

name: Compilation

### System Info

- `transformers` version: 4.40.0.dev0
- Platform: Linux-5.15.0-71-generic-x86_64-with-glibc2.31
- Python version: 3.10.9
- Huggingface_hub version: 0.21.4
- Safetensors version: 0.4.2
- Accelerate version: 0.28.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
from transformers.utils.fx import symbolic_trace
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained('mistralai/Mistral-7b-v0.1')
gm = symbolic_trace(model)
```
---------------------------------------------------------------------------

```
NotImplementedError                       Traceback (most recent call last)
Cell In[20], line 4
      2 from transformers import AutoModelForCausalLM
      3 model = AutoModelForCausalLM.from_pretrained('mistralai/Mistral-7b-v0.1')
----> 4 gm = symbolic_trace(model)

File ~/anaconda3/lib/python3.10/site-packages/transformers/utils/fx.py:1303, in symbolic_trace(model, input_names, disable_check, tracer_cls)
   1300 concrete_args = get_concrete_args(model, input_names)
   1302 if not disable_check:
-> 1303     check_if_model_is_supported(model)
   1305 # Tracing.
   1306 tracer = tracer_cls()

File ~/anaconda3/lib/python3.10/site-packages/transformers/utils/fx.py:1261, in check_if_model_is_supported(model)
   1259 if not is_model_supported(model):
   1260     supported_model_names = ", ".join(_SUPPORTED_MODELS)
-> 1261     raise NotImplementedError(
   1262         f"Model {model.__class__.__name__} is not supported yet, supported models: {supported_model_names}"
   1263     )

NotImplementedError: Model MistralForCausalLM is not supported yet, supported models: AlbertForMaskedLM, AlbertForMultipleChoice, AlbertForPreTraining, AlbertForQuestionAnswering, AlbertForSequenceClassification, AlbertForTokenClassification, AlbertModel, AltCLIPModel, AltCLIPTextModel, AltCLIPVisionModel, BartForCausalLM, BartForConditionalGeneration, BartForQuestionAnswering, BartForSequenceClassification, BartModel, BertForMaskedLM, BertForMultipleChoice, BertForNextSentencePrediction, BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification, BertForTokenClassification, BertLMHeadModel, BertModel, BlenderbotForCausalLM, BlenderbotForConditionalGeneration, BlenderbotModel, BlenderbotSmallForCausalLM, BlenderbotSmallForConditionalGeneration, BlenderbotSmallModel, BloomForCausalLM, BloomForQuestionAnswering, BloomForSequenceClassification, BloomForTokenClassification, BloomModel, CLIPForImageClassification, CLIPModel, CLIPTextModel, CLIPTextModelWithProjection, CLIPVisionModel, CLIPVisionModelWithProjection, CohereForCausalLM, CohereModel, ConvNextBackbone, ConvNextForImageClassification, ConvNextModel, DebertaForMaskedLM, DebertaForQuestionAnswering, DebertaForSequenceClassification, DebertaForTokenClassification, DebertaModel, DebertaV2ForMaskedLM, DebertaV2ForMultipleChoice, DebertaV2ForQuestionAnswering, DebertaV2ForSequenceClassification, DebertaV2ForTokenClassification, DebertaV2Model, Dinov2Backbone, Dinov2ForImageClassification, Dinov2Model, DistilBertForMaskedLM, DistilBertForMultipleChoice, DistilBertForQuestionAnswering, DistilBertForSequenceClassification, DistilBertForTokenClassification, DistilBertModel, DonutSwinModel, ElectraForCausalLM, ElectraForMaskedLM, ElectraForMultipleChoice, ElectraForPreTraining, ElectraForQuestionAnswering, ElectraForSequenceClassification, ElectraForTokenClassification, ElectraModel, GPT2DoubleHeadsModel, GPT2ForQuestionAnswering, GPT2ForSequenceClassification, GPT2ForTokenClassification, GPT2LMHeadModel, GPT2Model, GPTJForCausalLM, GPTJForQuestionAnswering, GPTJForSequenceClassification, GPTJModel, GPTNeoForCausalLM, GPTNeoForQuestionAnswering, GPTNeoForSequenceClassification, GPTNeoForTokenClassification, GPTNeoModel, GitVisionModel, HubertForCTC, HubertForSequenceClassification, HubertModel, LayoutLMForMaskedLM, LayoutLMForQuestionAnswering, LayoutLMForSequenceClassification, LayoutLMForTokenClassification, LayoutLMModel, LlamaForCausalLM, LlamaForQuestionAnswering, LlamaForSequenceClassification, LlamaModel, LxmertForPreTraining, LxmertForQuestionAnswering, LxmertModel, M2M100ForConditionalGeneration, M2M100Model, MBartForCausalLM, MBartForConditionalGeneration, MBartForQuestionAnswering, MBartForSequenceClassification, MBartModel, MT5ForConditionalGeneration, MT5ForQuestionAnswering, MT5ForSequenceClassification, MT5ForTokenClassification, MT5Model, MarianForCausalLM, MarianMTModel, MarianModel, MegatronBertForCausalLM, MegatronBertForMaskedLM, MegatronBertForMultipleChoice, MegatronBertForNextSentencePrediction, MegatronBertForPreTraining, MegatronBertForQuestionAnswering, MegatronBertForSequenceClassification, MegatronBertForTokenClassification, MegatronBertModel, MobileBertForMaskedLM, MobileBertForMultipleChoice, MobileBertForNextSentencePrediction, MobileBertForPreTraining, MobileBertForQuestionAnswering, MobileBertForSequenceClassification, MobileBertForTokenClassification, MobileBertModel, NezhaForMaskedLM, NezhaForMultipleChoice, NezhaForNextSentencePrediction, NezhaForPreTraining, NezhaForQuestionAnswering, NezhaForSequenceClassification, NezhaForTokenClassification, NezhaModel, OPTForCausalLM, OPTForQuestionAnswering, OPTForSequenceClassification, OPTModel, PLBartForCausalLM, PLBartForConditionalGeneration, PLBartForSequenceClassification, PLBartModel, PeftModelForCausalLM, PeftModelForSeq2SeqLM, PegasusForCausalLM, PegasusForConditionalGeneration, PegasusModel, ResNetBackbone, ResNetForImageClassification, ResNetModel, RobertaForCausalLM, RobertaForMaskedLM, RobertaForMultipleChoice, RobertaForQuestionAnswering, RobertaForSequenceClassification, RobertaForTokenClassification, RobertaModel, SegformerForImageClassification, SegformerForSemanticSegmentation, SegformerModel, Speech2Text2Decoder, Speech2Text2ForCausalLM, Speech2TextForConditionalGeneration, Speech2TextModel, SwinBackbone, SwinForImageClassification, SwinForMaskedImageModeling, SwinModel, T5ForConditionalGeneration, T5ForQuestionAnswering, T5ForSequenceClassification, T5ForTokenClassification, T5Model, TrOCRDecoder, TrOCRForCausalLM, ViTForImageClassification, ViTForMaskedImageModeling, ViTModel, Wav2Vec2ForCTC, Wav2Vec2ForMaskedLM, Wav2Vec2ForPreTraining, Wav2Vec2ForSequenceClassification, Wav2Vec2Model, XGLMForCausalLM, XGLMModel
```

### Expected behavior

Can we get fx support for mistral, Thanks!

Comments:
Hi @avishka22, thanks for opening this feature request! 

If you or anyone in the community would like to tackle enabling this for Mistral, we'd be happy to review the PR! 

I want to work on this