fix

Author: flyinglandlord

lightllm/models/gemma3/gemma3_visual.py

@@ -29,8 +29,8 @@ def load_model(self, 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.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)
 
@@ -43,7 +43,7 @@ def load_model(self, weight_dir):
     def load_hf_model(self, config, weight_dir):
         from transformers import AutoConfig, AutoProcessor, Gemma3ForConditionalGeneration
 
-        config = AutoConfig.from_pretrained(weight_dir, trust_remote_code=True)
+        # config = AutoConfig.from_pretrained(weight_dir, trust_remote_code=True)
         processor = AutoProcessor.from_pretrained(weight_dir)
         self.image_processor = processor.image_processor
 
@@ -79,6 +79,7 @@ def cuda(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
@@ -89,7 +90,7 @@ def _norm(x):
     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)
@@ -102,10 +103,14 @@ def forward(self, x):
         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.float32)
+        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"]
+        )
 
-        projected_vision_outputs = torch.matmul(normed_vision_outputs, self.projector_weights['model.mm_projector.linear'])
-        #print(projected_vision_outputs.type_as(x))
         return projected_vision_outputs.type_as(x)
 
     def encode(self, images: List[ImageItem]):
@@ -120,7 +125,6 @@ def encode(self, images: List[ImageItem]):
                 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"]
-                #print(t)
                 img_tensors.append(t)
             else:
                 raise Exception("Unsupport input types: {} for {}".format(type(img), img))

lightllm/models/gemma3/infer_struct.py

@@ -27,20 +27,36 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
             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_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)
+            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_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
+            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/layer_infer/post_layer_infer.py

@@ -6,7 +6,6 @@
 from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
 
 
-
 class Gemma3PostLayerInfer(LlamaPostLayerInfer):
     """ """
 
@@ -15,9 +14,10 @@ def __init__(self, network_config, mode):
         self.eps_ = 1e-6
         return
 
-    def gemma3_rmsnorm(self, input, weight, eps: float = 1e-6, out = None):
+    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:
@@ -26,7 +26,7 @@ def _inner_norm(x):
 
     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):
         # print('last_hidden_before_norm', input_embdings)
         last_input, token_num = self._slice_get_last_input(input_embdings, infer_state)
@@ -58,4 +58,4 @@ def token_forward(self, input_embdings, infer_state, layer_weight):
         )
         ans_logics[:, :] = gather_data.permute(1, 0)
         gather_data = None
-        return ans_logics
+        return ans_logics

lightllm/models/gemma3/layer_infer/pre_layer_infer.py

@@ -1,20 +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.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):
-        input_embedding = super().context_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)
+        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)
-