

TITLE
Make all Transformer models compatible with model parallelism

COMMENTS
27

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 4
rocket: 3
eyes: 2

BODY
Accelerate makes it easy to load a model on multiple GPUs with `device_map="auto"`. This in turn allows users to train model with naive model parallelism if they have several GPUs.

A problem that happens in Transformers, with model with heads (so not XxxModel but for instance XxxModelForSequenceClassification) is that the labels end up on a different device than the logits and there is a device mistmatch error.

Thankfully, there is an easy fix for that! #22535 shows how to fix this for T5 by just moving the labels to the same device as the logits they are compared to. This is a noop when the devices are the same, and fixes the issue if devices are different.

We would like help from the community to extend this to all models that support model parallelism, which are:
- [x] BART
- [x] BigBirdPegasus
- [x] BLIP2
- [x] BLOOM
- [x] BridgeTower
- [x] CamemBERT
- [x] CLIP
- [x] CLIPSeg
- [x] CodeGen
- [x] Data2Vec Text
- [x] Deit
- [x] ESM
- [x] GPT-2
- [x] GPT-Neo
- [x] GPT-NeoX
- [x] GPT-NeoX Japanese
- [x] GPT-J
- [x] GPT-San
- [x] JukeBox
- [x] Lilt
- [x] LLaMA (`LlamaForSequenceClassification` only)
- [x] Longformer
- [x] LongT5
- [x] Luke
- [x] M2M100
- [x] mBART
- [x] mT5
- [x] NLLB
- [x] OPT
- [x] Owl-ViT
- [x] Pix2Struct
- [x] PLBART
- [x] RoBERTa
- [x] RoBERTa PreLayerNorm
- [x] SwitchTransformer
- [x] T5
- [x] Vilt
- [x] ViT
- [x] ViT-Hybrid
- [x] Whisper
- [x] XLM-RoBERTa

If you would like to grab one of those models and apply the same fix as #22535 to all the model with heads, please leave a comment here!