[Model] Add support for gemma3 model (#825)

Co-authored-by: gushiqiao <[email protected]>
Co-authored-by: hiworldwzj <[email protected]>

Author: 3 people

docs/CN/source/models/supported_models.rst

@@ -79,6 +79,8 @@ lightllm 支持大多数的主流的开源大语言模型以及多模态模型
     -  :code:`--enable_multimodal`
   * - `Qwen2-VL <https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct>`_
     -  :code:`--enable_multimodal`
+  * - `Google Gemma3 <https://huggingface.co/google/gemma-3-12b-it>`_
+    -  :code:`--enable_multimodal`
 
 
 Reward模型

docs/EN/source/models/supported_models.rst

@@ -76,6 +76,8 @@ VLM
     -  :code:`--enable_multimodal`
   * - `Llava-13b <https://huggingface.co/liuhaotian/llava-v1.5-13b>`_
     -  :code:`--enable_multimodal`
+  * - `Google Gemma3 <https://huggingface.co/google/gemma-3-12b-it>`_
+    -  :code:`--enable_multimodal`
 
 
 Reward Model

lightllm/models/gemma3/__init__.py

@@ -0,0 +1,142 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import json
+import os
+from PIL import Image
+from typing import List, Union
+from safetensors import safe_open
+from io import BytesIO
+from lightllm.server.multimodal_params import MultimodalParams, ImageItem
+from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data
+from lightllm.utils.log_utils import init_logger
+
+
+logger = init_logger(__name__)
+
+
+class Gemma3VisionModel:
+    def __init__(self):
+        pass
+
+    def load_model(self, weight_dir):
+        config_file = os.path.join(weight_dir, "config.json")
+        config = json.load(open(config_file))
+
+        # for llava-v1.5-7b-hf model, should load config from transformers
+        if "text_config" in config:
+            self.load_hf_model(config, weight_dir)
+        else:
+            assert False, "only hf format model is supported for Gemma3"
+
+        self.patches_per_image = int(config["vision_config"]["image_size"] // config["vision_config"]["patch_size"])
+        self.tokens_per_side = int(config["mm_tokens_per_image"] ** 0.5)
+        self.kernel_size = self.patches_per_image // self.tokens_per_side
+        self.avg_pool = nn.AvgPool2d(kernel_size=self.kernel_size, stride=self.kernel_size)
+
+        self.vision_tower.requires_grad_(False)
+        self.device = torch.device("cpu")
+
+        assert "model.mm_projector.linear" in self.projector_weights
+        assert "model.mm_projector.norm" in self.projector_weights
+
+    def load_hf_model(self, config, weight_dir):
+        from transformers import AutoConfig, AutoProcessor, Gemma3ForConditionalGeneration
+
+        # config = AutoConfig.from_pretrained(weight_dir, trust_remote_code=True)
+        processor = AutoProcessor.from_pretrained(weight_dir)
+        self.image_processor = processor.image_processor
+
+        model = Gemma3ForConditionalGeneration.from_pretrained(
+            weight_dir,
+            torch_dtype=torch.float16,
+        )
+        self.vision_tower = model.vision_tower
+        model.multi_modal_projector = None
+        model.language_model = None
+
+        # load projector weights
+        self.projector_weights = {}
+        for f in os.listdir(weight_dir):
+            if f.endswith(".safetensors"):
+                d = safe_open(os.path.join(weight_dir, f), "pt", "cpu")
+                for k in d.keys():
+                    if "multi_modal_projector.mm_input_projection_weight" in k:
+                        self.projector_weights[
+                            k.replace("multi_modal_projector.mm_input_projection_weight", "model.mm_projector.linear")
+                        ] = d.get_tensor(k).to(torch.bfloat16)
+                    if "multi_modal_projector.mm_soft_emb_norm.weight" in k:
+                        self.projector_weights[
+                            k.replace("multi_modal_projector.mm_soft_emb_norm.weight", "model.mm_projector.norm")
+                        ] = d.get_tensor(k).to(torch.bfloat16)
+
+    def cuda(self):
+        self.vision_tower = self.vision_tower.cuda()
+        for k, v in self.projector_weights.items():
+            self.projector_weights[k] = v.cuda()
+        return self
+
+    def gemma3_rms_norm(self, input, weight, eps: float = 1e-6):
+        def _norm(x):
+            return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+        output = _norm(input.float())
+        # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
+        # See https://github.com/huggingface/transformers/pull/29402
+        output = output * (1.0 + weight.float())
+        return output.type_as(input)
+
+    # batch images infer
+    def forward(self, x):
+        x = x.to(torch.bfloat16).cuda()
+        x = self.vision_tower(x, output_hidden_states=True).last_hidden_state
+
+        batch_size, _, seq_length = x.shape
+
+        reshaped_vision_outputs = x.transpose(1, 2)
+        reshaped_vision_outputs = reshaped_vision_outputs.reshape(
+            batch_size, seq_length, self.patches_per_image, self.patches_per_image
+        )
+        reshaped_vision_outputs = reshaped_vision_outputs.contiguous()
+
+        pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs)
+        pooled_vision_outputs = pooled_vision_outputs.flatten(2)
+        pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2)
+
+        normed_vision_outputs = self.gemma3_rms_norm(
+            pooled_vision_outputs.float(), self.projector_weights["model.mm_projector.norm"]
+        ).to(torch.bfloat16)
+
+        projected_vision_outputs = torch.matmul(
+            normed_vision_outputs, self.projector_weights["model.mm_projector.linear"]
+        )
+
+        return projected_vision_outputs.type_as(x)
+
+    def encode(self, images: List[ImageItem]):
+        img_tensors = []
+        uuids = []
+        valid_id = 0
+        valid_ids = []
+
+        for i, img in enumerate(images):
+            if isinstance(img, ImageItem):
+                uuids.append(img.uuid)
+                image_data = read_shm(get_shm_name_data(img.uuid))
+                image_data = Image.open(BytesIO(image_data))
+                t = self.image_processor.preprocess(image_data, return_tensors="pt")["pixel_values"]
+                img_tensors.append(t)
+            else:
+                raise Exception("Unsupport input types: {} for {}".format(type(img), img))
+
+            cur_num = img_tensors[-1].shape[0]
+            valid_ids.append([valid_id, valid_id + cur_num])
+            valid_id += cur_num
+
+        if len(img_tensors) <= 0:
+            return None
+
+        img = torch.cat(img_tensors, dim=0)
+        all_img_embeds = self.forward(img)
+
+        return all_img_embeds, uuids, valid_ids

