Add review comments

Author: tirthasheshpatel

keras_nlp/layers/modeling/rotary_embedding.py

@@ -97,7 +97,7 @@ def _apply_rotary_pos_emb(self, tensor, cos_emb, sin_emb):
         return (tensor * cos_emb) + (half_rot_tensor * sin_emb)
 
     def _compute_cos_sin_embedding(self, x, rotary_dim, start_index):
-        freq_range = ops.cast(ops.arange(0, rotary_dim, 2), self.compute_dtype)
+        freq_range = ops.arange(0, rotary_dim, 2)
         freq_range = ops.cast(freq_range, self.compute_dtype)
         freq_range = freq_range / ops.cast(
             self.scaling_factor, self.compute_dtype

keras_nlp/models/mistral/mistral_attention.py

@@ -29,7 +29,6 @@ class CachedMistralAttention(keras.layers.Layer):
 
     def __init__(
         self,
-        *,
         num_query_heads,
         num_key_value_heads,
         rope_max_wavelength=10000,
@@ -124,6 +123,9 @@ def build(self, inputs_shape):
             dtype=self.compute_dtype,
         )
 
+        self._dot_product_equation = "bquh,bkuh->buqk"
+        self._combine_equation = "buqk,bkuh->bquh"
+
         self.built = True
 
     def call(
@@ -258,7 +260,7 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        attention_scores = ops.einsum("aecd,abcd->acbe", key, query)
+        attention_scores = ops.einsum(self._dot_product_equation, key, query)
 
         norm_factor = ops.sqrt(ops.cast(self._head_dim, self.compute_dtype))
 
@@ -268,7 +270,7 @@ def _compute_attention(self, query, key, value, attention_mask=None):
             attention_scores, attention_mask
         )
         attention_output = ops.einsum(
-            "acbe,aecd->abcd", attention_scores, value
+            self._combine_equation, attention_scores, value
         )
 
         return attention_output

keras_nlp/models/mistral/mistral_backbone.py

@@ -64,6 +64,7 @@ class MistralBackbone(Backbone):
             layers in each transformer decoder. Only `sliding_window` number of tokens
             are saved in the cache and used to generate the next token.
             Defaults to `512`.
+        dtype (str, optional): The dtype policy for the mistral model.
 
     Examples:
 
@@ -95,7 +96,6 @@ class MistralBackbone(Backbone):
 
     def __init__(
         self,
-        *,
         vocabulary_size,
         num_layers,
         num_query_heads,

keras_nlp/models/mistral/mistral_layer_norm.py

@@ -15,8 +15,9 @@
 from keras_nlp.backend import ops
 
 
-# TODO: Deprecate this in favor of `keras.layers.LayerNormalization` once
-#       Keras 2 support is removed.
+# TODO: Deprecate this in favor of
+# `keras.layers.LayerNormalization(rms_scaling=True)` once Keras 2 support is
+# removed.
 class MistralLayerNormalization(keras.layers.Layer):
     """A normalization layer for Mistral that implements RMS normalization."""
 

keras_nlp/models/mistral/mistral_transformer_decoder.py

@@ -31,7 +31,6 @@ class MistralTransformerDecoder(keras.layers.Layer):
 
     def __init__(
         self,
-        *,
         intermediate_dim,
         num_query_heads,
         num_key_value_heads,

tools/checkpoint_conversion/convert_mistral_checkpoints.py

@@ -13,16 +13,298 @@
 # limitations under the License.
 import json
 import pathlib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+from typing import Tuple
 
 import torch
+from torch import nn
 
 from keras_nlp.models import MistralBackbone
 
-from .scripts.mistral_torch import ModelArgs
-from .scripts.mistral_torch import Transformer as TorchTransformer
-
 MODEL_PATH = pathlib.Path("mistral-7B-v0.1")
 
+# Torch model taken from:
+# https://github.com/mistralai/mistral-src/blob/147c4e68279b90eb61b19bdea44e16f5539d5a5d/one_file_ref.py
+
+
+@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
+
+
+def repeat_kv(keys: torch.Tensor, values: torch.Tensor, repeats: int):
+    keys = torch.repeat_interleave(keys, repeats=repeats, dim=2)
+    values = torch.repeat_interleave(values, repeats=repeats, dim=2)
+    return keys, values
+
+
+def _reshape_for_broadcast(
+    freqs_cis: torch.Tensor, x: torch.Tensor
+) -> torch.Tensor:
+    """
+    freqs_cis: complex - (seq_len, head_dim / 2)
+    x: complex - (bsz, seq_len, head_dim / 2)
+    """
+    ndim = x.ndim
+    assert 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1]), (
+        freqs_cis.shape,
+        (x.shape[1], x.shape[-1]),
+    )
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = _reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+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.scale = self.args.head_dim**-0.5
+
+        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,
+            ),
+            dtype=torch.float16,
+        )
+        self.cache_v = torch.empty(
+            (
+                args.max_batch_size,
+                args.sliding_window,
+                self.n_kv_heads,
+                self.args.head_dim,
+            ),
+            dtype=torch.float16,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        positions: torch.Tensor,
+        mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        bsz, seqlen, _ = x.shape
+
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        xq = xq.view(bsz, seqlen, self.n_heads, self.args.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_kv_heads, self.args.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_kv_heads, self.args.head_dim)
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
+
+        # The cache is a rotating buffer
+        scatter_pos = (positions[-self.sliding_window :] % self.sliding_window)[
+            None, :, None, None
+        ]
+        scatter_pos = scatter_pos.repeat(
+            bsz, 1, self.n_kv_heads, self.args.head_dim
+        )
+        self.cache_k[:bsz].scatter_(
+            dim=1,
+            index=scatter_pos,
+            src=xk[:, -self.sliding_window :].to(self.cache_k.dtype),
+        )
+        self.cache_v[:bsz].scatter_(
+            dim=1,
+            index=scatter_pos,
+            src=xv[:, -self.sliding_window :].to(self.cache_v.dtype),
+        )
+
+        if positions.shape[0] > 1:
+            # prefill
+            key, value = repeat_kv(xk, xv, self.repeats)
+        else:
+            cur_pos = positions[-1].item() + 1
+            key, value = repeat_kv(
+                self.cache_k[:bsz, :cur_pos, ...].to(xk.dtype),
+                self.cache_v[:bsz, :cur_pos, ...].to(xv.dtype),
+                self.repeats,
+            )
+
+        query = xq.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        # scores : [bsz, n_heads, seqlen | 1, seqlen]
+        scores = torch.matmul(query, key.transpose(2, 3)) * self.scale
+
+        if mask is not None:
+            scores += mask[None, None, ...]
+
+        scores = scores.float()
+        scores = nn.functional.softmax(scores, dim=-1).type_as(query)
+        output = torch.matmul(
+            scores, value
+        )  # (bs, n_local_heads, slen, head_dim)
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        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 RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class TransformerBlock(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
+
+
+def precompute_freqs_cis(
+    dim: int, end: int, theta: float = 10000.0
+) -> torch.Tensor:
+    freqs = 1.0 / (
+        theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
+    )
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    return torch.polar(torch.ones_like(freqs), freqs)  # complex64
+
+
+class TorchTransformer(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(
+            [TransformerBlock(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 = precompute_freqs_cis(self.args.head_dim, 128_000)
+
+    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()
+
+    @staticmethod
+    def from_folder(
+        folder: Path, max_batch_size: int = 1, device="cpu", dtype=torch.float16
+    ):
+        with open(folder / "params.json", "r") as f:
+            model_args = ModelArgs(**json.loads(f.read()))
+        model_args.max_batch_size = max_batch_size
+        model = TorchTransformer(model_args).to(device=device, dtype=dtype)
+        loaded = torch.load(folder / "consolidated.00.pth")
+        model.load_state_dict(loaded)
+        return model
+
 
 def port_weights(
     model_k3: MistralBackbone, model_torch: TorchTransformer, params: ModelArgs