models/llama.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from models.norm import RMSNorm
from models.rope import precompute_freqs_cis, apply_rotary_emb
import bitsandbytes as bnb
import math


class NormalLinear(nn.Linear):
    def reset_parameters(self) -> None:
        pass


class BnbInt8Linear(bnb.nn.Linear8bitLt):
    def __init__(self, *args, **kwargs):
        super().__init__(has_fp16_weights=False, threshold=6.0, *args, **kwargs)

    def reset_parameters(self) -> None:
        pass


def get_linear_layer(use_int8):
    if use_int8:
        return BnbInt8Linear
    return NormalLinear


class WordEmbedding(nn.Module):
    def __init__(self, args):
        super(WordEmbedding, self).__init__()
        self.embedding = nn.Embedding(args.vocab_size, args.emb_size)

    def forward(self, src):
        emb = self.embedding(src)
        return emb


class MultiHeadedAttention(nn.Module):
    def __init__(self, args, hidden_size, heads_num, attention_head_size, has_bias=True, use_int8=True):
        super(MultiHeadedAttention, self).__init__()
        self.heads_num = heads_num

        self.per_head_size = attention_head_size
        self.inner_hidden_size = heads_num * attention_head_size

        Linear = get_linear_layer(use_int8)
        self.linear_layers = nn.ModuleList(
            [Linear(hidden_size, self.inner_hidden_size, bias=has_bias) for _ in range(3)]
        )

        self.final_linear = Linear(self.inner_hidden_size, hidden_size, bias=has_bias)

        # add cache to reduce compute source.
        self.cache_k = torch.zeros(
            (args.batch_size, args.seq_length, self.heads_num, self.per_head_size)
        )
        self.cache_v = torch.zeros(
            (args.batch_size, args.seq_length, self.heads_num, self.per_head_size)
        )

    def forward(self, key, value, query, start_pos, continue_exsample, mask, freqs_cis):
        batch_size, seq_length, _ = query.size()
        heads_num = self.heads_num
        per_head_size = self.per_head_size
        query, key, value = [l(x).view(batch_size, -1, heads_num, per_head_size) \
                             for l, x in zip(self.linear_layers, (query, key, value))]
        query, key = apply_rotary_emb(query, key, freqs_cis=freqs_cis)
        if self.cache_k.device != key.device:
            self.cache_k = self.cache_k.to(key)
        if self.cache_v.device != value.device:
            self.cache_v = self.cache_v.to(value)

        self.cache_k[continue_exsample, start_pos: start_pos + seq_length] = key
        self.cache_v[continue_exsample, start_pos: start_pos + seq_length] = value

        key = self.cache_k[continue_exsample, : start_pos + seq_length]
        value = self.cache_v[continue_exsample, : start_pos + seq_length]

        query, key, value = [x.transpose(1, 2) for x in (query, key, value)]

        scores = torch.matmul(query, key.transpose(-2, -1))
        scores = scores / math.sqrt(float(per_head_size))
        if mask is not None:
            scores += mask
        # probs = nn.Softmax(dim=-1)(scores)
        probs = F.softmax(scores.float(), dim=-1).type_as(query)
        output = torch.matmul(probs, value).transpose(1, 2).\
            contiguous().view(batch_size, seq_length, -1)
        return self.final_linear(output)


class GatedFeedForward(nn.Module):
    def __init__(self, hidden_size, feedforward_size, has_bias=True, use_int8=True):
        super(GatedFeedForward, self).__init__()
        Linear = get_linear_layer(use_int8)
        self.linear_gate = Linear(hidden_size, feedforward_size, bias=has_bias)
        self.linear_1 = Linear(hidden_size, feedforward_size, bias=has_bias)
        self.linear_2 = Linear(feedforward_size, hidden_size, bias=has_bias)
        self.act = F.silu

    def forward(self, x):
        # gate = self.act(self.linear_gate(x))
        gate = self.act(self.linear_gate(x)).type_as(x)
        inter_linear = self.linear_1(x)
        inter = gate * inter_linear
        output = self.linear_2(inter)
        return output


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

        if hasattr(args, "attention_head_size"):
            attention_head_size = args.attention_head_size
        else:
            attention_head_size = args.hidden_size // args.heads_num

        has_bias = bool(1 - args.remove_transformer_bias)
        # Multi-head Attention
        self.self_attn = MultiHeadedAttention(
            args, args.hidden_size, args.heads_num, attention_head_size, has_bias=has_bias,
            use_int8=args.use_int8
        )

        # FFN
        self.feed_forward = GatedFeedForward(
            args.hidden_size, args.feedforward_size, has_bias, use_int8=args.use_int8
        )

        self.layer_norm_1 = RMSNorm(args.hidden_size)
        self.layer_norm_2 = RMSNorm(args.hidden_size)

    def forward(self, hidden, start_pos, continue_exsample, mask, freqs_cis=None):
        inter = self.layer_norm_1(hidden)
        inter = self.self_attn(inter, inter, inter, start_pos, continue_exsample, mask, freqs_cis)
        hidden = hidden + inter
        output = self.layer_norm_2(hidden)
        output = self.feed_forward(output) + hidden
        return output


class TransformerEncoder(nn.Module):
    def __init__(self, args):
        super(TransformerEncoder, self).__init__()
        self.mask = args.mask
        self.layers_num = args.layers_num

        self.transformer = nn.ModuleList(
            [TransformerLayer(args) for _ in range(self.layers_num)]
        )

        self.layer_norm = RMSNorm(args.hidden_size)
        self.freqs_cis = precompute_freqs_cis(args.hidden_size // args.heads_num, args.max_seq_length * 2)

    def forward(self, emb, start_pos, continue_exsample):
        batch_size, seq_length, _ = emb.size()
        mask = None
        if seq_length > 1:
            mask = torch.ones(seq_length, seq_length, device=emb.device)
            mask = torch.tril(mask)
            mask = (1.0 - mask) * -10000
            mask = mask.repeat(batch_size, 1, 1, 1)

        hidden = emb
        freqs_cis = self.freqs_cis[start_pos: start_pos + seq_length].to(hidden.device)

        for i in range(self.layers_num):
            hidden = self.transformer[i](hidden, start_pos, continue_exsample, mask, freqs_cis=freqs_cis)
        return self.layer_norm(hidden)


class LmOutput(nn.Module):
    def __init__(self, args):
        super(LmOutput, self).__init__()
        # update: lm output not use int8
        Linear = get_linear_layer(False)
        self.lm = Linear(args.hidden_size, args.vocab_size, bias=False)

    def forward(self, x):
        return self.lm(x[:, -1, :])


class LLaMa(nn.Module):
    def __init__(self, args):
        super(LLaMa, self).__init__()
        self.embedding = WordEmbedding(args)
        self.encoder = TransformerEncoder(args)
        self.target = LmOutput(args)

    #@torch.inference_mode()
    def forward(self, src, start_pos, continue_exsample):
        emb = self.embedding(src)
        output = self.encoder(emb, start_pos, continue_exsample)
        output = self.target(output)
        return output