lightllm/models/gemma3/gemma3_visual.py

@@ -0,0 +1,62 @@
+import torch
+import numpy as np
+from lightllm.common.basemodel import InferStateInfo
+from lightllm.common.req_manager import ReqManager
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+
+class Gemma3InferStateInfo(LlamaInferStateInfo):
+    def __init__(self):
+        super().__init__()
+        self.position_cos_global = None
+        self.position_sin_global = None
+        self.position_sin_local = None
+        self.position_cos_local = None
+
+    def init_some_extra_state(self, model, input_ids: torch.Tensor):
+        if self.is_prefill:
+            b_seq_len_numpy = self.b_seq_len.cpu().numpy()
+            self.max_seq_len = b_seq_len_numpy.max()
+            b_ready_cache_len_numpy = self.b_ready_cache_len.cpu().numpy()
+            position_ids = torch.from_numpy(
+                np.concatenate(
+                    [np.arange(b_ready_cache_len_numpy[i], b_seq_len_numpy[i]) for i in range(len(b_seq_len_numpy))],
+                    axis=0,
+                )
+            ).cuda()
+            self.position_cos = torch.index_select(model._cos_cached, 0, position_ids).view(position_ids.shape[0], -1)
+            self.position_sin = torch.index_select(model._sin_cached, 0, position_ids).view(position_ids.shape[0], -1)
+
+            self.position_cos_local = torch.index_select(model._cos_cached_local, 0, position_ids).view(
+                position_ids.shape[0], -1
+            )
+            self.position_sin_local = torch.index_select(model._sin_cached_local, 0, position_ids).view(
+                position_ids.shape[0], -1
+            )
+
+            self.position_cos_global = torch.index_select(model._cos_cached_global, 0, position_ids).view(
+                position_ids.shape[0], -1
+            )
+            self.position_sin_global = torch.index_select(model._sin_cached_global, 0, position_ids).view(
+                position_ids.shape[0], -1
+            )
+            position_ids = None
+        else:
+            position_ids = self.b_seq_len - 1
+            self.position_cos = torch.index_select(model._cos_cached, 0, position_ids).view(self.b_seq_len.shape[0], -1)
+            self.position_sin = torch.index_select(model._sin_cached, 0, position_ids).view(self.b_seq_len.shape[0], -1)
+
+            self.position_cos_local = torch.index_select(model._cos_cached_local, 0, position_ids).view(
+                self.b_seq_len.shape[0], -1
+            )
+            self.position_sin_local = torch.index_select(model._sin_cached_local, 0, position_ids).view(
+                self.b_seq_len.shape[0], -1
+            )
+
+            self.position_cos_global = torch.index_select(model._cos_cached_global, 0, position_ids).view(
+                self.b_seq_len.shape[0], -1
+            )
+            self.position_sin_global = torch.index_select(model._sin_cached_global, 0, position_ids).view(
+                self.b_seq_len.shape[0], -1
+            )
+        return

