LLaMA 13B does not forward or generate properly after being converted to HuggingFace checkpoints
Author: RiverGao
Pull request: no

### System Info

## System info
- Transformer Version: 4.28.0.dev0
- Platform: Ubuntu 18.04.3 LTS (GNU/Linux 4.15.0-194-generic x86_64)
- GPU: Nvidia Tesla V100 SXM2 32GB
- CUDA version: 11.0
- PyTorch version: 1.12.1

## Summary
The LLaMA 13B model converted to HuggingFace format seems not able to generate any text, which may be due to NaN attentions produced in the last layer.

## Bug description
- Downloaded LLaMA weights (7B and 13B) from facebook, and checked their generation by `llama/example.py` provided with them ✅
- Used `transformers/src/transformers/models/llama/convert_llama_weights_to_hf.py` to convert 7B and 13B models into huggingface format ✅
- Loaded the 13B model by `LlamaForCausalLM.from_pretrained(model_path, device_map= "balanced", load_in_8bit=True)` and try to generate text by `model.generate` and ` tokenizer.batch_decode`, but the model simply outputs the given input and nothing else. ❌
- Loaded the 13B model by `LlamaModel.from_pretrained(model_path, device_map="balanced", load_in_8bit=True)` and try to obtain its attentions by `attentions = model(**encoded_input, output_attentions=True).attentions`, and found that all the attention weights in the last layer (layer 39) are NaN's. ❌

## Trying to debug
- Loaded the 7B model by `LlamaForCausalLM.from_pretrained(model_path, device_map= "balanced", load_in_8bit=True)`, and checked its generation by `model.generate` and ` tokenizer.batch_decode` ✅
- Loaded the 7B model by `LlamaModel.from_pretrained(model_path, device_map="balanced", load_in_8bit=True)` and try to obtain its attentions by `attentions = model(**encoded_input, output_attentions=True).attentions`, and the attentions are fine ✅
- I guess there may be corruptions in the model parameters. So I printed the parameters in layer 38 and 39,  `model.norm.weight`, and `lm_head.weight` of the 13B model, but did not find any trace.

### Who can help?

@ArthurZucker @sgugger

### Information

- [x] 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

Step 1: Convert LLaMA weights to HuggingFace
```
python src/transformers/models/llama/convert_llama_weights_to_hf.py \
    --input_dir downloaded_llama_weights --model_size 13B --output_dir llama-hf
```
Step 2: Try to generate text
```
import torch
from transformers import LlamaTokenizer, LlamaForCausalLM

model_path = f'llama-hf/13B'

sentences = [
    'Hello!',
    'Translate this sentence into German: I love baseball.',
    'Tell me about New York.'
]

# model initialization
tokenizer = LlamaTokenizer.from_pretrained(model_path)
model = LlamaForCausalLM.from_pretrained(model_path, device_map= "balanced", load_in_8bit=True)

# iterate over sentences
for sentence in sentences:
    print('-----\nInput:\n' + sentence + '\n-----')
    encoded_input = tokenizer(sentence, return_tensors='pt').to('cuda')
    generate_ids = model.generate(encoded_input.input_ids, max_length=256)
    output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    print('-----\nOutput:\n' + output + '\n-----')
```

Step 3: Try to obtain model attentions
```
# model initialization
tokenizer = LlamaTokenizer.from_pretrained(model_path)
model = LlamaModel.from_pretrained(model_path, device_map= "balanced", load_in_8bit=True, torch_dtype=torch.float32)

n_layer = model.config.num_hidden_layers

# iterate over sentences
attention_per_sentence = [[] for j in range(n_layer)]  # n_layer * n_sent * (n_head, max_sent_len, max_sent_len)
for sentence in sentences:
    encoded_input = tokenizer(sentence, return_tensors='pt')
    attentions = model(**encoded_input, output_attentions=True).attentions  # n_layer * shape (1, n_head, L, L)

    for lyr in range(n_layer):
        attn_array = attentions[lyr].detach().cpu().numpy()
        assert not np.isnan(np.sum(attn_array)), f"layer {lyr} has NaN attentions"
```


### Expected behavior

- In step 1, everything works fine.
- In step 2, the outputs are the same as inputs, no token is generated
- In step 3, assertion fails in layer 39: "AssertionError: layer 39 has NaN attentions"

