[gemma_2b] support (#346)

to be confirmed: 
in triton attention operator for context and decode,
because the head_dim == 256,
so "assert Lk in {16, 32, 64, 128}" -> "assert Lk in {16, 32, 64, 128,
256}"
if it will affect the performance?

Author: ANDgate99

lightllm/models/gemma_2b/__init__.py

@@ -0,0 +1,52 @@
+import torch
+import torch.distributed as dist
+import numpy as np
+
+from lightllm.common.basemodel.splitfuse_infer_struct import SplitFuseInferStateInfo
+from lightllm.models.gemma_2b.layer_weights.pre_and_post_layer_weight import Gemma_2bPreAndPostLayerWeight
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+from lightllm.common.basemodel import PreLayerInferTpl
+from lightllm.utils.infer_utils import mark_cost_time
+
+
+class Gemma_2bPreLayerInfer(PreLayerInferTpl):
+    """ """
+
+    def __init__(self, tp_rank, world_size, network_config, mode):
+        super().__init__(tp_rank, world_size, network_config, mode)
+        tp_vob_ids = np.linspace(0, network_config["vocab_size"], self.world_size_ + 1, dtype=np.int64)
+        self.vob_start_id_, self.vob_end_id_ = int(tp_vob_ids[self.tp_rank_]), int(tp_vob_ids[self.tp_rank_ + 1])
+        self.normfactor = network_config["hidden_size"]**0.5
+        return
+
+    def _norm(self, input, infer_state, layer_weight : Gemma_2bPreAndPostLayerWeight) -> torch.Tensor:
+        return input * self.normfactor
+
+    @mark_cost_time("pre context forward")
+    def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: Gemma_2bPreAndPostLayerWeight):
+        input_mask = torch.logical_or(self.vob_start_id_ > input_ids, input_ids >= self.vob_end_id_)
+        tmp_input_ids = input_ids - self.vob_start_id_
+        tmp_input_ids[input_mask] = 0
+        input_embdings = torch.embedding(layer_weight.wte_weight_, tmp_input_ids, padding_idx=-1)
+        input_embdings[input_mask] = 0.0
+        if self.world_size_ > 1:
+            dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False)
+        input_embdings = self._norm(input_embdings, infer_state, layer_weight)
+        return input_embdings
+
+    def token_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: Gemma_2bPreAndPostLayerWeight):
+        input_mask = torch.logical_or(self.vob_start_id_ > input_ids, input_ids >= self.vob_end_id_)
+        tmp_input_ids = input_ids - self.vob_start_id_
+        tmp_input_ids[input_mask] = 0
+        input_embdings = torch.embedding(layer_weight.wte_weight_, tmp_input_ids, padding_idx=-1)
+        input_embdings[input_mask] = 0.0
+        if self.world_size_ > 1:
+            dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False)
+        input_embdings = self._norm(input_embdings, infer_state, layer_weight)
+        return input_embdings
+
+    # @mark_cost_time("splitfuse forward")
+    def splitfuse_forward(
+        self, input_ids, infer_state: SplitFuseInferStateInfo, layer_weight: Gemma_2bPreAndPostLayerWeight
+    ):
+        return self.token_forward(input_ids, infer_state, layer_weight)

lightllm/models/gemma_2b/layer_infer/__init__.py

@@ -0,0 +1,32 @@
+import torch
+import torch.functional as F
+import torch.distributed as dist
+import numpy as np
+from typing import Tuple
+from functools import partial
+import triton
+
+from lightllm.models.gemma_2b.layer_weights.transformer_layer_weight import Gemma_2bTransformerLayerWeight
+from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
+from lightllm.models.gemma_2b.triton_kernel.gelu_and_mul import gelu_and_mul_fwd
+
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+
+class Gemma_2bTransformerLayerInfer(LlamaTransformerLayerInfer):
+    """ """
+
+    def __init__(self, layer_num, tp_rank, world_size, network_config, mode=[]):
+        super().__init__(layer_num, tp_rank, world_size, network_config, mode)
+        self.tp_k_head_num_ = network_config["num_key_value_heads"] # [SYM] always == 1
+        self.tp_v_head_num_ = network_config["num_key_value_heads"]
+        return
+
+    def _ffn(self, input, infer_state: LlamaInferStateInfo, layer_weight: Gemma_2bTransformerLayerWeight) -> torch.Tensor:
+        up_gate_out = torch.mm(input.view(-1, self.embed_dim_), layer_weight.gate_up_proj)
+        ffn1_out = gelu_and_mul_fwd(up_gate_out)
+        input = None
+        up_gate_out = None
+        ffn2_out = torch.mm(ffn1_out, layer_weight.down_proj)
+        ffn1_out = None
+        return ffn2_out

