Training model from checkpoint with adamw_torch_fused
Author: bhperry
Pull request: no

### System Info

- `transformers` version: 4.33.1
- Platform: Linux-5.4.253-167.359.amzn2.x86_64-x86_64-with-glibc2.35
- Python version: 3.9.18
- Huggingface_hub version: 0.17.1
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.0.1+cu117 (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?: Yes
- Using distributed or parallel set-up in script?: No


### Who can help?

_No response_

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

1. Create a transformers Trainer and TrainingArguments with ` optim="adamw_torch_fused"`
2. To quickly reproduce, set `save_strategy="steps"` and `save_steps=1`
3. Run training script until at least the first step is completed, then cancel
4. Run training again with `trainer.train(resume_from_checkpoint=True)`

```Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:48<00:00, 24.08s/it]
trainable params: 4,194,304 || all params: 6,742,609,920 || trainable%: 0.06220594176090199
  0%|                                                                                                                                                                                                                                                                                                                                                        | 0/38523 [00:32<?, ?it/s]
Traceback (most recent call last):                                                                                                                                                                                                                                                                                                                           | 0/77049 [00:00<?, ?it/s]
  File "/home/jovyan/workspace/llm/llm/training/training.py", line 150, in <module>
    tune()
  File "/home/jovyan/workspace/llm/llm/training/training.py", line 136, in tune
    trainer.train(resume_from_checkpoint=RESUME_FROM_CHECKPOINT and has_checkpoint(output_dir))
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/transformers/trainer.py", line 1553, in train
    return inner_training_loop(
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/accelerate/utils/memory.py", line 136, in decorator
    return function(batch_size, *args, **kwargs)
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/transformers/trainer.py", line 1914, in _inner_training_loop
    self.optimizer.step()
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/torch/optim/lr_scheduler.py", line 69, in wrapper
    return wrapped(*args, **kwargs)
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/accelerate/optimizer.py", line 145, in step
    self.optimizer.step(closure)
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/torch/optim/lr_scheduler.py", line 69, in wrapper
    return wrapped(*args, **kwargs)
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/torch/optim/optimizer.py", line 280, in wrapper
    out = func(*args, **kwargs)
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/torch/optim/adamw.py", line 171, in step
    adamw(
  File "/opt/saturncloud/envs/saturn/lib/python3.9/site-packages/torch/optim/adamw.py", line 321, in adamw
    func(
  File "/opt/saturncloud/envs/saturn/lib/python3.9/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_)
```

### Expected behavior

Expected to be able to resume training when using adamw_torch_fused optimizer. This works fine with adamw_torch optimizer.

Comments:
This thread seems potentially relevant https://discuss.pytorch.org/t/loading-a-model-runtimeerror-expected-all-tensors-to-be-on-the-same-device-but-found-at-least-two-devices-cuda-0-and-cpu/143897

Guessing there's just something missing when re-loading the fused optimizer.

I'm pretty sure this is expected, but will ping @pacman100 for his insight! 

> This thread seems potentially relevant https://discuss.pytorch.org/t/loading-a-model-runtimeerror-expected-all-tensors-to-be-on-the-same-device-but-found-at-least-two-devices-cuda-0-and-cpu/143897

> Guessing there's just something missing when re-loading the fused optimizer.

Hello, yes, that seems to be the case. Maybe, try passing the model loaded on the correct device to the Trainer as per the discussion in the above link you shared.

Model is loaded to the correct device, my understanding from the thread is that the optimizer is not loaded correctly. Thought the trainer would be able to figure that out based on the TrainingArguments, but I guess the answer is to load the optimizer manually when using fused.

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.