Upload Calf

f8ad943 verified almost 2 years ago

19.5 kB

import math
import struct
import inspect
import time

from .LMConfig import LMConfig
from typing import Any, Optional, Tuple
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast


class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float):
        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


def precompute_pos_cis(dim: int, end: int, theta: float = 10000.0):
    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
    pos_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return pos_cis


def apply_rotary_emb(xq, xk, pos_cis):
    def unite_shape(pos_cis, x):
        ndim = x.ndim
        assert 0 <= 1 < ndim
        assert pos_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 pos_cis.view(*shape)

    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))
    pos_cis = unite_shape(pos_cis, xq_)
    xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
    bs, slen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
    )


class Attention(nn.Module):
    def __init__(self, args: LMConfig):
        super().__init__()
        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
        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
        self.k_cache, self.v_cache = None, None
        self.attn_dropout = nn.Dropout(args.dropout)
        self.resid_dropout = nn.Dropout(args.dropout)
        self.dropout = args.dropout
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn

        # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
        mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
        mask = torch.triu(mask, diagonal=1)
        self.register_buffer("mask", mask, persistent=False)

    def forward(self, x: torch.Tensor, pos_cis: torch.Tensor, kv_cache=False):
        bsz, seqlen, _ = x.shape

        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, pos_cis)

        # 更高效的kv_cache实现
        if kv_cache and self.eval():
            if seqlen == 1 and all(cache is not None for cache in (self.k_cache, self.v_cache)):
                xk = torch.cat((self.k_cache, xk), dim=1)
                xv = torch.cat((self.v_cache, xv), dim=1)
            self.k_cache, self.v_cache = xk, xv

        xk = repeat_kv(xk, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
        xv = repeat_kv(xv, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)

        xq = xq.transpose(1, 2)
        xk = xk.transpose(1, 2)
        xv = xv.transpose(1, 2)

        if self.flash and seqlen != 1:
            output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None,
                                                                      dropout_p=self.dropout if self.training else 0.0,
                                                                      is_causal=True)
        else:
            scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
            scores = scores + self.mask[:, :, :seqlen, :seqlen]  # (bs, n_local_heads, seqlen, cache_len + seqlen)
            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
            scores = self.attn_dropout(scores)
            output = torch.matmul(scores, xv)  # (bs, n_local_heads, seqlen, head_dim)

        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)

        output = self.wo(output)
        output = self.resid_dropout(output)
        return output


class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = 4 * dim
            hidden_dim = int(2 * hidden_dim / 3)
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class MoEGate(nn.Module):
    def __init__(self, config: LMConfig):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts

        self.scoring_func = config.scoring_func
        self.alpha = config.aux_loss_alpha
        self.seq_aux = config.seq_aux

        self.norm_topk_prob = config.norm_topk_prob
        self.gating_dim = config.dim
        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        import torch.nn.init as init
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape

        hidden_states = hidden_states.view(-1, h)
        logits = F.linear(hidden_states, self.weight, None)
        if self.scoring_func == 'softmax':
            scores = logits.softmax(dim=-1)
        else:
            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')

        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

        if self.top_k > 1 and self.norm_topk_prob:
            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
            topk_weight = topk_weight / denominator

        if self.training and self.alpha > 0.0:
            scores_for_aux = scores
            aux_topk = self.top_k
            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
            if self.seq_aux:
                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
                ce.scatter_add_(1, topk_idx_for_aux_loss,
                                torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
                    seq_len * aux_topk / self.n_routed_experts)
                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
            else:
                mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
                ce = mask_ce.float().mean(0)
                Pi = scores_for_aux.mean(0)
                fi = ce * self.n_routed_experts
                aux_loss = (Pi * fi).sum() * self.alpha
        else:
            aux_loss = None
        return topk_idx, topk_weight, aux_loss