Comments:
@RiverGao  - could you retry running these steps with the most recent version of transformers? Llama was under a lot of development and has only be officially added to the library in the most recent release - v4.28.0 (4.28.0.dev is some commit between 4.27 and 4.28). This ensures that you have all the most recent updates to the model. 

I've encountered the same problem(specifically step 3) with version 4.29.0.dev0. 
A quick fix would be not quantizing the final 2 layers by manually passing the `qunatization_config`:

```python
from transformers.utils.quantization_config import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig.from_dict({
    'load_in_8bit': True, 'llm_int8_skip_modules': ['model.layers.39', 'model.layers.38', 'lm_head']}, False)

llama_13B = './models/13B_hf'
tokenizer = LlamaTokenizer.from_pretrained(llama_13B)
model = LlamaForCausalLM.from_pretrained(llama_13B, torch_dtype=torch.float16, 
                                         quantization_config=quantization_config, load_in_8bit=False, device_map='auto')

```

@aliwalker Thank you for the quick fix!

cc @younesbelkada 

@aliwalker @amyeroberts I have retried with updated transformers (version 4.29.0dev0), and the result is:
- When using `BitsAndBytesConfig` mentioned above, the model is able to generate normal text, but it still produces NaN's in layer 39
- When `load_in_8bit` is set to `False` for all layers, no NaN value is produced.
So, in conclusion, it may be a problem with the 8-bit quantization, instead of a problem with the converting script.

Hi @RiverGao 
This is probably related to the fact that you are using a V100. Could you share with us your `bitsandbytes` version? Or alternatively update `bitsandbytes` and let us know if you still face the issue

```bash
pip install --upgrade bitsandbytes
```

@younesbelkada Thank you for your helpful information, I updated `bitsandbytes` from 0.37.2 to 0.38.1, and LLaMA 13B produced no NaN attention when loaded with `load_in_8bit=True`. However, when applied with [Alpaca-LoRA tuned weights](https://github.com/tloen/alpaca-lora), the model produces NaN's in layer 39 again if it is loaded in int8.

Hi @RiverGao 
Interesting, how you apply the LoRA weights on that model? Can you share a repro script?

@younesbelkada I applied the LoRA weights following using `export_hf_checkpoint.py` from [their repo](https://github.com/tloen/alpaca-lora) by
`$ BASE_MODEL=path/to/converted/llama python export_hf_checkpoint.py`:
```
# export_hf_checkpoint.py

import os

import torch
import transformers
from peft import PeftModel
from transformers import LlamaForCausalLM, LlamaTokenizer  # noqa: F402

BASE_MODEL = os.environ.get("BASE_MODEL", None)
assert (
    BASE_MODEL
), "Please specify a value for BASE_MODEL environment variable, e.g. `export BASE_MODEL=huggyllama/llama-7b`"  # noqa: E501

tokenizer = LlamaTokenizer.from_pretrained(BASE_MODEL)

base_model = LlamaForCausalLM.from_pretrained(
    BASE_MODEL,
    load_in_8bit=False,
    torch_dtype=torch.float16,
    device_map={"": "cpu"},
)

first_weight = base_model.model.layers[0].self_attn.q_proj.weight
first_weight_old = first_weight.clone()

lora_model = PeftModel.from_pretrained(
    base_model,
    "Angainor/alpaca-lora-13b",
    device_map={"": "cpu"},
    torch_dtype=torch.float16,
)

lora_weight = lora_model.base_model.model.model.layers[
    0
].self_attn.q_proj.weight

assert torch.allclose(first_weight_old, first_weight)

# merge weights - new merging method from peft
lora_model = lora_model.merge_and_unload()

lora_model.train(False)

# did we do anything?
assert not torch.allclose(first_weight_old, first_weight)

lora_model_sd = lora_model.state_dict()
deloreanized_sd = {
    k.replace("base_model.model.", ""): v
    for k, v in lora_model_sd.items()
    if "lora" not in k
}

LlamaForCausalLM.save_pretrained(
    base_model, "../alpaca-hf/13B", state_dict=deloreanized_sd, max_shard_size="400MB"
)
```

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.