lightllm/models/gemma3/infer_struct.py

@@ -0,0 +1,60 @@
+import numpy as np
+import torch
+
+from lightllm.distributed.communication_op import all_gather
+from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
+from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
+
+
+class Gemma3PostLayerInfer(LlamaPostLayerInfer):
+    """ """
+
+    def __init__(self, network_config, mode):
+        super().__init__(network_config, mode)
+        self.eps_ = 1e-6
+        return
+
+    def gemma3_rmsnorm(self, input, weight, eps: float = 1e-6, out=None):
+        def _inner_norm(x):
+            return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+        output = _inner_norm(input.float())
+        output = output * (1.0 + weight.float())
+        if out is not None:
+            out = output.to(out.dtype)
+        return output
+
+    def _norm(self, input, infer_state, layer_weight) -> torch.Tensor:
+        return self.gemma3_rmsnorm(input, layer_weight.final_norm_weight_, eps=self.eps_)
+
+    def token_forward(self, input_embdings, infer_state, layer_weight):
+        last_input, token_num = self._slice_get_last_input(input_embdings, infer_state)
+        input_embdings_dtype = input_embdings.dtype
+        last_input = self._norm(last_input.float(), infer_state, layer_weight).to(torch.bfloat16)
+        last_input = last_input.permute(1, 0).view(-1, token_num)
+        logic_batch = self.alloc_tensor(
+            (layer_weight.lm_head_weight_.shape[0], last_input.shape[1]), dtype=last_input.dtype
+        )
+        torch.mm(layer_weight.lm_head_weight_.to(last_input.dtype), last_input, out=logic_batch)
+        last_input = None
+        if self.tp_world_size_ == 1:
+            gather_data = logic_batch
+        else:
+            gather_data = self.alloc_tensor((self.vocab_size_, token_num), dtype=input_embdings_dtype)
+            split_indexes = np.linspace(0, self.vocab_size_, self.tp_world_size_ + 1, dtype=np.int64)
+            all_gather(
+                [gather_data[split_indexes[i] : split_indexes[i + 1], :] for i in range(self.tp_world_size_)],
+                logic_batch,
+                group=infer_state.dist_group,
+                async_op=False,
+            )
+        logic_batch = None
+        ans_logics = self.alloc_tensor(
+            (token_num, self.vocab_size_),
+            dtype=torch.float32,
+            is_graph_out=True,
+            microbatch_index=infer_state.microbatch_index,
+        )
+        ans_logics[:, :] = gather_data.permute(1, 0)
+        gather_data = None
+        return ans_logics

