bf16 with DeepSpeed stage 3 with CPU offload breaks LLaMA 13b+ training
TL;DR deepspeed stage 3 with cpu offload and bf16 breaks llama 13b+ when fine-tuning. The loss starts high and then immediately drops to 0 after the first step and learning rate stays 0 the entire time.

### System Info

- `transformers` version: 4.30.2
- Platform: Linux-3.10.0-1160.92.1.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.2
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 1.12.1+cu113 (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?: Using 8 A100s in a training script.
- Using distributed or parallel set-up in script?: Using deepspeed stage 3 with CPU offload

### Who can help?

@sgugger @pacman100 

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

I'm running a really straightforward, bare-bones fine-tuning script to train llama 13b. The problem is, if I turn bf16 on, I run into the following problem:

```
  0%|          | 0/82 [00:00<?, ?it/s]
  1%|          | 1/82 [01:39<2:14:49, 99.87s/it]
                                                
{'loss': 8.2056, 'learning_rate': 0.0, 'epoch': 0.0}

  1%|          | 1/82 [01:39<2:14:49, 99.87s/it]
  2%|▏         | 2/82 [02:55<1:54:16, 85.70s/it]
                                                
{'loss': 0.0, 'learning_rate': 0.0, 'epoch': 0.0}

  2%|▏         | 2/82 [02:55<1:54:16, 85.70s/it]
  4%|▎         | 3/82 [04:10<1:46:23, 80.81s/it]
                                                
{'loss': 0.0, 'learning_rate': 0.0, 'epoch': 0.01}

  4%|▎         | 3/82 [04:10<1:46:23, 80.81s/it]
  5%|▍         | 4/82 [05:26<1:42:33, 78.90s/it]
                                                
{'loss': 0.0, 'learning_rate': 0.0, 'epoch': 0.01}

  5%|▍         | 4/82 [05:26<1:42:33, 78.90s/it]
  6%|▌         | 5/82 [06:40<1:38:56, 77.10s/it]
                                                
{'loss': 0.0, 'learning_rate': 0.0, 'epoch': 0.01}
```

This continues for the remainder of the training, with the loss and learning rate never changing. By the end, the model outputs gibberish.

Here is my launch script:

```bash
deepspeed train.py \
    --model_name_or_path /home/ashaw8/compute/models/$MODEL_NAME \
    --dataset_path datasets/$TOPIC/$MODEL_NAME \
    --run_name $RUN_NAME \
    --bf16 True \
    --output_dir $OUTPUT_DIR \
    --num_train_epochs 3 \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --evaluation_strategy "no" \
    --save_strategy "no" \
    --logging_strategy "steps" \
    --logging_steps 1 \
    --learning_rate 5e-6 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --deepspeed ds_config.json \
    --max_grad_norm 1.0 \
    --tf32 False \
    --report_to wandb
```

If I set bf16 to false, everything returns to normal and the training works fine, but then I cannot train the large models (e.g. 65b) because I can't reduce the model size with bf16. As far as I can tell, this issue has not been documented anywhere else.

A related issue seems to be documented here where similar problems occurred with fp16. I guess large losses were preventing the optimizer from stepping? The optimizer must be stepping in my case though, if the model is outputting gibberish by the end.

https://github.com/huggingface/transformers/issues/14531

Here is my actual python training script.

```[python]
from dataclasses import dataclass, field
from typing import Dict, Optional, Union
from transformers import (
    TrainingArguments as HfTrainingArguments,
    HfArgumentParser,
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    DataCollatorForLanguageModeling,
)
from datasets import Dataset


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(
        default="/mnt/pccfs2/backed_up/models/llama/hf/llama-7b-hf/"
    )


@dataclass
class DataArguments:
    dataset_path: str = field(
        default="/mnt/pccfs2/backed_up/alexshaw/media-training/datasets/dataset",
        metadata={"help": "Path to the training data."},
    )


@dataclass
class TrainingArguments(HfTrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    model_max_length: int = field(
        default=512,
        metadata={
            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    dataloader_num_workers: int = field(default=32)


if __name__ == "__main__":
    parser = HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)  # type: ignore
    )
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    config = AutoConfig.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
    )

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
        use_fast=False,
    )

    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        config=config,
        cache_dir=training_args.cache_dir,
    )

    dataset = Dataset.load_from_disk(data_args.dataset_path)

    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,  # type: ignore
        data_collator=data_collator,
    )

    trainer.train()
    trainer.save_state()
    trainer.save_model()
```

Additionally, here is my `ds_config.json`

```json
{
  "bf16": {
    "enabled": "auto"
  },
  "fp16": {
    "enabled": "auto"
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": "auto",
      "weight_decay": "auto"
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "total_num_steps": "auto",
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "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": 5,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}
```

### Expected behavior

The model should start with a higher loss and gradually decrease throughout training. The learning rate should rise to 5e-6 in a few steps.
Can you provide a minimal reproducer as I don't have access to  `/mnt/pccfs2/backed_up/alexshaw/media-training/datasets/dataset`? A reproducer should be minimal and run without having us spend time changing and debugging things.



<|||||>Okay, I built this repo that reproduces the issue with as few dependencies as possible.

https://github.com/alexgshaw/simple-trainer

You should be able to clone the repo, pip install the requirements.txt and run bash `train.sh`

When I ran it, it reproduced the error I described above.

Note that I am running with Python version: 3.8.2 and cuda 11.4<|||||>I am running into a similar issue, are there any updates?

I'm using python 3.11 and cuda 11.8