Repository: pytorch
Issue: RuntimeError: "_amp_foreach_non_finite_check_and_unscale_cuda" not implemented for 'BFloat16'
Labels: triaged,module: amp (automated mixed precision),module: mta

### 🐛 Describe the bug

i use pytorch==2.3.0 and peft to train llama3 8b , when i run my code, its raise error like:

```text 
    torch._amp_foreach_non_finite_check_and_unscale_(
RuntimeError: "_amp_foreach_non_finite_check_and_unscale_cuda" not implemented for 'BFloat16'
```
full of my training code like:

```python 
import torch

from datasets import Dataset
import pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer

from peft import LoraConfig, TaskType, get_peft_model

json_path = r"D:/codes/llm_about/self-llm/zzzzz_train/Qwen15B05B/short_name_10k.json"
df = pd.read_json(json_path)
ds = Dataset.from_pandas(df)

pretrained_model = "D:/codes/nlp_about/pretrained_model/hfl_llama-3-chinese-8b"

tokenizer = AutoTokenizer.from_pretrained(pretrained_model, use_fast=False,
                                          trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

print("tokenizer about: ", tokenizer.pad_token, tokenizer.pad_token_id, tokenizer.eos_token_id)


def process_func(example):
    MAX_LENGTH = 256  
    input_ids, attention_mask, labels = [], [], []
    instruction = tokenizer(
        f"
user
\n\n{example['instruction'] + ':' + example['input']}

assistant
\n\n",
        add_special_tokens=False)  
    response = tokenizer(f"{example['output']}
", add_special_tokens=False)
    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
    attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1] 
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
    if len(input_ids) > MAX_LENGTH: 
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }


tokenized_id = ds.map(process_func, remove_columns=ds.column_names)


model = AutoModelForCausalLM.from_pretrained(pretrained_model,
                                             device_map="auto", torch_dtype=torch.bfloat16)

model.enable_input_require_grads() 

print("model.dtype=", model.dtype)

config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    target_modules=["q_proj", "k_proj", "v_proj"],
    inference_mode=False, 
    r=8,  
    lora_alpha=32,  
    lora_dropout=0.1
)
# print("config=", config)

model = get_peft_model(model, config)
model.print_trainable_parameters()
model.enable_input_require_grads()

args = TrainingArguments(
    output_dir="./output/llama3",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    logging_steps=10,
    num_train_epochs=16,
    save_steps=300,
    learning_rate=1e-4,
    save_on_each_node=True,
    gradient_checkpointing=True,
    fp16=True,

)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_id,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
)

trainer.train()


```


and the all output in terminal like:

```text
C:\ProgramData\miniconda3\envs\llama\python.exe D:/codes/llm_about/self-llm/zzzzz_train/llama38B/finetune_llama3_8b.py
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
tokenizer about:  
 128001 128001
Map: 100%|██████████| 10000/10000 [00:03<00:00, 3302.79 examples/s]
Loading checkpoint shards: 100%|██████████| 4/4 [00:08<00:00,  2.21s/it]
model.dtype= torch.bfloat16
trainable params: 4,718,592 || all params: 8,034,979,840 || trainable%: 0.0587
C:\ProgramData\miniconda3\envs\llama\lib\site-packages\accelerate\accelerator.py:446: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches']). Please pass an `accelerate.DataLoaderConfiguration` instead: 
dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False)
  warnings.warn(
  0%|          | 0/10000 [00:00<?, ?it/s]`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...
C:\ProgramData\miniconda3\envs\llama\lib\site-packages\torch\utils\checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.
  warnings.warn(
C:\ProgramData\miniconda3\envs\llama\lib\site-packages\transformers\models\llama\modeling_llama.py:728: UserWarning: 1Torch was not compiled with flash attention. (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\transformers\cuda\sdp_utils.cpp:455.)
  attn_output = torch.nn.functional.scaled_dot_product_attention(
Traceback (most recent call last):
  File "D:\codes\llm_about\self-llm\zzzzz_train\llama38B\finetune_llama3_8b.py", line 96, in <module>
    trainer.train()
  File "C:\ProgramData\miniconda3\envs\llama\lib\site-packages\transformers\trainer.py", line 1539, in train
    return inner_training_loop(
  File "C:\ProgramData\miniconda3\envs\llama\lib\site-packages\transformers\trainer.py", line 1911, in _inner_training_loop
    self.accelerator.clip_grad_norm_(
  File "C:\ProgramData\miniconda3\envs\llama\lib\site-packages\accelerate\accelerator.py", line 2269, in clip_grad_norm_
    self.unscale_gradients()
  File "C:\ProgramData\miniconda3\envs\llama\lib\site-packages\accelerate\accelerator.py", line 2219, in unscale_gradients
    self.scaler.unscale_(opt)
  File "C:\ProgramData\miniconda3\envs\llama\lib\site-packages\torch\amp\grad_scaler.py", line 337, in unscale_
    optimizer_state["found_inf_per_device"] = self._unscale_grads_(
  File "C:\ProgramData\miniconda3\envs\llama\lib\site-packages\torch\amp\grad_scaler.py", line 278, in _unscale_grads_
    torch._amp_foreach_non_finite_check_and_unscale_(
RuntimeError: "_amp_foreach_non_finite_check_and_unscale_cuda" not implemented for 'BFloat16'
  0%|          | 0/10000 [00:02<?, ?it/s]


```

i didnot find any thing about this error in issue, help me plz, and thanks a lot.


### Versions

Collecting environment information...
PyTorch version: 2.3.0
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Pro
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.14 | packaged by Anaconda, Inc. | (main, May  6 2024, 19:44:50) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22631-SP0
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 555.85
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\bin\cudnn_ops64_9.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:










Revision=

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.3.0
[pip3] torchaudio==2.3.0
[pip3] torchvision==0.18.0
[conda] blas                      1.0                         mkl    defaults
[conda] mkl                       2021.4.0                 pypi_0    pypi
[conda] mkl-service               2.4.0           py310h2bbff1b_0    defaults
[conda] mkl_fft                   1.3.1           py310ha0764ea_0    defaults
[conda] mkl_random                1.2.2           py310h4ed8f06_0    defaults
[conda] numpy                     1.24.3          py310hdc03b94_0    defaults
[conda] numpy-base                1.24.3          py310h3caf3d7_0    defaults
[conda] pytorch                   2.3.0           py3.10_cuda12.1_cudnn8_0    pytorch
[conda] pytorch-cuda              12.1                 hde6ce7c_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.3.0                    pypi_0    pypi
[conda] torchvision               0.18.0                   pypi_0    pypi

Process finished with exit code 0

cc @mcarilli @ptrblck @leslie-fang-intel @jgong5 @crcrpar @janeyx99