lightllm/models/gemma_2b/layer_infer/pre_layer_infer.py

@@ -0,0 +1,28 @@
+import torch
+import numpy as np
+from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
+
+
+class Gemma_2bPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
+    def __init__(self, tp_rank, world_size, data_type, network_config, mode):
+        super().__init__(tp_rank, world_size, data_type, network_config, mode)
+        return
+
+    def load_hf_weights(self, weights):
+        vob_size = self.network_config_["vocab_size"]
+        split_indexes = np.linspace(0, vob_size, self.world_size_ + 1, dtype=np.int64)
+        split_start = split_indexes[self.tp_rank_]
+        split_end = split_indexes[self.tp_rank_ + 1]
+        if "model.embed_tokens.weight" in weights:
+            # print(weights['model.embed_tokens.weight'].shape)
+            self.wte_weight_ = self._cuda(weights["model.embed_tokens.weight"][split_start:split_end, :])
+        if "lm_head.weight" in weights:
+            # print(weights['lm_head.weight'].shape)
+            self.lm_head_weight_ = self._cuda(weights["lm_head.weight"][split_start:split_end, :])
+        else:
+            self.lm_head_weight_ = self.wte_weight_
+        if "model.norm.weight" in weights:
+            self.final_norm_weight_ = self._cuda(weights["model.norm.weight"])
+            self.final_norm_weight_ = self.final_norm_weight_ + 1
+
+        return

lightllm/models/gemma_2b/layer_infer/transformer_layer_infer.py

@@ -0,0 +1,74 @@
+import torch
+import math
+import numpy as np
+from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight
+
+
+class Gemma_2bTransformerLayerWeight(LlamaTransformerLayerWeight):
+    def __init__(self, layer_num, tp_rank, world_size, data_type, network_config, mode=[]):
+        super().__init__(layer_num, tp_rank, world_size, data_type, network_config, mode)
+        return
+
+    def _load_qkvo_weights(self, weights):
+        # input layernorm params
+        if f"model.layers.{self.layer_num_}.input_layernorm.weight" in weights:
+            self.att_norm_weight_ = self._cuda(weights[f"model.layers.{self.layer_num_}.input_layernorm.weight"])
+            self.att_norm_weight_ += 1
+
+        n_embed = self.network_config_["hidden_size"]
+        q_split_n_embed = n_embed // self.world_size_
+
+        # q k v weights for llama
+        if f"model.layers.{self.layer_num_}.self_attn.q_proj.weight" in weights:
+            self.q_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.q_proj.weight"]
+            self.q_weight_ = self.q_weight_[q_split_n_embed * self.tp_rank_ : q_split_n_embed * (self.tp_rank_ + 1), :]
+            self.q_weight_ = self._cuda(self.q_weight_.transpose(0, 1))
+
+        if f"model.layers.{self.layer_num_}.self_attn.k_proj.weight" in weights:
+            k_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.k_proj.weight"]
+            self.k_weight_ = k_weight_.transpose(0, 1)
+
+        if f"model.layers.{self.layer_num_}.self_attn.v_proj.weight" in weights:
+            v_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.v_proj.weight"]
+            self.v_weight_ = v_weight_.transpose(0, 1)
+
+        # attention output dense params
+        if f"model.layers.{self.layer_num_}.self_attn.o_proj.weight" in weights:
+            self.o_weight_ = weights[f"model.layers.{self.layer_num_}.self_attn.o_proj.weight"]
+            self.o_weight_ = self.o_weight_[:, q_split_n_embed * self.tp_rank_ : q_split_n_embed * (self.tp_rank_ + 1)]
+            self.o_weight_ = self._cuda(self.o_weight_.transpose(0, 1))
+
+        self._try_cat_to(["k_weight_", "v_weight_"], "kv_weight_", cat_dim=1)
+
+        return
+
+    def _load_ffn_weights(self, weights):
+        if f"model.layers.{self.layer_num_}.post_attention_layernorm.weight" in weights:
+            self.ffn_norm_weight_ = self._cuda(
+                weights[f"model.layers.{self.layer_num_}.post_attention_layernorm.weight"]
+            )
+            self.ffn_norm_weight_ += 1
+
+        inter_size = self.network_config_["intermediate_size"]
+        split_inter_size = inter_size // self.world_size_
+
+        if f"model.layers.{self.layer_num_}.mlp.up_proj.weight" in weights:
+            up_proj = weights[f"model.layers.{self.layer_num_}.mlp.up_proj.weight"][
+                split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1), :
+            ]
+            self.up_proj = up_proj.transpose(0, 1)
+
+        if f"model.layers.{self.layer_num_}.mlp.gate_proj.weight" in weights:
+            gate_proj = weights[f"model.layers.{self.layer_num_}.mlp.gate_proj.weight"][
+                split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1), :
+            ]
+            self.gate_proj = gate_proj.transpose(0, 1)
+
+        self._try_cat_to(["gate_proj", "up_proj"], "gate_up_proj", cat_dim=1)
+
+        if f"model.layers.{self.layer_num_}.mlp.down_proj.weight" in weights:
+            self.down_proj = weights[f"model.layers.{self.layer_num_}.mlp.down_proj.weight"][
+                :, split_inter_size * self.tp_rank_ : split_inter_size * (self.tp_rank_ + 1)
+            ]
+            self.down_proj = self._cuda(self.down_proj.transpose(0, 1))
+        return

