feat: add c4ai

Author: yunqian

lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_cohere_template.py

@@ -0,0 +1,202 @@
+from functools import partial
+from typing import Tuple
+
+import torch
+import torch.distributed as dist
+
+from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl
+from lightllm.utils.infer_utils import mark_cost_time
+
+from ...infer_struct import InferStateInfo
+from ...splitfuse_infer_struct import SplitFuseInferStateInfo
+from ..transformer_layer_infer import TransformerLayerInfer
+
+
+class TransformerLayerCohereInferTpl(TransformerLayerInferTpl):
+    """ """
+
+    def __init__(self, layer_num, tp_rank, world_size, network_config, mode):
+        super().__init__(layer_num, tp_rank, world_size, network_config, mode)
+
+        self.use_qk_norm_ = self.network_config_.get("use_qk_norm", False)
+        return
+
+    def _att_norm(
+        self, input, infer_state: InferStateInfo, layer_weight
+    ) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    def _q_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    def _k_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    def _bind_norm(
+        self, input, infer_state: InferStateInfo, layer_weight
+    ) -> torch.Tensor:
+        self._att_norm = partial(TransformerLayerCohereInferTpl._q_norm, self)
+        self._q_norm = partial(TransformerLayerCohereInferTpl._k_norm, self)
+        self._k_norm = partial(TransformerLayerCohereInferTpl._att_norm, self)
+
+    def _rotary_emb_fwd(self, q, kv, position_cos, position_sin):
+        raise Exception("need to impl")
+
+    def _bind_rotary_emb_fwd(self):
+        raise Exception("need to impl")
+
+    def _get_qkv(
+        self, input, cache_kv, infer_state: InferStateInfo, layer_weight
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        q = torch.mm(input.view(-1, self.embed_dim_), layer_weight.q_weight_)
+        torch.mm(
+            input.view(-1, self.embed_dim_),
+            layer_weight.kv_weight_,
+            out=cache_kv.view(
+                -1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_
+            ),
+        )
+        if self.use_qk_norm_:
+            q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
+            k = cache_kv[:, 0 : self.tp_k_head_num_, :]
+            q = self._q_norm(q, infer_state, layer_weight)
+            cache_kv[:, 0 : self.tp_k_head_num_, :] = self._k_norm(
+                k, infer_state, layer_weight
+            )
+        self._rotary_emb_fwd(
+            q, cache_kv, infer_state.position_cos, infer_state.position_sin
+        )
+        return q, cache_kv
+
+    def _context_attention_kernel(
+        self, q, kv, infer_state: InferStateInfo, layer_weight, out=None
+    ) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    def _token_attention_kernel(
+        self, q, infer_state: InferStateInfo, layer_weight, out=None
+    ) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    def _splitfuse_attention_kernel(
+        self, q, infer_state: SplitFuseInferStateInfo, layer_weight, out=None
+    ) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    def _get_o(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    def _ffn(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
+        raise Exception("need to impl")
+
+    @mark_cost_time(
+        "trans context flash forward time cost"
+    )  # dont to remove this, will make performence down, did not know why
+    def _context_attention(
+        self, input_embding, infer_state: InferStateInfo, layer_weight
+    ):
+        cache_kv = self._pre_cache_kv(infer_state, layer_weight)
+        q, cache_kv = self._get_qkv(input_embding, cache_kv, infer_state, layer_weight)
+        self._post_cache_kv(cache_kv, infer_state, layer_weight)
+        o = self._context_attention_kernel(q, cache_kv, infer_state, layer_weight)
+        q = None
+        o = self._get_o(o, infer_state, layer_weight)
+        if self.world_size_ > 1:
+            dist.all_reduce(o, op=dist.ReduceOp.SUM, async_op=False)
+        infer_state._attn_out = o
+        return
+
+    @mark_cost_time(
+        "trans context ffn forward time cost"
+    )  # dont to remove this, will make performence down, did not know why
+    def _context_ffn(self, input_embdings, infer_state: InferStateInfo, layer_weight):
+        ffn_out = self._ffn(input_embdings, infer_state, layer_weight)
+        if self.world_size_ > 1:
+            dist.all_reduce(ffn_out, op=dist.ReduceOp.SUM, async_op=False)
+        infer_state._ffn_out = ffn_out
+        return
+
+    # this impl dont to use @mark_cost_time
+    def _token_attention(
+        self, input_embding, infer_state: InferStateInfo, layer_weight
+    ):
+        cache_kv = self._pre_cache_kv(infer_state, layer_weight)
+        q, cache_kv = self._get_qkv(input_embding, cache_kv, infer_state, layer_weight)
+        self._post_cache_kv(cache_kv, infer_state, layer_weight)
+        o = self._token_attention_kernel(q, infer_state, layer_weight)
+        q = None
+        o = self._get_o(o, infer_state, layer_weight)
+        if self.world_size_ > 1:
+            dist.all_reduce(o, op=dist.ReduceOp.SUM, async_op=False)
+        infer_state._attn_out = o
+        return
+
+    # this impl dont to use @mark_cost_time
+    def _token_ffn(self, input_embdings, infer_state: InferStateInfo, layer_weight):
+        ffn_out = self._ffn(input_embdings, infer_state, layer_weight)
+        if self.world_size_ > 1:
+            dist.all_reduce(ffn_out, op=dist.ReduceOp.SUM, async_op=False)
+        infer_state._ffn_out = ffn_out
+        return
+
+    # @mark_cost_time("trans context flash forward time cost")  # dont to remove this, will make performence down, did not know why
+    def _splitfuse_attention(
+        self, input_embding, infer_state: SplitFuseInferStateInfo, layer_weight
+    ):
+        cache_kv = self._pre_cache_kv(infer_state, layer_weight)
+        q, cache_kv = self._get_qkv(input_embding, cache_kv, infer_state, layer_weight)
+        self._post_cache_kv(cache_kv, infer_state, layer_weight)
+        o = self._splitfuse_attention_kernel(q, infer_state, layer_weight)
+        q = None
+        o = self._get_o(o, infer_state, layer_weight)
+        if self.world_size_ > 1:
+            dist.all_reduce(o, op=dist.ReduceOp.SUM, async_op=False)
+        infer_state._attn_out = o
+        return
+
+    # @mark_cost_time("trans context ffn forward time cost")  # dont to remove this, will make performence down, did not know why
+    def _splitfuse_ffn(
+        self, input_embdings, infer_state: SplitFuseInferStateInfo, layer_weight
+    ):
+        ffn_out = self._ffn(input_embdings, infer_state, layer_weight)
+        if self.world_size_ > 1:
+            dist.all_reduce(ffn_out, op=dist.ReduceOp.SUM, async_op=False)
+        infer_state._ffn_out = ffn_out
+        return
+
+    def _cohere_residual(self, input_embdings, infer_state: InferStateInfo):
+        # emb_addr = input_embdings.data_ptr()
+        # attn_out_addr = infer_state._attn_out.data_ptr()
+        # ffn_addr = infer_state._ffn_out.data_ptr()
+        # assert emb_addr != attn_out_addr
+        # assert emb_addr != ffn_addr
+        # assert attn_out_addr != ffn_addr
+        input_embdings.add_(
+            infer_state._attn_out.view(-1, self.embed_dim_)
+            + infer_state._ffn_out.view(-1, self.embed_dim_)
+        )
+
+    def context_forward(
+        self, input_embdings, infer_state: InferStateInfo, layer_weight
+    ):
+        input1 = self._att_norm(input_embdings, infer_state, layer_weight)
+        self._context_attention(input1, infer_state, layer_weight=layer_weight)
+        self._context_ffn(input1, infer_state, layer_weight)
+        self._cohere_residual(input_embdings, infer_state)
+        return input_embdings
+
+    def token_forward(self, input_embdings, infer_state: InferStateInfo, layer_weight):
+        input1 = self._att_norm(input_embdings, infer_state, layer_weight)
+        self._token_attention(input1, infer_state, layer_weight=layer_weight)
+        self._token_ffn(input1, infer_state, layer_weight)
+        self._cohere_residual(input_embdings, infer_state)
+        return input_embdings
+
+    def splitfuse_forward(
+        self, input_embdings, infer_state: SplitFuseInferStateInfo, layer_weight
+    ):
+        input1 = self._att_norm(input_embdings, infer_state, layer_weight)
+        self._splitfuse_attention(input1, infer_state, layer_weight=layer_weight)
+        self._splitfuse_ffn(input1, infer_state, layer_weight)
+        self._cohere_residual(input_embdings, infer_state)
+        return input_embdings

lightllm/models/cohere/__init__.py

@@ -0,0 +1,8 @@
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+
+class CohereInferStateInfo(LlamaInferStateInfo):
+    def __init__(self):
+        super().__init__()
+        self._attn_out = None
+        self._ffn_out = None

lightllm/models/cohere/infer_struct.py

@@ -0,0 +1,138 @@
+import torch
+import torch.distributed as dist
+import numpy as np
+
+from lightllm.models.cohere.infer_struct import CohereInferStateInfo
+from lightllm.models.cohere.layer_weights.pre_and_post_layer_weight import CoherePreAndPostLayerWeight
+from lightllm.models.cohere.triton_kernels.layernorm import layernorm_forward, multi_head_layernorm_forward
+from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight
+from lightllm.common.basemodel.splitfuse_infer_struct import SplitFuseInferStateInfo
+
+from einops import rearrange
+from lightllm.common.basemodel import PostLayerInferTpl
+
+
+class CoherePostLayerInfer(PostLayerInferTpl):
+    def __init__(self, tp_rank, world_size, network_config, mode):
+        super().__init__(tp_rank, world_size, network_config, mode)
+        self.eps_ = network_config["layer_norm_eps"]
+        self.vocab_size_ = network_config["vocab_size"]
+        self.embed_dim_ = network_config["n_embed"]
+        self.logits_scale = network_config["logit_scale"]
+        return
+
+    def _norm(self, input, infer_state, layer_weight: CoherePreAndPostLayerWeight) -> torch.Tensor:
+        return layernorm_forward(input, layer_weight.final_norm_weight_, eps=self.eps_)
+
+    def _slice_get_last_input(self, input_embdings, infer_state: CohereInferStateInfo):
+        if infer_state.is_splitfuse:
+            # for SplitFuse
+            batch_size = infer_state.batch_size
+            last_input = torch.empty(
+                (batch_size, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype
+            )
+            tmp_ = torch.cat(
+                [
+                    torch.ones(infer_state.decode_req_num, dtype=torch.int32, device="cuda"),
+                    infer_state.prefill_b_seq_len - infer_state.prefill_b_split_ready_cache_len,
+                ],
+                dim=0,
+            )
+            last_index = torch.cumsum(tmp_, dim=0, dtype=torch.long) - 1
+            last_input[:, :] = input_embdings[last_index, :]
+            return last_input, batch_size
+
+        if infer_state.is_prefill and infer_state.is_token_healing:
+            batch_size = infer_state.batch_size
+            b_seq_len_numpy = (infer_state.b_seq_len - infer_state.b_ready_cache_len).detach().cpu().numpy()
+            select_index = []
+            start_index = 0
+            select_token_num = 0
+            for cur_len in b_seq_len_numpy:
+                if cur_len == 1:
+                    select_index.append(start_index + cur_len - 1)
+                    start_index += cur_len
+                    select_token_num += 1
+                else:
+                    select_index.append(start_index + cur_len - 2)
+                    select_index.append(start_index + cur_len - 1)
+                    start_index += cur_len
+                    select_token_num += 2
+
+            last_index = torch.tensor(select_index, dtype=torch.long, device=input_embdings.device)
+            last_input = torch.empty(
+                (select_token_num, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype
+            )
+
+            last_input[:, :] = input_embdings[last_index, :]
+            return last_input, select_token_num
+
+        if not infer_state.is_splitfuse and infer_state.is_prefill and not infer_state.return_all_prompt_logics:
+            batch_size = infer_state.batch_size
+            last_input = torch.empty(
+                (batch_size, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype
+            )
+            last_index = (
+                torch.cumsum(infer_state.b_seq_len - infer_state.b_ready_cache_len, dim=0, dtype=torch.long) - 1
+            )
+            last_input[:, :] = input_embdings[last_index, :]
+            return last_input, batch_size
+
+        if not infer_state.is_splitfuse and infer_state.is_prefill and infer_state.return_all_prompt_logics:
+            total_tokens = infer_state.total_token_num
+            return input_embdings, total_tokens
+
+        if not infer_state.is_splitfuse and not infer_state.is_prefill:
+            batch_size = infer_state.batch_size
+            return input_embdings[-batch_size:, :], batch_size
+
+        assert False, "Error State"
+
+    def soft_max(self, data):
+        return torch.softmax(data.permute(1, 0).float(), dim=-1)
+
+    def token_forward(
+        self,
+        input_embdings,
+        infer_state: CohereInferStateInfo,
+        layer_weight: CoherePreAndPostLayerWeight,
+        return_logics=False,
+    ):
+        last_input, token_num = self._slice_get_last_input(input_embdings, infer_state)
+        input_embdings_dtype = input_embdings.dtype
+        input_embdings = None
+        last_input = self._norm(last_input, infer_state, layer_weight)
+        last_input = rearrange(last_input, "batch embed_dim -> embed_dim batch").contiguous().reshape(-1, token_num)
+        logic_batch = torch.mm(layer_weight.lm_head_weight_, last_input)
+
+        last_input = None
+        if self.world_size_ == 1:
+            gather_data = logic_batch
+        else:
+            gather_data = torch.empty(
+                (self.vocab_size_, token_num), device=logic_batch.device, dtype=input_embdings_dtype
+            )
+            split_indexes = np.linspace(0, self.vocab_size_, self.world_size_ + 1, dtype=np.int64)
+            dist.all_gather(
+                [gather_data[split_indexes[i] : split_indexes[i + 1], :] for i in range(self.world_size_)],
+                logic_batch,
+                group=None,
+                async_op=False,
+            )
+        gather_data = gather_data * self.logits_scale
+        logic_batch = None
+
+        if not return_logics:
+            prob_out = self.soft_max(gather_data)
+            gather_data = None
+            return prob_out
+        else:
+            ans_logics = gather_data.permute(1, 0).float()
+            gather_data = None
+            return ans_logics
+
+    # @mark_cost_time("splitfuse post forward")
+    def splitfuse_forward(
+        self, input_embdings, infer_state: SplitFuseInferStateInfo, layer_weight: BaseLayerWeight, return_logics=False
+    ):
+        return self.token_forward(input_embdings, infer_state, layer_weight, return_logics=return_logics)