Extend existing LLMs way beyond the original training length with constant memory usage, and without retraining
__version__ = "0.2.4.dev"

from transformers import AutoTokenizer

from .attention_sink_kv_cache import AttentionSinkKVCache
from .models import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    FalconForCausalLM,
    FalconForQuestionAnswering,
    FalconForSequenceClassification,
    FalconForTokenClassification,
    FalconModel,
    FalconPreTrainedModel,
    GPTJForCausalLM,
    GPTJForQuestionAnswering,
    GPTJForSequenceClassification,
    GPTJModel,
    GPTJPreTrainedModel,
    GPTNeoXForCausalLM,
    GPTNeoXForQuestionAnswering,
    GPTNeoXForSequenceClassification,
    GPTNeoXForTokenClassification,
    GPTNeoXModel,
    GPTNeoXPreTrainedModel,
    LlamaForCausalLM,
    LlamaForSequenceClassification,
    LlamaModel,
    MistralForCausalLM,
    MistralForSequenceClassification,
    MistralModel,
    MptForCausalLM,
    MptForQuestionAnswering,
    MptForSequenceClassification,
    MptForTokenClassification,
    MptModel,
    MptPreTrainedModel,
)

attention_sinks/__init__.py
tomaarsen/attention_sinks
"""
Adapted from https://github.com/mit-han-lab/streaming-llm
"""

from dataclasses import dataclass

import torch


def slice1d(x, start, end):
    return x[:, start:end, ...]


def slice2d(x, start, end):
    return x[:, :, start:end, ...]


def slice3d(x, start, end):
    return x[:, :, :, start:end, ...]


DIM_TO_SLICE = {
    1: slice1d,
    2: slice2d,
    3: slice3d,
}


@dataclass
class AttentionSinkKVCache:
    attention_sink_size: int = 4
    attention_sink_window_size: int = 1020
    k_seq_dim: int = 2
    v_seq_dim: int = 2

    def __post_init__(self):
        self.cache_size = self.attention_sink_size + self.attention_sink_window_size
        self.k_slice = DIM_TO_SLICE[self.k_seq_dim]
        self.v_slice = DIM_TO_SLICE[self.v_seq_dim]

    def __call__(self, past_key_values):
        if past_key_values is None:
            return None
        seq_len = past_key_values[0][0].size(self.k_seq_dim)
        if seq_len <= self.cache_size:
            return past_key_values
        return [
            [
                torch.cat(
                    [
                        self.k_slice(k, 0, self.attention_sink_size),
                        self.k_slice(k, seq_len - self.attention_sink_window_size, seq_len),
                    ],
                    dim=self.k_seq_dim,
                ),
                torch.cat(
                    [
                        self.v_slice(v, 0, self.attention_sink_size),
                        self.v_slice(v, seq_len - self.attention_sink_window_size, seq_len),
                    ],
                    dim=self.v_seq_dim,
                ),
            ]
            for k, v in past_key_values
        ]

    def evict_for_space(self, past_key_values, num_coming):
        if past_key_values is None:
            return None
        seq_len = past_key_values[0][0].size(self.k_seq_dim)
        if seq_len + num_coming <= self.cache_size:
            return past_key_values
        return [
            [
                torch.cat(
                    [
                        self.k_slice(k, 0, self.attention_sink_size),
                        self.k_slice(
                            k,
                            seq_len - self.attention_sink_window_size + num_coming,
                            seq_len,
                        ),
                    ],
                    dim=self.k_seq_dim,
                ),
                torch.cat(
                    [
                        self.v_slice(v, 0, self.attention_sink_size),
                        self.v_slice(
                            v,
                            seq_len - self.attention_sink_window_size + num_coming,
                            seq_len,
                        ),
                    ],
                    dim=self.v_seq_dim,
                ),
            ]
            for k, v in past_key_values
        ]

    def evict_range(self, past_key_values, start, end):
        if past_key_values is None:
            return None
        seq_len = past_key_values[0][0].size(self.k_seq_dim)
        assert start <= end and end <= seq_len
        return [
            [
                torch.cat(
                    [
                        self.k_slice(k, 0, start),
                        self.k_slice(k, end, seq_len),
                    ],
                    dim=self.k_seq_dim,
                ),
                torch.cat(
                    [
                        self.v_slice(v, 0, start),
                        self.v_slice(v, end, seq_len),
                    ],
                    dim=self.v_seq_dim,
                ),
            ]
            for k, v in past_key_values
        ]

attention_sinks/attention_sink_kv_cache.py
tomaarsen/attention_sinks
from typing import Any, Dict

import torch
from transformers.generation.utils import GenerationMixin as TGenerationMixin
from transformers.utils import ModelOutput


class GenerationMixin(TGenerationMixin):
    """
    This GenerationMixin must be overridden to prevent the `attention_mask`
    from extending beyond the window size.
    """

    def _update_model_kwargs_for_generation(
        self,
        outputs: ModelOutput,
        model_kwargs: Dict[str, Any],
        is_encoder_decoder: bool = False,
        standardize_cache_format: bool = False,
    ) -> Dict[str, Any]:
        # update past_key_values
        model_kwargs["past_key_values"] = self._extract_past_from_model_output(
            outputs, standardize_cache_format=standardize_cache_format
        )
        if getattr(outputs, "state", None) is not None:
            model_kwargs["state"] = outputs.state

        # update token_type_ids with last value
        if "token_type_ids" in model_kwargs:
            token_type_ids = model_kwargs["token_type_ids"]
            model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1)

        if not is_encoder_decoder:
            # update attention mask
            if "attention_mask" in model_kwargs:
                attention_mask = model_kwargs["attention_mask"]
                # Only this `if`-statement is changed, it's required to stop the attention_mask from extending itself too far
                if model_kwargs["attention_mask"].size(-1) == model_kwargs["past_key_values"][0][0].size(2):
                    model_kwargs["attention_mask"] = torch.cat(
                        [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
                    )
        else:
            # update decoder attention mask
            if "decoder_attention_mask" in model_kwargs:
                decoder_attention_mask = model_kwargs["decoder_attention_mask"]
                model_kwargs["decoder_attention_mask"] = torch.cat(
                    [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))],
                    dim=-1,
                )

        return model_kwargs