lightllm/models/gemma3/layer_infer/__init__.py

@@ -0,0 +1,84 @@
+import torch
+from lightllm.common.basemodel.triton_kernel.multimodal_emb import multimodal_emb
+from lightllm.distributed.communication_op import all_reduce
+from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
+from lightllm.server.embed_cache.utils import bytes2tensor, get_shm_name_embed, read_shm
+
+
+class Gemma3PreLayerInfer(LlamaMultimodalPreLayerInfer):
+    def __init__(self, network_config, mode):
+        super().__init__(network_config, mode)
+        self.embed_scale = torch.tensor(network_config["hidden_size"] ** 0.5, dtype=torch.float32)
+        self.boi_token_index: int = 255_999
+        self.eoi_token_index: int = 256_000
+        return
+
+    def context_forward(self, input_ids, infer_state, layer_weight):
+        img_weight = []
+        img_start_token_ids = []
+        img_token_lens = []
+        img_start_loc = 0
+        img_start_locs = []
+        device = layer_weight.wte_weight_.device
+        dtype = layer_weight.wte_weight_.dtype
+        hidden_size = layer_weight.wte_weight_.shape[1]
+        weight_mask = torch.zeros((len(input_ids)), dtype=torch.float32, device=device)
+
+        scale = self.embed_scale
+        for idx, input_id in enumerate(input_ids):
+            if input_id == self.boi_token_index:
+                weight_mask[idx] = scale
+                scale = 1.0
+            elif input_id == self.eoi_token_index:
+                scale = self.embed_scale
+                weight_mask[idx] = scale
+            else:
+                weight_mask[idx] = scale
+
+        for batch_id, p in enumerate(infer_state.multimodal_params):
+            for img in p["images"]:
+                # skip the same image
+                if img["token_id"] in img_start_token_ids:
+                    continue
+                # pull the img_embeds by uid from shm
+                data = read_shm(get_shm_name_embed(img["uuid"]))
+                img_weight.append(bytes2tensor(data).cuda().reshape(img["token_num"], -1))
+                img_start_token_ids.append(img["token_id"])
+                img_token_lens.append(img["token_num"])
+                img_start_locs.append(img_start_loc)
+                img_start_loc += img["token_num"]
+        out = torch.zeros((len(input_ids), hidden_size), dtype=dtype, device=device)
+        if len(img_weight) > 0:
+            img_weight = torch.cat(img_weight, dim=0).to(device=device, dtype=dtype)
+        else:
+            img_weight = torch.empty((0, hidden_size), device=device, dtype=dtype)
+        assert img_weight.shape[1] == hidden_size, (
+            f"Dimension mismatch: text weight dimension is {hidden_size}, "
+            f"but image weight dimension is {img_weight.shape[1]}"
+        )
+        # each tp will fill the img embeds, should divide by world_size
+        img_weight = img_weight / self.tp_world_size_
+        img_start_token_ids = torch.Tensor(img_start_token_ids).to(device=device, dtype=torch.long)
+        img_token_lens = torch.Tensor(img_token_lens).to(device=device, dtype=torch.long)
+        img_start_locs = torch.Tensor(img_start_locs).to(device=device, dtype=torch.long)
+
+        multimodal_emb(
+            out,
+            input_ids,
+            layer_weight.wte_weight_,
+            img_weight,
+            img_token_lens,
+            img_start_token_ids,
+            img_start_locs,
+            self.vob_start_id_,
+            self.vob_end_id_,
+        )
+        input_dtype = out.dtype
+        if self.tp_world_size_ > 1:
+            all_reduce(out, group=infer_state.dist_group, op=torch.dist.ReduceOp.SUM, async_op=False)
+        return (out.float() * weight_mask.unsqueeze(1).float()).to(input_dtype)
+
+    def token_forward(self, input_ids, infer_state, layer_weight):
+        input_embedding = super().token_forward(input_ids, infer_state, layer_weight)
+        input_dtype = input_embedding.dtype
+        return (input_embedding.float() * self.embed_scale.to(input_embedding.device).float()).to(input_dtype)