bf16 with DeepSpeed stage 3 with CPU offload breaks LLaMA 13b+ training
Author: alexgshaw
Pull request: no

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.

Comments:
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

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/main/CONTRIBUTING.md) are likely to be ignored.

try zero stage3 change to zero stage2