attention_sinks/generation/utils.py
tomaarsen/attention_sinks
import types
from typing import Callable, Optional

from transformers import PreTrainedModel
from transformers.utils import logging

from attention_sinks.attention_sink_kv_cache import AttentionSinkKVCache
from attention_sinks.generation.utils import GenerationMixin

logger = logging.get_logger(__name__)

MODEL_NAME_MAPPING = {
    "llama": "LlamaModel",
    "falcon": "FalconModel",
    "mpt": "MptModel",
    "gpt_neox": "GPTNeoXModel",
    "gptj": "GPTJModel",
    "mistral": "MistralModel",
}
ATTENTION_NAME_MAPPING = {
    "llama": "LlamaAttention",
    "falcon": "FalconAttention",
    "mpt": "MptAttention",
    "gpt_neox": "GPTNeoXAttention",
    "gptj": "GPTJAttention",
    "mistral": "MistralAttention",
}
KV_DIM_MAPPING = {
    "llama": (2, 2),
    "falcon": (2, 2),
    "mpt": (2, 2),
    "gpt_neox": (2, 2),
    "gptj": (2, 2),
    "mistral": (2, 2),
}


class InjectAttentionSinksMixin(GenerationMixin):
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        # Separate Attention Sink kwargs from regular kwargs
        attention_sink_kwargs = {key: value for key, value in kwargs.items() if key.startswith("attention_sink")}
        for key in attention_sink_kwargs:
            kwargs.pop(key)

        model = super().from_pretrained(
            pretrained_model_name_or_path,
            *model_args,
            **kwargs,
        )
        model_type = model.config.model_type
        if model_type not in MODEL_NAME_MAPPING:
            raise NotImplementedError(
                f"`attention_sinks` does not support models with the `{model_type}` architecture at this time."
            )

        # Enable position shifting attention
        call_count = cls._inject_pos_shift_attention(model)
        if call_count is not None:
            logger.warn(
                f"[Attention Sinks] Injected Position Shifting into {call_count} attention class{'es' if call_count != 1 else ''}."
            )

        # Inject the Attention Sink KV Cache to the model
        call_count = cls._inject_attention_sink_kv_cache(model, **attention_sink_kwargs)
        logger.warn(
            f"[Attention Sinks] Injected Attention Sink KV Cache into {call_count} model class{'es' if call_count != 1 else ''}."
        )

        return model

    @classmethod
    def _inject_pos_shift_attention(cls, model: PreTrainedModel) -> Optional[int]:
        model_type = model.config.model_type

        from attention_sinks.models import (
            falcon_pos_shift_attention_forward,
            gpt_neox_pos_shift_attention_forward,
            gptj_pos_shift_attention_forward,
            llama_pos_shift_attention_forward,
            mistral_pos_shift_attention_forward,
        )

        ATTENTION_FORWARD_MAPPING = {
            "llama": llama_pos_shift_attention_forward,
            "falcon": falcon_pos_shift_attention_forward,
            "mpt": None,
            "gpt_neox": gpt_neox_pos_shift_attention_forward,
            "gptj": gptj_pos_shift_attention_forward,
            "mistral": mistral_pos_shift_attention_forward,
        }

        # Not all models require updated attention forwards
        if ATTENTION_FORWARD_MAPPING[model_type] is None:
            return

        def overwrite_forward(module) -> None:
            module.forward = types.MethodType(ATTENTION_FORWARD_MAPPING[model_type], module)

        return cls._call_modules_by_name(model, ATTENTION_NAME_MAPPING[model_type], overwrite_forward)

    @classmethod
    def _inject_attention_sink_kv_cache(cls, model, **attention_sink_kwargs) -> int:
        model_type = model.config.model_type
        attention_sink_kwargs["k_seq_dim"], attention_sink_kwargs["v_seq_dim"] = KV_DIM_MAPPING[model_type]

        def overwrite_forward(module):
            # Create the new cache
            module.attention_sink_kv_cache = AttentionSinkKVCache(**attention_sink_kwargs)

            # Keep track of the old forward method, we need it in the wrapped one
            old_forward = module.forward

            # Wrap the forward by overriding the past_key_values using the cache
            def wrapped_forward(self, *args, **kwargs):
                outputs = old_forward(*args, **kwargs)
                outputs.past_key_values = self.attention_sink_kv_cache(outputs.past_key_values)
                return outputs

            module.forward = types.MethodType(wrapped_forward, module)

        return cls._call_modules_by_name(model, MODEL_NAME_MAPPING[model_type], overwrite_forward)

    @classmethod
    def _call_modules_by_name(cls, module, target_name: str, func: Callable) -> int:
        if module.__class__.__name__ == target_name:
            func(module)
            return 1

        return sum(cls._call_modules_by_name(module, target_name, func) for module in module.children())

attention_sinks/inject_mixin.py
tomaarsen/attention_sinks
from .auto import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
)
from .falcon import (
    FalconForCausalLM,
    FalconForQuestionAnswering,
    FalconForSequenceClassification,
    FalconForTokenClassification,
    FalconModel,
    FalconPreTrainedModel,
    falcon_pos_shift_attention_forward,
)
from .gpt_neox import (
    GPTNeoXForCausalLM,
    GPTNeoXForQuestionAnswering,
    GPTNeoXForSequenceClassification,
    GPTNeoXForTokenClassification,
    GPTNeoXModel,
    GPTNeoXPreTrainedModel,
    gpt_neox_pos_shift_attention_forward,
)
from .gptj import (
    GPTJForCausalLM,
    GPTJForQuestionAnswering,
    GPTJForSequenceClassification,
    GPTJModel,
    GPTJPreTrainedModel,
    gptj_pos_shift_attention_forward,
)
from .llama import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, llama_pos_shift_attention_forward
from .mistral import (
    MistralForCausalLM,
    MistralForSequenceClassification,
    MistralModel,
    mistral_pos_shift_attention_forward,
)
from .mpt import (
    MptForCausalLM,
    MptForQuestionAnswering,
    MptForSequenceClassification,
    MptForTokenClassification,
    MptModel,
    MptPreTrainedModel,
)