class MOEFeedForward(nn.Module):
    def __init__(self, config: LMConfig):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            FeedForward(
                dim=config.dim,
                hidden_dim=config.hidden_dim,
                multiple_of=config.multiple_of,
                dropout=config.dropout,
            )
            for _ in range(config.n_routed_experts)
        ])

        self.gate = MoEGate(config)
        if config.n_shared_experts is not None:
            self.shared_experts = FeedForward(
                dim=config.dim,
                hidden_dim=config.hidden_dim,
                multiple_of=config.multiple_of,
                dropout=config.dropout,
            )

    def forward(self, x):
        identity = x
        orig_shape = x.shape
        bsz, seq_len, _ = x.shape

        # 使用门控机制选择专家
        topk_idx, topk_weight, aux_loss = self.gate(x)

        x = x.view(-1, x.shape[-1])
        flat_topk_idx = topk_idx.view(-1)

        if self.training:
            # 训练模式下，重复输入数据
            x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
            y = torch.empty_like(x, dtype=torch.float16)
            for i, expert in enumerate(self.experts):
                y[flat_topk_idx == i] = expert(x[flat_topk_idx == i])
            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
            y = y.view(*orig_shape)
        else:
            # 推理模式下，只选择最优专家
            y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)

        if self.config.n_shared_experts is not None:
            y = y + self.shared_experts(identity)

        return y

    @torch.no_grad()
    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
        token_idxs = idxs // self.config.num_experts_per_tok
        # 例如当tokens_per_expert=[6, 15, 20, 26, 33, 38, 46, 52]
        # 当token_idxs=[3, 7, 19, 21, 24, 25,  4,  5,  6, 10, 11, 12...]
        # 意味着当token_idxs[:6] -> [3,  7, 19, 21, 24, 25,  4]位置的token都由专家0处理，token_idxs[6:15]位置的token都由专家1处理......
        for i, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
            if start_idx == end_idx:
                continue
            expert = self.experts[i]
            exp_token_idx = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idx]
            expert_out = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            # 使用 scatter_add_ 进行 sum 操作
            expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)

        return expert_cache


class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: LMConfig):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)

        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

        if args.use_moe:
            self.feed_forward = MOEFeedForward(args)
        else:
            self.feed_forward = FeedForward(
                dim=args.dim,
                hidden_dim=args.hidden_dim,
                multiple_of=args.multiple_of,
                dropout=args.dropout,
            )

    def forward(self, x, pos_cis, kv_cache=False):
        h = x + self.attention(self.attention_norm(x), pos_cis, kv_cache)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out


class VisionProj(nn.Module):
    def __init__(self, vision_out_dim=768, lm_dim=512, image_ids=[1, 2, 3, 4]):
        super().__init__()
        self.vision_out_dim = vision_out_dim
        self.lm_dim = lm_dim
        self.image_ids = image_ids
        self.vision_proj = nn.Sequential(
            nn.Linear(self.vision_out_dim, self.lm_dim),
        )

    def forward(self, image_encoders):
        vision_proj = self.vision_proj(image_encoders)
        return vision_proj


