Make datatype conversion consistent in LlamaRMSNorm between weight and hidden_states
Author: jaroslawnowak
Pull request: yes

# What does this PR do?

This PR makes the datatype conversion operate on the whole output of LlamaRMSNorm.compute. As it stands, the weight is initialized to the default datatype, which causes the return statement to coerce the entire output to the weight datatype, even if hidden_states datatype is different. This causes problems when finetuning/training the model with e.g. bf16 datatype. 

This is also causing the following error when running `model.generate`:

`RuntimeError: Expected attn_mask dtype to be bool or to match query dtype, but got attn_mask.dtype: c10::BFloat16 and  query.dtype: float instead.`

@ArthurZucker 

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR 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:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
> Thanks, this is a duplicate of #25421. Could you provide a complete reproducer for the issue you are mentioning?

Thanks for reviewing. This is the code to repro:

```
import numpy as np
import torch

import os
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
    logging,
)
from transformers.generation.configuration_utils import GenerationConfig
from peft import LoraConfig
from trl import SFTTrainer
from trl.trainer.utils import ConstantLengthDataset

base_model = "NousResearch/Llama-2-7b-chat-hf"
guanaco_dataset = "mlabonne/guanaco-llama2-1k"

dataset = load_dataset(guanaco_dataset, split="train")

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=False,
)

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=quant_config,
    device_map={"": 0},
)
model.config.use_cache = False
model.config.pretraining_tp = 1

tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

peft_params = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=8,
    bias="none",
    task_type="CAUSAL_LM",
)

training_params = TrainingArguments(
    output_dir="./results",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    optim="paged_adamw_32bit",
    save_steps=25,
    logging_steps=25,
    learning_rate=2e-4,
    weight_decay=0.001,
    fp16=False,
    bf16=True,
    max_grad_norm=0.3,
    max_steps=-1,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="constant",
    report_to="tensorboard"
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_params,
    dataset_text_field="text",
    max_seq_length=512,
    tokenizer=tokenizer,
    args=training_params,
    packing=False,
)

trainer.train()

prompt = "Who is the current president of the United States?"
inputs = tokenizer([f"

 {prompt} 
"], return_tensors="pt").to("cuda")

generated_ids = model.generate(
    input_ids=inputs['input_ids'], 
    max_new_tokens=250, 
    do_sample=True,
)
```


I think this is caused by `peft_module_casting_to_bf16` in trl/trainer/utils.py:661 explicitly not converting dtypes of `LayerNorm` to float16, it seems this is done to prevent division by 0 problems

There are specific branches for this in the modeling code. But TLDR this is not planned (see the other linked PR)! FYI @younesbelkada 

@ArthurZucker thank you for reviewing. In addition to the `RuntimeError` explained above, the issue is this: if the suggested bug fix is not implemented, the memory usage goes up by about 20% because the data is implicitly converted to float32. 

What is the next step: would you rather I close this PR or are you happy to merge as is?



`peft_module_casting_to_bf16` is important to guarantee training stability for FA2 models cc @pacman100 , can you elaborate more about this fix @jaroslawnowak ? 🙏 

> peft_module_casting_to_bf16 is important to guarantee training stability for FA2 models cc @pacman100

Hello, `peft_module_casting_to_bf16` was added to be inline with QLoRA implementation meant to further reduce GPU memory requirments: https://github.com/artidoro/qlora

I'm able to run the reproducible scr4ipt without any errors with the usage of `with torch.autocast(device_type="cuda", dtype=torch.bfloat16):` during generate as shown below:

```diff
import numpy as np
import torch

import os
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
    logging,
)
from transformers.generation.configuration_utils import GenerationConfig
from peft import LoraConfig
from trl import SFTTrainer
from trl.trainer.utils import ConstantLengthDataset

base_model = "NousResearch/Llama-2-7b-chat-hf"
guanaco_dataset = "mlabonne/guanaco-llama2-1k"

dataset = load_dataset(guanaco_dataset, split="train")

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=False,
)

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=quant_config,
    device_map={"": 0},
)
model.config.use_cache = False
model.config.pretraining_tp = 1

tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

peft_params = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=8,
    bias="none",
    task_type="CAUSAL_LM",
)

training_params = TrainingArguments(
    output_dir="./results",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    optim="paged_adamw_32bit",
    save_steps=25,
    logging_steps=25,
    learning_rate=2e-4,
    weight_decay=0.001,
    fp16=False,
    bf16=True,
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="constant",
    report_to="tensorboard",
    max_steps=50
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_params,
    dataset_text_field="text",
    max_seq_length=512,
    tokenizer=tokenizer,
    args=training_params,
    packing=False,
)

trainer.train()

prompt = "Who is the current president of the United States?"
inputs = tokenizer([f"

 {prompt} 
"], return_tensors="pt").to("cuda")


+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    generated_ids = model.generate(
        input_ids=inputs['input_ids'], 
        max_new_tokens=250, 
        do_sample=True,
    )
```

Output:
```
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████| 2/2 [00:06<00:00,  3.02s/it]
/raid/sourab/transformers/src/transformers/generation/configuration_utils.py:500: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.9` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`. This was detected when initializing the generation config instance, which means the corresponding file may hold incorrect parameterization and should be fixed.
  warnings.warn(
/raid/sourab/transformers/src/transformers/generation/configuration_utils.py:505: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.6` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`. This was detected when initializing the generation config instance, which means the corresponding file may hold incorrect parameterization and should be fixed.
  warnings.warn(
/raid/sourab/transformers/src/transformers/generation/configuration_utils.py:500: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.9` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.
  warnings.warn(
/raid/sourab/transformers/src/transformers/generation/configuration_utils.py:505: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.6` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`.
  warnings.warn(
Detected kernel version 5.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.
max_steps is given, it will override any value given in num_train_epochs
[2024-04-29 12:36:40,689] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
{'loss': 1.3877, 'grad_norm': 0.4296875, 'learning_rate': 0.0002, 'epoch': 0.1}                                                         
{'loss': 1.6284, 'grad_norm': 12.0, 'learning_rate': 0.0002, 'epoch': 0.2}                                                              
{'train_runtime': 364.1265, 'train_samples_per_second': 0.549, 'train_steps_per_second': 0.137, 'train_loss': 1.508057861328125, 'epoch': 0.2}
100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [06:04<00:00,  7.28s/it]

```

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.