Llama runs very slow in 4.38.2
Author: lhkhiem28
Pull request: no

### System Info

- `transformers` version: 4.38.2
- Platform: Linux-4.18.0-513.5.1.el8_9.x86_64-x86_64-with-glibc2.28
- Python version: 3.10.13
- Huggingface_hub version: 0.20.2
- Safetensors version: 0.4.1
- Accelerate version: 0.26.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+cu121 (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?: NVIDIA A100 GPUs
- Using distributed or parallel set-up in script?: "auto"

### Who can help?

@ArthurZucker @gante

### 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

pip install transformers==4.38.2

### Expected behavior

When I use transformers==4.36.2, Llama 2 functions well. However, after updating to 4.38.2, Llama runs very much slower (20 times). No error, no warning, just very slow running.

Please help to check the issue on a timely basis! Thanks!

Comments:
Hi @lhkhiem28, thanks for opening an issue! 

Please provide a minimal code snippet that can be used for benchmarking as well as exact running times to compare the two versions

Hi @amyeroberts , below are the code that I am using

inference.py
```
import os
import torch
import gc
from tqdm import tqdm
from torch.utils.data import DataLoader

from source.utils.help_funcs import seed_everything
from source.config import parse_args_llama
from source.models import load_model, llama_model_path
from source.datasets import load_dataset
from source.utils.evaluation import eval_funcs
from source.utils.help_funcs import collate_fn

import warnings; warnings.filterwarnings("ignore")

def main(args):

    seed = args.seed

    seed_everything(seed=seed)
    # print(args)

    # Step 2: Build Dataset
    test_dataset = load_dataset[args.dataset](data = args.data, split = 'test', task = args.task)
    test_loader = DataLoader(test_dataset, batch_size=args.eval_batch_size, drop_last=False, pin_memory=True, shuffle=False, collate_fn=collate_fn)

    # Step 3: Build Model
    args.llm_model_path = llama_model_path[args.llm_model_name]
    model = load_model[args.model_name](args=args)

    # Step 4. Evaluating
    model.eval()
    eval_output = []
    progress_bar_test = tqdm(range(len(test_loader)))
    for _, batch in enumerate(test_loader):
        with torch.no_grad():
            output = model.inference(batch)
            eval_output.append(output)

        progress_bar_test.update(1)

    # Step 5. Post-processing & Evaluating
    os.makedirs(f'{args.output_dir}/{args.dataset}', exist_ok=True)
    path = f'{args.output_dir}/{args.dataset}/{args.model_name}_{args.llm_model_name}_{args.data}_{args.task}.csv'
    if "prediction" in args.dataset:
        acc, f1 = eval_funcs[args.dataset](eval_output, path)
        print('Test Accuracy: {:.4f}'.format(acc))
        print('Test F1: {:.4f}'.format(f1))
    if "generation" in args.dataset:
        scores = eval_funcs[args.dataset](eval_output, path, model.tokenizer, args.max_new_tokens)
        print("BLEU-2: {:.4f} | BLEU-4: {:.4f} | ROUGE-1: {:.4f} | ROUGE-2: {:.4f} | ROUGE-L: {:.4f} | METEOR: {:.4f}".format(
            *scores
        ))

if __name__ == "__main__":

    args = parse_args_llama()

    main(args)
    # torch.cuda.empty_cache()
    # torch.cuda.reset_max_memory_allocated()
    # gc.collect()
```
baseline_llm.py
```
import contextlib
import torch
from torch.cuda.amp import autocast as autocast
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_int8_training,
)

# self.BOS = '

'
# self.EOS_USER = '
'
# self.EOS = '
'
# self.IGNORE_INDEX = -100

class BaselineLLM(torch.nn.Module):

    def __init__(
        self,
        args,
        **kwargs
    ):
        super().__init__()
        self.max_new_tokens = args.max_new_tokens
        if "llama" in args.llm_model_path:
            self.BOS = '

'
            self.EOS_USER = '
'
            self.EOS = '
'
            self.IGNORE_INDEX = -100
        if "gemma" in args.llm_model_path:
            self.BOS = '<bos><start_of_turn>user'
            self.EOS_USER = '<end_of_turn><start_of_turn>model'
            self.EOS = '
'
            self.IGNORE_INDEX = -100

        print(f'Loading {args.llm_model_path}')
        kwargs = {
            "max_memory": {0: '80GiB', 1: '80GiB'},
            "device_map": "auto",
            "revision": "main",
        }
        self.tokenizer = AutoTokenizer.from_pretrained(args.llm_model_path, use_fast=False, revision=kwargs["revision"])
        self.tokenizer.pad_token_id = 0
        self.tokenizer.padding_side = 'left'

        model = AutoModelForCausalLM.from_pretrained(
            args.llm_model_path,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            **kwargs
        )

        if args.llm_frozen == 'True':
            print(f"{args.llm_model_path} has been frozen!")
            for name, param in model.named_parameters():
                param.requires_grad = False
        else:
            print(f"{args.llm_model_path} has been factorized for training!")
            model = prepare_model_for_int8_training(model)

            lora_r: int = 8
            lora_alpha: int = 16
            lora_dropout: float = 0.05
            lora_target_modules = [
                "q_proj",
                "v_proj",
            ]
            config = LoraConfig(
                r=lora_r,
                lora_alpha=lora_alpha,
                target_modules=lora_target_modules,
                lora_dropout=lora_dropout,
                bias="none",
                task_type="CAUSAL_LM",
            )
            model = get_peft_model(model, config)

        self.model = model

        self.word_embedding = self.model.model.get_input_embeddings()

    @property
    def device(self):
        return list(self.parameters())[0].device

    def maybe_autocast(self, dtype=torch.bfloat16):
        # if on cpu, don't use autocast
        # if on gpu, use autocast with dtype if provided, otherwise use torch.float16
        enable_autocast = self.device != torch.device("cpu")

        if enable_autocast:
            return torch.cuda.amp.autocast(dtype=dtype)
        else:
            return contextlib.nullcontext()

    def forward(self, samples):
        # encode description, questions and labels
        questions = self.tokenizer(samples["question"], add_special_tokens=False)
        labels = self.tokenizer(samples["label"], add_special_tokens=False)

        # encode special tokens
        eos_tokens = self.tokenizer(self.EOS, add_special_tokens=False)
        eos_user_tokens = self.tokenizer(self.EOS_USER, add_special_tokens=False)
        bos_embeds = self.word_embedding(self.tokenizer(self.BOS, add_special_tokens=False, return_tensors='pt').input_ids[0])
        pad_embeds = self.word_embedding(torch.tensor(self.tokenizer.pad_token_id)).unsqueeze(0)

        batch_size = len(samples['id'])
        batch_inputs_embeds = []
        batch_attention_mask = []
        batch_label_input_ids = []
        for i in range(batch_size):
            # Add bos & eos token
            label_input_ids = labels.input_ids[i][:self.max_new_tokens] + eos_tokens.input_ids
            input_ids = questions.input_ids[i] + eos_user_tokens.input_ids + label_input_ids
            inputs_embeds = self.word_embedding(torch.tensor(input_ids).to(self.model.device))
            inputs_embeds = torch.cat([bos_embeds, inputs_embeds], dim=0)

            batch_inputs_embeds.append(inputs_embeds)
            batch_attention_mask.append([1] * inputs_embeds.shape[0])
            label_input_ids = [self.IGNORE_INDEX] * (inputs_embeds.shape[0]-len(label_input_ids)) + label_input_ids
            batch_label_input_ids.append(label_input_ids)

        # pad inputs_embeds
        max_length = max([x.shape[0] for x in batch_inputs_embeds])
        for i in range(batch_size):
            pad_length = max_length-batch_inputs_embeds[i].shape[0]
            batch_inputs_embeds[i] = torch.cat([pad_embeds.repeat(pad_length, 1), batch_inputs_embeds[i]])
            batch_attention_mask[i] = [0]*pad_length + batch_attention_mask[i]
            batch_label_input_ids[i] = [self.IGNORE_INDEX] * pad_length+batch_label_input_ids[i]

        inputs_embeds = torch.stack(batch_inputs_embeds, dim=0).to(self.model.device)
        attention_mask = torch.tensor(batch_attention_mask).to(self.model.device)
        label_input_ids = torch.tensor(batch_label_input_ids).to(self.model.device)

        with self.maybe_autocast():
            outputs = self.model(
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                return_dict=True,
                labels=label_input_ids,
            )

        return outputs.loss

    def inference(self, samples):

        # encode description and questions
        questions = self.tokenizer(samples["question"], add_special_tokens=False)

        # encode special tokens
        eos_user_tokens = self.tokenizer(self.EOS_USER, add_special_tokens=False)
        bos_embeds = self.word_embedding(self.tokenizer(self.BOS, add_special_tokens=False, return_tensors='pt').input_ids[0])
        pad_embeds = self.word_embedding(torch.tensor(self.tokenizer.pad_token_id)).unsqueeze(0)

        batch_size = len(samples['id'])
        batch_inputs_embeds = []
        batch_attention_mask = []
        for i in range(batch_size):
            # Add bos & eos token
            input_ids = questions.input_ids[i] + eos_user_tokens.input_ids
            inputs_embeds = self.word_embedding(torch.tensor(input_ids).to(self.model.device))
            inputs_embeds = torch.cat([bos_embeds, inputs_embeds], dim=0)
            batch_inputs_embeds.append(inputs_embeds)
            batch_attention_mask.append([1] * inputs_embeds.shape[0])

        # pad inputs_embeds
        max_length = max([x.shape[0] for x in batch_inputs_embeds])
        for i in range(batch_size):
            pad_length = max_length - batch_inputs_embeds[i].shape[0]
            batch_inputs_embeds[i] = torch.cat([pad_embeds.repeat(pad_length, 1), batch_inputs_embeds[i]])
            batch_attention_mask[i] = [0]*pad_length + batch_attention_mask[i]

        inputs_embeds = torch.stack(batch_inputs_embeds, dim=0).to(self.model.device)
        attention_mask = torch.tensor(batch_attention_mask).to(self.model.device)

        with self.maybe_autocast():
            outputs = self.model.generate(
                inputs_embeds=inputs_embeds,
                max_new_tokens=self.max_new_tokens,
                attention_mask=attention_mask,
                # do_sample=True,
                use_cache=True  # IMPORTANT!
            )
        pred = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)

        return {'id': samples['id'],
                'pred': [p.strip() for p in pred],
                'label': samples['label'],
                'question': samples['question'],
        }

    def print_trainable_params(self):
        trainable_params = 0
        all_param = 0

        for _, param in self.named_parameters():
            num_params = param.numel()

            all_param += num_params
            if param.requires_grad:
                trainable_params += num_params

        return trainable_params, all_param
```

with transformers 4.36.2, the running time is ~5s/iter
with transformers 4.38.2 (and the same code), the running time is ~100s/iter

Follow-up information:
with transformers-4.39.0.dev0 (and the same code), the running time is ~40s/iter

Hope you can help to identify the issue and make Llama 2 in 4.38.2 run as fast as 4.36.2. Thanks a lot!

@lhkhiem28 This isn't a benchmarking script, nor is it minimal. 

For example: 
* `BaselineLLM` doesn't appear to be used at all in `inference.py` obviously - there's a call to model.inference, but it's not clear how this model is loaded
* the model `BaselineLLM` has many custom operations, which are outside of the world of transformers
* we don't have access to `collate_fn` so can't prepare the dataset
* "`source`" as a library is completely unclear
* there's a bunch of commented-out code which should just be removed
* we don't have access to `parse_args_llama` so don't know that values of the args, in particular, we cannot just copy-paste this code and run it ourselves, nor do we have access to the dataset, model or tokenizer. 

We get many issues and pull requests per day, in order for us to be able to address them in a timely manner, we need you to help us help you. Any issues which would required lots of additional work on our part to get working and try to infer the environment to reproduce will not be prioritised and may not be addressed at all. 

The script should:
* Be minimal i.e. the fewest lines of code necessary to replicate
* Be reproducible - we can copy-paste the code and run it out of the box without having to modify it

We understand you may not be able to share checkpoints, but it should be possible to share equivalent public checkpoints to reproduce the issue

@lhkhiem28 There is a similar issue open, let's focus all discussion on the same issue: https://github.com/huggingface/transformers/issues/29412 🤗