

TITLE
DeepSpeed and HF trainer return hugely different losses and perplexity

COMMENTS
10

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
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. -->
