feat: add Gemma4 VL support

aiak_megatron/megatron/core/transformer/transformer_config.py

@@ -3,7 +3,7 @@
 """tranformer config"""
 
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Callable, List, Optional, Tuple, Union
 
 import torch
@@ -530,6 +530,51 @@ class TransformerConfig(ModelParallelConfig):
     Using fp32 or fp64 can improve stability especially when the number of experts is large (e.g. finegrained-moe).
     None means no changes for dtype."""
 
+    final_logit_softcapping: Optional[float] = None
+    """If set, apply ``cap * tanh(logits / cap)`` to the language-model output logits.
+    Used by Gemma4 (cap=30.0). None disables the softcap."""
+
+    use_layer_scalar: bool = False
+    """If True, each transformer layer learns a scalar gate (init 1.0) applied to the
+    layer output before the residual add. Used by Gemma4."""
+
+    partial_rotary_factor: float = 1.0
+    """Fraction of the head dim that receives RoPE rotation. Gemma4 uses 0.25.
+    A value of 1.0 (default) keeps full RoPE and is a no-op for existing models."""
+
+    sliding_window: Optional[int] = None
+    """Window size used by sliding-window attention layers. None means no sliding mask."""
+
+    rotary_base_sliding: Optional[int] = None
+    """Optional separate RoPE base for sliding-window attention layers.
+    If None, sliding layers share ``rotary_base`` with global layers."""
+
+    layer_pattern: list = field(default_factory=list)
+    """Per-layer attention type as a list of strings ('sliding' / 'global'),
+    length == num_layers. Empty list (default) means all layers use the same
+    config-level head_dim / num_query_groups (no hybrid behavior)."""
+
+    per_layer_kv_channels: dict = field(default_factory=dict)
+    """Optional per-layer-type override of ``kv_channels`` (a.k.a. head_dim).
+    Keyed by layer-type string ('sliding' / 'global'). Empty dict (default)
+    means use ``kv_channels`` for every layer."""
+
+    per_layer_num_query_groups: dict = field(default_factory=dict)
+    """Optional per-layer-type override of ``num_query_groups``.
+    Empty dict (default) means use ``num_query_groups`` for every layer."""
+
+    attention_k_eq_v: bool = False
+    """If True, Gemma4 global layers use K=V tying according to ``kv_tied_layers``.
+    Existing models keep the default False behavior."""
+
+    kv_tied_layers: list = field(default_factory=list)
+    """Layer indices for which K and V projections share the same weight tensor
+    (used by Gemma4 global layers). Empty list (default) means no tying."""
+
+    scale_emb_by_sqrt_hidden: bool = False
+    """If True, multiply token embeddings by ``sqrt(hidden_size)`` before the
+    transformer (Gemma family convention)."""
+
     def __post_init__(self):
         """Python dataclass method that is used to modify attributes after initialization.
         See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more

aiak_training_llm/data/chat_templete.py

@@ -13,10 +13,11 @@
 import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Type, Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Set, Tuple, Type, Union
 
 from aiak_training_llm.utils.constants import DataRoles
-from .mm_plugin import MMPlugin, Qwen2VLPlugin
+
+from .mm_plugin import Gemma4VLPlugin, MMPlugin, Qwen2VLPlugin
 
 
 if TYPE_CHECKING:
@@ -459,3 +460,15 @@ def get_support_templates() -> List[str]:
     format_user=StringFormatter(slots=["<｜User｜>{{content}}<｜Assistant｜>"]),
     format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
 )
+
+_register_chat_template(
+    name="gemma4",
+    format_user=StringFormatter(
+        slots=["<|turn>user\n{{content}}<turn|>\n<|turn>model\n"]
+    ),
+    format_assistant=StringFormatter(slots=["{{content}}<turn|>\n"]),
+    format_system=StringFormatter(slots=["<|turn>system\n{{content}}<turn|>\n"]),
+    format_separator=EmptyFormatter(slots=[""]),
+    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
+    mm_plugin=Gemma4VLPlugin(image_token="<|image|>", video_token=None),
+)

