Errors when using FSDP's cpu-offload to train the OPT model
Author: young-chao
Pull request: no

### System Info

### Environment
ubuntu 20.0.4
python 3.9.10
torch=1.12.0+cu113
transformers==4.26.0

### Who can help?

@sgugger @ArthurZucker 

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

### command
torchrun --nproc_per_node=4 [run_clm.py](https://github.com/huggingface/transformers/blob/v4.26.0/examples/pytorch/language-modeling/run_clm.py) \\
--model_name_or_path "facebook/opt-1.3b" \\
--per_device_train_batch_size 1 \\
--output_dir "/workspace/workfile/Downloads/output" \\
--overwrite_output_dir --fp16 --do_train \\
--max_train_samples 500 --num_train_epochs 1 \\
--dataset_name "wikitext" --dataset_config "wikitext-2-v1" \\
--fsdp "full_shard offload auto_wrap" \\
--fsdp_transformer_layer_cls_to_wrap "OPTDecoderLayer"  \\
--optim "adamw_torch" --block_size 1024

### Expected behavior

Traceback (most recent call last):
  File "/workspace/workfile/Projects/fsdp_test/run_clm.py", line 608, in <module>
    main()
  File "/workspace/workfile/Projects/fsdp_test/run_clm.py", line 558, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1543, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 1791, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 2539, in training_step
    loss = self.compute_loss(model, inputs)
  File "/usr/local/lib/python3.9/site-packages/transformers/trainer.py", line 2571, in compute_loss
    outputs = model(**inputs)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2280, in forward
    outputs = self.module(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 476, in forward
    return self.module(*inputs, **kwinputs)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/opt/modeling_opt.py", line 932, in forward
    outputs = self.model.decoder(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/transformers/models/opt/modeling_opt.py", line 697, in forward
    layer_outputs = decoder_layer(
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2244, in forward
    self._lazy_init()
  File "/usr/local/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 1454, in _lazy_init
    self._init_param_attributes(p)
  File "/usr/local/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 1527, in _init_param_attributes
    assert p.device == torch.device("cpu"), (
AssertionError: Expected param to be on CPU when cpu_offloading is enabled. If CPU offloading is enabled correctly, you may be accidentally moving the model to CUDA after FSDP initialization.

Comments:
cc @pacman100 

Hello @young-chao,

Could you please update the torch version to 1.13 and check if that fixes the issue?
 

Also, cpu offload seems to hang indefinitely with mixed_precision, will raise issue with pytorch repo on this.  FSDP CPU Offload without mixed precision is working on 1.13 and later versions 

> Also, cpu offload seems to hang indefinitely with mixed_precision, will raise issue with pytorch repo on this. FSDP CPU Offload without mixed precision is working on 1.13 and later versions

Thank you for your reply，I updated the torch version to 1.13 and  this error disappeared.

Hello @young-chao, happy that the issue is resolved. Please feel free to close the issue.