Finetuning OPT models with 8bit and LoRA on QA tasks leads to NAN weight in `model.qa_outputs`
Author: shiningrain
Pull request: no

### System Info

- `transformers` version: 4.32.1
- Platform: Linux-6.5.0-28-generic-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.23.1
- Safetensors version: 0.4.3
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (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?: Yes
- Using distributed or parallel set-up in script?: Distribute on 2 or 3 GPUs

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I have faced `nan` and `inf` weights when finetuning OPT model (2.7b or 125m) using `load_in_8bit=True` on a QA dataset `squad` (referring to [notebook 1](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/question_answering.ipynb) and [notebook2](https://colab.research.google.com/drive/1jCkpikz0J2o20FBQmYmAGdiKmJGOMo-o)).
My script is shown as follows.

In the first batch of training, after the backward of AdamW optimizer in line 1917 (i.e., `self.optimizer.step()`) of `/site-packages/transformers/trainer.py` , the weight and bias of `model.base_model.model.qa_outputs.modules_to_save.default` will become NaN/Inf.
I have increased `eps` of the optimizer (following this [discussion](https://discuss.pytorch.org/t/adam-half-precision-nans/1765)), but nothing happens.
I am not sure whether this is caused by a problem in my implementation or a bug in the implementation for 8-bit training.
This problem disappears after disabling `load_in_8bit`.

Thanks a lot!

```
from datasets import load_dataset
from transformers import AutoTokenizer
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ["WANDB_LOG_MODEL"] = "checkpoint"
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
from peft import (  # noqa: E402
    LoraConfig,
    get_peft_model,
)
import transformers
from datasets import load_dataset


model_dir="facebook/opt-125m"
output_dir='./tmp_model-opt125-QA'

squad_v2 = True
lora_r=16
lora_alpha=32
lora_dropout=0.05#https://opendelta.readthedocs.io/en/latest/modules/deltas.html
lora_target_modules=[
        "q_proj",
        "v_proj",
]
batch_size = 16
max_length = 384 # The maximum length of a feature (question and context)
doc_stride = 128 # The authorized overlap between two part of the context when splitting it is needed.
lr=2e-4
epochs=3


datasets = load_dataset("squad_v2" if squad_v2 else "squad")
model = AutoModelForQuestionAnswering.from_pretrained(
    model_dir,
    load_in_8bit=True,
    device_map='auto',
)
tokenizer = AutoTokenizer.from_pretrained(model_dir,padding_side="left")
assert isinstance(tokenizer, transformers.PreTrainedTokenizerFast)
model.gradient_checkpointing_enable()  # reduce number of stored activations
model.enable_input_require_grads()

pad_on_right = tokenizer.padding_side == "right"
def prepare_train_features(examples):
    # Some of the questions have lots of whitespace on the left, which is not useful and will make the
    # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
    # left whitespace
    examples["question"] = [q.lstrip() for q in examples["question"]]

    # Tokenize our examples with truncation and padding, but keep the overflows using a stride. This results
    # in one example possible giving several features when a context is long, each of those features having a
    # context that overlaps a bit the context of the previous feature.
    tokenized_examples = tokenizer(
        examples["question" if pad_on_right else "context"],
        examples["context" if pad_on_right else "question"],
        truncation="only_second" if pad_on_right else "only_first",
        max_length=max_length,
        stride=doc_stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

    # Since one example might give us several features if it has a long context, we need a map from a feature to
    # its corresponding example. This key gives us just that.
    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
    # The offset mappings will give us a map from token to character position in the original context. This will
    # help us compute the start_positions and end_positions.
    offset_mapping = tokenized_examples.pop("offset_mapping")

    # Let's label those examples!
    tokenized_examples["start_positions"] = []
    tokenized_examples["end_positions"] = []

    for i, offsets in enumerate(offset_mapping):
        # We will label impossible answers with the index of the CLS token.
        input_ids = tokenized_examples["input_ids"][i]
        # cls_index = input_ids.index(tokenizer.cls_token_id)
        if tokenizer.cls_token_id in input_ids:
            cls_index = input_ids.index(tokenizer.cls_token_id)
        elif tokenizer.bos_token_id in input_ids:
            cls_index = input_ids.index(tokenizer.bos_token_id)
        else:
            cls_index = 0

        # Grab the sequence corresponding to that example (to know what is the context and what is the question).
        sequence_ids = tokenized_examples.sequence_ids(i)

        # One example can give several spans, this is the index of the example containing this span of text.
        sample_index = sample_mapping[i]
        answers = examples["answers"][sample_index]
        # If no answers are given, set the cls_index as answer.
        if len(answers["answer_start"]) == 0:
            tokenized_examples["start_positions"].append(cls_index)
            tokenized_examples["end_positions"].append(cls_index)
        else:
            # Start/end character index of the answer in the text.
            start_char = answers["answer_start"][0]
            end_char = start_char + len(answers["text"][0])

            # Start token index of the current span in the text.
            token_start_index = 0
            while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
                token_start_index += 1

            # End token index of the current span in the text.
            token_end_index = len(input_ids) - 1
            while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
                token_end_index -= 1

            # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).
            if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
            else:
                # Otherwise move the token_start_index and token_end_index to the two ends of the answer.
                # Note: we could go after the last offset if the answer is the last word (edge case).
                while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
                    token_start_index += 1
                tokenized_examples["start_positions"].append(token_start_index - 1)
                while offsets[token_end_index][1] >= end_char:
                    token_end_index -= 1
                tokenized_examples["end_positions"].append(token_end_index + 1)

    return tokenized_examples

config = LoraConfig(
    r=lora_r,
    lora_alpha=lora_alpha,
    target_modules=lora_target_modules,
    lora_dropout=lora_dropout,
    bias="none",
    task_type='QUESTION_ANS'
)
model = get_peft_model(model, config)
model.print_trainable_parameters()

tokenized_datasets = datasets.map(prepare_train_features, batched=True, remove_columns=datasets["train"].column_names,batch_size=batch_size, num_proc=4)

args = transformers.TrainingArguments(
    output_dir=output_dir,
    evaluation_strategy = "epoch",
    learning_rate=lr,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=epochs,
    weight_decay=0.01,
)
from transformers import default_data_collator

data_collator = default_data_collator
trainer = transformers.Trainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)
model.config.use_cache = False  
trainer.train()
trainer.save_model(output_dir)

```



### Expected behavior

It is expected to be trained without `nan` or `inf`.

Comments:
Hi @shiningrain , can you try the following:

Update `bitsandbytes`: `pip install bitsandbytes -U`
And try using `load_in_4bit=True` instead of `load_in_8bit=True` to check if it works


> Hi @shiningrain , can you try the following:
> 
> Update `bitsandbytes`: `pip install bitsandbytes -U` And try using `load_in_4bit=True` instead of `load_in_8bit=True` to check if it works

I have just tried. After updating `bitsandbytes` and using `load_in_4bit` , 
```
model = AutoModelForQuestionAnswering.from_pretrained(
    model_dir,
    load_in_4bit=True,
    device_map='auto',
)
```
the inf and nan are still in the weight and bias. it looks like this solution didn't work on this problem.
![image](https://github.com/huggingface/transformers/assets/46860123/865844d7-57be-4c09-8170-ffd21b1d644a)





Hi @shiningrain, firstly in the above solution, **'the `load_in_4bit` and `load_in_8bit` arguments are deprecated and will be removed in the future versions'**.

The supported format is by using a `BitsAndBytesConfig` instance and putting the `load_in_8bit` arg there as follows:
```
from transformers import  BitsAndBytesConfig
.
.
.
model = AutoModelForQuestionAnswering.from_pretrained(
    model_dir,
    quantization_config=BitsAndBytesConfig(load_in_8bit=True),
    device_map='auto',
)
```
And have you tried the code on a single GPU 'cause in most such issues, problems occured on multi-GPU setup and as that would clarify if the issue exits on quantization side or in the gpu parallel training?

Hi @RUFFY-369. Thank you for your advice, and so sorry for the delay. 

1) I have tried to use `quantization_config=BitsAndBytesConfig(load_in_8bit=True)# load_in_4bit=True` in my code but the `inf` still exists after the first `self.optimizer.step()` in the trainer.

2) I have also tried to use only one GPU, as shown in the setting `os.environ["CUDA_VISIBLE_DEVICES"]="0,1"` -> `os.environ["CUDA_VISIBLE_DEVICES"]="0"`, but there is still a `inf` in bias/weights.

3) When I remove the quantization configuration in the code, it has normal bias and weights in the model (on both single GPU and multiple GPUs).
I think **this may be a bug related to the quantization** in the trainer.

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.