aiak_training_llm/data/mm_plugin.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Type, TypedDict, Union
 
 import numpy as np
+import torch
 from PIL import Image
 from PIL.Image import Image as ImageObject
 from typing_extensions import override
@@ -16,6 +17,7 @@
     import torch
 
     from transformers.image_processing_utils import BaseImageProcessor
+    from transformers.processing_utils import ProcessorMixin
 
     class EncodedImage(TypedDict):
         """Encoded image type."""
@@ -292,3 +294,178 @@ def get_mm_inputs(
     ) -> Dict[str, Union[List[int], "torch.Tensor"]]:
         self._validate_input(images, videos)
         return self._get_mm_inputs(images, videos, processor)
+
+
+class Gemma4VLPlugin(MMPlugin):
+    """Gemma4-VL passthrough plugin: process_messages returns (messages, mm_inputs)
+    where mm_inputs follows the existing OV2 flattened-patch contract.
+
+    Cross-module contract (downstream consumers depend on these invariants):
+    - ``pixel_values`` shape ``[total_imgs_in_batch, P, D]`` — FLATTENED, NOT ``[B, ...]``.
+      Indexing by batch index will silently return wrong tensor.
+    - ``image_grid_thw`` is synthesized from HF ``image_position_ids`` as
+      ``[num_images, 3]`` rows of ``[1, H_p, W_p]``.
+    - Text-only batches OMIT multimodal keys entirely (no zero-shape sentinel).
+      All downstream consumers MUST guard with ``in mm_inputs``.
+    """
+
+    @staticmethod
+    def _flatten_gemma4_image_outputs(
+        image_outputs: dict[str, "torch.Tensor"],
+    ) -> dict[str, Union["torch.Tensor", list["torch.Tensor"]]]:
+        pixel_values = image_outputs["pixel_values"]
+        image_position_ids = image_outputs["image_position_ids"]
+        num_soft_tokens_per_image = image_outputs["num_soft_tokens_per_image"]
+
+        valid_mask = (image_position_ids != -1).all(dim=-1)
+        flat_pixel_values = pixel_values[valid_mask]
+
+        image_grid_rows: list[list[int]] = []
+        patch_positions: list[torch.Tensor] = []
+        for image_idx in range(image_position_ids.shape[0]):
+            valid_positions = image_position_ids[image_idx][valid_mask[image_idx]].to(dtype=torch.int64)
+            if valid_positions.numel() == 0:
+                raise ValueError(f"Gemma4 image {image_idx} has no valid patch positions.")
+
+            width = int(valid_positions[:, 0].max().item()) + 1
+            height = int(valid_positions[:, 1].max().item()) + 1
+            patch_count = int(valid_positions.shape[0])
+            if patch_count != height * width:
+                raise ValueError(
+                    "Gemma4 image patch positions are not a dense single-frame grid: "
+                    f"image_idx={image_idx}, patch_count={patch_count}, height={height}, width={width}."
+                )
+
+            image_grid_rows.append([1, height, width])
+            patch_positions.append(
+                torch.stack(
+                    (
+                        torch.zeros(patch_count, dtype=torch.int64, device=valid_positions.device),
+                        valid_positions[:, 1],
+                        valid_positions[:, 0],
+                    ),
+                    dim=-1,
+                )
+            )
+
+        image_grid_thw = torch.tensor(
+            image_grid_rows,
+            dtype=torch.int32,
+            device=image_position_ids.device,
+        )
+
+        return {
+            "pixel_values": flat_pixel_values,
+            "image_grid_thw": image_grid_thw,
+            "patch_positions": patch_positions,
+            "image_position_ids": image_position_ids,
+            "num_soft_tokens_per_image": num_soft_tokens_per_image,
+        }
+
+    def _build_gemma4_mm_inputs(
+        self,
+        images: Sequence["ImageInput"],
+        processor: Optional["ProcessorMixin"],
+    ) -> tuple[Optional[list["ImageObject"]], dict[str, Union["torch.Tensor", list["torch.Tensor"]]]]:
+        regularized_images = self._regularize_images(images) if len(images) != 0 else None
+        if regularized_images is None:
+            return None, {}
+
+        image_outputs = processor.image_processor(regularized_images, return_tensors="pt")
+        return regularized_images, self._flatten_gemma4_image_outputs(dict(image_outputs))
+
+    def _expand_image_placeholders(
+        self,
+        messages: Sequence[dict[str, str]],
+        num_soft_tokens_per_image: Sequence[int],
+        processor: Optional["ProcessorMixin"],
+    ) -> list[dict[str, str]]:
+        messages = deepcopy(messages)
+        actual_num_images = len(num_soft_tokens_per_image)
+
+        image_placeholder_count = sum(message["content"].count(Placeholder.IMAGE) for message in messages)
+        if actual_num_images > 0 and image_placeholder_count != actual_num_images:
+            for message in messages:
+                message["content"] = message["content"].replace(Placeholder.IMAGE, "")
+
+            first_user_msg = None
+            for message in messages:
+                if message.get("role") == "user":
+                    first_user_msg = message
+                    break
+
+            if first_user_msg is None:
+                raise ValueError("Cannot rebuild Gemma4 image placeholders: no user message found.")
+
+            image_placeholders = "\n".join([Placeholder.IMAGE] * actual_num_images)
+            user_content = first_user_msg["content"].lstrip("\n")
+            first_user_msg["content"] = f"{image_placeholders}\n{user_content}"
+
+        image_idx = 0
+        for message in messages:
+            content = message["content"]
+            while Placeholder.IMAGE in content:
+                if image_idx >= actual_num_images:
+                    raise ValueError(
+                        f"The number of {Placeholder.IMAGE} tokens is greater than available images."
+                    )
+
+                n_soft_tokens = int(num_soft_tokens_per_image[image_idx])
+                replacement = (
+                    f"{processor.boi_token}{self.image_token * n_soft_tokens}{processor.eoi_token}"
+                )
+                content = content.replace(Placeholder.IMAGE, replacement, 1)
+                image_idx += 1
+
+            if Placeholder.VIDEO in content:
+                raise ValueError("Gemma4-VL video placeholders are not supported in this OV2 path yet.")
+
+            message["content"] = content
+
+        if image_idx != actual_num_images:
+            raise ValueError(
+                f"The number of images ({actual_num_images}) does not match expanded placeholders ({image_idx})."
+            )
+
+        return messages
+
+    @override
+    def _preprocess_image(self, image: "ImageObject", **kwargs) -> "ImageObject":
+        return super()._preprocess_image(image, **kwargs)
+
+    @override
+    def process_messages(
+        self,
+        messages: Sequence[dict[str, str]],
+        images: Sequence["ImageInput"],
+        videos: Sequence["VideoInput"],
+        processor: Optional["ProcessorMixin"],
+    ) -> tuple[list[dict[str, str]], dict[str, "torch.Tensor"]]:
+        self._validate_input(images, videos)
+        _regularized_images, mm_inputs = self._build_gemma4_mm_inputs(images, processor)
+
+        if "num_soft_tokens_per_image" in mm_inputs:
+            messages = self._expand_image_placeholders(
+                messages,
+                mm_inputs["num_soft_tokens_per_image"],
+                processor,
+            )
+        else:
+            messages = list(messages)
+
+        return messages, dict(mm_inputs)
+
+    @override
+    def get_mm_inputs(
+        self,
+        images: Sequence["ImageInput"],
+        videos: Sequence["VideoInput"],
+        imglens: Sequence[int],
+        vidlens: Sequence[int],
+        seqlens: Sequence[int],
+        processor: Optional["ProcessorMixin"],
+    ) -> dict[str, Union[list[int], "torch.Tensor"]]:
+        self._validate_input(images, videos)
+        del imglens, vidlens, seqlens
+        _regularized_images, mm_inputs = self._build_gemma4_mm_inputs(images, processor)
+        return dict(mm_inputs)