Enable multi-device support for DPT
Author: OmarManzoor
Pull request: yes

# What does this PR do?
Adds multi device for for Dpt.  I tested this on a kaggle notebook with two T4 gpus.
<img width="1209" alt="Screenshot 2024-05-27 at 8 11 03 PM" src="https://github.com/huggingface/transformers/assets/17495884/71817575-64b3-45ff-a1bc-1acb6aab0226">


<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Towards #29786


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@amyeroberts @ArthurZucker 

Comments:
I am not sure what happened but the same test that passed previously is now failing.

![Screenshot 2024-05-28 at 12 46 24 PM](https://github.com/huggingface/transformers/assets/17495884/b8f89c66-7c1a-4949-8268-d75cbc897710)


cc @sunmarc

Hi @OmarManzoor, not sure what it is happening. However, usually, to make it work on multi device, you need to put some modules in  `_no_split_modules`. In this case , the test is not passing because the embedding size it too big for the first gpu (more than half the size of the model). In `test_model_parallelism`, you could try making it work with
`new_model = model_class.from_pretrained(tmp_dir, device_map="sequential", max_memory=max_memory) `, by changing device_map="auto" to "sequential" ? This should fix the first assert error. I will soon create a PR to do that since the test is not super well designed for unbalanced small models. 

> In `test_model_parallelism`, you could try making it work with `new_model = model_class.from_pretrained(tmp_dir, device_map="sequential", max_memory=max_memory) `, by changing device_map="auto" to "sequential" ?

Hi @SunMarc,
If I do that won't it actually change this test generically for all models which is supposed to test for device_map="auto"? Or should I override this test only for this particular model?

> Hi @SunMarc,
> If I do that won't it actually change this test generically for all models which is supposed to test for device_map="auto"? Or should I override this test only for this particular model?

The plan is to switch to `sequential` in the future. I need to check on our CI which tests will fail and fix it. 
Right now, for this test, there is no way to pass it if we go with device_map="auto" because of the unbalanced model. We can potentially skip the test but it would be great if we can still enable the multi-gpu nevertheless. 

> > Hi @SunMarc,
> > If I do that won't it actually change this test generically for all models which is supposed to test for device_map="auto"? Or should I override this test only for this particular model?
> 
> The plan is to switch to `sequential` in the future. I need to check on our CI which tests will fail and fix it. Right now, for this test, there is no way to pass it if we go with device_map="auto" because of the unbalanced model. We can potentially skip the test but it would be great if we can still enable the multi-gpu nevertheless.

Then maybe I should mark this particular test to be skipped for DPT?

> Then maybe I should mark this particular test to be skipped for DPT?

Sure, but right now it doesn't work with multi gpu. When I changed to `sequential`, I was hitting a device mismatch issue. So that would be great to fix this before merging this ! 

So I tried adding
`_no_split_modules = ["DPTViTEmbeddings", "DPTViTSelfAttention"]`
to DPT but now the other tests are failing as well and the test_model_parallelism returns the same error
```python
def test_model_parallelism(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
    
        for model_class in self.all_model_classes:
            if model_class._no_split_modules is None:
                continue
    
            inputs_dict_class = self._prepare_for_class(inputs_dict, model_class)
            model = model_class(config).eval()
            model = model.to(torch_device)
    
            torch.manual_seed(0)
            base_output = model(**inputs_dict_class)
    
            model_size = compute_module_sizes(model)[""]
            # We test several splits of sizes to make sure it works.
            max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents[1:]]
            with tempfile.TemporaryDirectory() as tmp_dir:
                model.cpu().save_pretrained(tmp_dir)
    
                for max_size in max_gpu_sizes:
                    max_memory = {0: max_size, 1: model_size * 2, "cpu": model_size * 2}
                    new_model = model_class.from_pretrained(tmp_dir, device_map="sequential", max_memory=max_memory)
                    # Making sure part of the model will actually end up offloaded
>                   self.assertSetEqual(set(new_model.hf_device_map.values()), {0, 1})
E                   AssertionError: Items in the second set but not the first:
E                   0

tests/test_modeling_common.py:3161: AssertionError
```

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.