BertForSequenceClassification does not support 'device_map':"auto" yet
Author: goodaytar
Pull request: no
Labels: name: Good Second Issue

### System Info

I have trained a model and am now trying to load and quantise it but getting the error: 

BertForSequenceClassification does not support 'device_map':"auto" yet

Code for loading is simply:
   ` model = AutoModelForSequenceClassification.from_pretrained(model_dir, device_map='auto', load_in_8bit=True)`

Help would be greatly appreciated!

Thanks,

Lee

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

model = AutoModelForSequenceClassification.from_pretrained(model_dir, device_map='auto', load_in_8bit=True)

### Expected behavior

The model would load and be usable.

Comments:
Hi @goodaytar, thanks for raising this issue! 

Yes, the BERT models don't support the use of `device_map=xxx` yet. In the full error message, you should have seen: 

```
BertForSequenceClassification not support `device_map="auto"`. To implement support, the model class needs to implement the `_no_split_modules` attribute.
```

In order to enable this the `_no_split_modules` attribute needs to be implemented for the model. If you or anyone else in the community would like to open a PR to add this, we'd be very happy to review! 

Thanks for the reply Amy. If you could give me a little bit more info on what needs adding, I'd be happy to.

Get Outlook for Android<https://aka.ms/AAb9ysg>
________________________________
From: amyeroberts ***@***.***>
Sent: Friday, August 4, 2023 12:26:00 PM
To: huggingface/transformers ***@***.***>
Cc: goodaytar ***@***.***>; Mention ***@***.***>
Subject: Re: [huggingface/transformers] BertForSequenceClassification does not support 'device_map':"auto" yet (Issue #25296)


Hi @goodaytar<https://github.com/goodaytar>, thanks for raising this issue!

Yes, the BERT models don't support the use of device_map=xxx yet. In the full error message, you should have seen:

BertForSequenceClassification not support `device_map="auto"`. To implement support, the model class needs to implement the `_no_split_modules` attribute.


In order to enable this the _no_split_modules attribute needs to be implemented for the model. If you or anyone else in the community would like to open a PR to add this, we'd be very happy to review!

—
Reply to this email directly, view it on GitHub<https://github.com/huggingface/transformers/issues/25296#issuecomment-1665455659>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/APRZ52NZV7QNND7TR7TMJDTXTTL4RANCNFSM6AAAAAA3DCB7KY>.
You are receiving this because you were mentioned.Message ID: ***@***.***>


In order to know how to properly place the model onto difference devices, the models need to have `_no_split_modules` implemented in their `PreTrainedModel` class e.g. [like here for Roberta](https://github.com/huggingface/transformers/blob/fdaef3368b7495f6d3f26739fece0ee370fa7ce6/src/transformers/models/roberta/modeling_roberta.py#L596).

For some modules, it's necessary to place all of the weights on the same device e.g. like [`Pix2StructVisionLayer` for Pix2Struct](https://github.com/huggingface/transformers/blob/fdaef3368b7495f6d3f26739fece0ee370fa7ce6/src/transformers/models/pix2struct/modeling_pix2struct.py#L552).

In order to add, it'll be a case of iterating to find the modules that should be split or not. Once implemented, the [accelerate tests should be run and pass](https://github.com/huggingface/transformers/blob/fdaef3368b7495f6d3f26739fece0ee370fa7ce6/src/transformers/models/pix2struct/modeling_pix2struct.py#L552). This should be tested with 1 and 2 GPUs. 


And how do I find the modules that should be split or not?

@goodaytar You'll need to experiment with the model to find out which modules should be split. I suggest starting with an empty list and looking at similar models to see how they set `_no_split_modules`. 

You can inspect where the layers are allocated by using `infer_auto_device_map`:

```python
device_map = infer_auto_device_map(model, no_split_module_classes=[])
```

The modules that can be added will be the layers defined in the modeling file e.g. `"BertEmbeddings"`

Once set, you can try running the accelerate tests (with GPUs!) to confirm the mapping works. If not, then inspect the device map. 

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.

Hi @amyeroberts, I would like to add the `'device_map': "auto"` functionality to BERT Models!

@tanaymeh Great! From next week, I'll be off for a few weeks. Please ping @younesbelkada for review in that time. 

@tanaymeh that would be really great, in  few words, you just need to make sure to add the module names that contain any skip connection to avoid potential device mismatch issues
Check for instance what has been done for RoBERTa here: https://github.com/huggingface/transformers/blob/main/src/transformers/models/roberta/modeling_roberta.py#L596

> @tanaymeh that would be really great, in a few words, you just need to make sure to add the module names that contain any skip connection to avoid potential device mismatch issues Check for instance what has been done for RoBERTa here: [`main`/src/transformers/models/roberta/modeling_roberta.py#L596](https://github.com/huggingface/transformers/blob/main/src/transformers/models/roberta/modeling_roberta.py?rgh-link-date=2023-09-13T10%3A58%3A02Z#L596)

That makes sense @younesbelkada! Will create a PR for this.
One question: Will the CI tests on Github also test my implementation of `device_map` (with 1 and 2 GPUs) every time I push a commit?

Hi @tanaymeh , 
Thanks, will look into it! 
The CI will not directly test it, we run "slow" tests every 24h on GPUs that will run those tests

@younesbelkada 
Hi Younes
Could you make it work for xlm_roberta_xl too?
Thanks 
Regards
Dragan

@younesbelkada  Any updates? We can't wait to use this great feature.

@Hambaobao I am working on the PR for this feature but waiting for a revert from @younesbelkada!

any update on this issue? or anyone fixed it?


Any update on this issue ? 

ValueError: SiglipVisionModel does not support `device_map='auto'`

Same for Siglip?

@lucasjinreal There are many models which don't yet have this enabled. I've opened a feature request to add this for vision and multimodal models which could have this added: #29786