attention_sinks/models/__init__.py
tomaarsen/attention_sinks
from .modeling_auto import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
)

attention_sinks/models/auto/__init__.py
tomaarsen/attention_sinks
from transformers import AutoModel as TAutoModel
from transformers import AutoModelForCausalLM as TAutoModelForCausalLM
from transformers import AutoModelForQuestionAnswering as TAutoModelForQuestionAnswering
from transformers import AutoModelForSequenceClassification as TAutoModelForSequenceClassification
from transformers import AutoModelForTokenClassification as TAutoModelForTokenClassification

from attention_sinks.inject_mixin import InjectAttentionSinksMixin


class AutoModel(InjectAttentionSinksMixin, TAutoModel):
    pass


class AutoModelForCausalLM(InjectAttentionSinksMixin, TAutoModelForCausalLM):
    pass


class AutoModelForSequenceClassification(InjectAttentionSinksMixin, TAutoModelForSequenceClassification):
    pass


class AutoModelForQuestionAnswering(InjectAttentionSinksMixin, TAutoModelForQuestionAnswering):
    pass


class AutoModelForTokenClassification(InjectAttentionSinksMixin, TAutoModelForTokenClassification):
    pass

attention_sinks/models/auto/modeling_auto.py
tomaarsen/attention_sinks
from .modeling_falcon import (
    FalconForCausalLM,
    FalconForQuestionAnswering,
    FalconForSequenceClassification,
    FalconForTokenClassification,
    FalconModel,
    FalconPreTrainedModel,
)
from .pos_shift import falcon_pos_shift_attention_forward

attention_sinks/models/falcon/__init__.py
tomaarsen/attention_sinks
from transformers import FalconForCausalLM as TFalconForCausalLM
from transformers import FalconForQuestionAnswering as TFalconForQuestionAnswering
from transformers import FalconForSequenceClassification as TFalconForSequenceClassification
from transformers import FalconForTokenClassification as TFalconForTokenClassification
from transformers import FalconModel as TFalconModel
from transformers import FalconPreTrainedModel as TFalconPreTrainedModel

from attention_sinks.inject_mixin import InjectAttentionSinksMixin


class FalconPreTrainedModel(InjectAttentionSinksMixin, TFalconPreTrainedModel):
    pass


class FalconModel(FalconPreTrainedModel, TFalconModel):
    pass


class FalconForCausalLM(FalconPreTrainedModel, TFalconForCausalLM):
    pass


class FalconForSequenceClassification(FalconPreTrainedModel, TFalconForSequenceClassification):
    pass


class FalconForTokenClassification(FalconPreTrainedModel, TFalconForTokenClassification):
    pass


class FalconForQuestionAnswering(FalconPreTrainedModel, TFalconForQuestionAnswering):
    pass

attention_sinks/models/falcon/modeling_falcon.py
tomaarsen/attention_sinks
import math
from typing import Optional, Tuple

import torch
import torch.nn.functional as F
import torch.utils.checkpoint

__all__ = ["falcon_pos_shift_attention_forward"]


def falcon_pos_shift_attention_forward(
    self,
    hidden_states: torch.Tensor,
    alibi: Optional[torch.Tensor],
    attention_mask: torch.Tensor,
    position_ids: Optional[torch.LongTensor] = None,
    layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    head_mask: Optional[torch.Tensor] = None,
    use_cache: bool = False,
    output_attentions: bool = False,
    padding_mask: Optional[torch.LongTensor] = None,
):
    fused_qkv = self.query_key_value(hidden_states)  # [batch_size, seq_length, 3 x hidden_size]
    num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
    # 3 x [batch_size, seq_length, num_heads, head_dim]
    (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)

    batch_size, query_length, _, _ = query_layer.shape

    query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, query_length, self.head_dim)
    key_layer = key_layer.transpose(1, 2).reshape(
        batch_size * num_kv_heads,
        query_length,
        self.head_dim,
    )
    value_layer = value_layer.transpose(1, 2).reshape(batch_size * num_kv_heads, query_length, self.head_dim)

    past_kv_length = 0 if layer_past is None else layer_past[0].shape[1]
    # CHANGED:
    # query_layer, key_layer = self.maybe_rotary(query_layer, key_layer, past_kv_length, position_ids)
    query_layer_copy = query_layer.clone()
    query_layer, _ = self.maybe_rotary(query_layer, query_layer_copy, past_kv_length, position_ids)
    # ########

    if layer_past is not None:
        past_key, past_value = layer_past
        # concatenate along seq_length dimension:
        #  - key: [batch_size * self.num_heads, kv_length, head_dim]
        #  - value: [batch_size * self.num_heads, kv_length, head_dim]
        key_layer = torch.cat((past_key, key_layer), dim=1)
        value_layer = torch.cat((past_value, value_layer), dim=1)

    _, kv_length, _ = key_layer.shape
    if use_cache:
        present = (key_layer, value_layer)
    else:
        present = None

    # CHANGED:
    key_layer_copy = key_layer.clone()
    key_position_ids = torch.arange(kv_length, device=position_ids.device).unsqueeze(0)
    _, key_layer = self.maybe_rotary(key_layer_copy, key_layer, 0, key_position_ids)
    # ########
    float_min = torch.finfo(query_layer.dtype).min
    attention_mask_float = (attention_mask * 1.0).masked_fill(attention_mask, float_min).to(query_layer.dtype)

    query_layer_ = query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
    key_layer_ = key_layer.reshape(batch_size, num_kv_heads, -1, self.head_dim)
    value_layer_ = value_layer.reshape(batch_size, num_kv_heads, -1, self.head_dim)

    if alibi is None:
        if hasattr(F, "scaled_dot_product_attention") and not output_attentions:
            # TODO: deprecate this once we add FA2 support in Falcon
            # query_layer_.shape, key_layer_.shape, value_layer_.shape
            attn_output = F.scaled_dot_product_attention(
                query_layer_, key_layer_, value_layer_, attention_mask_float, 0.0, is_causal=False
            )
            attention_scores = None
        else:
            attention_scores = query_layer_ @ key_layer_.transpose(-1, -2)
            attention_scores /= math.sqrt(self.head_dim)

            attention_scores = F.softmax(attention_scores + attention_mask_float, dim=-1, dtype=hidden_states.dtype)
            attn_output = attention_scores @ value_layer_

        attn_output = attn_output.view(batch_size, self.num_heads, query_length, self.head_dim)
        attn_output = attn_output.permute(0, 2, 1, 3)
        attn_output = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)

        output_tensor = self.dense(attn_output)

        if output_attentions:
            return output_tensor, present, attention_scores
        else:
            return output_tensor, present

    else:
        matmul_result = query_layer_ @ key_layer_.transpose(-1, -2)

        # change view to [batch_size, num_heads, q_length, kv_length]
        attention_scores = matmul_result.view(batch_size, self.num_heads, query_length, kv_length)

        # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]
        input_dtype = attention_scores.dtype
        # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`
        if input_dtype == torch.float16 or input_dtype == torch.bfloat16:
            attention_scores = attention_scores.to(torch.float32)
        # Matt (HF) note: We could possibly use F.scaled_dot_product_attention here too, by
        # adding (alibi * self.inv_norm_factor) to attention_mask_float. I think this would be mathematically
        # equivalent and more performant, but there might be a numerical difference. If you're reading this
        # and you'd like to experiment and maybe file a PR, feel free!
        attention_logits = attention_scores + alibi.view(batch_size, self.num_heads, 1, -1)
        attention_logits *= self.inv_norm_factor
        attention_probs = F.softmax(attention_logits + attention_mask_float, dim=-1, dtype=hidden_states.dtype)
        # [batch_size, num_heads, q_length, kv_length]
        attention_probs = self.attention_dropout(attention_probs)

        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        # change view [batch_size, num_heads, q_length, kv_length]
        attention_probs_reshaped = attention_probs.view(batch_size, self.num_heads, query_length, kv_length)

        # matmul: [batch_size * num_heads, q_length, head_dim]
        context_layer = (attention_probs_reshaped @ value_layer_).flatten(0, 1)

        # change view [batch_size, q_length, num_heads * head_dim]
        context_layer = self._merge_heads(context_layer)

        output_tensor = self.dense(context_layer)

        if output_attentions:
            return output_tensor, present, attention_probs
        else:
            return output_tensor, present

attention_sinks/models/falcon/pos_shift.py
tomaarsen/attention_sinks
from .modeling_gpt_neox import (
    GPTNeoXForCausalLM,
    GPTNeoXForQuestionAnswering,
    GPTNeoXForSequenceClassification,
    GPTNeoXForTokenClassification,
    GPTNeoXModel,
    GPTNeoXPreTrainedModel,
)
from .pos_shift import gpt_neox_pos_shift_attention_forward

attention_sinks/models/gpt_neox/__init__.py
tomaarsen/attention_sinks
from transformers import GPTNeoXForCausalLM as TGPTNeoXForCausalLM
from transformers import GPTNeoXForQuestionAnswering as TGPTNeoXForQuestionAnswering
from transformers import GPTNeoXForSequenceClassification as TGPTNeoXForSequenceClassification
from transformers import GPTNeoXForTokenClassification as TGPTNeoXForTokenClassification
from transformers import GPTNeoXModel as TGPTNeoXModel
from transformers import GPTNeoXPreTrainedModel as TGPTNeoXPreTrainedModel

from attention_sinks.inject_mixin import InjectAttentionSinksMixin


class GPTNeoXPreTrainedModel(InjectAttentionSinksMixin, TGPTNeoXPreTrainedModel):
    pass


class GPTNeoXModel(GPTNeoXPreTrainedModel, TGPTNeoXModel):
    pass


class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, TGPTNeoXForCausalLM):
    pass


class GPTNeoXForSequenceClassification(GPTNeoXPreTrainedModel, TGPTNeoXForSequenceClassification):
    pass


class GPTNeoXForTokenClassification(GPTNeoXPreTrainedModel, TGPTNeoXForTokenClassification):
    pass


class GPTNeoXForQuestionAnswering(GPTNeoXPreTrainedModel, TGPTNeoXForQuestionAnswering):
    pass

attention_sinks/models/gpt_neox/modeling_gpt_neox.py
tomaarsen/attention_sinks
"""
Adapted from https://github.com/mit-han-lab/streaming-llm
"""


from typing import Optional, Tuple

import torch
import torch.utils.checkpoint
from transformers.models.gpt_neox.modeling_gpt_neox import rotate_half

__all__ = ["gpt_neox_pos_shift_attention_forward"]


def apply_rotary_pos_emb_single(x, cos, sin, position_ids):
    gather_indices = position_ids[:, None, :, None]  # [bs, 1, seq_len, 1]
    gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3])
    cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
    sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
    x_embed = (x * cos) + (rotate_half(x) * sin)
    return x_embed


def gpt_neox_pos_shift_attention_forward(
    self,
    hidden_states: torch.FloatTensor,
    attention_mask: torch.FloatTensor,
    position_ids: torch.LongTensor,
    head_mask: Optional[torch.FloatTensor] = None,
    layer_past: Optional[Tuple[torch.Tensor]] = None,
    use_cache: Optional[bool] = False,
    output_attentions: Optional[bool] = False,
):
    has_layer_past = layer_past is not None

    # Compute QKV
    # Attention heads [batch, seq_len, hidden_size]
    #   --> [batch, seq_len, (np * 3 * head_size)]
    qkv = self.query_key_value(hidden_states)

    # [batch, seq_len, (num_heads * 3 * head_size)]
    #   --> [batch, seq_len, num_heads, 3 * head_size]
    new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
    qkv = qkv.view(*new_qkv_shape)

    # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
    query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
    key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
    value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)

    # Compute rotary embeddings on rotary_ndims
    query_rot = query[..., : self.rotary_ndims]
    query_pass = query[..., self.rotary_ndims :]

    # Compute token offset for rotary embeddings (when decoding)
    seq_len = key.shape[-2]
    if has_layer_past:
        seq_len += layer_past[0].shape[-2]
    cos, sin = self.rotary_emb(value, seq_len=seq_len)
    query = apply_rotary_pos_emb_single(query_rot, cos, sin, position_ids)
    query = torch.cat((query, query_pass), dim=-1)

    # Cache QKV values
    if has_layer_past:
        past_key = layer_past[0]
        past_value = layer_past[1]
        key = torch.cat((past_key, key), dim=-2)
        value = torch.cat((past_value, value), dim=-2)

    present = (key, value) if use_cache else None

    key_rot = key[..., : self.rotary_ndims]
    key_pass = key[..., self.rotary_ndims :]
    key_position_ids = torch.arange(seq_len, device=position_ids.device).unsqueeze(0)
    key = apply_rotary_pos_emb_single(key_rot, cos, sin, key_position_ids)
    key = torch.cat((key, key_pass), dim=-1)

    # Compute attention
    attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

    # Reshape outputs
    attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size)
    attn_output = self.dense(attn_output)

    outputs = (attn_output, present)
    if output_attentions:
        outputs += (attn_weights,)

    return outputs

attention_sinks/models/gpt_neox/pos_shift.py
tomaarsen/attention_sinks
from .modeling_gptj import (
    GPTJForCausalLM,
    GPTJForQuestionAnswering,
    GPTJForSequenceClassification,
    GPTJModel,
    GPTJPreTrainedModel,
)
from .pos_shift import gptj_pos_shift_attention_forward

attention_sinks/models/gptj/__init__.py
tomaarsen/attention_sinks
from transformers import GPTJForCausalLM as TGPTJForCausalLM
from transformers import GPTJForQuestionAnswering as TGPTJForQuestionAnswering
from transformers import GPTJForSequenceClassification as TGPTJForSequenceClassification
from transformers import GPTJModel as TGPTJModel
from transformers import GPTJPreTrainedModel as TGPTJPreTrainedModel

from attention_sinks.inject_mixin import InjectAttentionSinksMixin


class GPTJPreTrainedModel(InjectAttentionSinksMixin, TGPTJPreTrainedModel):
    pass


class GPTJModel(GPTJPreTrainedModel, TGPTJModel):
    pass


class GPTJForCausalLM(GPTJPreTrainedModel, TGPTJForCausalLM):
    pass


class GPTJForSequenceClassification(GPTJPreTrainedModel, TGPTJForSequenceClassification):
    pass


class GPTJForQuestionAnswering(GPTJPreTrainedModel, TGPTJForQuestionAnswering):
    pass

attention_sinks/models/gptj/modeling_gptj.py
tomaarsen/attention_sinks
"""
Adapted from https://github.com/mit-han-lab/streaming-llm
"""


from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from transformers.models.gptj.modeling_gptj import get_embed_positions, rotate_every_two
from transformers.utils import is_torch_fx_proxy

__all__ = ["gptj_pos_shift_attention_forward"]


def apply_rotary_pos_emb(tensor: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
    sin = torch.repeat_interleave(sin[:, None, :, :], 2, 3)
    cos = torch.repeat_interleave(cos[:, None, :, :], 2, 3)
    return (tensor * cos) + (rotate_every_two(tensor) * sin)


def gptj_pos_shift_attention_forward(
    self,
    hidden_states: torch.FloatTensor,
    layer_past: Optional[Tuple[torch.Tensor]] = None,
    attention_mask: Optional[torch.FloatTensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    head_mask: Optional[torch.FloatTensor] = None,
    use_cache: Optional[bool] = False,
    output_attentions: Optional[bool] = False,
) -> Union[
    Tuple[torch.Tensor, Tuple[torch.Tensor]],
    Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
]:
    query = self.q_proj(hidden_states)
    key = self.k_proj(hidden_states)
    value = self.v_proj(hidden_states)

    query = self._split_heads(query, self.num_attention_heads, self.head_dim, False)
    key = self._split_heads(key, self.num_attention_heads, self.head_dim, False)
    value = self._split_heads(value, self.num_attention_heads, self.head_dim, False)

    if is_torch_fx_proxy(position_ids) or torch.jit.is_tracing():
        # The logic to conditionally copy to GPU could not be traced, so we do this
        # every time in the torch.fx case
        embed_positions = get_embed_positions(self.embed_positions, position_ids)
    else:
        embed_positions = self._get_embed_positions(position_ids)

    repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1])
    sincos = torch.gather(embed_positions, 1, repeated_position_ids)
    sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)

    if self.rotary_dim is not None:
        # [bsz, num_attn_heads, seq_len, head_dim]
        q_rot = query[:, :, :, : self.rotary_dim]
        q_pass = query[:, :, :, self.rotary_dim :]
        q_rot = apply_rotary_pos_emb(q_rot, sin, cos)
        query = torch.cat([q_rot, q_pass], dim=-1)
    else:
        query = apply_rotary_pos_emb(query, sin, cos)

    if layer_past is not None:
        past_key = layer_past[0]
        past_value = layer_past[1]
        key = torch.cat((past_key, key), dim=-2)
        value = torch.cat((past_value, value), dim=-2)

    if use_cache is True:
        present = (key, value)
    else:
        present = None

    key_position_ids = torch.arange(key.shape[-2], device=position_ids.device).unsqueeze(0)
    repeated_key_position_ids = key_position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1])
    sincos = torch.gather(embed_positions, 1, repeated_key_position_ids)
    sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)
    if self.rotary_dim is not None:
        # [bsz, num_attn_heads, seq_len, head_dim]
        k_rot = key[:, :, :, : self.rotary_dim]
        k_pass = key[:, :, :, self.rotary_dim :]
        k_rot = apply_rotary_pos_emb(k_rot, sin, cos)
        key = torch.cat([k_rot, k_pass], dim=-1)
    else:
        key = apply_rotary_pos_emb(key, sin, cos)

    # compute self-attention: V x Softmax(QK^T)
    attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

    attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
    attn_output = self.out_proj(attn_output)
    attn_output = self.resid_dropout(attn_output)

    outputs = (attn_output, present)
    if output_attentions:
        outputs += (attn_weights,)

    return outputs  # a, present, (attentions)

attention_sinks/models/gptj/pos_shift.py
tomaarsen/attention_sinks
from .modeling_llama import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel
from .pos_shift import llama_pos_shift_attention_forward

attention_sinks/models/llama/__init__.py
tomaarsen/attention_sinks
from transformers import LlamaForCausalLM as TLlamaForCausalLM
from transformers import LlamaForSequenceClassification as TLlamaForSequenceClassification
from transformers import LlamaModel as TLlamaModel
from transformers import LlamaPreTrainedModel as TLlamaPreTrainedModel

from attention_sinks.inject_mixin import InjectAttentionSinksMixin


class LlamaPreTrainedModel(InjectAttentionSinksMixin, TLlamaPreTrainedModel):
    pass


class LlamaModel(LlamaPreTrainedModel, TLlamaModel):
    pass


class LlamaForCausalLM(LlamaPreTrainedModel, TLlamaForCausalLM):
    pass


class LlamaForSequenceClassification(LlamaPreTrainedModel, TLlamaForSequenceClassification):
    pass

attention_sinks/models/llama/modeling_llama.py
tomaarsen/attention_sinks
"""
Adapted from https://github.com/mit-han-lab/streaming-llm
"""


import math
from typing import Optional, Tuple

import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch import nn
from transformers.models.llama.modeling_llama import repeat_kv, rotate_half

__all__ = ["llama_pos_shift_attention_forward"]


def apply_rotary_pos_emb_single(x, cos, sin, position_ids):
    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    x_embed = (x * cos) + (rotate_half(x) * sin)
    return x_embed


def llama_pos_shift_attention_forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Tuple[torch.Tensor]] = None,
    output_attentions: bool = False,
    use_cache: bool = False,
    padding_mask: Optional[torch.LongTensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
    bsz, q_len, _ = hidden_states.size()

    if self.config.pretraining_tp > 1:
        key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
        query_slices = self.q_proj.weight.split((self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0)
        key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
        value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)

        query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
        query_states = torch.cat(query_states, dim=-1)

        key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
        key_states = torch.cat(key_states, dim=-1)

        value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
        value_states = torch.cat(value_states, dim=-1)

    else:
        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

    kv_seq_len = key_states.shape[-2]
    if past_key_value is not None:
        kv_seq_len += past_key_value[0].shape[-2]
    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
    ### Shift Pos: query pos is min(cache_size, idx)
    # query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
    query_states = apply_rotary_pos_emb_single(query_states, cos, sin, position_ids)
    ###

    if past_key_value is not None:
        # reuse k, v, self_attention
        key_states = torch.cat([past_key_value[0], key_states], dim=2)
        value_states = torch.cat([past_key_value[1], value_states], dim=2)

    past_key_value = (key_states, value_states) if use_cache else None

    ### Shift Pos: key pos is the pos in cache
    key_position_ids = torch.arange(kv_seq_len, device=position_ids.device).unsqueeze(0)
    key_states = apply_rotary_pos_emb_single(key_states, cos, sin, key_position_ids)
    ###

    # repeat k/v heads if n_kv_heads < n_heads
    key_states = repeat_kv(key_states, self.num_key_value_groups)
    value_states = repeat_kv(value_states, self.num_key_value_groups)

    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

    if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
        raise ValueError(
            f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
            f" {attn_weights.size()}"
        )

    if attention_mask is not None:
        if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
            raise ValueError(
                f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
            )
        attn_weights = attn_weights + attention_mask

    # upcast attention to fp32
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    attn_output = torch.matmul(attn_weights, value_states)

    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
        raise ValueError(
            f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
            f" {attn_output.size()}"
        )

    attn_output = attn_output.transpose(1, 2).contiguous()
    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

    if self.config.pretraining_tp > 1:
        attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
        o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
        attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
    else:
        attn_output = self.o_proj(attn_output)

    if not output_attentions:
        attn_weights = None

    return attn_output, attn_weights, past_key_value

attention_sinks/models/llama/pos_shift.py
tomaarsen/attention_sinks
from .modeling_mistral import (
    MistralForCausalLM,
    MistralForSequenceClassification,
    MistralModel,
)
from .pos_shift import mistral_pos_shift_attention_forward

attention_sinks/models/mistral/__init__.py
tomaarsen/attention_sinks
from transformers import MistralForCausalLM as TMistralForCausalLM
from transformers import MistralForSequenceClassification as TMistralForSequenceClassification
from transformers import MistralModel as TMistralModel
from transformers import MistralPreTrainedModel as TMistralPreTrainedModel

from attention_sinks.inject_mixin import InjectAttentionSinksMixin


class MistralPreTrainedModel(InjectAttentionSinksMixin, TMistralPreTrainedModel):
    pass


class MistralModel(MistralPreTrainedModel, TMistralModel):
    pass


class MistralForCausalLM(MistralPreTrainedModel, TMistralForCausalLM):
    pass


class MistralForSequenceClassification(MistralPreTrainedModel, TMistralForSequenceClassification):
    pass

attention_sinks/models/mistral/modeling_mistral.py
tomaarsen/attention_sinks
import math
from typing import Optional, Tuple

import torch
import torch.utils.checkpoint
from torch import nn
from transformers.models.mistral.modeling_mistral import repeat_kv, rotate_half

__all__ = ["mistral_pos_shift_attention_forward"]


def apply_rotary_pos_emb_single(x, cos, sin, position_ids):
    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    x_embed = (x * cos) + (rotate_half(x) * sin)
    return x_embed


def mistral_pos_shift_attention_forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Tuple[torch.Tensor]] = None,
    output_attentions: bool = False,
    use_cache: bool = False,
    padding_mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
    bsz, q_len, _ = hidden_states.size()

    query_states = self.q_proj(hidden_states)
    key_states = self.k_proj(hidden_states)
    value_states = self.v_proj(hidden_states)

    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

    kv_seq_len = key_states.shape[-2]
    if past_key_value is not None:
        kv_seq_len += past_key_value[0].shape[-2]
    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
    ### Shift Pos: query pos is min(cache_size, idx)
    # query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
    query_states = apply_rotary_pos_emb_single(query_states, cos, sin, position_ids)
    ###

    if past_key_value is not None:
        # reuse k, v, self_attention
        key_states = torch.cat([past_key_value[0], key_states], dim=2)
        value_states = torch.cat([past_key_value[1], value_states], dim=2)

    past_key_value = (key_states, value_states) if use_cache else None

    ### Shift Pos: key pos is the pos in cache
    key_position_ids = torch.arange(kv_seq_len, device=position_ids.device).unsqueeze(0)
    key_states = apply_rotary_pos_emb_single(key_states, cos, sin, key_position_ids)
    ###

    # repeat k/v heads if n_kv_heads < n_heads
    key_states = repeat_kv(key_states, self.num_key_value_groups)
    value_states = repeat_kv(value_states, self.num_key_value_groups)

    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

    if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
        raise ValueError(
            f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
            f" {attn_weights.size()}"
        )

    if attention_mask is not None:
        if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
            raise ValueError(
                f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
            )

        attn_weights = attn_weights + attention_mask

    # upcast attention to fp32
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    attn_output = torch.matmul(attn_weights, value_states)

    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
        raise ValueError(
            f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
            f" {attn_output.size()}"
        )

    attn_output = attn_output.transpose(1, 2).contiguous()
    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

    attn_output = self.o_proj(attn_output)

    if not output_attentions:
        attn_weights = None

    return attn_output, attn_weights, past_key_value

attention_sinks/models/mistral/pos_shift.py
tomaarsen/attention_sinks
from .modeling_mpt import (
    MptForCausalLM,
    MptForQuestionAnswering,
    MptForSequenceClassification,
    MptForTokenClassification,
    MptModel,
    MptPreTrainedModel,
)

attention_sinks/models/mpt/__init__.py
tomaarsen/attention_sinks
from transformers import MptForCausalLM as TMptForCausalLM
from transformers import MptForQuestionAnswering as TMptForQuestionAnswering
from transformers import MptForSequenceClassification as TMptForSequenceClassification
from transformers import MptForTokenClassification as TMptForTokenClassification
from transformers import MptModel as TMptModel
from transformers import MptPreTrainedModel as TMptPreTrainedModel

from attention_sinks.inject_mixin import InjectAttentionSinksMixin


class MptPreTrainedModel(InjectAttentionSinksMixin, TMptPreTrainedModel):
    pass


class MptModel(MptPreTrainedModel, TMptModel):
    pass


class MptForCausalLM(MptPreTrainedModel, TMptForCausalLM):
    pass


class MptForSequenceClassification(MptPreTrainedModel, TMptForSequenceClassification):
    pass


class MptForTokenClassification(MptPreTrainedModel, TMptForTokenClassification):
    pass


class MptForQuestionAnswering(MptPreTrainedModel, TMptForQuestionAnswering):
    pass

attention_sinks/models/mpt/modeling_mpt.py
tomaarsen/attention_sinks
"""
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


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()
    # Which experiment to run?
    parser.add_argument(
        "--experiment", choices=["attention_sinks", "transformers", "windowed"], default="attention_sinks"
    )

    # Model args
    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=8192)

    # Where to log
    parser.add_argument("--output_dir", type=str, default="benchmark/outputs")
    parser.add_argument("--overwrite", action="store_true")

    # Window size for windowed and attention_sinks
    parser.add_argument("--window_size", type=int, default=1024)

    # Attention Sinks-only settings
    # Attention Sink window size is calculated with args.window_size - args.attention_sink_size
    parser.add_argument("--attention_sink_size", type=int, default=4)

    args = parser.parse_args()

    # Initialize the model, either via transformers or via attention_sinks
    if args.experiment == "transformers":
        from transformers import AutoModelForCausalLM
    else:
        from attention_sinks import AutoModelForCausalLM
    kwargs = {}
    if args.experiment == "attention_sinks":
        kwargs = {
            "attention_sink_size": args.attention_sink_size,
            "attention_sink_window_size": args.window_size - args.attention_sink_size,  # default: 1020
        }
    elif args.experiment == "windowed":
        kwargs = {
            "attention_sink_size": 0,
            "attention_sink_window_size": args.window_size,
        }
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name_or_path,
        revision=args.revision,
        trust_remote_code=bool(args.trust_remote_code),
        torch_dtype=torch.float16,
        device_map="auto",
        **kwargs,
    )
    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()

benchmark/perplexity.py
tomaarsen/attention_sinks
"""
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)
            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 or "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="benchmark/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

    plt.show()


if __name__ == "__main__":
    main()

benchmark/plot_perplexity.py
tomaarsen/attention_sinks
import argparse
from pathlib import Path

import torch
from transformers import (
    AutoTokenizer,
    GenerationConfig,
    PreTrainedModel,
    PreTrainedTokenizer,
)
from utils import FileStreamer


@torch.no_grad()
def endless_generate(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizer,
    text: str,
    log_file: str,
    min_new_tokens: int = 1000,
    max_new_tokens: int = 10000,
) -> str:
    streamer = FileStreamer(tokenizer, log_file)
    input_ids = tokenizer.encode(text, return_tensors="pt").to(model.device)
    generated_tokens = model.generate(
        input_ids,
        generation_config=GenerationConfig(
            # use_cache=True is required, the rest can be changed up.
            use_cache=True,
            min_new_tokens=min_new_tokens,
            max_new_tokens=max_new_tokens,
            penalty_alpha=0.6,
            top_k=4,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        ),
        streamer=streamer,
    )
    # Decode the final generated text
    return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)


def main():
    parser = argparse.ArgumentParser()

    # Which experiment to run?
    parser.add_argument(
        "--experiment", choices=["attention_sinks", "transformers", "windowed"], default="attention_sinks"
    )

    # Model args
    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")

    # Input text
    parser.add_argument("--text", type=str, default="Vaswani et al. (2017) introduced the Transformers")

    # Generation args
    parser.add_argument("--min_new_tokens", type=int, default=1000)
    parser.add_argument("--max_new_tokens", type=int, default=10000)

    # Where to log
    parser.add_argument("--log_file", type=str, default=None)

    # Window size for windowed and attention_sinks
    parser.add_argument("--window_size", type=int, default=512)

    # Attention Sinks-only settings
    # Attention Sink window size is calculated with args.window_size - args.attention_sink_size
    parser.add_argument("--attention_sink_size", type=int, default=4)

    args = parser.parse_args()

    # Initialize the model, either via transformers or via attention_sinks
    if args.experiment == "transformers":
        from transformers import AutoModelForCausalLM
    else:
        from attention_sinks import AutoModelForCausalLM
    kwargs = {}
    if args.experiment == "attention_sinks":
        kwargs = {
            "attention_sink_size": args.attention_sink_size,
            "attention_sink_window_size": args.window_size - args.attention_sink_size,  # default: 1020
        }
    elif args.experiment == "windowed":
        kwargs = {
            "attention_sink_size": 0,
            "attention_sink_window_size": args.window_size,
        }
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name_or_path,
        revision=args.revision,
        trust_remote_code=bool(args.trust_remote_code),
        torch_dtype=torch.float16,
        device_map="auto",
        **kwargs,
    )
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=bool(args.trust_remote_code))
    tokenizer.pad_token_id = tokenizer.eos_token_id

    log_file = args.log_file or Path("demo") / "endless_logs" / args.experiment / f"{args.model_name_or_path}.txt"
    endless_generate(
        model,
        tokenizer,
        args.text,
        log_file=log_file,
        min_new_tokens=args.min_new_tokens,
        max_new_tokens=args.max_new_tokens,
    )


if __name__ == "__main__":
    main()

demo/endless_generation.py
tomaarsen/attention_sinks
import argparse
from pathlib import Path
from typing import Any, Dict, List

import torch
from datasets import Dataset, load_dataset
from transformers import (
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizer,
)
from utils import FileStreamer


def create_prompts(samples: Dict[str, List[Any]]) -> Dict[str, Any]:
    return {"prompt": [prompt for prompts in samples["prompt"] for prompt in prompts]}


@torch.no_grad()
def greedy_generate(
    model: PreTrainedModel, tokenizer: PreTrainedTokenizer, dataset: Dataset, log_file: str, max_new_tokens: int = 1000
):
    streamer = FileStreamer(tokenizer, log_file)
    past_key_values = None
    new_line_tokens = tokenizer("\n\n", return_tensors="pt", add_special_tokens=False).input_ids

    for prompt_index, prompt in enumerate(dataset["prompt"]):
        # Use the chat template initially, as it adds the system prompt if the model has one, and then use [INST] and [/INST]
        if prompt_index:
            prompt = f"[INST] {prompt} [/INST]"
        else:
            prompt = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False)
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
        input_ids = input_ids.to(model.device)

        streamer.put(input_ids)
        for _ in range(max_new_tokens):
            outputs = model(input_ids, past_key_values=past_key_values, use_cache=True)
            past_key_values = outputs.past_key_values
            pred_token_idx = outputs.logits[:, -1, :].argmax(dim=-1).unsqueeze(1)
            streamer.put(pred_token_idx)
            input_ids = pred_token_idx

            if pred_token_idx == tokenizer.eos_token_id:
                break

        streamer.put(new_line_tokens)


def main():
    parser = argparse.ArgumentParser()

    # Which experiment to run?
    parser.add_argument(
        "--experiment", choices=["attention_sinks", "transformers", "windowed"], default="attention_sinks"
    )

    # Model args
    parser.add_argument("--model_name_or_path", type=str, default="mistralai/Mistral-7B-Instruct-v0.1")
    parser.add_argument("--revision", type=str, default="main")
    parser.add_argument("--trust_remote_code", action="store_true")

    # Dataset args, not recommended to change:
    parser.add_argument("--dataset_name", type=str, default="HuggingFaceH4/mt_bench_prompts")

    # Where to log
    parser.add_argument("--log_file", type=str, default=None)

    # Window size for windowed and attention_sinks
    parser.add_argument("--window_size", type=int, default=1024)

    # Attention Sinks-only settings
    # Attention Sink window size is calculated with args.window_size - args.attention_sink_size
    parser.add_argument("--attention_sink_size", type=int, default=4)

    args = parser.parse_args()

    # Initialize the model, either via transformers or via attention_sinks
    if args.experiment == "transformers":
        from transformers import AutoModelForCausalLM
    else:
        from attention_sinks import AutoModelForCausalLM
    kwargs = {}
    if args.experiment == "attention_sinks":
        kwargs = {
            "attention_sink_size": args.attention_sink_size,
            "attention_sink_window_size": args.window_size - args.attention_sink_size,  # default: 1020
        }
    elif args.experiment == "windowed":
        kwargs = {
            "attention_sink_size": 0,
            "attention_sink_window_size": args.window_size,
        }
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name_or_path,
        revision=args.revision,
        trust_remote_code=bool(args.trust_remote_code),
        torch_dtype=torch.float16,
        device_map="auto",
        **kwargs,
    )
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=bool(args.trust_remote_code))
    tokenizer.pad_token_id = tokenizer.eos_token_id

    # Set up the dataset
    dataset = load_dataset(args.dataset_name, split="train")
    dataset = dataset.map(create_prompts, batched=True, remove_columns=dataset.column_names)

    log_file = args.log_file or Path("demo") / "streaming_logs" / args.experiment / f"{args.model_name_or_path}.txt"
    greedy_generate(model, tokenizer, dataset, log_file=log_file)


if __name__ == "__main__":
    main()

demo/streaming.py
tomaarsen/attention_sinks
from pathlib import Path

from transformers import AutoTokenizer, TextStreamer


class FileStreamer(TextStreamer):
    def __init__(self, tokenizer: AutoTokenizer, log_file: str, skip_prompt: bool = False, **decode_kwargs):
        super().__init__(tokenizer, skip_prompt, **decode_kwargs)
        self.log_file = Path(log_file)
        self.log_file.parent.mkdir(parents=True, exist_ok=True)

    def on_finalized_text(self, text: str, stream_end: bool = False):
        """Prints the new text to stdout and writes it to a file. If the stream is ending, also prints a newline."""
        super().on_finalized_text(text, stream_end=stream_end)
        with open(self.log_file, "a+", encoding="utf8") as f:
            print(text, flush=True, end="" if not stream_end else None, file=f)

demo/utils.py
tomaarsen/attention_sinks