Multi-worker HF training using trainer API in torch-xla result in too many graph compilations after saving checkpoint (transformers>=4.35)
Author: jeffhataws
Pull request: no

### System Info

transformers>=4.35
Neuron SDK 2.15 with torch-neuronx 1.13

### Who can help?

@muellerz  

### Information

- [ ] The official example scripts
- [X] 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

(Duplicate of https://github.com/aws-neuron/aws-neuron-sdk/issues/813)

I followed [PyTorch Neuron for Trainium Hugging Face BERT MRPC task finetuning using Hugging Face Trainer API](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/tutorials/training/finetune_hftrainer.html#torch-hf-bert-finetune) to fine-tune BERT. I ran run_2w.sh and see the following behavior where it runs normally until the first checkpoint is saved, but then starts doing compilation for every step (I changed save_steps option in run_2w.sh to 10 steps in order to trigger the issue faster):
```
[INFO|trainer.py:1712] 2024-01-09 17:04:08,045 >> ***** Running training *****                                                                                 
[INFO|trainer.py:1713] 2024-01-09 17:04:08,045 >>   Num examples = 1,840                                                                                       
[INFO|trainer.py:1714] 2024-01-09 17:04:08,045 >>   Num Epochs = 5                                                                                             
[INFO|trainer.py:1715] 2024-01-09 17:04:08,045 >>   Instantaneous batch size per device = 8                                                                    
[INFO|trainer.py:1718] 2024-01-09 17:04:08,045 >>   Total train batch size (w. parallel, distributed & accumulation) = 16                                      
[INFO|trainer.py:1719] 2024-01-09 17:04:08,045 >>   Gradient Accumulation steps = 1                                                                            
[INFO|trainer.py:1720] 2024-01-09 17:04:08,045 >>   Total optimization steps = 1,150                                                                           [INFO|trainer.py:1721] 2024-01-09 17:04:08,045 >>   Number of trainable parameters = 109,483,778                                                               
  0%|          | 0/1150 [00:00<?, ?it/s]2024-01-09 17:04:08.000173:  140637  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache          
2024-01-09 17:04:08.000175:  140637  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_1650
6334326618155050+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                              
  0%|          | 1/1150 [00:00<04:53,  3.92it/s]2024-01-09 17:04:09.000508:  140742  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache  2024-01-09 17:04:09.000603:  140742  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_2044
823947559839528+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                               
  0%|          | 2/1150 [00:02<29:23,  1.54s/it]2024-01-09 17:04:13.000328:  140780  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache  
2024-01-09 17:04:13.000442:  140780  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_7850
734058944619683+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                               
  1%|          | 10/1150 [00:09<08:40,  2.19it/s][INFO|trainer.py:2859] 2024-01-09 17:04:17,051 >> Saving model checkpoint to /tmp/mrpc/tmp-checkpoint-10 

(Done saving checkpoint, then compilation every step below)
    
2024-01-09 17:04:17.000789:  141260  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache                                                  
2024-01-09 17:04:17.000873:  141260  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_2523
922307180626946+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                               
2024-01-09 17:04:20.000215:  141270  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache                                                  
2024-01-09 17:04:20.000216:  141270  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_6208
462474369064908+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                               
2024-01-09 17:04:21.000202:  141279  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache                                                  
2024-01-09 17:04:21.000282:  141279  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_1498
3430005009285767+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                              
2024-01-09 17:04:23.000265:  141288  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache                                                  
2024-01-09 17:04:23.000266:  141288  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_3356
031905174227108+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                               
2024-01-09 17:04:24.000025:  141297  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache                                                  
2024-01-09 17:04:24.000104:  141297  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_5950
234423484734321+abb26765/model.neff. Exiting with a successfully compiled graph.                                                                               
2024-01-09 17:04:26.000063:  141306  INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache                                                  
2024-01-09 17:04:26.000064:  141306  INFO ||NEURON_CC_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.11.0.35+4f5279863/MODULE_1050
0036830841255848+abb26765/model.neff. Exiting with a successfully compiled graph.  

(Compilation repeated many times, and eventually run out of device memory in Neuron runtime)
```                             

This issue starts in transformers version 4.35. 

### Expected behavior

We should see training completes normally with few torch-xla complations.

Comments: