Resume from checkpoint on fused AdamW raises device errors
Author: jleem-iso
Pull request: no

### System Info

Hi, I'm trying to resume a very minimal T5 model after some pre-training.

* transformers==4.30.1
* accelerate==0.20.3
* datasets==2.12.0

I'm using `Trainer`'s built-in `resume_from_checkpoint` argument, and I get the following error message:

```RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument state_steps in method wrapper_CUDA___fused_adamw_)```

The model trains fine, so I don't think there's anything wrong with the model training code per se.  

### Who can help?

@sgugger

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```
model = T5ForConditionalGeneration(hf_model_config)
tokenizer = load_tokenizer(tokenizer_path, max_length=max_position_embeddings)

training_args = TrainingArguments(
        ...
        bf16=True,
        optim='adamw_torch_fused',  
        ...
    )
trainer = Trainer(
        model=model,
        ...
)
trainer.train(resume_from_checkpoint=True)
```

### Expected behavior

Resuming from my latest checkpoint

Comments:
Hi @ideasbyjin, thanks for raising this issue. 

In order for us to be able to help, we'll need a minimal code snippet for us to be able to reproduce the error.

Could you provide us with some more information on the running environment: run `transformers-cli env` in the terminal and copy-paste the output? 

Have you run tried training and resuming from checkpoint with a different optimizer than `adamw_torch_fused`? Was it successful?

Hi @amyeroberts! Yep, if I use `adamw_torch` then it seems to train & resume perfectly OK, so I take it it's the fused implementation that's raising issues. 

```
- `transformers` version: 4.30.1
- Platform: Linux-5.15.0-1019-aws-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: DDP
```

Also, I can't provide a more comprehensive example as it's private at this stage, but the gist is that it's running the `T5ForConditionalGeneration` model. 

Again, with standard `adamw_torch` it seems to train/resume with 0 issues. When running with `adamw_torch_fused` , it clearly detects the correct checkpoint and loads the optimizer state (see below) but I can't really pin down where there might be device discrepancies that's upsetting the optimizer

```
Currently training with a batch size of: 256
***** Running training *****
  Num examples = 14,805,069
  Num Epochs = 1
  Instantaneous batch size per device = 256
  Total train batch size (w. parallel, distributed & accumulation) = 256
  Gradient Accumulation steps = 1
  Total optimization steps = 100
  Number of trainable parameters = 8,311,296
  Continuing training from checkpoint, will skip to saved global_step
  Continuing training from epoch 0
  Continuing training from global step 20
  Will skip the first 0 epochs then the first 20 batches in the first epoch.
  0%|                                                                                                                                                 | 0/100 [00:00<?, ?it/s]
```

@ideasbyjin OK, interesting. Could you try updating the pytorch version? 

There were some know issues with fused AdamW and fp16 - #22144 - which _should_ have been resolved in 2.0.1.

Thanks @amyeroberts , no avail on using PyTorch 2.0.1 though.

```
- `transformers` version: 4.30.1
- Platform: Linux-5.15.0-1019-aws-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.16.3
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: DDP
```

```
Will skip the first 0 epochs then the first 20 batches in the first epoch.
  0%|                                                                                                                                                                                                                 | 0/100 [00:00<?, ?it/s]Traceback (most recent call last):

    trainer.train(resume_from_checkpoint=True)

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/transformers/trainer.py", line 1645, in train
    return inner_training_loop(

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/transformers/trainer.py", line 2007, in _inner_training_loop
    self.optimizer.step()

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/accelerate/optimizer.py", line 140, in step
    self.optimizer.step(closure)

  File "/.../conda/envs/pytorch2.0.1/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 69, in wrapper
    return wrapped(*args, **kwargs)

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/torch/optim/optimizer.py", line 280, in wrapper
    out = func(*args, **kwargs)

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
    ret = func(self, *args, **kwargs)

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/torch/optim/adamw.py", line 171, in step
    adamw(

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/torch/optim/adamw.py", line 321, in adamw
    func(

  File "/.../.conda/envs/pytorch2.0.1/lib/python3.10/site-packages/torch/optim/adamw.py", line 615, in _fused_adamw
    torch._fused_adamw_(

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument state_steps in method wrapper_CUDA___fused_adamw_)
```

Right, I think I discovered the issue, and it's a bit to do with both PyTorch and HF. I'll explain:

This first starts with how the HF Trainer reloads an optimizer from a checkpoint,
https://github.com/huggingface/transformers/blob/v4.30.1/src/transformers/trainer.py#L2542

Since the optimizer's states are loaded onto `cpu` (in the case of a single worker, even with multiple GPUs I think?) then when you come to spinning up the fused AdamW optimizer,
https://github.com/pytorch/pytorch/blob/v2.0.1/torch/optim/adamw.py#L614

`device_state_dict` has its Tensors in `cpu` even though all others are in `cuda`, so it raises the error!

I found that either deliberately loading the optimizer states into `cuda` from the `Trainer`, or modifying the `torch.optim.AdamW` code to shift everything to `cuda` did the trick, though I feel like the fix on HF's end is a bit more elegant. 

Perhaps there's argument for changing the `map_location` of the optimizer states, especially in a scenario where we have multiple GPUs on one worker? I'll leave this to your judgment on how to navigate/fix though.