[BUG] Very high loss when using DeepSpeed with CPU offloading for versions>=4.36.0. 
Author: pacman100
Pull request: no

### System Info

- `transformers` version: 4.37.0.dev0
- Platform: Linux-5.4.0-166-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.20.2
- Safetensors version: 0.4.0
- Accelerate version: 0.25.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.2+cu121 (True)
- Tensorflow version (GPU?): 2.15.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?

_No response_

### 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

**Describe the bug**
Very high loss when using DeepSpeed with CPU offloading for versions>=4.36.0 
After spending 4 hours found that the cause is this PR: https://github.com/huggingface/transformers/pull/27709 and reverting this PR results in expected training loss curves.

**To Reproduce**
Steps to reproduce the behavior:
1. `run_glue.py` transformers example
```
cd transformers
export CUDA_VISISBLE_DEVICES=0,1
export TASK_NAME=mrpc
```
2. ds_config.json
```
{
  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },

  "bf16": {
    "enabled": "auto"
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },

  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 2000,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}
```
3. Launch command:
```
torchrun --nnodes 1 --nproc-per-node 2 ./examples/pytorch/text-classification/run_glue.py --model_name_or_path bert-base-cased  --task_name $TASK_NAME  --do_train  --do_eval  --max_seq_length 128  --per_device_train_batch_size 16  --learning_rate 5e-5  --num_train_epochs 3  --output_dir /tmp/$TASK_NAME/ --overwrite_output_dir --deepspeed ds_config.json --lr_scheduler_type cosine --save_strategy "epoch" --evaluation_strategy "epoch" --logging_steps 1
```
4. See very high loss values:
```
/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py:1330: UserWarning: The torch.cuda.*DtypeTensor constructors are no longer recommended. It's best to use methods such as torch.tensor(data, dtype=*, device='cuda') to create tensors. (Triggered internally at ../torch/csrc/tensor/python_tensor.cpp:83.)
  total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py:1330: UserWarning: The torch.cuda.*DtypeTensor constructors are no longer recommended. It's best to use methods such as torch.tensor(data, dtype=*, device='cuda') to create tensors. (Triggered internally at ../torch/csrc/tensor/python_tensor.cpp:83.)
  total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
{'loss': 18.693, 'learning_rate': 4.999896350176413e-05, 'epoch': 0.01}                                                                                                      
{'loss': 14.7233, 'learning_rate': 4.999585409300281e-05, 'epoch': 0.02}                                                                                                     
{'loss': 7.5012, 'learning_rate': 4.999067203154777e-05, 'epoch': 0.03}                                                                                                      
{'loss': 7.0167, 'learning_rate': 4.9983417747094816e-05, 'epoch': 0.03}                                                                                                     
{'loss': 5.6229, 'learning_rate': 4.9974091841168195e-05, 'epoch': 0.04}                                                                                                     
{'loss': 8.2446, 'learning_rate': 4.99626950870707e-05, 'epoch': 0.05}                                                                                                       
{'loss': 4.521, 'learning_rate': 4.994922842981958e-05, 'epoch': 0.06}                                                                                                       
{'loss': 2.8402, 'learning_rate': 4.9933692986068165e-05, 'epoch': 0.07}                                                                                                     
{'loss': 3.8474, 'learning_rate': 4.991609004401324e-05, 'epoch': 0.08}                                                                                                      
{'loss': 3.3916, 'learning_rate': 4.9896421063288286e-05, 'epoch': 0.09}                                                                                                     
{'loss': 3.2036, 'learning_rate': 4.98746876748424e-05, 'epoch': 0.1}                                                                                                        
{'loss': 2.467, 'learning_rate': 4.985089168080509e-05, 'epoch': 0.1}                                                                                                        
{'loss': 2.9375, 'learning_rate': 4.982503505433683e-05, 'epoch': 0.11}                                                                                                      
{'loss': 2.2244, 'learning_rate': 4.979711993946543e-05, 'epoch': 0.12}                                                                                                      
{'loss': 3.5073, 'learning_rate': 4.976714865090827e-05, 'epoch': 0.13}                                                                                                      
{'loss': 4.1485, 'learning_rate': 4.973512367388038e-05, 'epoch': 0.14}                                                                                                      
{'loss': 3.1678, 'learning_rate': 4.970104766388832e-05, 'epoch': 0.15}                                                                                                      
{'loss': 2.0759, 'learning_rate': 4.966492344651005e-05, 'epoch': 0.16}                                                                                                      
{'loss': 4.6716, 'learning_rate': 4.962675401716056e-05, 'epoch': 0.17}                                                                                                      
{'loss': 2.0029, 'learning_rate': 4.958654254084355e-05, 'epoch': 0.17}                                                                                                      
{'loss': 2.9068, 'learning_rate': 4.9544292351888966e-05, 'epoch': 0.18}                                                                                                     
{'loss': 2.8601, 'learning_rate': 4.95000069536765e-05, 'epoch': 0.19}                                                                                                       
{'loss': 2.4923, 'learning_rate': 4.9453690018345144e-05, 'epoch': 0.2}                                                                                                      
{'loss': 2.4786, 'learning_rate': 4.9405345386488614e-05, 'epoch': 0.21}                                                                                                     
{'loss': 2.1235, 'learning_rate': 4.9354977066836986e-05, 'epoch': 0.22}                                                                                                     
{'loss': 1.6675, 'learning_rate': 4.930258923592418e-05, 'epoch': 0.23}                                                                                                      
{'loss': 2.2374, 'learning_rate': 4.924818623774178e-05, 'epoch': 0.23}                                                                                                      
{'loss': 3.5783, 'learning_rate': 4.9191772583378705e-05, 'epoch': 0.24}                                                                                                     
{'loss': 3.2699, 'learning_rate': 4.91333529506472e-05, 'epoch': 0.25}                                                                                                       
{'loss': 2.0839, 'learning_rate': 4.907293218369499e-05, 'epoch': 0.26}                                                                                                      
{'loss': 2.7985, 'learning_rate': 4.901051529260352e-05, 'epoch': 0.27}                                                                                                      
{'loss': 1.7308, 'learning_rate': 4.89461074529726e-05, 'epoch': 0.28}                                                                                                       
{'loss': 1.1548, 'learning_rate': 4.88797140054912e-05, 'epoch': 0.29}                                                                                                       
{'loss': 2.1195, 'learning_rate': 4.8811340455494624e-05, 'epoch': 0.3}                                                                                                      
{'loss': 3.4952, 'learning_rate': 4.874099247250798e-05, 'epoch': 0.3}                                                                                                       
{'loss': 2.2034, 'learning_rate': 4.8668675889776095e-05, 'epoch': 0.31}                                                                                                     
{'loss': 1.9979, 'learning_rate': 4.85943967037798e-05, 'epoch': 0.32}                                                                                                       
{'loss': 1.1768, 'learning_rate': 4.851816107373871e-05, 'epoch': 0.33}                                                                                                      
{'loss': 1.3446, 'learning_rate': 4.843997532110051e-05, 'epoch': 0.34}                                                                                                      
{'loss': 2.8533, 'learning_rate': 4.835984592901678e-05, 'epoch': 0.35}                                                                                                      
{'loss': 1.32, 'learning_rate': 4.82777795418054e-05, 'epoch': 0.36}                                                                                                         
{'loss': 2.8985, 'learning_rate': 4.819378296439961e-05, 'epoch': 0.37}                                                                                                      
{'loss': 2.1229, 'learning_rate': 4.8107863161783773e-05, 'epoch': 0.37}                                                                                                     
{'loss': 2.2255, 'learning_rate': 4.8020027258415764e-05, 'epoch': 0.38}                                                                                                     
{'loss': 2.0219, 'learning_rate': 4.793028253763633e-05, 'epoch': 0.39}                                                                                                      
 13%|█████████████████▌                                                 
```
5. When reverting PR https://github.com/huggingface/transformers/pull/27709 or changing the offloading device from `cpu` to `none`, the training happens properly:
```
{'loss': 0.6931, 'learning_rate': 4.999896350176413e-05, 'epoch': 0.01}                                                                                                      
{'loss': 0.6833, 'learning_rate': 4.999585409300281e-05, 'epoch': 0.02}                                                                                                      
{'loss': 0.6762, 'learning_rate': 4.999067203154777e-05, 'epoch': 0.03}                                                                                                      
{'loss': 0.6624, 'learning_rate': 4.9983417747094816e-05, 'epoch': 0.03}                                                                                                     
{'loss': 0.653, 'learning_rate': 4.9974091841168195e-05, 'epoch': 0.04}                                                                                                      
{'loss': 0.7168, 'learning_rate': 4.99626950870707e-05, 'epoch': 0.05}                                                                                                       
{'loss': 0.6631, 'learning_rate': 4.994922842981958e-05, 'epoch': 0.06}                                                                                                      
{'loss': 0.6233, 'learning_rate': 4.9933692986068165e-05, 'epoch': 0.07}                                                                                                     
{'loss': 0.6324, 'learning_rate': 4.991609004401324e-05, 'epoch': 0.08}                                                                                                      
{'loss': 0.677, 'learning_rate': 4.9896421063288286e-05, 'epoch': 0.09}                                                                                                      
{'loss': 0.5409, 'learning_rate': 4.98746876748424e-05, 'epoch': 0.1}                                                                                                        
{'loss': 0.6643, 'learning_rate': 4.985089168080509e-05, 'epoch': 0.1}                                                                                                       
{'loss': 0.7062, 'learning_rate': 4.982503505433683e-05, 'epoch': 0.11}                                                                                                      
{'loss': 0.6542, 'learning_rate': 4.979711993946543e-05, 'epoch': 0.12}                                                                                                      
{'loss': 0.6364, 'learning_rate': 4.976714865090827e-05, 'epoch': 0.13}                                                                                                      
{'loss': 0.661, 'learning_rate': 4.973512367388038e-05, 'epoch': 0.14}                                                                                                       
{'loss': 0.6675, 'learning_rate': 4.970104766388832e-05, 'epoch': 0.15}                                                                                                      
{'loss': 0.6261, 'learning_rate': 4.966492344651005e-05, 'epoch': 0.16}                                                                                                      
{'loss': 0.654, 'learning_rate': 4.962675401716056e-05, 'epoch': 0.17}                                                                                                       
{'loss': 0.6515, 'learning_rate': 4.958654254084355e-05, 'epoch': 0.17}                                                                                                      
{'loss': 0.6065, 'learning_rate': 4.9544292351888966e-05, 'epoch': 0.18}                                                                                                     
{'loss': 0.5913, 'learning_rate': 4.95000069536765e-05, 'epoch': 0.19}                                                                                                       
{'loss': 0.6339, 'learning_rate': 4.9453690018345144e-05, 'epoch': 0.2}                                                                                                      
{'loss': 0.6514, 'learning_rate': 4.9405345386488614e-05, 'epoch': 0.21}                                                                                                     
{'loss': 0.6394, 'learning_rate': 4.9354977066836986e-05, 'epoch': 0.22}                                                                                                     
{'loss': 0.652, 'learning_rate': 4.930258923592418e-05, 'epoch': 0.23}                                                                                                       
{'loss': 0.6087, 'learning_rate': 4.924818623774178e-05, 'epoch': 0.23}                                                                                                      
{'loss': 0.6242, 'learning_rate': 4.9191772583378705e-05, 'epoch': 0.24}                                                                                                     
{'loss': 0.534, 'learning_rate': 4.91333529506472e-05, 'epoch': 0.25}                                                                                                        
{'loss': 0.6013, 'learning_rate': 4.907293218369499e-05, 'epoch': 0.26}                                                                                                      
{'loss': 0.6088, 'learning_rate': 4.901051529260352e-05, 'epoch': 0.27}                                                                                                      
{'loss': 0.5327, 'learning_rate': 4.89461074529726e-05, 'epoch': 0.28}                                                                                                       
{'loss': 0.5884, 'learning_rate': 4.88797140054912e-05, 'epoch': 0.29}                                                                                                       
{'loss': 0.6137, 'learning_rate': 4.8811340455494624e-05, 'epoch': 0.3}                                                                                                      
{'loss': 0.5845, 'learning_rate': 4.874099247250798e-05, 'epoch': 0.3}                                                                                                       
{'loss': 0.5141, 'learning_rate': 4.8668675889776095e-05, 'epoch': 0.31}                                                                                                     
{'loss': 0.5223, 'learning_rate': 4.85943967037798e-05, 'epoch': 0.32}                                                                                                       
{'loss': 0.5016, 'learning_rate': 4.851816107373871e-05, 'epoch': 0.33}                                                                                                      
{'loss': 0.5055, 'learning_rate': 4.843997532110051e-05, 'epoch': 0.34}                                                                                                      
{'loss': 0.5048, 'learning_rate': 4.835984592901678e-05, 'epoch': 0.35}                                                                                                      
{'loss': 0.4792, 'learning_rate': 4.82777795418054e-05, 'epoch': 0.36}                                                                                                       
{'loss': 0.6927, 'learning_rate': 4.819378296439961e-05, 'epoch': 0.37}                                                                                                      
{'loss': 0.6419, 'learning_rate': 4.8107863161783773e-05, 'epoch': 0.37}                                                                                                     
{'loss': 0.5379, 'learning_rate': 4.8020027258415764e-05, 'epoch': 0.38}                                                                                                     
{'loss': 0.4458, 'learning_rate': 4.793028253763633e-05, 'epoch': 0.39}                                                                                                      
{'loss': 0.5525, 'learning_rate': 4.783863644106502e-05, 'epoch': 0.4}                                                                                                       
{'loss': 0.4353, 'learning_rate': 4.7745096567983256e-05,
```

**Expected behavior**
Same training loss curves with and without CPU offloading

**ds_report output**
```
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch']
torch version .................... 2.1.2+cu121
deepspeed install path ........... ['/raid/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.12.6, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 2.1, cuda 12.1
shared memory (/dev/shm) size .... 251.77 GB
```

**System info (please complete the following information):**
 - GPU count and types: two A100s
 - Python version 3.10.13
 - Transformers version: 4.36.2
 - Accelerate version: 0.25.0
 
 
cc @ArthurZucker as you have better insights wrt the changes done in PR https://github.com/huggingface/transformers/pull/27709.
 

### Expected behavior

Same training loss curves with and without CPU offloading

Comments:
Closing as #28245 fixed it

These are brutally expensive bugs for your users.  What kind of test cases could be implemented to prevent this from happening?

We added a new test that you can find here: https://github.com/huggingface/transformers/pull/28245/files#diff-712bbefe24deb3653111486b471859bd0f2b4d61b87a81845189da697edde29d 

@ArthurZucker thank you :pray: