mistral.py

import torch
from torch import nn
from dataclasses import dataclass
from pathlib import Path

import math
import json
from typing import Optional, Tuple, List
from sentencepiece import SentencePieceProcessor
from llama import (
    repeat_kv_head,
    apply_rotary_embeddings,
    get_complex_rotary_matrix,
    RMSNorm,
)
from einops import rearrange


@dataclass
class ModelArgs:
    dim: int
    n_layers: int
    head_dim: int
    hidden_dim: int
    n_heads: int
    n_kv_heads: int
    sliding_window: int
    norm_eps: float
    vocab_size: int

    max_batch_size: int = 0


class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args

        self.n_heads: int = args.n_heads
        self.n_kv_heads: int = args.n_kv_heads

        self.repeats = self.n_heads // self.n_kv_heads
        self.sliding_window = self.args.sliding_window

        self.wq = nn.Linear(args.dim, args.n_heads * args.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, args.n_kv_heads * args.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, args.n_kv_heads * args.head_dim, bias=False)
        self.wo = nn.Linear(args.n_heads * args.head_dim, args.dim, bias=False)
        self.cache_k = torch.empty(
            (
                args.max_batch_size,
                args.sliding_window,
                self.n_kv_heads,
                self.args.head_dim,
            )
        ).cuda()
        self.cache_v = torch.empty(
            (
                args.max_batch_size,
                args.sliding_window,
                self.n_kv_heads,
                self.args.head_dim,
            ),
        ).cuda()

    def update_cache(self, key, value, positions):
        batch_size, _, _, _ = key.shape
        scatter_pos = (positions[-self.sliding_window :] % self.sliding_window)[
            None, :, None, None
        ]
        scatter_pos = scatter_pos.repeat(
            batch_size, 1, self.n_kv_heads, self.args.head_dim
        )
        self.cache_k[:batch_size].scatter_(
            dim=1, index=scatter_pos.cuda(), src=key[:, -self.sliding_window :]
        )
        self.cache_v[:batch_size].scatter_(
            dim=1, index=scatter_pos.cuda(), src=value[:, -self.sliding_window :]
        )

        if positions.shape[0] > 1:
            # prefill
            key, value = repeat_kv_head(key, value, self.repeats)
        else:
            cur_pos = positions[-1].item() + 1
            key, value = repeat_kv_head(
                self.cache_k[:batch_size, :cur_pos, ...],
                self.cache_v[:batch_size, :cur_pos, ...],
                self.repeats,
            )
        return key, value

    def forward(
        self,
        x: torch.Tensor,
        complex_rotary_matrix: torch.Tensor,
        positions: torch.Tensor,
        mask: Optional[torch.Tensor],
    ) -> torch.Tensor:

        query = rearrange(
            self.wq(x),
            "B seq_len (n_head h_dim) -> B seq_len n_head h_dim",
            h_dim=self.args.head_dim,
        )
        key = rearrange(
            self.wk(x),
            "B seq_len (n_head h_dim) -> B seq_len n_head h_dim",
            h_dim=self.args.head_dim,
        )

        value = rearrange(
            self.wv(x),
            "B seq_len (n_head h_dim) -> B seq_len n_head h_dim",
            h_dim=self.args.head_dim,
        )

        query, key = apply_rotary_embeddings(
            query, key, complex_rotary_matrix, x.device
        )
        key, value = self.update_cache(key, value, positions)

        key = rearrange(key, "B seq_len n_head head_dim -> B n_head seq_len head_dim")
        value = rearrange(
            value, "B seq_len n_head head_dim -> B n_head seq_len head_dim"
        )
        query = rearrange(
            query, "B seq_len n_head head_dim -> B n_head seq_len head_dim"
        )

        scores = torch.einsum("bhqd,bhkd -> bhqk", query, key) / math.sqrt(
            self.args.head_dim
        )
        if mask is not None:
            scores += mask[None, None, ...]

        scores = nn.functional.softmax(scores.float(), dim=-1).type_as(query)

        output = torch.einsum("bhqk,bhkv -> bhqv", scores, value)

        output = rearrange(
            output, " B n_head seq_len head_dim -> B seq_len n_head head_dim"
        )
        output = rearrange(
            output, "B seq_len n_head head_dim -> B seq_len (n_head head_dim)"
        )

        return self.wo(output)


class FeedForward(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.w1 = nn.Linear(args.dim, args.hidden_dim, bias=False)
        self.w2 = nn.Linear(args.hidden_dim, args.dim, bias=False)
        self.w3 = nn.Linear(args.dim, args.hidden_dim, bias=False)

    def forward(self, x) -> torch.Tensor:
        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))


class EncoderBlock(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.attention = Attention(args)
        self.feed_forward = FeedForward(args=args)
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.args = args

    def forward(
        self,
        x: torch.Tensor,
        freqs_cis: torch.Tensor,
        positions: torch.Tensor,
        mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        r = self.attention.forward(self.attention_norm(x), freqs_cis, positions, mask)
        h = x + r
        r = self.feed_forward.forward(self.ffn_norm(h))
        out = h + r
        return out


class MistralTransformer(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.vocab_size = args.vocab_size
        self.n_layers = args.n_layers
        assert self.vocab_size > 0

        self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)

        self.layers = torch.nn.ModuleList(
            [EncoderBlock(args=args) for _ in range(args.n_layers)]
        )

        self.norm = RMSNorm(args.dim, eps=args.norm_eps)

        self.output = nn.Linear(args.dim, args.vocab_size, bias=False)

        self.freqs_cis = get_complex_rotary_matrix(self.args.head_dim, 128_000).to(
            self.args.device
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ):
        h = self.tok_embeddings(input_ids)
        freqs_cis = self.freqs_cis[positions]

        mask: Optional[torch.Tensor] = None
        if input_ids.shape[1] > 1:
            seqlen = input_ids.shape[1]
            tensor = torch.full(
                (seqlen, seqlen),
                dtype=h.dtype,
                fill_value=1,
                device=h.device,
            )
            mask = torch.tril(tensor, diagonal=0).to(h.dtype)
            # make the mask banded to account for sliding window
            mask = torch.triu(mask, diagonal=-self.args.sliding_window)
            mask = torch.log(mask)

        for layer in self.layers:
            h = layer(h, freqs_cis, positions, mask)

        return self.output(self.norm(h)).float()


class Mistral:
    def __init__(
        self,
        max_batch_size: int,
        device: str,
    ):

        self.device = device
        self.max_batch_size = max_batch_size
        self.tokenizer = SentencePieceProcessor()
        self.model = None

    def from_pretrained(self, checkpoints_dir: Path, tokenizer_path: Path):
        with open(checkpoints_dir / "params.json", "r") as f:
            self.args = ModelArgs(
                **json.loads(f.read()),
            )
        if self.device == "cuda":
            torch.set_default_tensor_type(torch.cuda.HalfTensor)
        else:
            torch.set_default_tensor_type(torch.BFloat16Tensor)
        self.args.device = self.device
        self.args.max_batch_size = self.max_batch_size
        self.tokenizer.load(str(tokenizer_path))
        self.args.vocab_size = self.tokenizer.vocab_size()
        state_dict = torch.load(checkpoints_dir / "consolidated.00.pth")
        self.model = MistralTransformer(self.args).to(device=self.device)
        self.model.load_state_dict(state_dict)

    @torch.no_grad()
    def text_completion(self, prompts: List[str], max_gen_len: int):
        encoded_prompts = [
            [self.tokenizer.bos_id(), *self.tokenizer.encdode(prompt)]
            for prompt in prompts
        ]
        prompt_lens = [len(x) for x in encoded_prompts]
        min_prompt_len = min(prompt_lens)
        max_prompt_len = max(prompt_lens)
        pad_id = self.tokenizer.pad_id()
        input_tokens = torch.full(
            (len(prompts), max_prompt_len),
            pad_id,
            dtype=torch.long,
            device=self.args.device,
        )
        for i, encoded in enumerate(encoded_prompts):
            input_tokens[i, : len(encoded)] = torch.tensor(encoded).to(input_tokens)
        input_mask = input_tokens != pad_id

        # pre-fill
        positions = torch.arange(0, min_prompt_len).to(self.args.device)
        logits = self.model.forward(input_tokens[:, :min_prompt_len], positions)
        logprobs = nn.functional.log_softmax(logits, dim=-1)

        # decode
        generated = []
        all_logprobs = [
            logprobs[:, :-1, :]
            .gather(2, input_tokens[:, 1:min_prompt_len, None])
            .squeeze(-1),
        ]
        cur_pos = min_prompt_len
        for _ in range(max_gen_len):
            next_token = torch.argmax(logprobs[:, -1, :], dim=-1)

            if cur_pos < input_mask.shape[1]:
                next_token = torch.where(
                    input_mask[:, cur_pos], input_tokens[:, cur_pos], next_token
                )
            all_logprobs.append(
                logprobs[:, -1, :].gather(1, next_token[:, None]),
            )
            generated.append(next_token[:, None])
            logits = self.model.forward(
                next_token[:, None], torch.LongTensor([cur_pos]).to(next_token)
            )
            logprobs = nn.functional.log_softmax(logits, dim=-1)
            cur_pos += 1

        all_logprobs = torch.cat(all_logprobs, 1)
        res = []
        if max_gen_len > 0:
            generated = torch.cat(generated, 1)

            for i, x in enumerate(encoded_prompts):
                res.append(
                    self.tokenizer.decode(x[:min_prompt_len] + generated[i].tolist())
                )
        return all_logprobs, res