Errors while training apple/mobilevit-xx-small on image-classification example with and without deepspeed
Author: prathikr
Pull request: no

### System Info

- transformers installed from source
- python 3.8
- ZeRO-Stage-1

### Who can help?

@amyeroberts @NielsRogge @JingyaHuang 

### 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

python -m torch.distributed.launch --nproc_per_node=8 ~/transformers/examples/pytorch/image-classification/run_image_classification.py --model_name_or_path apple/mobilevit-xx-small --dataset_name beans --overwrite_output_dir --output_dir ./outputs/ --remove_unused_columns False --do_train --do_eval --learning_rate 2e-5 --num_train_epochs 50 --per_device_train_batch_size 8  --per_device_eval_batch_size 8 --logging_strategy steps --logging_steps 10 --evaluation_strategy epoch --seed 1337 --fp16 True --report_to none --ignore_mismatched_sizes True

AttributeError: 'MobileViTImageProcessor' object has no attribute 'image_mean'

python -m torch.distributed.launch --nproc_per_node=8 ~/transformers/examples/pytorch/image-classification/run_image_classification.py --model_name_or_path apple/mobilevit-xx-small --dataset_name beans --overwrite_output_dir --output_dir ./outputs/ --remove_unused_columns False --do_train --do_eval --learning_rate 2e-5 --num_train_epochs 50 --per_device_train_batch_size 8  --per_device_eval_batch_size 8 --logging_strategy steps --logging_steps 10 --evaluation_strategy epoch --seed 1337 --fp16 True --report_to none  --ignore_mismatched_sizes True --deepspeed ~/zero_stage_1.json

AttributeError: 'MobileViTConfig' object has no attribute 'hidden_size' 

### Expected behavior

I expect both examples to train with the deepspeed-enabled run completing faster than baseline. Currently, both scenarios error out. Thank you in advance for the assistance.

Comments:
cc @amyeroberts and @alaradirik 

See #21221 - the example scripts aren't meant to work out-of-the-box for any model, as MobileViT for instance doesn't normalizes the images with a mean and std, so one needs to comment out the normalization line.

@NielsRogge thank you, I can confirm that commenting out the normalization line resolves this issue. How about the DeepSpeed compatibility issue? @JingyaHuang mentioned to me that MobileViT has more of a CNN architecture rather than transformer so it may not work at all with DeepSpeed. Is this accurate? Or is there another workaround you can provide? Thank you in advance.

Hi @prathikr, as your previous assumption of the issue's root was `hidden_sizes`, I suspected it could come from the lack of [support for variable hidden_size in ONNX Runtime's graph optimization](https://github.com/microsoft/onnxruntime/blob/81120e9e8b377567daa00d55614c902f35b2ae8f/onnxruntime/python/tools/transformers/optimizer.py#L145) other than a problem from the DeepSpeed(but it could be, I am not aware of this).

@JingyaHuang I don't think so because this issue arises when running without ONNX Runtime as well

@NielsRogge @JingyaHuang any updates on the deepspeed issue for MobileViT? 

What's the reason you'd like to use Deepspeed?

Also, please provide a full stacktrace

@NielsRogge, most Microsoft internal training pipelines including AzureML leverage DeepSpeed since it provides better training speed and smaller memory footprint. When we evaluate any Hugging Face models, we always try to integrate both ORT and DeepSpeed to maximize training speed.

Traceback (most recent call last):
  File "/home/prathikrao/transformers/examples/pytorch/image-classification/run_image_classification.py", line 392, in <module>
    main()
  File "/home/prathikrao/transformers/examples/pytorch/image-classification/run_image_classification.py", line 366, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/prathikrao/transformers/src/transformers/trainer.py", line 1547, in train
    return inner_training_loop(
  File "/home/prathikrao/transformers/src/transformers/trainer.py", line 1616, in _inner_training_loop
    deepspeed_engine, optimizer, lr_scheduler = deepspeed_init(
  File "/home/prathikrao/transformers/src/transformers/deepspeed.py", line 312, in deepspeed_init
    hf_deepspeed_config.trainer_config_finalize(args, model, num_training_steps)
  File "/home/prathikrao/transformers/src/transformers/deepspeed.py", line 174, in trainer_config_finalize
    hidden_size = model.config.hidden_size
  File "/home/prathikrao/transformers/src/transformers/configuration_utils.py", line 260, in __getattribute__
    return super().__getattribute__(key)
AttributeError: 'MobileViTConfig' object has no attribute 'hidden_size'

@NielsRogge above is the full stacktrace. I believe I'm seeing this because MobileViT has an attribute named [hidden_sizes](https://github.com/huggingface/transformers/blob/7119bb052a3f492b9af3afe4f3f13132445eba6e/src/transformers/models/mobilevit/configuration_mobilevit.py#L63) which is a list (different from ViT, for example, which has [hidden_size](https://github.com/huggingface/transformers/blob/7119bb052a3f492b9af3afe4f3f13132445eba6e/src/transformers/models/vit/configuration_vit.py#L47) as an int). Not sure if this is an architectural difference that would make Deepspeed incompatible or if there is some workaround for this.

I'll cc @stas00 here as he's our Deepspeed expert. Several vision models indeed have a `hidden_sizes` attribute in their config (as a list of integers), as vision models often consist of several stages, each stage having its own dimensionality (unlike models like ViT which use the same hidden size for each Transformer block).

Yes, of course, I can help here. Thank you for pinging me, @NielsRogge 

OK, here is what's happening. The issue isn't of deepspeed but of its integration into transformers.

Here we pull `model.config.hidden_size` to create the most efficient for the model setup 

https://github.com/huggingface/transformers/blob/5b67ab9924cf7587b39b59eb0bf0abd3d099e8b9/src/transformers/deepspeed.py#L174-L175

now you're saying not all models have it.

So here are some suggestions at how to overcome this problem, while keeping the optimization as close as possible to the best.

1. check if the model has `config.hidden_size` and  if it doesn't and we have these 2 settings in the incoming ds_config:
```
            zero_optimization.stage3_prefetch_bucket_size = "auto"
            zero_optimization.stage3_param_persistence_threshold = "auto"
```
assert about it and then the user can replace `auto` with the value they think works the best and run again.

2. check if the model has `config.hidden_size` and if it doesn't and we have these 2 settings in the incoming ds_config:
```
            zero_optimization.stage3_prefetch_bucket_size = "auto"
            zero_optimization.stage3_param_persistence_threshold = "auto"
```
check that it has `hidden_sizes` and use the one with the highest value and use that instead of `config.hidden_size` 

Should we try option #2?

Please try this PR and let me know if it fixes the problem https://github.com/huggingface/transformers/pull/21504

Thank you!

p.s. I'm making an assumption that the largest hidden size is the most optimal choice, I could be wrong here. 

Thank you @stas00, I can confirm this solves the issue. Not sure about the performance but it at least runs with this fix.

Thanks a lot for testing, @prathikr. Will merge this asap.