llama3.efficient.np.py

from __future__ import annotations

import math
import sys
import time
from typing import TypeVar, Generic, Optional

import numpy as np
import warnings

from config import ModelArgs
from tokenizer import Tokenizer
from utils import load_parameters

Shape = TypeVar("Shape")
warnings.filterwarnings(
    "ignore",
    category=RuntimeWarning,
    message=".*encountered in matmul.*",
)


class Array(np.ndarray, Generic[Shape]): ...


def softmax(x):
    exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))  # for float16 precision
    return exp_x / exp_x.sum(axis=-1, keepdims=True)


def rmsnorm(x, w, eps):
    rms = np.sqrt((x ** 2).mean(-1, keepdims=True) + eps)
    return (x / rms) * w


def silu(x):
    return x * (1 / (1 + np.exp(-x)))


def compute_cos_sin_cache(head_dim: int, max_seq_len: int, base: int = 10000):
    inv_freq: Array["HD//2"] = 1.0 / (base ** (np.arange(0, head_dim, 2)[: (head_dim // 2)] / head_dim))
    t: Array["M"] = np.arange(max_seq_len)
    freqs: Array["M, HD//2"] = np.outer(t, inv_freq).astype(np.float32)

    freqs_cos: Array["1, M, 1, HD//2"] = np.expand_dims(np.cos(freqs), axis=(0, 2))
    freqs_sin: Array["1, M, 1, HD//2"] = np.expand_dims(np.sin(freqs), axis=(0, 2))

    return freqs_cos, freqs_sin


def repeat_kv(x: Array["B, L, KVHN, HD"], n_rep: int):
    if n_rep == 1:
        return x
    z: Array["B, L, QHN, HD"] = np.repeat(x, n_rep, axis=2)
    return z


class FeedForward:
    def __init__(self, up_weight: Array["FD, D"], gate_weight: Array["FD, D"], down_weight: Array["D, FD"]):
        self.up_weight = up_weight
        self.gate_weight = gate_weight
        self.down_weight = down_weight

    def __call__(self, x: Array["B, L or 1, D"]):
        # FD = 2 * 4 * D / 3
        x: Array["B, L or 1, FD"] = silu(x @ self.gate_weight) * (x @ self.up_weight)
        x: Array["B, L or 1, D"] = x @ self.down_weight
        return x


class Attention:
    def __init__(self, q_weight: Array["D, D"], k_weight: Array["D, D"], v_weight: Array["D, D"],
                 o_weight: Array["D, D"], args: ModelArgs):
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        assert args.n_heads % self.n_kv_heads == 0

        self.n_local_heads = args.n_heads
        self.n_local_kv_heads = self.n_kv_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

        # (dim, dim) e.g., (288, 288)
        self.q_weight = q_weight
        self.k_weight = k_weight
        self.v_weight = v_weight
        self.o_weight = o_weight

        # (B, L, heads, head_dim) e.g., (1, 256, 6, 48)
        self.cache_k = np.zeros((args.max_batch_size, args.max_seq_len, self.n_local_kv_heads, self.head_dim), dtype=np.float32)
        self.cache_v = np.zeros((args.max_batch_size, args.max_seq_len, self.n_local_kv_heads, self.head_dim), dtype=np.float32)

    def __call__(self, x: Array["B, L or 1, D"], start_pos: int, mask: Optional[Array["L, L"]],
                 freqs: Array["L or 1, HD//2, 2, 2]"]):
        B, L, _ = x.shape

        # Linear of QKV
        q: Array["B, L or 1, D"] = x @ self.q_weight
        k: Array["B, L or 1, D"] = x @ self.k_weight
        v: Array["B, L or 1, D"] = x @ self.v_weight

        # RoPE
        q: Array["B, L or 1, QHN, HD//2, 2, 1"] = q.reshape(B, L, self.n_local_heads, self.head_dim // 2, 2, 1)
        q: Array["B, L or 1, QHN, HD//2, 2"] = np.sum(q * freqs, axis=-2)
        q: Array["B, L or 1, QHN, HD"] = q.reshape(B, L, self.n_local_heads, self.head_dim)

        k: Array["B, L or 1, KVHN, HD//2, 2, 1"] = k.reshape(B, L, self.n_local_kv_heads, self.head_dim // 2, 2, 1)
        k: Array["B, L or 1, KVHN, HD//2, 2"] = np.sum(k * freqs, axis=-2)
        k: Array["B, L or 1, KVHN, HD"] = k.reshape(B, L, self.n_local_kv_heads, self.head_dim)

        v: Array["B, L or 1, KVHN, HD"] = v.reshape(B, L, self.n_local_kv_heads, self.head_dim)

        # Update KV Cache
        end_pos = start_pos + L
        self.cache_k[:B, start_pos:end_pos] = k
        self.cache_v[:B, start_pos:end_pos] = v

        # GQA
        k: Array["B, L, HN, HD"] = repeat_kv(self.cache_k[:B, :end_pos], self.n_rep)
        v: Array["B, L, HN, HD"] = repeat_kv(self.cache_v[:B, :end_pos], self.n_rep)

        # Scaled Dot-Product Attention
        q: Array["B, HN, L or 1, HD"] = q.transpose(0, 2, 1, 3)
        k: Array["B, HN, L, HD"] = k.transpose(0, 2, 1, 3)
        v: Array["B, HN, L, HD"] = v.transpose(0, 2, 1, 3)

        attention: Array["B, HN, L or 1, L"] = q @ k.transpose(0, 1, 3, 2) / math.sqrt(self.head_dim)
        if mask is not None:
            attention = attention + mask
        attention = softmax(attention)
        output: Array["B, HN, L or 1, HD"] = attention @ v
        output: Array["B, L or 1, D"] = output.transpose(0, 2, 1, 3).reshape(B, L, -1)

        # Linear after Concat
        output: Array["B, L or 1, D"] = output @ self.o_weight

        return output


class TransformerBlock:
    def __init__(self, weight: dict, layer_id: int, args: ModelArgs):
        layer_id_str = f"model.layers.{layer_id}"
        self.attention = Attention(
            weight.get(f"{layer_id_str}.self_attn.q_proj.weight").T.copy(),
            weight.get(f"{layer_id_str}.self_attn.k_proj.weight").T.copy(),
            weight.get(f"{layer_id_str}.self_attn.v_proj.weight").T.copy(),
            weight.get(f"{layer_id_str}.self_attn.o_proj.weight").T.copy(),
            args
        )
        self.feed_forward = FeedForward(
            weight.get(f"{layer_id_str}.mlp.up_proj.weight").T.copy(),
            weight.get(f"{layer_id_str}.mlp.gate_proj.weight").T.copy(),
            weight.get(f"{layer_id_str}.mlp.down_proj.weight").T.copy(),
        )
        self.input_norm_weight = weight.get(f"{layer_id_str}.input_layernorm.weight").T.copy()
        self.post_norm_weight = weight.get(f"{layer_id_str}.post_attention_layernorm.weight").T.copy()

    def __call__(self, x: Array["B, L or 1, D"], start_pos: int, mask: Array["L, L"],
                 freqs: Array["L or 1, HD//2, 2, 2"]):
        # RMS Norm
        norm_x = rmsnorm(x, self.input_norm_weight, args.norm_eps)

        # Grouped Multi-Head Attention
        z = x + self.attention(norm_x, start_pos, mask, freqs)

        # RMS Norm
        norm_z = rmsnorm(z, self.post_norm_weight, args.norm_eps)

        # Feed Forward with SwiGLU
        o = z + self.feed_forward(norm_z)

        return o


class Llama:
    def __init__(self, model_path: str, args: ModelArgs):
        self.args = args

        weight = load_parameters(model_path)

        # (vocab_size, dim) e.g., (32000, 288)
        self.embedding: Array["VS, D"] = weight.get("model.embed_tokens.weight")

        self.layers = []
        for layer_id in range(args.n_layers):  # n_layers=6
            self.layers.append(TransformerBlock(weight, layer_id, args))

        self.norm_weight = weight.get("model.norm.weight")
        self.norm_eps = args.norm_eps
        self.lm_head_weight: Array["D, VS"] = weight.get("lm_head.weight").T.copy()

        # RoPE #1
        self.freqs_cos, self.freqs_sin = compute_cos_sin_cache(args.dim // args.n_heads, args.max_seq_len)
        self.freqs1 = np.stack([self.freqs_cos, -self.freqs_sin], axis=-1)
        self.freqs2 = np.stack([self.freqs_sin, self.freqs_cos], axis=-1)
        self.freqs = np.stack([self.freqs1, self.freqs2], axis=-1)

        del weight

    def __call__(self, input_ids: Array["B, L"], start_pos: int):
        # `mask` is generated only once at the beginning.
        # e.g., L = 3, start_pos = 2
        # array([[  0.,   0.,   0., -inf, -inf],
        #        [  0.,   0.,   0.,   0., -inf],
        #        [  0.,   0.,   0.,   0.,   0.]])
        _, L = input_ids.shape
        freqs: Array["1, L or 1, HD//2"] = self.freqs[0:1, start_pos: start_pos + L]
        mask: Array["L, L"] = None
        if L > 1:
            mask = np.full((L, L), float("-inf"))
            mask = np.triu(mask, k=1).astype(np.float32)
            mask = np.concatenate([np.zeros((L, start_pos)), mask], axis=1)

        # Embeddings
        h: Array["B, L or 1, D"] = self.embedding[input_ids]

        # N x Transformer Layers
        for i in range(len(self.layers)):
            h: Array["B, L or 1, D"] = self.layers[i](h, start_pos, mask, freqs)

        # RMS Norm
        h: Array["B, L or 1, D"] = rmsnorm(h, self.norm_weight, self.norm_eps)

        # Linear
        # Only forward the output from the last position.
        # (B, 1, dim) @ (dim, vocab_size) = (B, 1, vocab_size) e.g., (B, 1, 32000)
        logit: Array["B, VS"] = h[:, -1, :] @ self.lm_head_weight

        # without softmax
        next_id: Array["B, 1"] = logit.argmax(-1, keepdims=True)

        return next_id

    def generate(self, input_ids: Array["B, L"], max_new_tokens: int):
        _, L = input_ids.shape
        for i, curr_pos in enumerate(range(L, max_new_tokens)):
            if i == 0:  # Prefill Phase
                inputs = input_ids
                pos = 0
            else:  # Decode Phase
                inputs = next_id
                pos = curr_pos
            next_id = self(inputs, pos)
            yield next_id


if __name__ == '__main__':
    args = ModelArgs()

    tokenizer = Tokenizer("./tokenizer.model.np")
    model = Llama("./stories15M.model.npz", args)

    if len(sys.argv) == 1:
        prompts = ["I have a dream", "Who is she?"]
    else:
        prompts = [sys.argv[1]]

    input_ids = np.array([tokenizer.encode(p) for p in prompts])
    start = time.time()
    B, L = input_ids.shape
    outputs = []
    for id in model.generate(input_ids, args.max_new_tokens):
        L += 1
        output_id = id
        if any([o in [tokenizer.eos_id, tokenizer.bos_id] for o in output_id]):
            break
        outputs.append(output_id)
    outputs = np.stack(outputs, axis=-1).squeeze()
    for b in range(B):
        print('---')
        print(prompts[b], end="")
        print(tokenizer.decode(outputs[b].tolist()))
    elapsed = time.time() - start
    print(f"---\nToken count: {B * L}, elapsed: {elapsed:.2f}s, {round(B * L / elapsed)} tokens/s")