Support symbolic tracing for NeoX models
Author: norabelrose
Pull request: no
Labels: name: Good Second Issue

name: Feature request

### Feature request

Currently `transformers.utils.fx.symbolic_trace` fails when passed any NeoX model, and I'd like to fix that:
```
NotImplementedError: Model GPTNeoXForCausalLM 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, CLIPModel, CLIPTextModel, 
CLIPTextModelWithProjection, CLIPVisionModel, CLIPVisionModelWithProjection, ConvNextBackbone, 
ConvNextForImageClassification, ConvNextModel, DebertaForMaskedLM, DebertaForQuestionAnswering, 
DebertaForSequenceClassification, DebertaForTokenClassification, DebertaModel, DebertaV2ForMaskedLM, 
DebertaV2ForMultipleChoice, DebertaV2ForQuestionAnswering, DebertaV2ForSequenceClassification, 
DebertaV2ForTokenClassification, DebertaV2Model, 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, LxmertForPreTraining, 
LxmertForQuestionAnswering, LxmertModel, M2M100ForConditionalGeneration, M2M100Model, MBartForCausalLM, 
MBartForConditionalGeneration, MBartForQuestionAnswering, MBartForSequenceClassification, MBartModel, 
MT5ForConditionalGeneration, 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, T5Model, TrOCRDecoder, TrOCRForCausalLM, 
ViTForImageClassification, ViTForMaskedImageModeling, ViTModel, Wav2Vec2ForCTC, Wav2Vec2ForMaskedLM, 
Wav2Vec2ForPreTraining, Wav2Vec2ForSequenceClassification, Wav2Vec2Model, XGLMForCausalLM, XGLMModel
```

### Motivation

The main motivation for this is to enable graph rewriting with the EleutherAI Pythia model suite. Graph rewriting has various interpretability use-cases and the Pythia suite was designed for interpretability research.

### Your contribution

I plan to implement a PR for this soon unless there's some major blocker for it.

Comments:
Hi @norabelrose 

That would be very nice if you are able to enable this 🤗 . Thanks in advance.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I'd love to contribute if possible!

> I'd love to contribute if possible!

Sorry I haven't gotten around to doing this, I was sort of hoping my related PR #24941 could get merged first. It's been stuck on some silly style issue in the CI. You can actually extract a computational graph easily by writing a [simple custom backend](https://pytorch.org/docs/stable/dynamo/custom-backends.html) for Torch Dynamo, and it'd be nice to be able to use either Dynamo or symbolic tracing.