Generate: SinkCache can handle iterative prompts
Author: gante
Pull request: yes

# What does this PR do?

Fixes the case where `SinkCache` is used in a chat bot, receiving new prompts after giving an answer. Fix developed with @tomaarsen 

Here's an example of a script that works after this PR:
```py
from transformers import AutoTokenizer, SinkCache, AutoModelForCausalLM, TextStreamer
import torch
from datasets import load_dataset

# Loading the model & tokenizer
model_id = "HuggingFaceH4/zephyr-7b-beta"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Loading the prompts to simulate user interactions
prompt_dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train")
prompts = [prompt for prompts in prompt_dataset["prompt"] for prompt in prompts]

# Prepare generation settings
cache = SinkCache(window_length=1024, num_sink_tokens=4)
streamer = TextStreamer(tokenizer)

input_ids = torch.tensor([], device=model.device, dtype=torch.int)
for prompt in prompts:
    # Tokenize the prompt with the correct chat template
    chat = [{"role": "user", "content": prompt}]
    input_ids = torch.cat((input_ids, tokenizer.apply_chat_template(chat, return_tensors="pt", add_generation_prompt=True).to(model.device)), dim=1)
    # input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(model.device)

    # Perform the generation
    gen_out = model.generate(input_ids, do_sample=False, max_new_tokens=100, past_key_values=cache, use_cache=True, streamer=streamer)

    # input_ids = torch.cat((input_ids, gen_out), dim=1)
    input_ids = gen_out

    # If desired, decode the output from this prompt
    decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
```

Comments:
Some additional results @gante 

