NeoX underperforming on A100
### System Info

- `transformers` version: 4.25.1
- Platform: Linux-5.15.0-50-generic-x86_64-with-debian-buster-sid
- Python version: 3.7.13
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.12.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?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

When running EleutherAI/pythia-12b-deduped on A100, generation slows down from 15tok/s to 3tok/s when sending 1024 tokens as input. This behaviour does not occur using A6000 with the same setup (14.04 tok/s, 1024 tok input). Ruled out "machine" issue as this happens on multiple servers on multiple locations (runpod, vast.ai, google...)

Does any of you have an idea what the root cause of this issue could be? Running pure python with no accelerate or anything that could "interfere" with the speed of HF. I am running a standard KoboldAI application.

Tagging: @ArthurZucker @younesbelkada @StellaAthena 

### Information

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

### Tasks

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

### Reproduction

1. Clone https://github.com/henk717/KoboldAI + install dependencies.
2. Start the KoboldAI using model EleutherAI/pythia-12b-deduped
3. Let it generate 100 tokens until it reaches 1024 tokens.
On A100, this would slowly start to slow down token generation. On A6000, this has no effect.

### Expected behavior

Expected behaviour would be that the token generation would be same, regardless of amount inserted.
Are you using a 40 GB or 80 GB A100? A 40 GB A100 has slightly less VRAM than an A6000 (48 GB), and it’s possible that you’re tripping a failsafe such as CPU-offload to avoid an OOM error.

12 billion params * 3 Bytes per parameter = 36 GB of VRAM. My general rule of thumb is to add 20% overhead space for a full context length of 2048, which would push the model over the 40 GB limit.

One way to test this theory would be to try running it with LLM.int8. Another would be to carefully monitor GPU and CPU usage during inference.<|||||>This might be the issue indeed. It's automatically offloading to CPU but not generating any overflow warnings, making it difficult to pinpoint where the issue would lie. Going for the 80Gb did indeed solve the issue.