DeepSpeed and HF trainer return hugely different losses and perplexity
I have been trying to pre-train GP2 models with HF Trainer and Deepspeed, but have noticed large differences between HF trainer's final loss and perplexity vs. that of Deepspeed trainer.

For the GPT-2 (100M) model on Wikitext-2-raw dataset on 4 A100 80GB GPU, with the same batchsize=32 per GPU:

HF trainer returns:
```
{'loss': 5.5565, 'learning_rate': 3.6842105263157892e-06, 'epoch': 46.32}       
{'loss': 5.5723, 'learning_rate': 3.1578947368421056e-06, 'epoch': 46.84}       
{'loss': 5.5063, 'learning_rate': 2.631578947368421e-06, 'epoch': 47.37}        
{'loss': 5.5367, 'learning_rate': 2.105263157894737e-06, 'epoch': 47.89}        
{'loss': 5.5109, 'learning_rate': 1.5789473684210528e-06, 'epoch': 48.42}       
{'loss': 5.5516, 'learning_rate': 1.0526315789473685e-06, 'epoch': 48.95}       
{'loss': 5.5457, 'learning_rate': 5.263157894736843e-07, 'epoch': 49.47}        
{'loss': 5.5739, 'learning_rate': 0.0, 'epoch': 50.0}                           
100%|#########################################| 950/950 [27:00<00:00,  1.28s/it][INFO|trainer.py:1391] 2021-09-20 13:44:34,292 >> 
***** train metrics *****
  epoch                    =       50.0
  train_loss               =      6.118
  train_runtime            = 0:27:16.45
  train_samples            =       2318
  train_samples_per_second =     70.824
  train_steps_per_second   =      0.581
09/20/2021 13:44:39 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:2209] 2021-09-20 13:44:39,362 >> ***** Running Evaluation *****
[INFO|trainer.py:2211] 2021-09-20 13:44:39,362 >>   Num examples = 240
[INFO|trainer.py:2214] 2021-09-20 13:44:39,362 >>   Batch size = 8
100%|#############################################| 8/8 [00:00<00:00,  8.73it/s]
***** eval metrics *****
  epoch                   =       50.0
  eval_loss               =     5.9233
  eval_runtime            = 0:00:00.92
  eval_samples            =        240
  eval_samples_per_second =    260.484
  eval_steps_per_second   =      8.683
  perplexity              =   373.6332
```
While the Deepspeed enabled trainer return:
```
{'loss': 23.5148, 'learning_rate': 2.631578947368421e-06, 'epoch': 47.37}       
{'loss': 23.2578, 'learning_rate': 2.105263157894737e-06, 'epoch': 47.89}       
{'loss': 23.187, 'learning_rate': 1.5789473684210528e-06, 'epoch': 48.42}       
{'loss': 23.3219, 'learning_rate': 1.0526315789473685e-06, 'epoch': 48.95}      
{'loss': 23.1348, 'learning_rate': 5.263157894736843e-07, 'epoch': 49.47}       
{'loss': 23.293, 'learning_rate': 0.0, 'epoch': 50.0}                           
100%|#########################################| 950/950 [35:25<00:00,  1.65s/it][INFO|trainer.py:1391] 2021-09-20 07:29:19,207 >> 
***** train metrics *****
  epoch                    =       50.0
  train_loss               =    31.7806
  train_runtime            = 0:35:25.78
  train_samples            =       2318
  train_samples_per_second =     54.521
  train_steps_per_second   =      0.447
09/20/2021 07:29:29 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:2209] 2021-09-20 07:29:29,273 >> ***** Running Evaluation *****
[INFO|trainer.py:2211] 2021-09-20 07:29:29,273 >>   Num examples = 240
[INFO|trainer.py:2214] 2021-09-20 07:29:29,273 >>   Batch size = 8
100%|#############################################| 8/8 [00:01<00:00,  6.59it/s]
***** eval metrics *****
  epoch                   =            50.0
  eval_loss               =          21.939
  eval_runtime            =      0:00:01.40
  eval_samples            =             240
  eval_samples_per_second =         171.321
  eval_steps_per_second   =           5.711
  perplexity              = 3372914502.3404
```

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.11.0.dev0
- Platform: Linux-5.4.0-72-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.9.0+cu111 (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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
- Deepspeed version: '0.5.3'
- 
### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:
- gpt2: @patrickvonplaten, @LysandreJik

Library:

- deepspeed: @stas00

 -->

## Information

GPT2

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Slightly modified clm [script](https://github.com/vinhngx/transformers/blob/deepspeed-test/examples/pytorch/language-modeling/my_run_clm.py) to allow training from scratch and gradient checkpointing
2. Native trainer launch:
```
LOG_DIR = "./models/gpt2-small-a100"
!rm -rf $LOG_DIR
cmd = """python3  -m torch.distributed.launch --nproc_per_node=4 run_clm.py \
    --model_name_or_path gpt2 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --do_eval \
    --eval_steps=10 \
    --logging_steps=10 \
    --save_steps=200 \
    --save_total_limit=1 \
    --fp16=true \
    --per_device_train_batch_size=32 \
    --output_dir {} \
    --num_train_epochs=50 \
    --overwrite_output_dir
""".format(LOG_DIR)
!$cmd
```` 
3. Deepspeed laucher
```
LOG_DIR = "./models/gpt2-deepspeed-a100"
!rm -rf $LOG_DIR
cmd = """deepspeed --num_gpus=4 run_clm.py \
    --model_name_or_path gpt2 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
    --do_eval \
    --eval_steps=10 \
    --logging_steps=10 \
    --save_steps=200 \
    --fp16=true \
    --per_device_train_batch_size=32\
    --output_dir {} \
    --save_total_limit=1 \
    --num_train_epochs=50 \
    --overwrite_output_dir=true \
    --deepspeed=deepspeed-gp2-A100.config.json
""".format(LOG_DIR)
! $cmd
```

With the following deepspeed config file:
```
%%writefile deepspeed-gp2-A100.config.json
{
    "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_fp16_weights_on_model_save": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 100,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

Matching loss and perplexity between HF & deepspeed trainer.

<!-- A clear and concise description of what you would expect to happen. -->

Very interesting report, @vinhngx. Perfect details too.

The 2 setups look very similar (identical). I don't see anything standing out.

I don't currently have an access to a similar hardware setup. I'm curious what happens if you disable gradient checkpointing? (and reduce the bs to fit the higher memory usage). Surely the diverging results if any should appear very quickly - no need to run for more than a few minutes.

Also, I realize that all work I have done was using DS optimizer and scheduler. It was only recently that the Deepspeed team said it's ok to use both as external ones. So perhaps that new ability has some issues? Does it make a large positive difference (making the results similar) if you enable the DS opt+sched by adding:

```
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
```

finally there is the mixed precision section, which could slightly impact things as well.  Usually I use the following:
```
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
```

Finally if you could attach both log files (as a file attachment) - which might help uncover something I'm missing.

Thank you!<|||||>Thanks @stas00 . I carry out a quick experiment on a DGX-1V with 8xV100 32GB, disabling gradient checkpointing this time.  BS is reduced to 4 for both HF trainer and DeepSpeed. 
The stark difference remains. This is for the smallest GPT model after 1 epoch on Wikitext-2.

DeepSpeed trainer always starts with a huge loss (228 vs. 9.8). Full log and notebook for reproduction is [here](https://github.com/vinhngx/transformers/blob/deepspeed-test/examples/pytorch/language-modeling/GPT-small-V100-debug.ipynb)

Will test other suggestions.

HF trainer:
```
***** train metrics *****
  epoch                    =        1.0
  train_loss               =     8.7321
  train_runtime            = 0:00:38.80
  train_samples            =       2318
  train_samples_per_second =     59.737
  train_steps_per_second   =      1.881
09/21/2021 00:59:20 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:2209] 2021-09-21 00:59:20,591 >> ***** Running Evaluation *****
[INFO|trainer.py:2211] 2021-09-21 00:59:20,592 >>   Num examples = 240
[INFO|trainer.py:2214] 2021-09-21 00:59:20,592 >>   Batch size = 8
100%|#############################################| 4/4 [00:01<00:00,  3.54it/s]
***** eval metrics *****
  epoch                   =        1.0
  eval_loss               =     7.9691
  eval_runtime            = 0:00:01.15
  eval_samples            =        240
  eval_samples_per_second =    208.103
  eval_steps_per_second   =      3.468
  perplexity              =  2890.3225
```

DeepSpeed:
```***** train metrics *****
  epoch                    =        1.0
  train_loss               =     80.659
  train_runtime            = 0:01:34.65
  train_samples            =       2318
  train_samples_per_second =     24.489
  train_steps_per_second   =      0.771
09/21/2021 01:02:05 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:2209] 2021-09-21 01:02:05,291 >> ***** Running Evaluation *****
[INFO|trainer.py:2211] 2021-09-21 01:02:05,291 >>   Num examples = 240
[INFO|trainer.py:2214] 2021-09-21 01:02:05,291 >>   Batch size = 8
100%|#############################################| 4/4 [00:01<00:00,  3.51it/s]
***** eval metrics *****
  epoch                   =                   1.0
  eval_loss               =               45.3922
  eval_runtime            =            0:00:01.60
  eval_samples            =                   240
  eval_samples_per_second =               149.585
  eval_steps_per_second   =                 2.493
  perplexity              = 5.171071334294747e+19
```<|||||>Adding some more experiments, which don't seem to help DeepSpeed convergence:

- No amp: https://github.com/vinhngx/transformers/blob/deepspeed-test/examples/pytorch/language-modeling/GPT-small-V100-debug-noAMP.ipynb
- Using DeepSpeed optimizer+LR scheduler:  https://github.com/vinhngx/transformers/blob/deepspeed-test/examples/pytorch/language-modeling/GPT-small-V100-debug-DS-optimizer.ipynb<|||||>Added a quick test for fairscale. It behaves well.

https://github.com/vinhngx/transformers/blob/deepspeed-test/examples/pytorch/language-modeling/GPT-small-V100-debug-fairscale.ipynb<|||||>Thank you for the code and logs and the experiments I asked you to run, @vinhngx! This helps a lot to see what's being run. But I'm yet to see the culprit.

@samyam, any chance you could have a look at one of the log files linked above. @vinhngx has done a ton of different variations and the problem is the same in all of them. And that's Deepspeed's training first step's loss is an order of magnitude larger than the same training w/o Deepspeed (i.e. plain HF Trainer). Thank you!

<|||||>@stas00, @vinhngx  this is very strange. The configs I looked at above have ZeRO Stage 3 enabled. Could you please try i) without any ZeRO, and ii) with ZeRO Stage 2, and see if the difference still persists? I wonder if this model is hitting a corner case in ZeRO Stage 3 where the communication hooks are not triggered properly.<|||||>@samyam good guess. With zero-2, the loss looks inline with the native HF trainer, as well as HF trainer+fairscale.

```
***** train metrics *****
  epoch                    =        1.0
  train_loss               =     8.5525
  train_runtime            = 0:00:34.08
  train_samples            =       2318
  train_samples_per_second =     67.998
  train_steps_per_second   =      2.141
09/23/2021 04:11:04 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:2209] 2021-09-23 04:11:04,653 >> ***** Running Evaluation *****
[INFO|trainer.py:2211] 2021-09-23 04:11:04,653 >>   Num examples = 240
[INFO|trainer.py:2214] 2021-09-23 04:11:04,653 >>   Batch size = 8
100%|#############################################| 4/4 [00:00<00:00,  7.36it/s]
***** eval metrics *****
  epoch                   =        1.0
  eval_loss               =     7.2266
  eval_runtime            = 0:00:00.56
  eval_samples            =        240
  eval_samples_per_second =    422.895
  eval_steps_per_second   =      7.048
  perplexity              =  1375.4861
```
https://github.com/vinhngx/transformers/blob/deepspeed-test/examples/pytorch/language-modeling/GPT-small-V100-debug-DS-optimizer-zero2.ipynb<|||||>@vinhngx, probably the best way to proceed is to repost this issue at https://github.com/microsoft/DeepSpeed since as you have discovered with @samyam's help this is something to do with deepspeed core, rather than HF integration.<|||||>Alright. Submitted a bug for DeepSpeed here: https://github.com/microsoft/DeepSpeed/issues/1402

Thanks for looking into this.<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.