![image](https://github.com/huggingface/transformers/assets/37621491/36b8ef6c-c5fa-418a-acc4-6193c4abc2c3)

* `transformers_old` refers to before #27700 (thanks @kevinhu for the change!)
* `transformers_56be5e80` refers to commit 56be5e80, i.e. near main
* `transformers_attn_sink_1024_4_pr-27907` refers to this PR, with `SinkCache(1024, 4)`

This experiment is with calling the models with individual tokens exclusively. Using `SinkCache` makes the memory usage linear at a very low cost in perplexity.

---

Additionally, I'm doing a test with calling Mistral-7B in this case with this PR using multiple tokens at once. I took indices 0 to 63k of a book from pg19, only kept 20% of all indices, and then fed the model with the tokens between subsequent indices. The running cache is also included. The NLL loss is then converted to perplexity.
**Note:** We can't compare this with the perplexities from the previous graph: we should only try and observe whether the model eventually increases in perplexity.

![image](https://github.com/huggingface/transformers/assets/37621491/6606dbbd-079c-4028-81df-2f9195a68a91)

The same script crashes on `main`. In this test, the perplexity stays constant, which is good. 
Edit: I have now continued with more tests:
![image](https://github.com/huggingface/transformers/assets/37621491/a9f47b15-55ae-4564-a1f7-82701ff21936)


* `transformers_multi_attn_sink_1024_4_pr-27907_left_cache`: This PR, with `SinkCache(1024, 4)`.
* `transformers_multi_attn_sink_1024_4_pr-27907_right_cache`: This PR, with `SinkCache(1024, 4)` & the change @ArthurZucker proposed regarding slicing the cache from the right.
* `transformers_multi_b31905d1`: `main` without a special SinkCache.

The perplexity diverges quite heavily between the SinkCache and non, which is not ideal. Perhaps this is indicative of some error/bug, or perhaps not. It's a bit hard to tell. Beyond that, the left and right cache implementations behave identically (unless I made some measuring mistake), which is a bit odd. I don't have 100% confidence in this fix anymore I'm afraid.

- Tom Aarsen

@ArthurZucker as per your suggestion, I've reworked the PR to avoid post hoc `attention_mask` slicing -- there is a new function to get the usable cache length, and that function is used to obtain `kv_seq_len`

@tomaarsen the rework seems to have resulted in a qualitative result upgrade (e.g. see the test case), so I suspect that I've inadvertently fixed a bug 👀 Would you be able to rerun your benchmarks for `SinkCache`? 

@gante I get `ValueError: Attention weights should be of size (1, 32, 5, 1027), but is torch.Size([1, 32, 5, 1024])` upon running my modified script. Do you get an error like this with the multi-step generation script from your PR?

@tomaarsen no, the script in the PR header runs endlessly without issues 🤔 LMK if you can find a reproducer

I have the same, that script works fine. Hmmm

Got a reproducer: change `max_new_tokens` in the script above to 512 👀  having a look!

@tomaarsen should be fixed

@gante Works great now!
![image](https://github.com/huggingface/transformers/assets/37621491/3ec53db0-49eb-483f-8611-ac01312c6a54)

Red is the baseline, I can only run it to about ~15k seq length until my PC completely freezes.

🙌 

@tomaarsen @gante @ArthurZucker  Hi, thanks for the work. Do you know how we can reproduce the benchmark test with the latest transformers repo implementation? Are there any scripts? Really appreciated!

@YJHMITWEB you can adapt the following scripts to your need. Note -- Llama + attention sinks is probably not working on the latest version, Mistral should be. We are reworking how cache works, things may be a bit rough :) Worst case scenario, use `v4.36`.

<details>
  <summary>Script to gather perplexity</summary>

  ```py
"""
Adapted from https://github.com/mit-han-lab/streaming-llm

Note: Although this script measures latency, it is not optimized whatsoever!
The latency is only tracked to see the impact of speed over time.

Usage:

python benchmark/perplexity.py --experiment attention_sinks
python benchmark/perplexity.py --experiment transformers
python benchmark/perplexity.py --experiment windowed
"""


import argparse
import itertools
import time
from collections import defaultdict
from pathlib import Path
from typing import Optional

import pandas as pd
import torch
from datasets import load_dataset
from torch.nn import CrossEntropyLoss
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM

def compute_perplexity(
    model,
    tokenizer,
    dataset,
    experiment: str,
    output_dir: str = "outputs",
    data_column: str = "text",
    num_samples: int = 1,
    num_tokens: Optional[int] = None,
    overwrite: bool = False,
) -> None:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / f"{experiment}.csv"

    if output_file.exists() and not overwrite:
        raise ValueError(
            f"The {output_file!r} output file already exists - if you really want to override it, then use `--overwrite`."
        )

    logs = defaultdict(list)
    loss_fn = CrossEntropyLoss(reduction="none")
    past_key_values = None
    num_processed_tokens = 0
    for text in itertools.islice(dataset, num_samples):
        encodings = tokenizer(text[data_column], return_tensors="pt")

        seq_len = encodings.input_ids.size(1)
        print(f"sequence length: {seq_len}")
        pbar = tqdm(range(0, seq_len - 1))

        for idx in pbar:
            start_t = time.time()
            input_ids = encodings.input_ids[:, idx : idx + 1].to(model.device)
            with torch.no_grad():
                outputs = model(input_ids, past_key_values=past_key_values, use_cache=True)
                logits = outputs.logits.view(-1, model.config.vocab_size)
                past_key_values = outputs.past_key_values
                label = encodings.input_ids[:, idx + 1 : idx + 2].to(logits.device).view(-1)
                neg_log_likelihood = loss_fn(logits, label)
                perplexity = neg_log_likelihood.exp()
            pbar.set_description(f"nll: {neg_log_likelihood.item():>5.2f}, ppl: {perplexity.item():>8.2f}")

            # Store data and save every 10 tokens
            logs["input_length"].append(idx + 1)
            logs["nll"].append(neg_log_likelihood.item())
            logs["ppl"].append(perplexity.item())
            logs["overall_ppl"].append(torch.tensor(logs["nll"]).mean().exp().item())
            logs["cuda_vram_allocated"].append(torch.cuda.memory_allocated(0) / 1024 / 1024 / 1024)  # in GB
            logs["latency"].append(time.time() - start_t)
            if num_processed_tokens % 10 == 0:
                try:
                    pd.DataFrame(logs).to_csv(output_file, index=False)
                except KeyboardInterrupt as ex:
                    # If there's a Keyboard Interrupt, still write the file, and then stop
                    pd.DataFrame(logs).to_csv(output_file, index=False)
                    raise ex

            num_processed_tokens += 1
            if num_tokens and num_processed_tokens >= num_tokens:
                return


def main():
    parser = argparse.ArgumentParser()
    # How to call this experiment?
    parser.add_argument(
        "--experiment", type=str, default="main"
    )

    # Model args
    # parser.add_argument("--model_name_or_path", type=str, default="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
    parser.add_argument("--model_name_or_path", type=str, default="meta-llama/Llama-2-7b-hf")
    parser.add_argument("--revision", type=str, default="main")
    parser.add_argument("--trust_remote_code", action="store_true")

    # Dataset args
    parser.add_argument("--dataset_name", type=str, default="emozilla/pg19-test")
    parser.add_argument("--data_column", type=str, default="text")
    parser.add_argument("--task", type=str, default=None)
    parser.add_argument("--split", type=str, default="test", choices=["validation", "test"])
    # parser.add_argument("--num_samples", type=int, default=1)
    parser.add_argument("--num_tokens", type=int, default=5000)
    parser.add_argument("--dtype", type=str, default="fp16")

    # Where to log
    parser.add_argument("--output_dir", type=str, default="/home/joao/joao_scripts/perplexity/outputs")
    parser.add_argument("--overwrite", action="store_true")

    args = parser.parse_args()

    if args.dtype == "fp16":
        dtype = torch.float16
    elif args.dtype == "fp32":
        dtype = torch.float32
    elif args.dtype == "bf16":
        dtype = torch.bfloat16
    else:
        raise ValueError(f"Unknown dtype: {args.dtype}")

    model = AutoModelForCausalLM.from_pretrained(
        args.model_name_or_path,
        revision=args.revision,
        trust_remote_code=bool(args.trust_remote_code),
        attn_implementation="eager",
        torch_dtype=dtype,
        device_map="auto",
    )
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=bool(args.trust_remote_code))

    # Set up the dataset
    dataset = load_dataset(args.dataset_name, args.task, split=args.split, streaming=True)

    compute_perplexity(
        model,
        tokenizer,
        dataset,
        args.experiment,
        output_dir=args.output_dir,
        data_column=args.data_column,
        num_samples=1,  # <- No support for more than one instance now
        num_tokens=args.num_tokens,
        overwrite=args.overwrite,
    )


if __name__ == "__main__":
    main()
  ```

</details>

<details>
  <summary>Script to plot perplexity (and other metrics)</summary>

  ```py
"""
First run `perplexity.py` to generate one or more `csv` files.
This script can plot those csv files.

Usage:
python benchmark/plot_perplexity.py
python benchmark/plot_perplexity.py --features perplexity latency --title "Log perplexity & latency of Llama 2 7B as a function of input lengths"
"""

import argparse
from pathlib import Path
from typing import List, Optional

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

FEATURE_DF_MAP = {
    "perplexity": "overall_ppl",
    "vram": "cuda_vram_allocated",
    "latency": "latency",
}
FEATURE_STYLE_MAP = {
    "perplexity": "-",
    "vram": "--",
    "latency": ":",
}
FEATURE_LABEL_MAP = {
    "perplexity": "Perplexity (log), lower is better",
    "vram": "CUDA VRAM Usage (GB), lower is better",
    "latency": "Time per token (sec), lower is better",
}


def plot(
    features: List[str],
    output_dir: str = "outputs",
    title: Optional[str] = None,
    perplexity_limit: Optional[float] = None,
    skip_first: int = 100,
):
    output_dir = Path(output_dir)

    fig, ax = plt.subplots()
    ax.set_xlabel("Input Sequence Length")

    for feature_i, feature in enumerate(features):
        # If we already plotted on this ax, make a new one
        if feature_i:
            ax = ax.twinx()

        for file in output_dir.glob("*.csv"):
            experiment = file.stem
            df = pd.read_csv(file)
            X = df["input_length"][skip_first:]
            Y = df[FEATURE_DF_MAP[feature]][skip_first:]
            if feature == "perplexity":
                Y = np.log(Y)
            if feature == "latency":
                poly = np.polyfit(X, Y, 20)
                poly_y = np.poly1d(poly)(X)
                ax.plot(X, poly_y, FEATURE_STYLE_MAP[feature], label=f"{experiment} {feature}")
            else:
                ax.plot(X, Y, FEATURE_STYLE_MAP[feature], label=f"{experiment} {feature}")

        ax.set_ylabel(FEATURE_LABEL_MAP[feature])
        if perplexity_limit and feature == "perplexity":
            ax.set_ylim(top=min(ax.get_ylim()[1], perplexity_limit))

        ax.legend(loc=[1, 2, 7][feature_i])  # upper right, upper left, center right

    ax.set_title(title.replace("\\n", "\n") if title else "Log perplexity as a function of input lengths")
    fig.tight_layout()

    return fig


def main():
    parser = argparse.ArgumentParser()
    # Where csv files have been logged
    parser.add_argument("--output_dir", type=str, default="/home/joao/joao_scripts/perplexity/outputs")
    parser.add_argument(
        "--features", choices=["perplexity", "vram", "latency"], nargs="+", default=["perplexity", "vram"]
    )
    parser.add_argument("--title", type=str, default=None)
    parser.add_argument("--log_perplexity_limit", type=float, default=5.0)
    # Perplexity starts a bit unstable, so we skip the start
    parser.add_argument("--skip_first", type=int, default=100)

    args = parser.parse_args()

    figure = plot(
        args.features,
        output_dir=args.output_dir,
        title=args.title,
        perplexity_limit=args.log_perplexity_limit,
        skip_first=args.skip_first,
    )

    # Add your own code here if you'd like to change the figure
    features = "_".join(args.features)
    save_path = f"/home/joao/joao_scripts/perplexity/outputs/plot_{features}.png"
    plt.savefig(save_path, dpi=600)
    print(f"plot saved to {save_path}")


if __name__ == "__main__":
    main()

  ```

</details>

Credits to @tomaarsen, who shared these with me :)

@gante @tomaarsen Thanks!