NameError: name 'save_offload_index' is not defined when use --model_revision sharded
Author: edchengg
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.20.0.dev0
- Platform: Linux-4.15.0-142-generic-x86_64-with-debian-stretch-sid
- Python version: 3.7.7
- Huggingface_hub version: 0.5.1
- PyTorch version (GPU?): 1.9.0+cu111 (True)
- Tensorflow version (GPU?): 2.9.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```


### Who can help?

@stas00 

### 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

`export BS=8; rm -r output_dir; PYTHONPATH=src USE_TF=0 deepspeed --num_gpus=4 run_translation.py --model_name_or_path google/mt5-xxl --output_dir output_dir --adam_eps 1e-06 --evaluation_strategy=steps --do_train --do_eval --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 500 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --per_device_eval_batch_size $BS --predict_with_generate --sortish_sampler --source_lang en --target_lang ro --dataset_name wmt16 --dataset_config "ro-en" --source_prefix "translate English to Romanian: " --val_max_target_length 128 --warmup_steps 50 --max_train_samples 500 --max_eval_samples 50 --deepspeed ds_zero3.json  --fp16 --model_revision sharded`


error:
```
[INFO|modeling_utils.py:2115] 2022-06-01 16:46:52,833 >> Detected DeepSpeed ZeRO-3: activating zero.init() for this model
[2022-06-01 16:47:17,303] [INFO] [partition_parameters.py:464:__exit__] finished initializing model with 12.92B parameters
Traceback (most recent call last):
  File "run_translation.py", line 654, in <module>
Traceback (most recent call last):
  File "run_translation.py", line 654, in <module>
    main()
  File "run_translation.py", line 377, in main
    main()
  File "run_translation.py", line 377, in main
    use_auth_token=True if model_args.use_auth_token else None,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/models/auto/auto_factory.py", line 446, in from_pretrained
    use_auth_token=True if model_args.use_auth_token else None,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/models/auto/auto_factory.py", line 446, in from_pretrained
    return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2179, in from_pretrained
    return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2179, in from_pretrained
    dtype=torch_dtype,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2397, in _load_pretrained_model
    dtype=torch_dtype,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2397, in _load_pretrained_model
    save_offload_index(offload_index, offload_folder)
NameError: name 'save_offload_index' is not defined
    save_offload_index(offload_index, offload_folder)
NameError: name 'save_offload_index' is not defined
Traceback (most recent call last):
Traceback (most recent call last):
  File "run_translation.py", line 654, in <module>
  File "run_translation.py", line 654, in <module>
        main()main()

  File "run_translation.py", line 377, in main
  File "run_translation.py", line 377, in main
    use_auth_token=True if model_args.use_auth_token else None,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/models/auto/auto_factory.py", line 446, in from_pretrained
    use_auth_token=True if model_args.use_auth_token else None,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/models/auto/auto_factory.py", line 446, in from_pretrained
    return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2179, in from_pretrained
    return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2179, in from_pretrained
    dtype=torch_dtype,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2397, in _load_pretrained_model
    dtype=torch_dtype,
  File "/srv/scratch/ychen3411/anaconda3/envs/unify_srl/lib/python3.7/site-packages/transformers/modeling_utils.py", line 2397, in _load_pretrained_model
    save_offload_index(offload_index, offload_folder)
NameError: name 'save_offload_index' is not defined
    save_offload_index(offload_index, offload_folder)
NameError: name 'save_offload_index' is not defined
[2022-06-01 16:53:23,841] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 81406
[2022-06-01 16:53:23,841] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 81407
[2022-06-01 16:53:23,841] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 81408
[2022-06-01 16:53:23,841] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 81409
[2022-06-01 16:53:23,841] [ERROR] [launch.py:184:sigkill_handler] ['/srv/scratch/ychen3411/anaconda3/envs/unify_srl/bin/python', '-u', 'run_translation.py', '--local_rank=3', '--model_name_or_path', 'google/mt5-xxl', '--output_dir', 'output_dir', '--adam_eps', '1e-06', '--evaluation_strategy=steps', '--do_train', '--do_eval', '--label_smoothing', '0.1', '--learning_rate', '3e-5', '--logging_first_step', '--logging_steps', '500', '--max_source_length', '128', '--max_target_length', '128', '--num_train_epochs', '1', '--overwrite_output_dir', '--per_device_train_batch_size', '8', '--per_device_eval_batch_size', '8', '--predict_with_generate', '--sortish_sampler', '--source_lang', 'en', '--target_lang', 'ro', '--dataset_name', 'wmt16', '--dataset_config', 'ro-en', '--source_prefix', 'translate English to Romanian: ', '--val_max_target_length', '128', '--warmup_steps', '50', '--max_train_samples', '500', '--max_eval_samples', '50', '--deepspeed', 'ds_zero3.json', '--fp16', '--model_revision', 'sharded'] exits with return code = 1`
```

### Expected behavior

```shell
I tried to run mt5-xxl (12b) on 4 gpus with deepspeed zero3 and sharded.
But I got the following error:

NameError: name 'save_offload_index' is not defined
```


Comments:
@sgugger, would you please kindly look at this - this looks related to the accelerate import - probably not checking if it was imported and using it anyway? 

https://github.com/huggingface/transformers/blob/58fb3c9f98877bf76efb03e376a5c92cf80f7952/src/transformers/modeling_utils.py#L75-L80

https://github.com/huggingface/transformers/blob/58fb3c9f98877bf76efb03e376a5c92cf80f7952/src/transformers/modeling_utils.py#L2397

Thank you!

Indeed, this needs some gating. Will send a fix shortly, thanks for flagging!