Make all Transformer models compatible with model parallelism

[sgugger]

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:

• BART
• BigBirdPegasus
• BLIP2
• BLOOM
• BridgeTower
• CamemBERT
• CLIP
• CLIPSeg
• CodeGen
• Data2Vec Text
• Deit
• ESM
• GPT-2
• GPT-Neo
• GPT-NeoX
• GPT-NeoX Japanese
• GPT-J
• GPT-San
• JukeBox
• Lilt
• LLaMA (LlamaForSequenceClassification only)
• Longformer
• LongT5
• Luke
• M2M100
• mBART
• mT5
• NLLB
• OPT
• Owl-ViT
• Pix2Struct
• PLBART
• RoBERTa
• RoBERTa PreLayerNorm
• SwitchTransformer
• T5
• Vilt
• ViT
• ViT-Hybrid
• Whisper
• 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!

[muaid-mughrabi]

I think I can help with this Issue :)

[iamarunbrahma]

I would like to work on this issue - BART model :)

[kausmeows]

Hi, I can take this up 🙌🏻

[zsc]

Indeed, this fix is required for BLOOM. main...zsc:transformers:main (my fix is hacky and not PR-ready. Just FYI)

[TerryCM]

Just to make sure does LlamaForCausalLM supports this feature already?(#22546 ) it seems that, still there are some errors when using device_map="auto" for this task.

[mollerup23]

Hi, I'd like to pick up the GPT-2 model!

[xssChauhan]

Hi! I am taking this up for LlamaForSequenceClassification.

[kooshi]

Just to make sure does LlamaForCausalLM supports this feature already?(#22546 ) it seems that, still there are some errors when using device_map="auto" for this task.

It does (#22329). I have started seeing similar errors to #22546, but only after updating my drivers from 525 to 530, similar to #22546 (comment)

(which is good news to me, I had no idea why that gpu started disappearing occasionally. It seems it can happen when that gpu is under any load, not just during training)

Edit: seems like the errors I was getting were actually caused by GPU sag. I haven't yet reproduced that exact error, but it has been reported elsewhere. It is certainly not consistent though.

[innat]

@younesbelkada @sgugger
Does this fix (moving label/logit to same device) supposed to work (model parallelism) for all models (listed above)? Or, a crucial step toward it? Also, this design fix is only for pytorch model and not for jax or tf?

[younesbelkada]

I think it is supposed to work for all models listed above, as long as you are loading your model with device_map=xxx. And yes this should be for Pytorch only, though I am not really aware of how model parallelism work on TF & Jax

[innat]

I think it is supposed to work for all models listed above, as long as you are loading your model with device_map=xxx

I tried with such fix here #22591 (comment) but sadly it didn't work out. Any catch?

[innat]

@sgugger
As the goal of this ticket is to enable model parallelism with easy fix, have the merged PR(s) checked on multi-gpu? I couldn't find any test script here #22663 regarding that .

[shahad-mahmud]

I would love to work with BridgeTower

[trantuantdt]

Hi. I would like to try with "Whisper"

[mollerup23]

I'd like to claim OPT model if no one else has picked it up.

[mayankagarwals]

Taking this up for the remaining GPT models

[jprivera44]

Hello, I just completed the GPT-J code. Just filling in the PR now.

[oscar-garzon]

Hello! I'd like to work in Whisper model

[abhigyan631]

Hi, is there any model on which I can work, please? Thanks.

[Tanmaypatil123]

Is there any remaining model on which I can work ? Thanks .

[JuheonChu]

@sgugger Hello, can I work on the JukeBox?

[elabongaatuo]

Hello @sgugger , I'd like to work on m2m100

[Batese2001]

@sgugger I would love to work on CodeGen if it is unclaimed

[katiele47]

Hi @sgugger I can work on Luke if it has not been taken

[VomV]

@sgugger I would like to work on SwitchTransformer, if not taken.

[sushmanthreddy]

@sgugger I think all transformers are covered, I have checked for others also...for example, switch transformers have parallelism implemented already. i think we can close this issue. The only pending models are clip,jukebox,owlvit, and Nllb , may be model parallelism is not applicable for some of there models

[sgugger]

Indeed, all models have been covered. Thanks a lot everyone!