lightllm/models/gemma_2b/layer_weights/__init__.py

@@ -0,0 +1,79 @@
+import os
+import json
+import torch
+
+from lightllm.common.basemodel import TpPartBaseModel
+from lightllm.models.gemma_2b.layer_weights.transformer_layer_weight import Gemma_2bTransformerLayerWeight
+from lightllm.models.gemma_2b.layer_weights.pre_and_post_layer_weight import Gemma_2bPreAndPostLayerWeight
+from lightllm.models.gemma_2b.layer_infer.pre_layer_infer import Gemma_2bPreLayerInfer
+from lightllm.models.gemma_2b.layer_infer.transformer_layer_infer import Gemma_2bTransformerLayerInfer
+from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+from lightllm.common.mem_utils import MemoryManager
+
+class Gemma_2bTpPartModel(TpPartBaseModel):
+    # weight class
+    pre_and_post_weight_class = Gemma_2bPreAndPostLayerWeight
+    transformer_weight_class = Gemma_2bTransformerLayerWeight
+
+    # infer class
+    pre_layer_infer_class = Gemma_2bPreLayerInfer
+    post_layer_infer_class = LlamaPostLayerInfer
+    transformer_layer_infer_class = Gemma_2bTransformerLayerInfer
+
+    # infer state class
+    infer_state_class = LlamaInferStateInfo
+
+    def __init__(self, kvargs):
+        super().__init__(kvargs)
+        return
+
+    def _init_config(self):
+        super()._init_config()
+        return
+
+    def _verify_params(self):
+        assert self.load_way in ["HF"], "gemma only supports HF format to load Now!"
+        # assert self.config["num_key_value_heads"] % self.world_size_ == 0
+        assert self.config["num_attention_heads"] % self.world_size_ == 0
+        return
+    
+    def _init_custom(self):
+        self._init_to_get_rotary()
+        return
+    
+    def _init_mem_manager(self):
+        self.mem_manager = MemoryManager(self.max_total_token_num,
+                                        dtype=torch.float16,
+                                        head_num=self.config["num_key_value_heads"], # [SYM] always == 1
+                                        head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
+                                        layer_num=self.config["num_hidden_layers"])       
+        return
+
+
+    def _init_to_get_rotary(self, default_base=10000):
+        if self.config.get("rope_scaling", {}) is None:
+            rope_scaling_factor = 1.0
+        else:
+            rope_scaling_factor = self.config.get("rope_scaling", {}).get("factor", 1.0)
+
+        base = self.config.get("rope_theta", float(default_base))
+
+        if "max_sequence_length" in self.config:
+            max_seq_len = self.config["max_sequence_length"]
+        else:
+            max_position_embeddings = self.config.get(
+                "max_position_embeddings",
+                2048 if base <= 10000.0 + 1e-5 else 16384
+            )
+            max_seq_len = max_position_embeddings * rope_scaling_factor
+
+        inv_freq = 1.0 / (base ** (torch.arange(0, self.head_dim_, 2, device="cpu", dtype=torch.float32) / self.head_dim_))
+        t = torch.arange(max_seq_len + 1024 * 64, device="cpu", dtype=torch.float32) / rope_scaling_factor
+        freqs = torch.outer(t, inv_freq)
+
+        self._cos_cached = torch.cos(freqs).to(torch.float16).cuda()
+        self._sin_cached = torch.sin(freqs).to(torch.float16).cuda()
+        return
+    

