port huggingface#3188 to gaudi backend

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

Author: sywangyi

backends/gaudi/server/text_generation_server/models/__init__.py

@@ -83,9 +83,6 @@
     from text_generation_server.models.custom_modeling.flash_neox_modeling import (
         FlashGPTNeoXForCausalLM,
     )
-    from text_generation_server.models.pali_gemma import (
-        PaliGemmaBatch,
-    )
     from text_generation_server.models.custom_modeling.flash_pali_gemma_modeling import (
         PaliGemmaForConditionalGeneration,
     )
@@ -153,7 +150,6 @@
     )
 
     VLM_BATCH_TYPES = {
-        PaliGemmaBatch,
         FlashVlmCausalLMBatch,
         FlashMllamaCausalLMBatch,
     }
@@ -635,6 +631,7 @@ def get_model(
                 default_dtype=torch.bfloat16,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
+                support_chunking=False,
             )
         elif model_type == BAICHUAN:
             return FlashCausalLM(
@@ -784,6 +781,8 @@ def get_model(
                 kv_cache_dtype=kv_cache_dtype,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
+                # TODO: Fix bug in rust image_text_replacement implementation
+                support_chunking=False,
             )
         elif model_type == QWEN2_5_VL:
             return FlashVlmCausalLM(
@@ -799,6 +798,8 @@ def get_model(
                 lora_adapter_ids=lora_adapter_ids,
                 config_class=Qwen2_5_VLConfig,
                 processor_class=Qwen2_5_VLProcessor,
+                # TODO: Fix bug in rust image_text_replacement implementation
+                support_chunking=False,
             )
         elif model_type == QWEN3:
             return FlashCausalLM(
@@ -824,6 +825,7 @@ def get_model(
                 default_dtype=torch.bfloat16,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
+                support_chunking=False,
             )
         elif model_type == IDEFICS2:
             return FlashVlmCausalLM(
@@ -868,7 +870,6 @@ def get_model(
                 default_dtype=torch.bfloat16,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
-                batch_class=PaliGemmaBatch,
             )
         elif model_type == LLAVA_NEXT:
             return FlashVlmCausalLM(

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_llava_next.py

@@ -163,111 +163,124 @@ def _merge_input_ids_with_image_features(
             )
         return inputs_embeds
 
-    def forward(
+    def get_vision_embeds(
         self,
-        input_ids: torch.Tensor,
-        position_ids: torch.Tensor,
-        cu_seqlen_prefill: Optional[torch.Tensor],
-        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
-        slots: torch.Tensor,
-        seqlen: Seqlen,
-        hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
-        lm_head_indices: Optional[torch.Tensor] = None,
-        pixel_values: torch.FloatTensor = None,
-        # Unused for this model
-        pixel_attention_mask=None,
-        image_sizes: Optional[torch.LongTensor] = None,
-        adapter_data: Optional[torch.Tensor] = None,
+        pixel_values: torch.FloatTensor,
+        pixel_attention_mask: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[torch.Tensor] = None,
         image_grid_thw: Optional[torch.LongTensor] = None,
     ):
-        inputs_embeds = self.text_model.embed_tokens(input_ids)
-        if pixel_values is not None and len(pixel_values) > 0:
-            # num_special_image_tokens = (input_ids == self.config.image_token_index).sum()
-            # assert num_special_image_tokens == len(pixel_values), f"Received {num_special_image_tokens} for {len(pixel_values)} images, this is invalid"
-            # 1. Extract the input embeddings
-
-            # 2. Merge text and images
-            num_images, num_patches, channels, height, width = pixel_values.shape
-            pixel_values = pixel_values.view(
-                num_images * num_patches, channels, height, width
-            )
-            image_features = self.vision_tower(pixel_values)
+        # num_special_image_tokens = (input_ids == self.config.image_token_index).sum()
+        # assert num_special_image_tokens == len(pixel_values), f"Received {num_special_image_tokens} for {len(pixel_values)} images, this is invalid"
+        # 1. Extract the input embeddings
+
+        # 2. Merge text and images
+        num_images, num_patches, channels, height, width = pixel_values.shape
+        pixel_values = pixel_values.view(
+            num_images * num_patches, channels, height, width
+        )
+        image_features = self.vision_tower(pixel_values)
 
-            # selected_image_feature = image_features.hidden_states[self.config.vision_feature_layer]
-            # Already done within the clip model
-            selected_image_feature = image_features.last_hidden_state
+        # selected_image_feature = image_features.hidden_states[self.config.vision_feature_layer]
+        # Already done within the clip model
+        selected_image_feature = image_features.last_hidden_state
 
-            if self.config.vision_feature_select_strategy == "default":
-                selected_image_feature = selected_image_feature[:, 1:]
-            elif self.config.vision_feature_select_strategy == "full":
-                selected_image_feature = selected_image_feature
-            else:
-                raise RuntimeError(
-                    f"Strategy `{self.config.vision_feature_select_strategy}` is not supported/valid."
-                )
+        if self.config.vision_feature_select_strategy == "default":
+            selected_image_feature = selected_image_feature[:, 1:]
+        elif self.config.vision_feature_select_strategy == "full":
+            selected_image_feature = selected_image_feature
+        else:
+            raise RuntimeError(
+                f"Strategy `{self.config.vision_feature_select_strategy}` is not supported/valid."
+            )
 
-            image_features = self.multi_modal_projector(selected_image_feature)
+        image_features = self.multi_modal_projector(selected_image_feature)
 
-            # split up image_features for each of the individual images
-            # hence we get a list of image_features, each of shape (5, num_patches, hidden_size)
-            # if we assume each image has 5 image features (base image + 4 patches)
-            split_sizes = [num_patches] * num_images
-            image_features = torch.split(image_features, split_sizes, dim=0)
+        # split up image_features for each of the individual images
+        # hence we get a list of image_features, each of shape (5, num_patches, hidden_size)
+        # if we assume each image has 5 image features (base image + 4 patches)
+        split_sizes = [num_patches] * num_images
+        image_features = torch.split(image_features, split_sizes, dim=0)
 
-            # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
-            height = width = (
-                self.config.vision_config.image_size
-                // self.config.vision_config.patch_size
-            )
+        # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
+        height = width = (
+            self.config.vision_config.image_size // self.config.vision_config.patch_size
+        )
 
-            new_image_features = []
-            for image_idx, image_feature in enumerate(image_features):
-                if image_feature.shape[0] > 1:
-                    base_image_feature = image_feature[0]
-                    image_feature = image_feature[1:]
-
-                    if height * width != base_image_feature.shape[0]:
-                        raise ValueError(
-                            "The number of patches is not consistent with the image size."
-                        )
-
-                    # Dimensions are intentionally swapped to be bug-compatible with
-                    # upstream: https://github.com/LLaVA-VL/LLaVA-NeXT/issues/59
-                    num_patch_width, num_patch_height = get_anyres_image_grid_shape(
-                        image_sizes[image_idx],
-                        self.config.image_grid_pinpoints,
-                        self.config.vision_config.image_size,
-                    )
-                    image_feature = image_feature.view(
-                        num_patch_height, num_patch_width, height, width, -1
+        new_image_features = []
+        for image_idx, image_feature in enumerate(image_features):
+            if image_feature.shape[0] > 1:
+                base_image_feature = image_feature[0]
+                image_feature = image_feature[1:]
+
+                if height * width != base_image_feature.shape[0]:
+                    raise ValueError(
+                        "The number of patches is not consistent with the image size."
                     )
-                    image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
-                    image_feature = image_feature.flatten(1, 2).flatten(2, 3)
-                    image_feature = unpad_image(image_feature, image_sizes[image_idx])
-                    image_feature = torch.cat(
-                        (
-                            image_feature,
-                            self.image_newline[:, None, None].expand(
-                                *image_feature.shape[:-1], 1
-                            ),
+
+                # Dimensions are intentionally swapped to be bug-compatible with
+                # upstream: https://github.com/LLaVA-VL/LLaVA-NeXT/issues/59
+                num_patch_width, num_patch_height = get_anyres_image_grid_shape(
+                    image_sizes[image_idx],
+                    self.config.image_grid_pinpoints,
+                    self.config.vision_config.image_size,
+                )
+                image_feature = image_feature.view(
+                    num_patch_height, num_patch_width, height, width, -1
+                )
+                image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
+                image_feature = image_feature.flatten(1, 2).flatten(2, 3)
+                image_feature = unpad_image(image_feature, image_sizes[image_idx])
+                image_feature = torch.cat(
+                    (
+                        image_feature,
+                        self.image_newline[:, None, None].expand(
+                            *image_feature.shape[:-1], 1
                         ),
-                        dim=-1,
-                    )
-                    image_feature = image_feature.flatten(1, 2).transpose(0, 1)
-                    image_feature = torch.cat(
-                        (base_image_feature, image_feature), dim=0
-                    )
-                else:
-                    image_feature = image_feature[0]
-                    image_feature = torch.cat(
-                        (image_feature, self.image_newline[None]), dim=0
-                    )
-                new_image_features.append(image_feature)
-            image_features = torch.stack(new_image_features, dim=0)
+                    ),
+                    dim=-1,
+                )
+                image_feature = image_feature.flatten(1, 2).transpose(0, 1)
+                image_feature = torch.cat((base_image_feature, image_feature), dim=0)
+            else:
+                image_feature = image_feature[0]
+                image_feature = torch.cat(
+                    (image_feature, self.image_newline[None]), dim=0
+                )
+            new_image_features.append(image_feature)
+        image_features = torch.stack(new_image_features, dim=0)
+        return image_features.view(-1, image_features.shape[-1])
+
+    def get_inputs_embeds(
+        self,
+        input_ids: torch.Tensor,
+        vision_embeds: torch.Tensor = None,
+        pixel_values: torch.FloatTensor = None,
+        image_sizes: Optional[torch.LongTensor] = None,
+    ):
+        inputs_embeds = self.text_model.embed_tokens(input_ids)
 
+        if vision_embeds is not None:
+            # When we generate, we don't want to replace the potential image_token_id that we generated by images
+            # that simply don't exist
             inputs_embeds = self._merge_input_ids_with_image_features(
-                input_ids, inputs_embeds, image_features
+                input_ids, inputs_embeds, vision_embeds
             )
+        return inputs_embeds
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        cu_seqlen_prefill: Optional[torch.Tensor],
+        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
+        slots: torch.Tensor,
+        seqlen: Seqlen,
+        hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
+        lm_head_indices: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        adapter_data: Optional[torch.Tensor] = None,
+    ):
 
         hidden_states = self.text_model.model(
             inputs_embeds=inputs_embeds,

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_pali_gemma_modeling.py

@@ -62,43 +62,54 @@ def __init__(self, prefix, config, weights):
         self.pad_token_id = (
             config.pad_token_id if config.pad_token_id is not None else -1
         )
+        self.dtype = weights.dtype
 
-    def forward(
+    def get_vision_embeds(
+        self,
+        pixel_values: torch.FloatTensor,
+        pixel_attention_mask: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[torch.Tensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+    ):
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        image_outputs = self.vision_tower(pixel_values)
+        last_hidden_state = self.post_vision_tower_layernorm(
+            image_outputs.last_hidden_state
+        )
+        image_features = self.multi_modal_projector(last_hidden_state)
+        image_features = image_features.view(-1, image_features.shape[-1])
+        return image_features
+
+    def get_inputs_embeds(
         self,
         input_ids: torch.Tensor,
+        vision_embeds: torch.Tensor = None,
+    ):
+        inputs_embeds = self.text_model.embed_tokens(input_ids)
+
+        if vision_embeds is not None:
+            mask = input_ids == self.config.image_token_index
+            inputs_embeds[mask] = vision_embeds
+
+        return inputs_embeds
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
         position_ids: torch.Tensor,
         cu_seqlen_prefill: Optional[torch.Tensor],
         kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
         slots: torch.Tensor,
         seqlen: Seqlen,
         hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
         lm_head_indices: Optional[torch.Tensor] = None,
-        pixel_values: torch.FloatTensor = None,
-        # Unused here
-        pixel_attention_mask: Optional[torch.BoolTensor] = None,
-        image_sizes: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
         adapter_data: Optional[torch.Tensor] = None,
-        image_grid_thw: Optional[torch.LongTensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        inputs_embeds = self.text_model.embed_tokens(input_ids)
         # TODO This is odd but apparently pali gemma position ids start at 1.
         if cu_seqlen_prefill is not None:
             position_ids += 1
 
-        if pixel_values is not None:
-            pixel_values = pixel_values.to(dtype=inputs_embeds.dtype)
-            image_outputs = self.vision_tower(pixel_values)
-            last_hidden_state = self.post_vision_tower_layernorm(
-                image_outputs.last_hidden_state
-            )
-            image_features = self.multi_modal_projector(last_hidden_state)
-
-            # mask where image or padding tokens
-            mask = input_ids == self.config.image_token_index
-
-            # insert image features into input embeddings
-            inputs_embeds[mask] = image_features.view(-1, image_features.shape[-1])
-
         hidden_states = self.text_model.model(
             inputs_embeds=inputs_embeds,
             position_ids=position_ids,