class Calf(PreTrainedModel):
    config_class = LMConfig
    last_loss: Optional[torch.Tensor]

    def __init__(self, params: LMConfig = None):
        super().__init__(params)
        if not params:
            params = LMConfig()
        self.params = params
        self.vocab_size = params.vocab_size
        self.n_layers = params.n_layers
        # image的特殊占位符，对应每张图切分成M个token，和get_img_process中的数量对应
        self.image_ids = params.image_ids

        self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
        self.dropout = nn.Dropout(params.dropout)
        self.layers = torch.nn.ModuleList()
        for layer_id in range(self.n_layers):
            self.layers.append(TransformerBlock(layer_id, params))
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
        self.tok_embeddings.weight = self.output.weight
        pos_cis = precompute_pos_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len)
        self.register_buffer("pos_cis", pos_cis, persistent=False)

        self.apply(self._init_weights)

        for pn, p in self.named_parameters():
            if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * params.n_layers))

        self.last_loss = None
        self.OUT = CausalLMOutputWithPast()
        self._no_split_modules = [name for name, _ in self.named_modules()]

        self.vision_proj = VisionProj(768, params.dim, self.image_ids)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    # VLM
    def count_vision_proj(self, tokens, h, image_encoders=None, seqlen=200):
        # 查找token中<image>片段的索引，为了替换做准备
        def find_indices(tokens, image_ids):
            image_ids_tensor = torch.tensor(image_ids).to(tokens.device)
            return [
                [i, i + len(image_ids) - 1]
                for batch_idx in range(tokens.size(0))
                for i in range(tokens.size(1) - len(image_ids) + 1)
                if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor)
            ] or None

        image_indices = find_indices(tokens, self.image_ids)

        # 如果此时有图像编码
        if image_encoders is not None:
            vision_proj = self.vision_proj(image_encoders)
            if image_indices is not None:
                # 创建一个新的张量来存储拼接后的结果
                new_h = []
                for i in range(h.size(0)):
                    before = h[i, :image_indices[i][0], :]
                    after = h[i, image_indices[i][1] + 1:, :]
                    # 拼接 before, vision_proj, after
                    new_h_i = torch.cat((before, vision_proj[i], after), dim=0)[:seqlen]
                    new_h.append(new_h_i)
                # 将所有拼接后的结果堆叠起来
                new_h = torch.stack(new_h, dim=0)
                return new_h

        return h

    def forward(self, tokens: Optional[torch.Tensor] = None, targets: Optional[torch.Tensor] = None,
                kv_cache=False, image_encoders=None, **keyargs):
        current_idx = 0
        if 'input_ids' in keyargs:
            tokens = keyargs['input_ids']
        if 'attention_mask' in keyargs:
            targets = keyargs['attention_mask']
        if 'current_idx' in keyargs:
            current_idx = int(keyargs['current_idx'])

        _bsz, seqlen = tokens.shape
        # language proj token
        h = self.tok_embeddings(tokens)
        h = self.dropout(h)
        # vision proj token
        h = self.count_vision_proj(tokens=tokens, h=h, image_encoders=image_encoders, seqlen=seqlen)

        pos_cis = self.pos_cis[current_idx:current_idx + seqlen]
        for idx, layer in enumerate(self.layers):
            h = layer(h, pos_cis, kv_cache)

        h = self.norm(h)

        if targets is not None:
            logits = self.output(h)
            self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=0,
                                             reduction='none')
        else:
            logits = self.output(h[:, [-1], :])
            self.last_loss = None

        self.OUT.__setitem__('logits', logits)
        self.OUT.__setitem__('last_loss', self.last_loss)
        return self.OUT

    @torch.inference_mode()
    def generate(self, idx, eos, max_new_tokens, temperature=0.7, top_k=8, stream=True, rp=1., kv_cache=True,
                 image_encoders=None):
        # rp: repetition_penalty
        index = idx.shape[1]
        init_inference = True
        while idx.shape[1] < max_new_tokens - 1:
            if init_inference or not kv_cache:
                inference_res, init_inference = self(idx, kv_cache=kv_cache, image_encoders=image_encoders), False
            else:
                inference_res = self(idx[:, -1:], kv_cache=kv_cache, current_idx=idx.shape[1] - 1)

            logits = inference_res.logits
            logits = logits[:, -1, :]

            for token in set(idx.tolist()[0]):
                logits[:, token] /= rp

            if temperature == 0.0:
                _, idx_next = torch.topk(logits, k=1, dim=-1)
            else:
                logits = logits / temperature
                if top_k is not None:
                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                    logits[logits < v[:, [-1]]] = -float('Inf')

                probs = F.softmax(logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1, generator=None)

            if idx_next == eos:
                break

            idx = torch.cat((idx, idx_next), dim=1)
            if stream:
                yield idx[:, index:]

        if not stream:
            yield idx[:, index:]

    @torch.inference_mode()
    def eval_answer(self, idx):
        idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:]
        inference_res = self(idx_cond)
        logits = inference_res.logits
        logits = logits[:, -1, :]
        return logits

	import math
	import struct
	import inspect
	import time

	from .LMConfig import LMConfig
	from typing import Any, Optional, Tuple
	import numpy as np
	import torch
	import torch.nn.functional as F
	from torch import nn
	from transformers import PreTrainedModel
	from transformers.modeling_outputs import CausalLMOutputWithPast


	class RMSNorm(torch.nn.Module):
	def __init__(self, dim: int, eps: float):
	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


	def precompute_pos_cis(dim: int, end: int, theta: float = 10000.0):
	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
	pos_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
	return pos_cis


	def apply_rotary_emb(xq, xk, pos_cis):
	def unite_shape(pos_cis, x):
	ndim = x.ndim
	assert 0 <= 1 < ndim
	assert pos_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 pos_cis.view(*shape)

	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))
	pos_cis = unite_shape(pos_cis, xq_)
	xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
	xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
	return xq_out.type_as(xq), xk_out.type_as(xk)


	def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
	"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
	bs, slen, n_kv_heads, head_dim = x.shape
	if n_rep == 1:
	return x
	return (
	x[:, :, :, None, :]
	.expand(bs, slen, n_kv_heads, n_rep, head_dim)
	.reshape(bs, slen, n_kv_heads * n_rep, head_dim)
	)


	class Attention(nn.Module):
	def __init__(self, args: LMConfig):
	super().__init__()
	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
	self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
	self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
	self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
	self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
	self.k_cache, self.v_cache = None, None
	self.attn_dropout = nn.Dropout(args.dropout)
	self.resid_dropout = nn.Dropout(args.dropout)
	self.dropout = args.dropout
	self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn

	# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
	mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
	mask = torch.triu(mask, diagonal=1)
	self.register_buffer("mask", mask, persistent=False)

	def forward(self, x: torch.Tensor, pos_cis: torch.Tensor, kv_cache=False):
	bsz, seqlen, _ = x.shape

	xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

	xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
	xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
	xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

	xq, xk = apply_rotary_emb(xq, xk, pos_cis)

	# 更高效的kv_cache实现
	if kv_cache and self.eval():
	if seqlen == 1 and all(cache is not None for cache in (self.k_cache, self.v_cache)):
	xk = torch.cat((self.k_cache, xk), dim=1)
	xv = torch.cat((self.v_cache, xv), dim=1)
	self.k_cache, self.v_cache = xk, xv

	xk = repeat_kv(xk, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)
	xv = repeat_kv(xv, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)

	xq = xq.transpose(1, 2)
	xk = xk.transpose(1, 2)
	xv = xv.transpose(1, 2)

	if self.flash and seqlen != 1:
	output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None,
	dropout_p=self.dropout if self.training else 0.0,
	is_causal=True)
	else:
	scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
	scores = scores + self.mask[:, :, :seqlen, :seqlen] # (bs, n_local_heads, seqlen, cache_len + seqlen)
	scores = F.softmax(scores.float(), dim=-1).type_as(xq)
	scores = self.attn_dropout(scores)
	output = torch.matmul(scores, xv) # (bs, n_local_heads, seqlen, head_dim)

	output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)

	output = self.wo(output)
	output = self.resid_dropout(output)
	return output


	class FeedForward(nn.Module):
	def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
	super().__init__()
	if hidden_dim is None:
	hidden_dim = 4 * dim
	hidden_dim = int(2 * hidden_dim / 3)
	hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
	self.w1 = nn.Linear(dim, hidden_dim, bias=False)
	self.w2 = nn.Linear(hidden_dim, dim, bias=False)
	self.w3 = nn.Linear(dim, hidden_dim, bias=False)
	self.dropout = nn.Dropout(dropout)

	def forward(self, x):
	return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))


	class MoEGate(nn.Module):
	def __init__(self, config: LMConfig):
	super().__init__()
	self.config = config
	self.top_k = config.num_experts_per_tok
	self.n_routed_experts = config.n_routed_experts

	self.scoring_func = config.scoring_func
	self.alpha = config.aux_loss_alpha
	self.seq_aux = config.seq_aux

	self.norm_topk_prob = config.norm_topk_prob
	self.gating_dim = config.dim
	self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
	self.reset_parameters()

	def reset_parameters(self) -> None:
	import torch.nn.init as init
	init.kaiming_uniform_(self.weight, a=math.sqrt(5))

	def forward(self, hidden_states):
	bsz, seq_len, h = hidden_states.shape

	hidden_states = hidden_states.view(-1, h)
	logits = F.linear(hidden_states, self.weight, None)
	if self.scoring_func == 'softmax':
	scores = logits.softmax(dim=-1)
	else:
	raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')

	topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

	if self.top_k > 1 and self.norm_topk_prob:
	denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
	topk_weight = topk_weight / denominator

	if self.training and self.alpha > 0.0:
	scores_for_aux = scores
	aux_topk = self.top_k
	topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
	if self.seq_aux:
	scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
	ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
	ce.scatter_add_(1, topk_idx_for_aux_loss,
	torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
	seq_len * aux_topk / self.n_routed_experts)
	aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
	else:
	mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
	ce = mask_ce.float().mean(0)
	Pi = scores_for_aux.mean(0)
	fi = ce * self.n_routed_experts
	aux_loss = (Pi * fi).sum() * self.alpha
	else:
	aux_loss = None
	return topk_idx, topk_weight, aux_loss


	class MOEFeedForward(nn.Module):
	def __init__(self, config: LMConfig):
	super().__init__()
	self.config = config
	self.experts = nn.ModuleList([
	FeedForward(
	dim=config.dim,
	hidden_dim=config.hidden_dim,
	multiple_of=config.multiple_of,
	dropout=config.dropout,
	)
	for _ in range(config.n_routed_experts)
	])

	self.gate = MoEGate(config)
	if config.n_shared_experts is not None:
	self.shared_experts = FeedForward(
	dim=config.dim,
	hidden_dim=config.hidden_dim,
	multiple_of=config.multiple_of,
	dropout=config.dropout,
	)

	def forward(self, x):
	identity = x
	orig_shape = x.shape
	bsz, seq_len, _ = x.shape

	# 使用门控机制选择专家
	topk_idx, topk_weight, aux_loss = self.gate(x)

	x = x.view(-1, x.shape[-1])
	flat_topk_idx = topk_idx.view(-1)

	if self.training:
	# 训练模式下，重复输入数据
	x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
	y = torch.empty_like(x, dtype=torch.float16)
	for i, expert in enumerate(self.experts):
	y[flat_topk_idx == i] = expert(x[flat_topk_idx == i])
	y = (y.view(topk_weight.shape, -1) topk_weight.unsqueeze(-1)).sum(dim=1)
	y = y.view(*orig_shape)
	else:
	# 推理模式下，只选择最优专家
	y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)

	if self.config.n_shared_experts is not None:
	y = y + self.shared_experts(identity)

	return y

	@torch.no_grad()
	def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
	expert_cache = torch.zeros_like(x)
	idxs = flat_expert_indices.argsort()
	tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
	token_idxs = idxs // self.config.num_experts_per_tok
	# 例如当tokens_per_expert=[6, 15, 20, 26, 33, 38, 46, 52]
	# 当token_idxs=[3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...]
	# 意味着当token_idxs[:6] -> [3, 7, 19, 21, 24, 25, 4]位置的token都由专家0处理，token_idxs[6:15]位置的token都由专家1处理......
	for i, end_idx in enumerate(tokens_per_expert):
	start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
	if start_idx == end_idx:
	continue
	expert = self.experts[i]
	exp_token_idx = token_idxs[start_idx:end_idx]
	expert_tokens = x[exp_token_idx]
	expert_out = expert(expert_tokens)
	expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
	# 使用 scatter_add_ 进行 sum 操作
	expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)

	return expert_cache


	class TransformerBlock(nn.Module):
	def __init__(self, layer_id: int, args: LMConfig):
	super().__init__()
	self.n_heads = args.n_heads
	self.dim = args.dim
	self.head_dim = args.dim // args.n_heads
	self.attention = Attention(args)

	self.layer_id = layer_id
	self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
	self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

	if args.use_moe:
	self.feed_forward = MOEFeedForward(args)
	else:
	self.feed_forward = FeedForward(
	dim=args.dim,
	hidden_dim=args.hidden_dim,
	multiple_of=args.multiple_of,
	dropout=args.dropout,
	)

	def forward(self, x, pos_cis, kv_cache=False):
	h = x + self.attention(self.attention_norm(x), pos_cis, kv_cache)
	out = h + self.feed_forward(self.ffn_norm(h))
	return out


	class VisionProj(nn.Module):
	def __init__(self, vision_out_dim=768, lm_dim=512, image_ids=[1, 2, 3, 4]):
	super().__init__()
	self.vision_out_dim = vision_out_dim
	self.lm_dim = lm_dim
	self.image_ids = image_ids
	self.vision_proj = nn.Sequential(
	nn.Linear(self.vision_out_dim, self.lm_dim),
	)

	def forward(self, image_encoders):
	vision_proj = self.vision_proj(image_encoders)
	return vision_proj


	class Calf(PreTrainedModel):
	config_class = LMConfig
	last_loss: Optional[torch.Tensor]

	def __init__(self, params: LMConfig = None):
	super().__init__(params)
	if not params:
	params = LMConfig()
	self.params = params
	self.vocab_size = params.vocab_size
	self.n_layers = params.n_layers
	# image的特殊占位符，对应每张图切分成M个token，和get_img_process中的数量对应
	self.image_ids = params.image_ids

	self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
	self.dropout = nn.Dropout(params.dropout)
	self.layers = torch.nn.ModuleList()
	for layer_id in range(self.n_layers):
	self.layers.append(TransformerBlock(layer_id, params))
	self.norm = RMSNorm(params.dim, eps=params.norm_eps)
	self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
	self.tok_embeddings.weight = self.output.weight
	pos_cis = precompute_pos_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len)
	self.register_buffer("pos_cis", pos_cis, persistent=False)

	self.apply(self._init_weights)

	for pn, p in self.named_parameters():
	if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
	torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * params.n_layers))

	self.last_loss = None
	self.OUT = CausalLMOutputWithPast()
	self._no_split_modules = [name for name, _ in self.named_modules()]

	self.vision_proj = VisionProj(768, params.dim, self.image_ids)

	def _init_weights(self, module):
	if isinstance(module, nn.Linear):
	torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
	if module.bias is not None:
	torch.nn.init.zeros_(module.bias)
	elif isinstance(module, nn.Embedding):
	torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

	# VLM
	def count_vision_proj(self, tokens, h, image_encoders=None, seqlen=200):
	# 查找token中<image>片段的索引，为了替换做准备
	def find_indices(tokens, image_ids):
	image_ids_tensor = torch.tensor(image_ids).to(tokens.device)
	return [
	[i, i + len(image_ids) - 1]
	for batch_idx in range(tokens.size(0))
	for i in range(tokens.size(1) - len(image_ids) + 1)
	if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor)
	] or None

	image_indices = find_indices(tokens, self.image_ids)

	# 如果此时有图像编码
	if image_encoders is not None:
	vision_proj = self.vision_proj(image_encoders)
	if image_indices is not None:
	# 创建一个新的张量来存储拼接后的结果
	new_h = []
	for i in range(h.size(0)):
	before = h[i, :image_indices[i][0], :]
	after = h[i, image_indices[i][1] + 1:, :]
	# 拼接 before, vision_proj, after
	new_h_i = torch.cat((before, vision_proj[i], after), dim=0)[:seqlen]
	new_h.append(new_h_i)
	# 将所有拼接后的结果堆叠起来
	new_h = torch.stack(new_h, dim=0)
	return new_h

	return h

	def forward(self, tokens: Optional[torch.Tensor] = None, targets: Optional[torch.Tensor] = None,
	kv_cache=False, image_encoders=None, **keyargs):
	current_idx = 0
	if 'input_ids' in keyargs:
	tokens = keyargs['input_ids']
	if 'attention_mask' in keyargs:
	targets = keyargs['attention_mask']
	if 'current_idx' in keyargs:
	current_idx = int(keyargs['current_idx'])

	_bsz, seqlen = tokens.shape
	# language proj token
	h = self.tok_embeddings(tokens)
	h = self.dropout(h)
	# vision proj token
	h = self.count_vision_proj(tokens=tokens, h=h, image_encoders=image_encoders, seqlen=seqlen)

	pos_cis = self.pos_cis[current_idx:current_idx + seqlen]
	for idx, layer in enumerate(self.layers):
	h = layer(h, pos_cis, kv_cache)

	h = self.norm(h)

	if targets is not None:
	logits = self.output(h)
	self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=0,
	reduction='none')
	else:
	logits = self.output(h[:, [-1], :])
	self.last_loss = None

	self.OUT.__setitem__('logits', logits)
	self.OUT.__setitem__('last_loss', self.last_loss)
	return self.OUT

	@torch.inference_mode()
	def generate(self, idx, eos, max_new_tokens, temperature=0.7, top_k=8, stream=True, rp=1., kv_cache=True,
	image_encoders=None):
	# rp: repetition_penalty
	index = idx.shape[1]
	init_inference = True
	while idx.shape[1] < max_new_tokens - 1:
	if init_inference or not kv_cache:
	inference_res, init_inference = self(idx, kv_cache=kv_cache, image_encoders=image_encoders), False
	else:
	inference_res = self(idx[:, -1:], kv_cache=kv_cache, current_idx=idx.shape[1] - 1)

	logits = inference_res.logits
	logits = logits[:, -1, :]

	for token in set(idx.tolist()[0]):
	logits[:, token] /= rp

	if temperature == 0.0:
	_, idx_next = torch.topk(logits, k=1, dim=-1)
	else:
	logits = logits / temperature
	if top_k is not None:
	v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
	logits[logits < v[:, [-1]]] = -float('Inf')

	probs = F.softmax(logits, dim=-1)
	idx_next = torch.multinomial(probs, num_samples=1, generator=None)

	if idx_next == eos:
	break

	idx = torch.cat((idx, idx_next), dim=1)
	if stream:
	yield idx[:, index:]

	if not stream:
	yield idx[:, index:]

	@torch.inference_mode()
	def eval_answer(self, idx):
	idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:]
	inference_res = self(idx_cond)
	logits = inference_res.logits
	logits = logits[:, -1, :]
	return logits