lightllm/models/gemma_2b/layer_weights/pre_and_post_layer_weight.py

@@ -0,0 +1,114 @@
+import torch
+import math
+import triton
+import triton.language as tl
+
+# copy from xformers impl.
+_kAlpha = math.sqrt(2.0 / math.pi)
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+@triton.jit
+def gelu(x):
+    """
+    GeLU_ activation - Gaussian error linear unit
+
+    .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf
+    """
+    return 0.5 * x * (1 + tanh(_kAlpha * (x + 0.044715 * x * x * x)))
+
+@triton.jit
+def _gelu_and_mul_kernel(
+    input_ptr,
+    stride_input_m,
+    stride_input_n,
+    stride_output_m,
+    stride_output_n,
+    size_m,
+    size_n,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    tid = tl.program_id(0)
+    input_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M)
+    output_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    pid = tl.program_id(1)
+    input_n_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    output_n_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    up_offsets = input_m_offsets[:, None] * stride_input_m + (input_n_offsets[None, :] + size_n) * stride_input_n
+    gate_offsets = input_m_offsets[:, None] * stride_input_m + input_n_offsets[None, :] * stride_input_n
+    res_offsets = output_m_offsets[:, None] * stride_output_m + output_n_offsets[None, :] * stride_output_n
+
+    up = tl.load(
+        input_ptr + up_offsets,
+        mask=(input_n_offsets < size_n)[None, :] * (input_m_offsets < size_m)[:, None],
+        other=0.0,
+    )
+    gate = tl.load(
+        input_ptr + gate_offsets,
+        mask=(input_n_offsets < size_n)[None, :] * (input_m_offsets < size_m)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+
+    gate = gelu(gate)
+    gate = gate.to(tl.float16)
+
+    tl.store(
+        input_ptr + res_offsets,
+        up * gate,
+        mask=(output_n_offsets < size_n)[None, :] * (output_m_offsets < size_m)[:, None],
+    )
+
+
+def gelu_and_mul_fwd(input):
+    stride_input_m = input.stride(0)
+    stride_input_n = input.stride(1)
+    stride_output_m = input.stride(0)
+    stride_output_n = input.stride(1)
+    size_m = input.shape[0]
+    size_n = input.shape[-1] // 2
+    BLOCK_M = 128
+    BLOCK_N = 128
+    grid = (
+        triton.cdiv(size_m, BLOCK_M),
+        triton.cdiv(size_n, BLOCK_N),
+    )
+    _gelu_and_mul_kernel[grid](
+        input,
+        stride_input_m,
+        stride_input_n,
+        stride_output_m,
+        stride_output_n,
+        size_m,
+        size_n,
+        BLOCK_M,
+        BLOCK_N,
+    )
+    return input[:, 0 : (input.shape[-1] // 2)]
+
+
+def torch_gelu_and_mul(input: torch.Tensor):
+    return torch.nn.functional.gelu(input[:, 0 : (input.shape[-1] // 2)]) * input[:, (input.shape[-1] // 2) :]
+
+
+def test_gelu_and_mul(M, N, dtype, device="cuda"):
+    # create data
+    X = torch.randn((M, N), dtype=dtype, device=device)
+
+    # run
+    y_tri = gelu_and_mul_fwd(X)
+    y_ref = torch_gelu_and_mul(X)
+
+    # compare
+    print("type:", y_tri.dtype, y_ref.dtype)
+    print("max delta:", torch.max(torch.abs(y_tri - y_ref)))
+    assert torch.allclose(y_tri, y_ref, atol=1e-2, rtol=0)
+    return
+
+
+# test_gelu_and_mul(16, 4096, torch.float16, device='cuda')