diff --git a/aiak_megatron/megatron/core/transformer/transformer_config.py b/aiak_megatron/megatron/core/transformer/transformer_config.py
index e41c3eeb..87f52bc7 100644
--- a/aiak_megatron/megatron/core/transformer/transformer_config.py
+++ b/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
diff --git a/aiak_training_llm/data/chat_templete.py b/aiak_training_llm/data/chat_templete.py
index 82262cc4..8eebdefe 100644
--- a/aiak_training_llm/data/chat_templete.py
+++ b/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),
+)
diff --git a/aiak_training_llm/data/mm_plugin.py b/aiak_training_llm/data/mm_plugin.py
index 6438e282..09a7d57d 100644
--- a/aiak_training_llm/data/mm_plugin.py
+++ b/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)
diff --git a/aiak_training_llm/data/multimodal/gemma4_vl_task_encoder.py b/aiak_training_llm/data/multimodal/gemma4_vl_task_encoder.py
new file mode 100644
index 00000000..0f47fa13
--- /dev/null
+++ b/aiak_training_llm/data/multimodal/gemma4_vl_task_encoder.py
@@ -0,0 +1,934 @@
+"""Gemma4VLTaskEncoder class.
+
+Verbatim mirror of :mod:`qwen2vl_task_encoder` for the Gemma4-VL family.
+The structural pipeline (mm_plugin → encode_multiturn → packing → batching) is
+identical; only Gemma4-specific surface differs:
+
+- Special-token strings (Gemma4 vocab is unrelated to Qwen):
+    * boi (vision_start)  -> ``<|image>``   (id=255999)
+    * image pad           -> ``<|image|>``  (id=258880)
+    * eoi (vision_end)    -> ``<image|>``   (id=258882)
+    * video               -> ``<|video|>``  (id=258884)
+- Spatial merge: Gemma4 uses ``image_processor.pooling_kernel_size=3`` instead of
+  Qwen's ``merge_size=2``; the existing ``getattr(..., 'merge_size', 2)`` lookup
+  reads the wrong attribute on Gemma4Processor so we resolve it via a small
+  helper (``_gemma4_spatial_merge_size``) that prefers ``pooling_kernel_size``.
+- ``smart_resize`` factor: Qwen uses 28 (= patch_size 14 × merge 2); Gemma4 uses
+  patch_size 16 × pooling 3 = 48. We pass ``size_factor=48`` through.
+
+We keep this as a fork rather than a subclass per user direction (avoids leaky
+class-attribute overrides on the parent Qwen module-level constants).
+
+CAVEAT — image placeholder wrapping mismatch (deferred fix, ack'd in P4.9):
+  Qwen2VL inline format is ``<|vision_start|><|image_pad|><|vision_end|>`` (3
+  tokens, fence + pad + fence). Gemma4 chat-template instead emits a single
+  ``<|image|>`` (no boi/eoi wrap). Our verbatim mirror keeps the 3-token wrap
+  using Gemma4 ids — correct for fence semantics but produces 1 extra image
+  placeholder vs the reference Gemma4 jinja path. Forward will run; lm_loss
+  numerics will be off until a future revision rewrites the
+  ``<image>``-substitution path to call ``processor.apply_chat_template``
+  directly. Tracked as a Phase-P4.9 known issue.
+"""
+
+import re
+import sys
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Optional, TypeVar, Union
+
+import numpy as np
+import torch
+from megatron.energon import CaptioningSample, SkipSample, VQASample
+from megatron.energon.flavors.base_dataset import (
+    BaseCoreDatasetFactory,
+    SavableDataset,
+)
+from megatron.energon.flavors.crude import CrudeWebdataset
+from megatron.energon.flavors.webdataset import VideoData
+from megatron.energon.metadataset.loader_interface import DatasetBlendMode
+from megatron.energon.task_encoder.base import stateless
+from megatron.energon.worker import WorkerConfig
+from megatron.energon.wrappers import (
+    BlendDataset,
+    EpochizeDataset,
+    LogSampleDataset,
+    ShuffleBufferDataset,
+)
+from megatron.energon.wrappers.repeat_dataset import RepeatDataset
+from qwen_vl_utils.vision_process import smart_nframes, smart_resize
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+from typing_extensions import override
+
+from aiak_training_llm.data.multimodal import MultiMixQASample
+from aiak_training_llm.data.multimodal.length_sort_dataset import LengthPoolSortDataset
+from aiak_training_llm.data.multimodal.packed_sort_dataset import PackedSeparateSortDataset
+from aiak_training_llm.utils import constants, get_chat_template
+from transformers import AutoProcessor
+
+from .task_encoder import ImageTaskBatchPacked, ImageTaskSample, ImageTaskSamplePacked, TaskEncoder
+
+
+T = TypeVar("T")
+V = TypeVar("V")
+T_sample = TypeVar("T_sample")
+T_encoded_sample = TypeVar("T_encoded_sample")
+T_raw_batch = TypeVar("T_raw_batch")
+T_batch = TypeVar("T_batch")
+
+
+IGNORE_INDEX = -100  # ID for labels that should be ignored.
+IMAGE_TOKEN = "<|image|>"
+VIDEO_TOKEN = "<|video|>"
+VISION_TAGS = ["<|image>", "<image|>"]
+IMAGE_TOKEN_WITH_TAGS = VISION_TAGS[0] + IMAGE_TOKEN + VISION_TAGS[1]
+VIDEO_TOKEN_WITH_TAGS = VISION_TAGS[0] + VIDEO_TOKEN + VISION_TAGS[1]
+SKIP_LOG_LIMIT = 5
+_SKIP_LOG_COUNTS: dict[str, int] = {}
+
+
+def _gemma4_spatial_merge_size(image_processor) -> int:
+    """Resolve Gemma4's spatial merge factor.
+
+    Gemma4Processor exposes ``pooling_kernel_size`` (default 3); the parent
+    Qwen lookup ``merge_size`` (default 2) returns the wrong fallback on Gemma4
+    and silently corrupts patch_position block layout. Prefer the Gemma4 name
+    when present.
+    """
+    val = getattr(image_processor, "pooling_kernel_size", None)
+    if val is not None:
+        return int(val)
+    return int(getattr(image_processor, "merge_size", 2))
+
+
+def skip_malformed_multimodal_sample(sample_key: str, signature: str, detail: str) -> None:
+    count = _SKIP_LOG_COUNTS.get(signature, 0) + 1
+    _SKIP_LOG_COUNTS[signature] = count
+
+    if count <= SKIP_LOG_LIMIT:
+        print(f"Skipping malformed multimodal sample {sample_key}: {detail}", file=sys.stderr)
+        if count == SKIP_LOG_LIMIT:
+            print(
+                f"Further '{signature}' skip logs will be suppressed for this worker.",
+                file=sys.stderr,
+            )
+
+    raise SkipSample(f"{sample_key}: {detail}")
+
+
+def get_stateless(fn: Callable[..., T_sample]) -> bool:
+    """Get whether a function is stateless."""
+    return getattr(fn, "__stateless__", False)
+
+
+def convert_positions_to_block_layout(
+    positions: torch.Tensor, t: int, h: int, w: int, spatial_merge_size: int = 2
+) -> torch.Tensor:
+    """
+    Convert patch positions from row-major order to 2x2 block layout.
+
+    This function reorders patch positions to match the 2x2 block arrangement
+    used by the image processor. Uses index-based reordering instead of reshape.
+
+    Args:
+        positions: Patch positions in row-major order, shape [t*h*w, 3]
+        t: temporal dimension
+        h: height (unmerged patch count)
+        w: width (unmerged patch count)
+        spatial_merge_size: size of spatial merge blocks (default: 2)
+
+    Returns:
+        torch.Tensor: Patch positions in 2x2 block order, same shape [t*h*w, 3]
+    """
+    sms = spatial_merge_size
+    if sms == 1:
+        return positions
+
+    device = positions.device
+    total_patches = t * h * w
+
+    # Generate row-major indices: [0, 1, 2, ..., t*h*w-1]
+    # Reshape to [t, h, w]
+    indices = torch.arange(total_patches, device=device).view(t, h, w)
+
+    # Calculate merged dimensions
+    h_merged = h // sms
+    w_merged = w // sms
+
+    # Reshape to [t, h_merged, sms, w_merged, sms]
+    indices = indices.view(t, h_merged, sms, w_merged, sms)
+
+    # Permute to [t, h_merged, w_merged, sms_h, sms_w] - 2x2 block order
+    indices = indices.permute(0, 1, 3, 2, 4).contiguous()
+
+    # Flatten to get the reordering indices
+    indices = indices.view(total_patches)
+
+    # Apply the reordering to positions
+    return positions[indices]
+
+
+@dataclass
+class Gemma4VLImageTaskSample(ImageTaskSample):
+    """An image task sample with a grid of tokens and their corresponding pixel values."""
+
+    image_grid_thw: torch.Tensor = None
+    video_grid_thw: torch.Tensor = None
+
+    def __init__(self, image_grid_thw: str, video_grid_thw=None, **kwargs):
+        super().__init__(**kwargs)
+        self.image_grid_thw = image_grid_thw
+        self.video_grid_thw = video_grid_thw
+
+
+@dataclass
+class Gemma4VLImageTaskSamplePacked(ImageTaskSamplePacked):
+    """An image task sample with a grid of tokens and their corresponding pixel values."""
+
+    image_grid_thw: torch.Tensor = None
+    video_grid_thw: torch.Tensor = None
+
+    def __init__(self, sample: ImageTaskSample, image_grid_thw: str, video_grid_thw=None):
+        super().__init__(**vars(sample))
+        self.image_grid_thw = image_grid_thw
+        self.video_grid_thw = video_grid_thw
+
+
+@dataclass
+class Gemma4VLImageTaskBatchPacked(ImageTaskBatchPacked):
+    """An image task sample with a grid of tokens and their corresponding pixel values."""
+
+    image_grid_thw: torch.Tensor = None
+    video_grid_thw: torch.Tensor = None
+
+    def __init__(self, sample: ImageTaskSample, image_grid_thw: str, video_grid_thw=None):
+        super().__init__(**vars(sample))
+        self.image_grid_thw = image_grid_thw
+        self.video_grid_thw = video_grid_thw
+
+
+class Gemma4VLTaskEncoder(TaskEncoder):
+    """A simple task encoder for VLMs."""
+
+    def __init__(self, args):
+        super().__init__()
+        if args.training_phase in ["sft"]:
+            self.chat_template = get_chat_template()
+        self.processor = AutoProcessor.from_pretrained(self.args.hf_tokenizer_path, trust_remote_code=True)
+
+        # image
+        self.min_pixels = args.min_pixels
+        self.max_pixels = args.max_pixels
+
+    def _normalize_image_backed_video_placeholders(
+        self,
+        messages: list[dict[str, str]],
+        image: list | None,
+        video: list | None,
+    ) -> list[dict[str, str]]:
+        """Rewrite image-backed <video> placeholders into image placeholders."""
+        has_images = image is not None and len(image) > 0
+        has_video = video is not None and (
+            (isinstance(video, list) and len(video) > 0) or (not isinstance(video, list))
+        )
+
+        if not has_images or has_video:
+            return messages
+
+        video_placeholder_count = sum(msg.get("content", "").count(constants.Placeholder.VIDEO) for msg in messages)
+        if video_placeholder_count == 0:
+            return messages
+
+        image_block = "\n".join([constants.Placeholder.IMAGE] * len(image))
+        for msg in messages:
+            if constants.Placeholder.VIDEO in msg.get("content", ""):
+                msg["content"] = msg["content"].replace(constants.Placeholder.VIDEO, image_block)
+
+        return messages
+
+    def _resize_image(self, image, size_factor=48):
+        resized_height, resized_width = smart_resize(
+            image.height,
+            image.width,
+            factor=size_factor,
+            min_pixels=self.min_pixels,
+            max_pixels=self.max_pixels,
+        )
+        image = image.resize((resized_width, resized_height))
+
+        return image
+
+    def _process(self, image, text):
+        """ " Process the data to get the model's input"""
+        inputs = self.processor(
+            text=text,
+            images=image,
+            padding=True,
+            return_tensors="pt",
+        )
+        input_ids = inputs["input_ids"][0]
+        attn_mask = inputs["attention_mask"][0].logical_not()
+        image_grid_thw = None
+        pixel = []
+        if image is not None:
+            if "image_grid_thw" in inputs:
+                image_grid_thw = inputs["image_grid_thw"]  # [t,h,w]
+                pixel = [inputs["pixel_values"]]  # [hw, 2*3*14*14]
+            else:
+                images = image if isinstance(image, list) else [image]
+                _, mm_inputs = self.chat_template.mm_plugin._build_gemma4_mm_inputs(images, self.processor)
+                image_grid_thw = mm_inputs["image_grid_thw"]
+                pixel = [mm_inputs["pixel_values"]]
+
+        target = input_ids.clone()
+        vision_start_id, img_pad_id, vision_end_id = self.tokenizer.convert_tokens_to_ids(
+            [VISION_TAGS[0], IMAGE_TOKEN, VISION_TAGS[1]]
+        )
+        target[target == vision_start_id] = IGNORE_INDEX
+        target[target == img_pad_id] = IGNORE_INDEX
+        target[target == vision_end_id] = IGNORE_INDEX
+
+        return input_ids, target, pixel, image_grid_thw, attn_mask
+
+    def process_sft_vqa(self, context, answer, image):
+        """process the data for sft vqa"""
+        text = self.processor.apply_chat_template(
+            [{"role": "user", "content": context}, {"role": "assistant", "content": answer}], tokenize=False
+        ).replace("<image>", IMAGE_TOKEN)
+        if text[-1] == "\n":
+            text = text[:-1]
+        input_ids, _, imgs, image_grid_thw, attn_mask = self._process(image, text)
+        target = torch.ones_like(input_ids) * IGNORE_INDEX
+        answer = self.tokenizer.tokenize(answer)
+        target[-len(answer) - 1 : -1] = torch.tensor(answer)
+
+        return input_ids, target, attn_mask, imgs, image_grid_thw
+
+    def compute_frame_timestamps(self, images, patch_positions, fps):
+        """Compute timestamps for unique frames from patch positions.
+
+        Args:
+            images: List of images
+            patch_positions: Tensor of shape (n, 3) with (t, h, w) coordinates,
+                             where t is the frame index
+            fps: Frames per second
+
+        Returns:
+            List of float timestamps for each unique frame
+        """
+        if patch_positions is None or len(patch_positions) == 0:
+            return []
+
+        t_values = patch_positions[:, 0]
+        unique_t = torch.unique(t_values)
+        timestamps = [float(t.item() / fps) for t in unique_t]
+
+        return timestamps
+
+    def _rewrap_vision_by_frame(
+        self,
+        messages: list[dict],
+        patch_positions: list[torch.Tensor],
+        timestamp_strings: list[str],
+    ) -> list[dict]:
+        """Rewrite vision blocks in message text from per-canvas to per-frame grouping.
+
+        After mm_plugin.process_messages, the message content has one vision block
+        per canvas image (e.g. 20 blocks of 144 pad tokens each):
+            VS PAD*144 VE \n VS PAD*144 VE \n ...
+
+        This rewrites to per-frame blocks with timestamps:
+            <t1> VS PAD*K1 VE \n <t2> VS PAD*K2 VE \n ...
+
+        Works at the string level BEFORE tokenization, avoiding error-prone
+        token-level merge/split.
+
+        Args:
+            messages: Message dicts with 'content' strings (after mm_plugin).
+            patch_positions: List of tensors in block layout, each [n_patches, 3].
+            timestamp_strings: List of timestamp strings, one per unique frame,
+                e.g. ["<0.0 seconds>", "<0.1 seconds>", ...].
+
+        Returns:
+            Modified messages list (same objects, content strings updated in-place).
+        """
+        # Compute per-frame merged token counts from block-layout patch positions.
+        # After block layout, every (sms*sms) consecutive patches = 1 merged token,
+        # and patches with the same t are contiguous.
+        sms = _gemma4_spatial_merge_size(self.processor.image_processor)
+        merge_unit = sms * sms
+        flat_positions = torch.cat(patch_positions, dim=0)
+        num_merged_tokens = len(flat_positions) // merge_unit
+        token_t_values = flat_positions[::merge_unit, 0].tolist()
+
+        frame_token_counts = []
+        idx = 0
+        while idx < num_merged_tokens:
+            current_t = token_t_values[idx]
+            j = idx + 1
+            while j < num_merged_tokens and token_t_values[j] == current_t:
+                j += 1
+            frame_token_counts.append(j - idx)
+            idx = j
+
+        # Build new vision string: for each frame, [timestamp] VS PAD*N VE \n
+        new_parts = []
+        for frame_i, count in enumerate(frame_token_counts):
+            if frame_i < len(timestamp_strings):
+                new_parts.append(timestamp_strings[frame_i])
+            new_parts.append(VISION_TAGS[0])
+            new_parts.append(IMAGE_TOKEN * count)
+            new_parts.append(VISION_TAGS[1])
+            new_parts.append("\n")
+        new_vision_str = "".join(new_parts)
+
+        # Find and replace the entire vision region in the message content.
+        # Original: VS...VE\nVS...VE\n...VS...VE\n\nDescribe...
+        # We replace from the first VS through the last VE + one \n,
+        # and our new_vision_str already ends with \n for each frame.
+        for msg in messages:
+            content = msg.get("content", "")
+            if not isinstance(content, str) or VISION_TAGS[0] not in content:
+                continue
+
+            first_vs = content.find(VISION_TAGS[0])
+            last_ve_end = content.rfind(VISION_TAGS[1]) + len(VISION_TAGS[1])
+            # Skip one \n after the last VE (our new string already provides it)
+            tail_start = last_ve_end
+            if tail_start < len(content) and content[tail_start] == "\n":
+                tail_start += 1
+
+            msg["content"] = content[:first_vs] + new_vision_str + content[tail_start:]
+            break  # only process the first message with vision content
+        return messages
+
+    def process_sft_qa(
+        self, messages: list, system: str, raw_video: list, raw_image: list, raw_patch_positions: list, **kwargs
+    ):
+        """process the data for sft qa"""
+        video_grid_thw = None
+        pixel_values_videos = []
+        image_grid_thw = None
+        pixel_values_images = []
+        video = []
+        image = []
+        patch_positions = []
+
+        has_image_inputs = raw_image is not None and len(raw_image) > 0
+        if has_image_inputs:
+            image = raw_image
+
+        if raw_patch_positions is not None:
+            for i in raw_patch_positions:
+                if i is not None:
+                    patch_positions.append(torch.tensor(i, dtype=torch.int64))
+
+        messages, mm_inputs = self.chat_template.mm_plugin.process_messages(
+            messages,
+            image,
+            video if raw_video is not None else [],
+            self.processor,
+        )
+
+        if has_image_inputs:
+            image_grid_thw = mm_inputs.get("image_grid_thw")
+            assert image_grid_thw is not None, (
+                "Missing `image_grid_thw` in multimodal inputs:\n"
+                f"- num_raw_image: {len(raw_image)}\n"
+                f"- mm_input_keys: {list(mm_inputs.keys())}\n"
+                f"- has_image_placeholder: {any('<image>' in msg.get('content', '') for msg in messages)}"
+            )
+            # get image patch num:
+            num_patches = 0
+            for img_idx in range(len(image_grid_thw)):
+                t_val, h_val, w_val = image_grid_thw[img_idx].tolist()
+                num_patches += h_val * w_val
+            pixel_values = mm_inputs.get("pixel_values")
+            assert pixel_values is not None, (
+                "Missing `pixel_values` in multimodal inputs:\n"
+                f"- num_raw_image: {len(raw_image)}\n"
+                f"- mm_input_keys: {list(mm_inputs.keys())}"
+            )
+            pixel_values_images = [pixel_values]
+            plugin_patch_positions = mm_inputs.get("patch_positions")
+            if len(patch_positions) == 0:
+                if plugin_patch_positions is not None:
+                    patch_positions = [pp.to(dtype=torch.int64) for pp in plugin_patch_positions]
+                else:
+                    # Fallback: derive row-major (t, h, w) coordinates from image_grid_thw.
+                    for img_idx in range(len(image_grid_thw)):
+                        t_val, h_val, w_val = image_grid_thw[img_idx].tolist()
+                        cur_num_patches = t_val * h_val * w_val
+                        h_coords = torch.arange(h_val, dtype=torch.int64).repeat_interleave(w_val).repeat(t_val)
+                        w_coords = torch.arange(w_val, dtype=torch.int64).repeat(h_val).repeat(t_val)
+                        t_coords = torch.zeros(cur_num_patches, dtype=torch.int64)
+                        patch_positions.append(torch.stack([t_coords, h_coords, w_coords], dim=1))
+            else:
+                image_grid_thw = torch.tensor([[len(image_grid_thw), image_grid_thw[0][1], image_grid_thw[0][2]]])
+                # Apply block layout conversion for temporal contiguity after spatial merge
+                flat_positions = torch.cat(patch_positions, dim=0)
+                t_v, h_v, w_v = int(image_grid_thw[0][0]), int(image_grid_thw[0][1]), int(image_grid_thw[0][2])
+                flat_positions = convert_positions_to_block_layout(
+                    flat_positions, t_v, h_v, w_v,
+                    spatial_merge_size=_gemma4_spatial_merge_size(self.processor.image_processor),
+                )
+                patch_positions = [flat_positions]
+            patch_positions_sum = sum(len(p) for p in patch_positions)
+            assert num_patches == patch_positions_sum, (
+                "num_patches mismatch:\n"
+                f"- num_patches: {num_patches}\n"
+                f"- patch_positions_sum: {patch_positions_sum}\n"
+                f"- len(image_grid_thw): {len(image_grid_thw)}\n"
+                f"- patch_positions len: {len(patch_positions)}\n"
+                f"- image_grid_thw[-1]: {image_grid_thw[-1]}\n"
+                f"- image_sizes: {[img.size for img in image]}"
+            )
+
+        # Compute timestamps and rewrite vision blocks by frame in the message text.
+        # This rewrites mm_plugin's per-canvas wrapping (VS PAD*N VE \n VS PAD*M VE \n ...)
+        # into per-frame wrapping with timestamps (<ts> VS PAD*K VE \n ...) at the STRING
+        # level, before tokenization — no token-level merge/split needed.
+        timestamp_strings = None
+        if kwargs is not None and "fps" in kwargs and len(patch_positions) > 0:
+            fps = kwargs["fps"][0] if isinstance(kwargs["fps"], list) else kwargs["fps"]
+            if fps is not None and fps > 0:
+                td = kwargs.get("timestamp_decimal", 1) or 1
+                pt_patch_position = torch.cat(patch_positions)
+                timestamps = self.compute_frame_timestamps(patch_positions, pt_patch_position, fps)
+                timestamps = [round(t, td) for t in timestamps]
+                timestamp_strings = [f"<{t:.{td}f} seconds>" for t in timestamps]
+
+        if timestamp_strings is not None and len(timestamp_strings) > 0 and len(patch_positions) > 0:
+            messages = self._rewrap_vision_by_frame(messages, patch_positions, timestamp_strings)
+        encode_pairs = self.chat_template.encode_multiturn(
+            tokenizer=self.tokenizer,
+            messages=messages,
+            system=system,
+        )
+        input_ids, target = [], []
+        for turn_idx, (source_ids, target_ids) in enumerate(encode_pairs):
+            input_ids += source_ids + target_ids
+            target += [IGNORE_INDEX] * len(source_ids) + target_ids
+        input_ids = torch.tensor(input_ids)
+        target = torch.tensor(target)
+        attn_mask = torch.zeros_like(input_ids).bool()
+
+        return (
+            input_ids,
+            target,
+            attn_mask,
+            pixel_values_images,
+            image_grid_thw,
+            pixel_values_videos,
+            video_grid_thw,
+            patch_positions,
+        )
+
+    def encode_vqa4packing(self, sample: VQASample) -> ImageTaskSample:
+        """Encode VQASample in Gemma4VL style."""
+        text = self.processor.apply_chat_template(
+            [{"role": "user", "content": sample.context}, {"role": "assistant", "content": sample.answers}],
+            tokenize=False,
+        ).replace("<image>", IMAGE_TOKEN)
+
+        if text[-1] == "\n":
+            text = text[:-1]
+            pass
+
+        input_ids, _, imgs, image_grid_thw, attn_mask = self._process(sample.image, text)
+        target = torch.ones_like(input_ids) * IGNORE_INDEX
+        answers = self.tokenizer.tokenize(sample.answers)
+        target[-len(answers) - 1 : -1] = torch.tensor(answers)
+        target[-1] = input_ids[-1]
+        # print(target[-1])
+
+        num_tiles = [len(image_grid_thw)]
+        if self.args.enable_discard_sample:
+            if len(input_ids) > self.args.seq_length:
+                skip_malformed_multimodal_sample(
+                    sample.__key__,
+                    "input_length_exceeds_seq_length",
+                    f"input length {len(input_ids)} exceeds seq_length={self.args.seq_length}",
+                )
+        else:
+            assert image_grid_thw.prod() / 4 <= self.args.seq_length, f"{sample.__key__} grid_thw: {image_grid_thw}"
+
+        return Gemma4VLImageTaskSample(
+            __key__=sample.__key__,
+            __restore_key__=sample.__restore_key__,
+            __subflavor__=None,
+            __subflavors__=sample.__subflavors__,
+            imgs=imgs,
+            image_grid_thw=image_grid_thw,
+            num_tiles=num_tiles,
+            tokens=input_ids,
+            labels=target,
+            attn_mask=attn_mask,
+            total_len=len(input_ids),
+        )
+
+    def encode_multi_vid_qa(self, sample: VQASample) -> ImageTaskSample:
+        """Encode sample in Gemma4VL style."""
+        if self.args.training_phase == constants.TrainingPhase.SFT:
+            input_ids, target, attn_mask, imgs, image_grid_thw, video, video_grid_thw = self.process_sft_qa(
+                sample.messages, sample.system, sample.video, None
+            )
+        else:
+            raise NotImplementedError(f"Unknown training phase {self.args.training_phase}")
+
+        if self.args.enable_discard_sample:
+            if len(input_ids) > self.args.seq_length:
+                skip_malformed_multimodal_sample(
+                    sample.__key__,
+                    "input_length_exceeds_seq_length",
+                    f"input length {len(input_ids)} exceeds seq_length={self.args.seq_length}",
+                )
+        else:
+            assert video_grid_thw.prod(dim=-1).sum() / 4 <= self.args.seq_length, (
+                f"{sample.__key__} grid_thw: {video_grid_thw}"
+            )
+
+        return Gemma4VLImageTaskSample(
+            __key__=sample.__key__,
+            __restore_key__=sample.__restore_key__,
+            __subflavor__=None,
+            __subflavors__=sample.__subflavors__,
+            imgs=imgs,
+            image_grid_thw=image_grid_thw,
+            pixel_values_videos=video,
+            video_grid_thw=video_grid_thw,
+            num_tiles=[len(video_grid_thw)],
+            tokens=input_ids,
+            labels=target,
+            attn_mask=attn_mask,
+            total_len=len(input_ids),
+        )
+
+    def encode_multi_mix_qa(self, sample: MultiMixQASample) -> ImageTaskSample:
+        """Encode sample in Gemma4VL style."""
+
+        if self.args.training_phase == constants.TrainingPhase.SFT:
+            num_tiles = []
+            kwargs = {}
+            if hasattr(sample, "fps") and sample.fps is not None:
+                kwargs["fps"] = sample.fps
+            if hasattr(sample, "timestamp_decimal") and sample.timestamp_decimal is not None:
+                kwargs["timestamp_decimal"] = sample.timestamp_decimal
+
+            def _remove_last_qa_round(messages: list[dict]) -> list[dict]:
+                assistant_idx = -1
+                for idx in range(len(messages) - 1, -1, -1):
+                    if messages[idx].get("role") == constants.DataRoles.ASSISTANT:
+                        assistant_idx = idx
+                        break
+
+                if assistant_idx == -1:
+                    return []
+
+                user_idx = -1
+                for idx in range(assistant_idx - 1, -1, -1):
+                    if messages[idx].get("role") == constants.DataRoles.USER:
+                        user_idx = idx
+                        break
+
+                if user_idx == -1:
+                    return []
+
+                return messages[:user_idx] + messages[assistant_idx + 1 :]
+
+            current_messages = self._normalize_image_backed_video_placeholders(
+                [dict(message) for message in sample.messages],
+                sample.image,
+                sample.video,
+            )
+            current_image = sample.image
+            current_video = sample.video
+            current_patch_positions = sample.patch_positions
+
+            while True:
+                (
+                    input_ids,
+                    target,
+                    attn_mask,
+                    imgs,
+                    image_grid_thw,
+                    pixel_values_videos,
+                    video_grid_thw,
+                    patch_positions,
+                ) = self.process_sft_qa(
+                    current_messages,
+                    sample.system,
+                    current_video,
+                    current_image,
+                    current_patch_positions,
+                    **kwargs,
+                )
+
+                if len(input_ids) <= self.args.seq_length:
+                    break
+
+                current_messages = _remove_last_qa_round(current_messages)
+                if len(current_messages) == 0:
+                    if self.args.enable_discard_sample:
+                        skip_malformed_multimodal_sample(
+                            sample.__key__,
+                            "qa_truncation_exhausted",
+                            (
+                                "sample has no QA rounds left after truncation "
+                                f"to fit seq_length={self.args.seq_length}"
+                            ),
+                        )
+                    raise AssertionError(
+                        "Sample has no QA rounds left after truncation to fit seq_length:\n"
+                        f"- sample: {sample.__key__}\n"
+                        f"- seq_length: {self.args.seq_length}"
+                    )
+
+                image_placeholder_count = sum(
+                    message.get("content", "").count(constants.Placeholder.IMAGE) for message in current_messages
+                )
+                video_placeholder_count = sum(
+                    message.get("content", "").count(constants.Placeholder.VIDEO) for message in current_messages
+                )
+
+                if current_image is not None:
+                    current_image = current_image[:image_placeholder_count]
+
+                if current_patch_positions is not None:
+                    current_patch_positions = current_patch_positions[:image_placeholder_count]
+
+                if isinstance(current_video, list):
+                    current_video = current_video[:video_placeholder_count]
+                elif current_video is not None and video_placeholder_count == 0:
+                    current_video = None
+
+            if video_grid_thw is not None:
+                num_tiles = [len(video_grid_thw)]
+            elif image_grid_thw is not None:
+                num_tiles = [len(image_grid_thw)]
+        else:
+            raise NotImplementedError(f"Unknown training phase {self.args.training_phase}")
+
+        assert len(input_ids) > 0, f"input_ids is empty in {sample.__key__}"
+
+        if self.args.enable_discard_sample:
+            if len(input_ids) > self.args.seq_length:
+                skip_malformed_multimodal_sample(
+                    sample.__key__,
+                    "input_length_exceeds_seq_length",
+                    f"input length {len(input_ids)} exceeds seq_length={self.args.seq_length}",
+                )
+        elif video_grid_thw is not None:
+            sms = _gemma4_spatial_merge_size(self.processor.image_processor)
+            assert video_grid_thw.prod(dim=-1).sum() / (sms * sms) <= self.args.seq_length, (
+                f"{sample.__key__} grid_thw: {video_grid_thw}"
+            )
+        elif image_grid_thw is not None:
+            sms = _gemma4_spatial_merge_size(self.processor.image_processor)
+            image_token_len = int(image_grid_thw.prod(dim=-1).sum().item() / (sms * sms))
+            assert image_token_len <= self.args.seq_length, (
+                "Image token length exceeds seq_length:\n"
+                f"- sample: {sample.__key__}\n"
+                f"- image_grid_thw: {image_grid_thw}\n"
+                f"- image_token_len: {image_token_len}\n"
+                f"- seq_length: {self.args.seq_length}"
+            )
+
+        return Gemma4VLImageTaskSample(
+            __key__=sample.__key__,
+            __restore_key__=sample.__restore_key__,
+            __subflavor__=None,
+            __subflavors__=sample.__subflavors__,
+            imgs=imgs,
+            image_grid_thw=image_grid_thw,
+            pixel_values_videos=pixel_values_videos,
+            video_grid_thw=video_grid_thw,
+            num_tiles=num_tiles,
+            tokens=input_ids,
+            labels=target,
+            attn_mask=attn_mask,
+            total_len=len(input_ids),
+            patch_positions=patch_positions,
+        )
+
+    def process_samples_grid(self, samples):
+        """concat grid_thw for image and video"""
+        image_grid_thw = [x.image_grid_thw for x in samples if x.image_grid_thw is not None]
+        video_grid_thw = [x.video_grid_thw for x in samples if x.video_grid_thw is not None]
+
+        if len(image_grid_thw) > 0:
+            image_grid_thw = torch.cat(image_grid_thw).to(dtype=torch.int32)
+        else:
+            image_grid_thw = None
+
+        if len(video_grid_thw) > 0:
+            video_grid_thw = torch.cat(video_grid_thw).to(dtype=torch.int32)
+        else:
+            video_grid_thw = None
+
+        return image_grid_thw, video_grid_thw
+
+    @override
+    @stateless
+    def pack_selected_samples(self, samples: list[Gemma4VLImageTaskSample]) -> list[Gemma4VLImageTaskSamplePacked]:
+        """Pack selected samples into one big sample."""
+        image_grid_thw, video_grid_thw = self.process_samples_grid(samples)
+        return Gemma4VLImageTaskSamplePacked(
+            super().pack_selected_samples(samples), image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw
+        )
+
+    @override
+    def batch(
+        self, samples: list[Union[Gemma4VLImageTaskSample, Gemma4VLImageTaskSamplePacked]]
+    ) -> Gemma4VLImageTaskBatchPacked:
+        """Batch samples together"""
+        image_grid_thw, video_grid_thw = self.process_samples_grid(samples)
+        return Gemma4VLImageTaskBatchPacked(
+            super().batch(samples), image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw
+        )
+
+    @override
+    def process_images(
+        self, samples: list[Union[Gemma4VLImageTaskSample, Gemma4VLImageTaskSamplePacked]]
+    ) -> torch.Tensor:
+        """ " Process the data to get the model's input"""
+        imgs = [img for s in samples if s.imgs is not None for img in s.imgs]
+        if len(imgs) > 0:
+            return torch.cat(imgs)
+        else:
+            return torch.tensor([[0]], dtype=torch.float32)
+
+    @override
+    def process_videos(
+        self, samples: list[Union[Gemma4VLImageTaskSample, Gemma4VLImageTaskSamplePacked]]
+    ) -> torch.Tensor:
+        """ " Process the data to get the model's input"""
+        pixel_values_videos = [
+            pixel_values_video
+            for s in samples
+            if s.pixel_values_videos is not None
+            for pixel_values_video in s.pixel_values_videos
+        ]
+        if len(pixel_values_videos) > 0:
+            return torch.cat(pixel_values_videos)
+        else:
+            return torch.tensor([[0]], dtype=torch.float32)
+
+    @override
+    def build_train_datasets(
+        self,
+        *,
+        datasets: list[tuple[BaseCoreDatasetFactory[T_sample], Union[float, int, None]]],
+        worker_config: WorkerConfig,
+        batch_size: Optional[int],
+        batch_drop_last: bool = False,
+        packing_buffer_size: Optional[int] = None,
+        virtual_epoch_length: int = 0,
+        shuffle_buffer_size: Optional[int] = None,
+        blend_mode: DatasetBlendMode = DatasetBlendMode.NONE,
+        repeat: bool = True,
+    ) -> SavableDataset[T_batch]:
+        """Combines train datasets to a single dataset."""
+
+        # Check if there's a CrudeWebdataset but no cookers
+        for dataset, _ in datasets:
+            if isinstance(dataset, CrudeWebdataset):
+                assert self.cookers, "CrudeWebdataset found, but no cookers registered."
+
+        global_workers = max(1, worker_config.num_workers) * worker_config.world_size
+        rotation_lengths = [len(dataset) for dataset, _ in datasets]
+        for i in range(1, len(rotation_lengths)):
+            rotation_lengths[i] += rotation_lengths[i - 1]
+        worker_rotation_offsets = [rotation_length % global_workers for rotation_length in [0] + rotation_lengths[:-1]]
+
+        if repeat:
+            inner_datasets = [
+                (
+                    RepeatDataset(
+                        dataset.build(worker_rotation_offset=worker_rotation_offset),
+                        worker_config=worker_config,
+                    ),
+                    1.0 if weight is None else float(weight),
+                )
+                for (dataset, weight), worker_rotation_offset in zip(datasets, worker_rotation_offsets)
+            ]
+        else:
+            assert blend_mode in (
+                DatasetBlendMode.NONE,
+                DatasetBlendMode.SAMPLE_REPETITIONS,
+            ) and all(isinstance(repetitions, int) for _dataset, repetitions in datasets), (
+                "If repeat is False, the datasets must be repeated with integer weights."
+            )
+            inner_datasets = [
+                (
+                    (
+                        dataset.build(worker_rotation_offset=worker_rotation_offset)
+                        if repetition is None or repetition == 1
+                        else RepeatDataset(
+                            dataset.build(worker_rotation_offset=worker_rotation_offset),
+                            repeats=int(repetition),
+                            worker_config=worker_config,
+                        )
+                    ),
+                    len(dataset) * (1 if repetition is None else int(repetition)),
+                )
+                for (dataset, repetition), worker_rotation_offset in zip(datasets, worker_rotation_offsets)
+            ]
+
+        if len(inner_datasets) > 1:
+            # The worker offset for each dataset is the cumsum of the dataset lengths, but modulo the
+            # global number of workers.
+            dataset = BlendDataset(
+                *inner_datasets,
+                worker_config=worker_config,
+            )
+        elif len(datasets) == 1:
+            dataset = inner_datasets[0][0]
+        else:
+            raise ValueError("No datasets given.")
+        if shuffle_buffer_size is not None and shuffle_buffer_size > 1:
+            dataset = ShuffleBufferDataset(
+                dataset,
+                size=shuffle_buffer_size,
+                worker_config=worker_config,
+            )
+        dataset = self.build_cook_crude_sample(dataset, worker_config=worker_config)
+        dataset = self.build_encode_sample(dataset, worker_config=worker_config)
+
+        # Insert pool sorting before entering BatchDataset
+        if getattr(self.args, "length_sort_pool_size", 0) and self.args.length_sort_pool_size > 0:
+
+            def sort_key_fn(s):
+                return getattr(s, "total_len", len(getattr(s, "tokens")))
+
+            sort_ascending = not getattr(self.args, "length_sort_desc", False)
+            sort_cls = (
+                PackedSeparateSortDataset
+                if getattr(self.args, "length_sort_separate_packed", False)
+                else LengthPoolSortDataset
+            )
+            dataset = sort_cls(
+                dataset,
+                pool_size=self.args.length_sort_pool_size,
+                key_fn=sort_key_fn,
+                ascending=sort_ascending,
+                worker_config=worker_config,
+                warmup_steps=getattr(self.args, "length_sort_warmup_steps", 0),
+                initial_pool_size=getattr(self.args, "length_sort_initial_pool_size", 10),
+            )
+        dataset = self.build_batch(
+            dataset,
+            batch_size=batch_size,
+            batch_drop_last=batch_drop_last,
+            packing_buffer_size=packing_buffer_size,
+            worker_config=worker_config,
+        )
+        if virtual_epoch_length > 0:
+            dataset = EpochizeDataset(
+                dataset,
+                length=virtual_epoch_length,
+                worker_config=worker_config,
+            )
+        if worker_config.should_log(level=1):
+            dataset = LogSampleDataset(dataset, mode="train", worker_config=worker_config)
+        return dataset
diff --git a/aiak_training_llm/models/__init__.py b/aiak_training_llm/models/__init__.py
index dbf99330..86d393e8 100644
--- a/aiak_training_llm/models/__init__.py
+++ b/aiak_training_llm/models/__init__.py
@@ -5,6 +5,7 @@
 from .qwen_vl import qwen2_vl_config, qwen2_vl_provider
 from .llava_onevision1_5 import llava_onevision1_5_config, llava_onevision1_5_provider
 from .llava_onevision2 import llava_onevision2_config, llava_onevision2_provider
+from .gemma4_vl import gemma4_vl_config, gemma4_vl_provider
 
 from .factory import (
     get_support_model_archs,
diff --git a/aiak_training_llm/models/gemma4_vl/__init__.py b/aiak_training_llm/models/gemma4_vl/__init__.py
new file mode 100644
index 00000000..f3faeeff
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/__init__.py
@@ -0,0 +1,7 @@
+"""Gemma4-VL model implementation."""
+
+from . import gemma4_vl_config, gemma4_vl_provider
+from .gemma4_vl_model import Gemma4VL
+
+
+__all__ = ["Gemma4VL", "gemma4_vl_config", "gemma4_vl_provider"]
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_adapter.py b/aiak_training_llm/models/gemma4_vl/gemma4_adapter.py
new file mode 100644
index 00000000..6d6a9e1d
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_adapter.py
@@ -0,0 +1,47 @@
+"""Gemma4 vision-to-LM adapter.
+
+Verbatim port of HF ``Gemma4MultimodalEmbedder`` (modeling_gemma4 lines 2023-2047).
+
+In the HF state_dict the adapter is stored under
+``model.embed_vision.{embedding_pre_projection_norm,embedding_projection}.*``.
+For Gemma4-26B-A4B-it specifically:
+- ``embedding_pre_projection_norm`` is a ``Gemma4RMSNorm`` with
+  ``with_scale=False`` (no learnable weight, ckpt has no entry for it).
+- ``embedding_projection`` is ``nn.Linear(vision.hidden_size=1152,
+  text.hidden_size=2816, bias=False)``; the only ckpt entry is
+  ``model.embed_vision.embedding_projection.weight``.
+
+So the entire adapter holds exactly ONE trainable tensor. Do not promote the
+RMSNorm to ``with_scale=True`` even though it would seem more "complete" — HF
+ckpt loading will fail on an unexpected ``embedding_pre_projection_norm.weight``
+key. The naming uses HF attribute names so converter mapping stays trivial.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .gemma4_vision_common import Gemma4RMSNorm
+
+
+class Gemma4Adapter(nn.Module):
+    def __init__(
+        self,
+        vision_hidden_size: int,
+        text_hidden_size: int,
+        rms_norm_eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.vision_hidden_size = vision_hidden_size
+        self.text_hidden_size = text_hidden_size
+        self.embedding_pre_projection_norm = Gemma4RMSNorm(
+            vision_hidden_size, eps=rms_norm_eps, with_scale=False
+        )
+        self.embedding_projection = nn.Linear(
+            vision_hidden_size, text_hidden_size, bias=False
+        )
+
+    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
+        embs_normed = self.embedding_pre_projection_norm(inputs_embeds)
+        return self.embedding_projection(embs_normed)
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_attention.py b/aiak_training_llm/models/gemma4_vl/gemma4_attention.py
new file mode 100644
index 00000000..54005081
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_attention.py
@@ -0,0 +1,258 @@
+"""Gemma4 self-attention with hybrid head_dim, K=V tying, and parameter-free V LayerNorm.
+
+This module subclasses Megatron's SelfAttention to implement Gemma4-specific
+behaviors WITHOUT modifying megatron/core/transformer/attention.py:
+
+  1. Hybrid per-layer head_dim and num_kv_heads (sliding vs global layers).
+     Implemented by cloning TransformerConfig and overriding kv_channels +
+     num_query_groups before super().__init__.
+
+  2. K=V tying on global layers (the K projection is also reused as V).
+     Implemented by rebuilding self.linear_qkv with reduced fan_out
+     (q + 1*kv instead of q + 2*kv) and overriding get_query_key_value_tensors
+     to perform single-KV split with value = key.
+
+  3. Parameter-free V LayerNorm (F.layer_norm with no weight/bias).
+     Applied at the end of get_query_key_value_tensors.
+
+  4. Dual-RoPE dispatch: when ``rotary_pos_emb`` is a ``dict`` mapping
+     layer-type strings ('sliding'|'global') to per-type RoPE tensors,
+     this layer selects the tensor matching its own ``layer_pattern[layer_number-1]``
+     entry before calling the parent forward. This lets ``Gemma4Model`` ship two
+     independent ``RotaryEmbedding`` instances (sliding theta=1e4 vs global
+     theta=1e6 + proportional zero-pad) through the unmodified ``TransformerBlock``
+     pass-through path. Pre-existing single-tensor callers are unaffected.
+
+  5. Sandwich post-attention RMSNorm: HF Gemma4DecoderLayer applies
+     ``post_attention_layernorm`` between ``self_attn(...)`` output and the
+     residual add (modeling_gemma4.py:1387). Megatron's stock pre-norm
+     ``TransformerLayer`` has no slot for this norm, so it is owned here and
+     applied after ``super().forward`` returns ``(output, bias)``. The state_dict
+     key becomes ``self_attention.post_attention_layernorm.weight`` and is
+     written by the HF→MG converter from HF ``post_attention_layernorm.weight``.
+"""
+
+import copy
+from typing import Optional
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.spec_utils import build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from aiak_training_llm.models.dispatch import multiacc_modules
+from aiak_training_llm.utils import is_te_min_version
+
+
+try:
+    from megatron.core.extensions.transformer_engine import SplitAlongDim
+except ImportError:
+    SplitAlongDim = None
+
+
+def _resolve_per_layer_overrides(
+    config: TransformerConfig, layer_number: int
+) -> TransformerConfig:
+    """Per-layer override of kv_channels / num_query_groups / window_size.
+
+    Reads ``config.layer_pattern[layer_number - 1]`` ('sliding' or 'global') and
+    selects per-layer kv_channels and num_query_groups from
+    ``config.per_layer_kv_channels`` / ``config.per_layer_num_query_groups``.
+
+    Critical contract with TE backend: also writes ``cloned.window_size``, which
+    ``TEDotProductAttention`` reads (megatron/core/extensions/transformer_engine.py
+    line 738) and forwards to TE FlashAttention as ``window_size=(left, right)``.
+    HF's ``sliding_window=1024`` ("1024 tokens visible") maps to TE's ``(1023, 0)``
+    because TE's window is inclusive: ``[i-1023, i+0]`` = exactly 1024 tokens.
+    Global layers get ``None`` (pure causal). If ``layer_pattern`` is empty, the
+    config is returned unchanged (no-op for non-Gemma4 models).
+
+    Caller MUST always receive the cloned config (never early-return after pattern
+    is set) — otherwise sliding layers silently fall back to global causal.
+    """
+    if not config.layer_pattern:
+        return config
+    idx = layer_number - 1
+    if idx < 0 or idx >= len(config.layer_pattern):
+        return config
+    layer_type = config.layer_pattern[idx]
+    new_kv_channels = (config.per_layer_kv_channels or {}).get(layer_type)
+    new_num_query_groups = (config.per_layer_num_query_groups or {}).get(layer_type)
+    sliding_w = getattr(config, "sliding_window", None)
+    cloned = copy.copy(config)
+    if new_kv_channels is not None:
+        cloned.kv_channels = new_kv_channels
+    if new_num_query_groups is not None:
+        cloned.num_query_groups = new_num_query_groups
+    cloned.window_size = (sliding_w - 1, 0) if (layer_type == "sliding" and sliding_w is not None) else None
+    return cloned
+
+
+class Gemma4SelfAttention(SelfAttention):
+    """Self-attention layer for Gemma4 with hybrid head_dim, K=V tying, and V LN.
+
+    Args mirror SelfAttention exactly. The class inspects `config.layer_pattern`
+    and `config.kv_tied_layers` to decide per-layer behavior.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: SelfAttentionSubmodules,
+        layer_number: int,
+        attn_mask_type: AttnMaskType = AttnMaskType.padding,
+        cp_comm_type: Optional[str] = None,
+    ):
+        effective_config = _resolve_per_layer_overrides(config, layer_number)
+        super().__init__(
+            config=effective_config,
+            submodules=submodules,
+            layer_number=layer_number,
+            attn_mask_type=attn_mask_type,
+            cp_comm_type=cp_comm_type,
+        )
+
+        self.kv_tied = (layer_number - 1) in (config.kv_tied_layers or [])
+
+        if self.kv_tied:
+            self.linear_qkv = build_module(
+                submodules.linear_qkv,
+                self.config.hidden_size,
+                self.query_projection_size + self.kv_projection_size,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=self.config.add_bias_linear or self.config.add_qkv_bias,
+                skip_bias_add=False,
+                is_expert=False,
+                tp_comm_buffer_name='qkv',
+            )
+
+        norm_cls = multiacc_modules.TENorm if is_te_min_version("1.9.0") else multiacc_modules.LocalNorm
+        self.post_attention_layernorm = norm_cls(
+            config=self.config,
+            hidden_size=self.config.hidden_size,
+            eps=self.config.layernorm_epsilon,
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        attn_mask_type=None,
+        key_value_states=None,
+        inference_params=None,
+        rotary_pos_emb=None,
+        rotary_pos_cos=None,
+        rotary_pos_sin=None,
+        attention_bias=None,
+        packed_seq_params=None,
+        sequence_len_offset=None,
+    ):
+        if isinstance(rotary_pos_emb, dict):
+            layer_type = self.config.layer_pattern[self.layer_number - 1]
+            try:
+                rotary_pos_emb = rotary_pos_emb[layer_type]
+            except KeyError as exc:
+                raise KeyError(
+                    f"Gemma4SelfAttention(layer_number={self.layer_number}): "
+                    f"layer_pattern entry '{layer_type}' missing from rotary_pos_emb "
+                    f"dict keys {list(rotary_pos_emb.keys())}."
+                ) from exc
+        output, bias = super().forward(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            attn_mask_type=attn_mask_type,
+            key_value_states=key_value_states,
+            inference_params=inference_params,
+            rotary_pos_emb=rotary_pos_emb,
+            rotary_pos_cos=rotary_pos_cos,
+            rotary_pos_sin=rotary_pos_sin,
+            attention_bias=attention_bias,
+            packed_seq_params=packed_seq_params,
+            sequence_len_offset=sequence_len_offset,
+        )
+        if bias is not None:
+            output = output + bias
+            bias = None
+        output = self.post_attention_layernorm(output)
+        return output, bias
+
+    def get_query_key_value_tensors(self, hidden_states, key_value_states=None):
+        mixed_qkv, _ = self.linear_qkv(hidden_states)
+
+        if self.kv_tied:
+            new_tensor_shape = mixed_qkv.size()[:-1] + (
+                self.num_query_groups_per_partition,
+                (
+                    (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 1)
+                    * self.hidden_size_per_attention_head
+                ),
+            )
+            mixed_qkv = mixed_qkv.view(*new_tensor_shape)
+
+            split_arg_list = [
+                (
+                    self.num_attention_heads_per_partition
+                    // self.num_query_groups_per_partition
+                    * self.hidden_size_per_attention_head
+                ),
+                self.hidden_size_per_attention_head,
+            ]
+
+            if SplitAlongDim is not None:
+                (query, key) = SplitAlongDim(mixed_qkv, 3, split_arg_list)
+            else:
+                (query, key) = torch.split(mixed_qkv, split_arg_list, dim=3)
+
+            # K=V: value shares the K tensor. No clone — downstream kernels (TE/flash-attn)
+            # treat K and V as read-only. Saves one tensor allocation per layer.
+            value = key
+        else:
+            new_tensor_shape = mixed_qkv.size()[:-1] + (
+                self.num_query_groups_per_partition,
+                (
+                    (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2)
+                    * self.hidden_size_per_attention_head
+                ),
+            )
+            mixed_qkv = mixed_qkv.view(*new_tensor_shape)
+
+            split_arg_list = [
+                (
+                    self.num_attention_heads_per_partition
+                    // self.num_query_groups_per_partition
+                    * self.hidden_size_per_attention_head
+                ),
+                self.hidden_size_per_attention_head,
+                self.hidden_size_per_attention_head,
+            ]
+
+            if SplitAlongDim is not None:
+                (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list)
+            else:
+                (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3)
+
+        query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head)
+
+        if self.q_layernorm is not None:
+            query = self.q_layernorm(query)
+
+        if self.k_layernorm is not None:
+            key = self.k_layernorm(key)
+
+        # Parameter-free V LayerNorm: normalize over the head_dim with weight=None bias=None.
+        # Equivalent to LayerNorm with weight=1 and bias=0 that is NOT learned.
+        # Applied AFTER any K=V tying so V normalization does not mutate K.
+        if self.kv_tied:
+            # Re-derive value as a normalized copy of key so K is untouched.
+            value = F.layer_norm(key, [self.hidden_size_per_attention_head])
+        else:
+            value = F.layer_norm(value, [self.hidden_size_per_attention_head])
+
+        if self.config.test_mode:
+            self.run_realtime_tests()
+
+        return query, key, value
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_mlp.py b/aiak_training_llm/models/gemma4_vl/gemma4_mlp.py
new file mode 100644
index 00000000..a3176ae1
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_mlp.py
@@ -0,0 +1,169 @@
+"""Gemma4 parallel dense + MoE FFN block.
+
+Each Gemma4 transformer layer runs TWO FFN branches in parallel from the same
+*raw post-attention residual* and sums their outputs. The HF reference
+(``transformers/models/gemma4/modeling_gemma4.py``,
+``Gemma4TextDecoderLayer.forward`` lines 1374-1409) reads::
+
+    residual = h_after_attn
+    hidden_states = pre_feedforward_layernorm(residual)
+    hidden_states = self.mlp(hidden_states)                       # dense
+    if enable_moe_block:
+        hidden_states_1 = post_feedforward_layernorm_1(hidden_states)
+        hidden_states_flat = residual.reshape(-1, H)
+        _, top_w, top_i = self.router(hidden_states_flat)         # router on RAW residual
+        hidden_states_2 = pre_feedforward_layernorm_2(hidden_states_flat)
+        hidden_states_2 = self.experts(hidden_states_2, top_i, top_w)
+        hidden_states_2 = post_feedforward_layernorm_2(hidden_states_2)
+        hidden_states = hidden_states_1 + hidden_states_2
+    hidden_states = post_feedforward_layernorm(hidden_states)     # done in Gemma4TransformerLayer
+    hidden_states = residual + hidden_states                      # done by Megatron mlp_bda
+
+Megatron's standard ``TransformerLayer`` always applies
+``self.pre_mlp_layernorm`` and feeds *that* to ``self.mlp``. To preserve the
+HF requirement that **dense, experts and router each see** the raw residual
+(through their own independent norms — or, for the router, no norm at all
+beyond its own internal RMSNorm), the Gemma4 layer spec sets
+``pre_mlp_layernorm = IdentityOp`` and this block holds three pre-merge norms:
+
+  - ``pre_feedforward_layernorm``    (in front of dense, with weight)
+  - ``post_feedforward_layernorm_1`` (after dense MLP, with weight)
+  - ``post_feedforward_layernorm_2`` (after experts, with weight)
+
+The fourth pre-merge norm (``pre_feedforward_layernorm_2``, the experts'
+input norm) lives on :class:`~aiak_training_llm.models.gemma4_vl.gemma4_moe_layer.Gemma4MoELayer`
+because :class:`MoELayer.forward` feeds the same tensor to both ``self.router``
+(which must see RAW residual) and ``self.token_dispatcher`` (which must see
+the normed residual). Splitting the two requires an override that lives next
+to that forward, not in this composer.
+
+The fifth norm ``post_feedforward_layernorm`` (post-merge, pre-residual-add)
+lives on :class:`~aiak_training_llm.models.gemma4_vl.gemma4_transformer_layer.Gemma4TransformerLayer`
+because Megatron's ``mlp_bda`` is what performs the residual add and there is
+no slot between ``self.mlp(...)`` and ``mlp_bda(...)`` that we can populate
+without subclassing the layer.
+
+Forward contract matches Megatron's ``MLP`` / ``MoELayer``::
+
+    forward(hidden_states) -> (output, bias_or_None)
+
+Bias handling: Gemma4 sets ``add_bias_linear=False`` everywhere, so both
+branches return ``bias=None``. We assert this and return ``None`` for the
+composed bias to refuse silently-wrong fallbacks.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+@dataclass
+class Gemma4ParallelDenseMoESubmodules:
+    """Submodule spec for :class:`Gemma4ParallelDenseMoE`."""
+
+    dense: Union[ModuleSpec, type] = None
+    moe: Union[ModuleSpec, type] = None
+    pre_feedforward_layernorm: Union[ModuleSpec, type] = None
+    post_feedforward_layernorm_1: Union[ModuleSpec, type] = None
+    post_feedforward_layernorm_2: Union[ModuleSpec, type] = None
+
+
+class Gemma4ParallelDenseMoE(MegatronModule):
+    """Parallel dense + MoE FFN composition for Gemma4.
+
+    Receives the raw post-attention residual (Megatron's ``pre_mlp_layernorm``
+    is forced to ``IdentityOp`` in the Gemma4 layer spec). The dense branch
+    norms locally; the MoE branch passes the raw residual straight through to
+    :class:`Gemma4MoELayer`, which owns the router-vs.-experts norm split.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: Gemma4ParallelDenseMoESubmodules,
+        layer_number: Optional[int] = None,
+    ):
+        super().__init__(config=config)
+        self.config = config
+        self.layer_number = layer_number
+
+        self.pre_feedforward_layernorm = build_module(
+            submodules.pre_feedforward_layernorm,
+            config=config,
+            hidden_size=config.hidden_size,
+            eps=config.layernorm_epsilon,
+        )
+        self.post_feedforward_layernorm_1 = build_module(
+            submodules.post_feedforward_layernorm_1,
+            config=config,
+            hidden_size=config.hidden_size,
+            eps=config.layernorm_epsilon,
+        )
+        self.post_feedforward_layernorm_2 = build_module(
+            submodules.post_feedforward_layernorm_2,
+            config=config,
+            hidden_size=config.hidden_size,
+            eps=config.layernorm_epsilon,
+        )
+
+        self.dense = build_module(submodules.dense, config=config)
+        self.moe = build_module(submodules.moe, config=config)
+        if layer_number is not None and hasattr(self.moe, "set_layer_number"):
+            self.moe.set_layer_number(layer_number)
+
+    def set_layer_number(self, layer_number: int):
+        self.layer_number = layer_number
+        if hasattr(self.moe, "set_layer_number"):
+            self.moe.set_layer_number(layer_number)
+
+    def forward(self, hidden_states: torch.Tensor):
+        residual = hidden_states
+
+        dense_in = self.pre_feedforward_layernorm(residual)
+        dense_out, dense_bias = self.dense(dense_in)
+        dense_out = self.post_feedforward_layernorm_1(dense_out)
+
+        moe_out, moe_bias = self.moe(residual)
+        moe_out = self.post_feedforward_layernorm_2(moe_out)
+
+        assert dense_bias is None, (
+            "Gemma4ParallelDenseMoE: dense branch returned non-None bias; "
+            "Gemma4 expects add_bias_linear=False everywhere."
+        )
+        assert moe_bias is None, (
+            "Gemma4ParallelDenseMoE: MoE branch returned non-None mlp_bias; "
+            "Gemma4 expects add_bias_linear=False everywhere."
+        )
+        return dense_out + moe_out, None
+
+    @property
+    def use_shared_expert(self) -> bool:
+        return getattr(self.moe, "use_shared_expert", False)
+
+    @property
+    def shared_expert_overlap(self) -> bool:
+        return getattr(self.moe, "shared_expert_overlap", False)
+
+    @property
+    def shared_experts(self):
+        return getattr(self.moe, "shared_experts", None)
+
+    @property
+    def router(self):
+        return self.moe.router
+
+    @property
+    def experts(self):
+        return self.moe.experts
+
+    @property
+    def local_experts(self):
+        return getattr(self.moe.experts, "local_experts", self.moe.experts)
+
+    @property
+    def token_dispatcher(self):
+        return self.moe.token_dispatcher
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_model.py b/aiak_training_llm/models/gemma4_vl/gemma4_model.py
new file mode 100644
index 00000000..6caa1c08
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_model.py
@@ -0,0 +1,296 @@
+"""Gemma4 LLM backbone wrapper (Megatron LanguageModule subclass).
+
+Mirrors :class:`aiak_training_llm.models.qwen.qwen_model.QwenModel` and adds three
+Gemma4-specific behaviors to ``forward``:
+
+  1. Embedding scaling: decoder_input *= sqrt(hidden_size). HF reference:
+     ``Gemma4TextScaledWordEmbedding.forward`` multiplies token embeddings by
+     ``embed_scale = config.hidden_size ** 0.5``.
+
+  2. Final logit soft-capping: ``logits = softcap * tanh(logits / softcap)``,
+     gated by ``config.final_logit_softcapping`` (=30.0 for Gemma4-26B-A4B).
+     Disabled when the value is None or <= 0.
+
+  3. Dual RoPE: when ``config.layer_pattern`` is non-empty (Gemma4 hybrid
+     attention), construct two independent rotary embeddings — sliding layers
+     get the stock :class:`RotaryEmbedding` (``head_dim=256``, ``theta=1e4``,
+     full rotation), global layers get :class:`Gemma4ProportionalRotaryEmbedding`
+     (``head_dim=512``, ``theta=1e6``, ``partial_rotary_factor=0.25`` with HF
+     proportional zero-padding). Both are evaluated each step and shipped as
+     a ``dict[str, Tensor]`` keyed by layer-type. The unmodified
+     :class:`TransformerBlock` passes the dict through verbatim;
+     :class:`Gemma4SelfAttention.forward` selects the matching tensor based on
+     its own ``self.layer_number`` and ``config.layer_pattern``.
+"""
+
+import math
+from collections import OrderedDict
+from typing import Literal, Optional
+
+import torch
+from megatron.core import InferenceParams, tensor_parallel
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType, ModelType
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_config import TransformerConfig
+from torch import Tensor
+
+from aiak_training_llm.models.gemma4_vl.gemma4_proportional_rotary import (
+    Gemma4ProportionalRotaryEmbedding,
+)
+from aiak_training_llm.models.qwen.qwen_model import _load_state_dict_hook_ignore_extra_state
+
+
+class Gemma4Model(LanguageModule):
+    """Gemma4 transformer language model backbone."""
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        transformer_layer_spec: ModuleSpec,
+        vocab_size: int,
+        max_sequence_length: int,
+        pre_process: bool = True,
+        post_process: bool = True,
+        fp16_lm_cross_entropy: bool = False,
+        parallel_output: bool = True,
+        share_embeddings_and_output_weights: bool = True,
+        position_embedding_type: Literal["learned_absolute", "rope"] = "rope",
+        rotary_percent: float = 1.0,
+        rotary_base: int = 10000,
+        rotary_base_sliding: Optional[int] = None,
+        rope_scaling: bool = False,
+        rope_scaling_factor: float = 8.0,
+        scatter_embedding_sequence_parallel: bool = True,
+        seq_len_interpolation_factor: Optional[float] = None,
+    ) -> None:
+        super().__init__(config=config)
+
+        if has_config_logger_enabled(config):
+            log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+        self.transformer_layer_spec: ModuleSpec = transformer_layer_spec
+        self.vocab_size = vocab_size
+        self.max_sequence_length = max_sequence_length
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
+        self.parallel_output = parallel_output
+        self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+        self.position_embedding_type = position_embedding_type
+
+        self.model_type = ModelType.encoder_or_decoder
+
+        self.max_position_embeddings = max_sequence_length
+        self.rotary_percent = rotary_percent
+        self.rotary_base = rotary_base
+        self.rotary_scaling = rope_scaling
+
+        # Gemma4 embedding scale factor: applied after embedding lookup.
+        # See HF Gemma4TextScaledWordEmbedding.embed_scale.
+        self._embed_scale = math.sqrt(self.config.hidden_size)
+
+        # Final logit soft-capping; off when <= 0 / None.
+        self._final_logit_softcapping = getattr(
+            self.config, "final_logit_softcapping", None
+        )
+
+        if self.pre_process:
+            self.embedding = LanguageModelEmbedding(
+                config=self.config,
+                vocab_size=self.vocab_size,
+                max_sequence_length=self.max_sequence_length,
+                position_embedding_type=position_embedding_type,
+                scatter_to_sequence_parallel=scatter_embedding_sequence_parallel,
+            )
+
+        if self.position_embedding_type == "rope" and not self.config.multi_latent_attention:
+            sliding_kv_channels = (
+                self.config.per_layer_kv_channels.get("sliding")
+                if self.config.per_layer_kv_channels
+                else None
+            ) or self.config.kv_channels
+            global_kv_channels = (
+                self.config.per_layer_kv_channels.get("global")
+                if self.config.per_layer_kv_channels
+                else None
+            ) or self.config.kv_channels
+
+            sliding_base = rotary_base_sliding if rotary_base_sliding is not None else rotary_base
+            global_base = rotary_base
+
+            partial_global = getattr(self.config, "partial_rotary_factor", 1.0) or 1.0
+
+            self.rotary_pos_emb_sliding = RotaryEmbedding(
+                kv_channels=sliding_kv_channels,
+                rotary_percent=rotary_percent,
+                rotary_interleaved=self.config.rotary_interleaved,
+                seq_len_interpolation_factor=seq_len_interpolation_factor,
+                rotary_base=sliding_base,
+                rope_scaling=rope_scaling,
+                rope_scaling_factor=rope_scaling_factor,
+                use_cpu_initialization=self.config.use_cpu_initialization,
+            )
+
+            if self.config.layer_pattern and "global" in self.config.layer_pattern:
+                self.rotary_pos_emb_global = Gemma4ProportionalRotaryEmbedding(
+                    head_dim=global_kv_channels,
+                    partial_rotary_factor=partial_global,
+                    rotary_base=global_base,
+                    rotary_interleaved=self.config.rotary_interleaved,
+                    seq_len_interpolation_factor=seq_len_interpolation_factor,
+                    use_cpu_initialization=self.config.use_cpu_initialization,
+                )
+            else:
+                self.rotary_pos_emb_global = None
+
+        self.rotary_pos_emb_cache = {}
+
+        self.decoder = TransformerBlock(
+            config=self.config,
+            spec=transformer_layer_spec,
+            pre_process=self.pre_process,
+            post_process=self.post_process,
+        )
+
+        if post_process:
+            if self.config.defer_embedding_wgrad_compute:
+                self.embedding_activation_buffer = []
+                self.grad_output_buffer = []
+            else:
+                self.embedding_activation_buffer = None
+                self.grad_output_buffer = None
+
+            self.output_layer = tensor_parallel.ColumnParallelLinear(
+                config.hidden_size,
+                self.vocab_size,
+                config=config,
+                init_method=config.init_method,
+                bias=False,
+                skip_bias_add=False,
+                gather_output=not self.parallel_output,
+                skip_weight_param_allocation=self.pre_process
+                and self.share_embeddings_and_output_weights,
+                embedding_activation_buffer=self.embedding_activation_buffer,
+                grad_output_buffer=self.grad_output_buffer,
+            )
+
+        if self.pre_process or self.post_process:
+            self.setup_embeddings_and_output_layer()
+
+        if has_config_logger_enabled(self.config):
+            log_config_to_disk(
+                self.config, self.state_dict(), prefix=f"{type(self).__name__}_init_ckpt"
+            )
+
+        self.register_load_state_dict_post_hook(_load_state_dict_hook_ignore_extra_state)
+
+    def set_input_tensor(self, input_tensor: Tensor) -> None:
+        """See :meth:`LanguageModule.set_input_tensor`."""
+        if not isinstance(input_tensor, list):
+            input_tensor = [input_tensor]
+        assert len(input_tensor) == 1, "input_tensor should only be length 1"
+        self.decoder.set_input_tensor(input_tensor[0])
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        position_ids: Tensor,
+        attention_mask: Tensor,
+        attn_mask_type: Optional[AttnMaskType] = None,
+        decoder_input: Tensor = None,
+        labels: Tensor = None,
+        rotary_pos_emb: Tensor = None,
+        inference_params: InferenceParams = None,
+        packed_seq_params: PackedSeqParams = None,
+        extra_block_kwargs: Optional[dict] = None,
+        runtime_gather_output: Optional[bool] = None,
+    ) -> Tensor:
+        if decoder_input is None:
+            if self.pre_process:
+                decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+                # Gemma4 embedding scaling: in-dtype multiply preserves the
+                # downstream casting policy (do NOT promote to fp32 here).
+                decoder_input = decoder_input * self._embed_scale
+            else:
+                decoder_input = None
+
+        if (
+            rotary_pos_emb is None
+            and self.position_embedding_type == "rope"
+            and not self.config.multi_latent_attention
+        ):
+            rotary_seq_len = self.rotary_pos_emb_sliding.get_rotary_seq_len(
+                inference_params, self.decoder, decoder_input, self.config, packed_seq_params
+            )
+            packed = (
+                packed_seq_params is not None and packed_seq_params.qkv_format == "thd"
+            )
+            sliding_emb = self.rotary_pos_emb_sliding(rotary_seq_len, packed_seq=packed)
+            if self.rotary_pos_emb_global is not None:
+                global_emb = self.rotary_pos_emb_global(rotary_seq_len, packed_seq=packed)
+                rotary_pos_emb = {"sliding": sliding_emb, "global": global_emb}
+            else:
+                rotary_pos_emb = sliding_emb
+
+        hidden_states = self.decoder(
+            hidden_states=decoder_input,
+            attention_mask=attention_mask,
+            attn_mask_type=attn_mask_type,
+            inference_params=inference_params,
+            rotary_pos_emb=rotary_pos_emb,
+            packed_seq_params=packed_seq_params,
+            **(extra_block_kwargs or {}),
+        )
+
+        if not self.post_process:
+            return hidden_states
+
+        output_weight = None
+        if self.share_embeddings_and_output_weights:
+            output_weight = self.shared_embedding_or_output_weight()
+
+        logits, _ = self.output_layer(
+            hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output
+        )
+
+        # Gemma4 final-logit soft-capping: logits = c * tanh(logits / c).
+        # Bounds extreme logits before cross-entropy. Off when c is None / <= 0.
+        if self._final_logit_softcapping and self._final_logit_softcapping > 0:
+            cap = float(self._final_logit_softcapping)
+            logits = cap * torch.tanh(logits / cap)
+
+        if has_config_logger_enabled(self.config):
+            payload = OrderedDict(
+                {
+                    "input_ids": input_ids,
+                    "position_ids": position_ids,
+                    "attention_mask": attention_mask,
+                    "decoder_input": decoder_input,
+                    "logits": logits,
+                }
+            )
+            log_config_to_disk(self.config, payload, prefix="input_and_logits")
+
+        if labels is None:
+            return logits.transpose(0, 1).contiguous()
+
+        loss = self.compute_language_model_loss(labels, logits)
+        return loss
+
+    def sharded_state_dict(
+        self, prefix: str = "", sharded_offsets: tuple = (), metadata: Optional[dict] = None
+    ) -> ShardedStateDict:
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+        output_layer_extra_state_key = f"{prefix}output_layer._extra_state"
+        output_extra_state = sharded_state_dict.pop(output_layer_extra_state_key, None)
+        assert not (
+            output_extra_state and output_extra_state.data
+        ), f"Expected output layer extra state to be empty, got: {output_extra_state}"
+        return sharded_state_dict
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_moe_layer.py b/aiak_training_llm/models/gemma4_vl/gemma4_moe_layer.py
new file mode 100644
index 00000000..2c1afa74
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_moe_layer.py
@@ -0,0 +1,121 @@
+"""Gemma4-VL MoE layer (subclass of Megatron's :class:`MoELayer`).
+
+Two reasons we subclass :class:`MoELayer`:
+
+1. **Router rebind.** Megatron's :class:`MoELayer.__init__` hard-codes
+   ``self.router = TopKRouter(config)`` and ignores anything you might pass
+   through :class:`MoESubmodules`. Gemma4 needs its own router with a
+   parameter-free RMSNorm + ``scale[H]`` + ``per_expert_scale[E]``
+   (see ``gemma4_router.py``), so we rebind ``self.router`` after the parent
+   has finished wiring the experts / dispatcher / shared experts.
+
+2. **Router input contract.** HF Gemma4 routes on the *raw post-attention
+   residual* but feeds the experts a separately-normed view of that same
+   residual (see ``modeling_gemma4.py`` lines 1394-1405)::
+
+       hidden_states_flat = residual.reshape(-1, H)              # raw
+       _, top_w, top_i  = self.router(hidden_states_flat)        # router on raw
+       hidden_states_2  = pre_feedforward_layernorm_2(hidden_states_flat)
+       hidden_states_2  = self.experts(hidden_states_2, top_i, top_w)
+
+   Megatron's :class:`MoELayer.forward` instead feeds the same tensor to both
+   ``self.router`` and ``self.token_dispatcher.token_permutation`` — making it
+   impossible to reproduce HF's split-input behavior from a wrapper that lives
+   *outside* :class:`MoELayer`. We therefore own ``pre_feedforward_layernorm_2``
+   here and override :meth:`forward` to thread raw vs. normed correctly.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+from megatron.core import tensor_parallel
+from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.training import utils
+
+from aiak_training_llm.models.gemma4_vl.gemma4_router import Gemma4Router
+
+
+@dataclass
+class Gemma4MoESubmodules(MoESubmodules):
+    """Extends :class:`MoESubmodules` with Gemma4's expert pre-norm slot."""
+
+    pre_feedforward_layernorm_2: Union[ModuleSpec, type] = None
+
+
+class Gemma4MoELayer(MoELayer):
+    """:class:`MoELayer` with Gemma4 router + split router/expert input norm."""
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: Gemma4MoESubmodules = None,
+        layer_number: int | None = None,
+    ) -> None:
+        super().__init__(config=config, submodules=submodules, layer_number=layer_number)
+        self.router = Gemma4Router(config=self.config)
+        if layer_number is not None:
+            self.router.set_layer_number(layer_number)
+
+        if submodules is None or submodules.pre_feedforward_layernorm_2 is None:
+            raise ValueError(
+                "Gemma4MoELayer requires submodules.pre_feedforward_layernorm_2 to be set; "
+                "the layer spec must wire this norm in."
+            )
+        self.pre_feedforward_layernorm_2 = build_module(
+            submodules.pre_feedforward_layernorm_2,
+            config=config,
+            hidden_size=config.hidden_size,
+            eps=config.layernorm_epsilon,
+        )
+
+    def forward(self, hidden_states: torch.Tensor):
+        if (
+            self.training
+            and self.config.tensor_model_parallel_size > 1
+            and not self.config.sequence_parallel
+        ):
+            raise ValueError(
+                "During training, performance may degrade if MoE and tensor parallelism "
+                "are enabled without also enabling sequence parallelism."
+            )
+
+        def custom_forward(hidden_states):
+            probs, routing_map = self.router(hidden_states)
+
+            expert_input = self.pre_feedforward_layernorm_2(hidden_states)
+            (dispatched_input, tokens_per_expert) = self.token_dispatcher.token_permutation(
+                expert_input, probs, routing_map
+            )
+
+            if self.config.enable_mem_monitor:
+                utils.monitor_max_memory_usage()
+                utils.monitor_max_dispatcher_tokens(tokens_per_expert)
+                if self.config.mem_monitor_log is not None:
+                    utils.write_monitor_data_to_file(
+                        self.config.mem_monitor_log, self.config.print_mem_monitor_interval
+                    )
+
+            expert_output, mlp_bias = self.experts(dispatched_input, tokens_per_expert)
+            output, mlp_bias = self.token_dispatcher.token_unpermutation(expert_output, mlp_bias)
+
+            if self.use_shared_expert and not self.shared_expert_overlap:
+                raise RuntimeError(
+                    "Gemma4MoELayer was built with shared experts enabled; HF Gemma4 has none. "
+                    "Check that moe_shared_expert_intermediate_size is unset (None)."
+                )
+            return output, mlp_bias
+
+        if self.moe_layer_recompute:
+            output, mlp_bias = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
+        else:
+            output, mlp_bias = custom_forward(hidden_states)
+
+        return output, mlp_bias
+
+
+__all__ = ["Gemma4MoELayer", "Gemma4MoESubmodules"]
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_proportional_rotary.py b/aiak_training_llm/models/gemma4_vl/gemma4_proportional_rotary.py
new file mode 100644
index 00000000..1a1e2143
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_proportional_rotary.py
@@ -0,0 +1,109 @@
+"""Gemma4 proportional RoPE — global-layer rotary with HF zero-padded inv_freq.
+
+HF Gemma4 global (full_attention) layers use ``rope_type='proportional'`` with
+``partial_rotary_factor=0.25`` and ``rope_theta=1e6``. This is **not** the same
+as Megatron's stock ``rotary_percent`` — the two differ in two ways that
+together push logits off by O(1e-3) on a real micro-batch:
+
+  1. Denominator. HF proportional uses the **full head_dim** as the denominator
+     when computing inv_freq exponents:
+
+         inv_freq_rotated[k] = 1 / base^( 2*k / head_dim )       for k in [0, rope_angles)
+
+     where ``rope_angles = int(rope_proportion * head_dim // 2)``. Megatron's
+     ``RotaryEmbedding`` instead uses ``dim = int(kv_channels * rotary_percent)``
+     as the denominator: ``inv_freq[k] = 1 / base^(2*k / dim)``. With
+     ``head_dim=512`` and ``partial_rotary_factor=0.25``, HF divides by 512 while
+     stock Megatron would divide by 128 — a 4× larger exponent — producing a
+     completely different frequency spectrum.
+
+  2. Padding vs. truncation. HF returns an inv_freq of length ``head_dim // 2``
+     where the first ``rope_angles`` entries are non-zero and the remainder is
+     zero-padded; the cat-(freqs, freqs) in the forward then produces a
+     ``head_dim``-long cos/sin pair whose tail is identity (cos=1, sin=0).
+     Megatron's stock partial-RoPE truncates instead: it returns inv_freq of
+     length ``dim // 2`` and the cos/sin pair is ``dim``-long, with the
+     attention head splitting Q/K into rotated + un-rotated sub-vectors. The
+     attention math is **incompatible** between the two without code changes
+     elsewhere; we match HF by zero-padding inside ``__init__``.
+
+This subclass therefore reimplements ``__init__``'s inv_freq computation:
+construct a length-``head_dim // 2`` buffer whose first ``rope_angles`` entries
+are HF-equivalent and the rest are zero. The parent ``forward`` is left intact
+because (a) ``cat((freqs, freqs), dim=-1)`` correctly emits a length-``head_dim``
+cos/sin pair, and (b) zeros in inv_freq produce cos=1, sin=0 — the desired
+identity rotation on the un-rotated tail.
+
+For sliding (``rope_type='default'`` with no ``partial_rotary_factor``) we
+**don't need this class** — the stock ``RotaryEmbedding(kv_channels=256,
+rotary_percent=1.0, rotary_base=1e4)`` is bit-for-bit equivalent to HF.
+
+References:
+  - HF ``_compute_proportional_rope_parameters`` in
+    ``transformers/src/transformers/modeling_rope_utils.py:187-254`` (commit c9de109).
+  - HF ``Gemma4TextRotaryEmbedding`` in
+    ``transformers/src/transformers/models/gemma4/modeling_gemma4.py:1046-1130``.
+"""
+
+from __future__ import annotations
+
+import torch
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+
+
+class Gemma4ProportionalRotaryEmbedding(RotaryEmbedding):
+    """Gemma4 global-layer RoPE with HF-equivalent zero-padded inv_freq.
+
+    Args:
+        head_dim: Full attention head dimension (e.g. 512 for Gemma4 global).
+            This is the **denominator** in the HF inv_freq exponent and the
+            length of the emitted cos/sin (= 2 * len(inv_freq)).
+        partial_rotary_factor: Fraction of the head_dim to actually rotate
+            (e.g. 0.25 for Gemma4 global → rotate first 64 angle pairs of 256).
+        rotary_base: RoPE base / theta (e.g. 1_000_000 for Gemma4 global).
+        rotary_interleaved: Forwarded to parent. Gemma4 uses False (cat layout).
+        seq_len_interpolation_factor: Forwarded to parent.
+        use_cpu_initialization: Forwarded to parent.
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        partial_rotary_factor: float,
+        rotary_base: float,
+        rotary_interleaved: bool = False,
+        seq_len_interpolation_factor: float | None = None,
+        use_cpu_initialization: bool = False,
+    ) -> None:
+        # Parent __init__ computes self.inv_freq using kv_channels + rotary_percent.
+        # We pass kv_channels=head_dim, rotary_percent=1.0 so the parent emits an
+        # inv_freq of length head_dim//2 with the **correct denominator** (head_dim).
+        # Then we zero out the tail to match HF proportional padding.
+        super().__init__(
+            kv_channels=head_dim,
+            rotary_percent=1.0,
+            rotary_interleaved=rotary_interleaved,
+            seq_len_interpolation_factor=seq_len_interpolation_factor,
+            rotary_base=rotary_base,
+            rope_scaling=False,
+            rope_scaling_factor=1.0,
+            use_cpu_initialization=use_cpu_initialization,
+        )
+
+        rope_angles = int(partial_rotary_factor * head_dim // 2)
+        if rope_angles < 0 or rope_angles > head_dim // 2:
+            raise ValueError(
+                f"Gemma4ProportionalRotaryEmbedding: rope_angles={rope_angles} out of "
+                f"range [0, {head_dim // 2}] for head_dim={head_dim}, "
+                f"partial_rotary_factor={partial_rotary_factor}."
+            )
+
+        # Zero-pad the inv_freq tail so cos[..., rope_angles:]=1 and sin[..., rope_angles:]=0
+        # in the parent forward. Use in-place fill on the registered buffer / parameter view.
+        # self.inv_freq is a Tensor (set by parent as a plain attribute, not a buffer).
+        if rope_angles < head_dim // 2:
+            with torch.no_grad():
+                self.inv_freq[rope_angles:].zero_()
+
+
+__all__ = ["Gemma4ProportionalRotaryEmbedding"]
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_router.py b/aiak_training_llm/models/gemma4_vl/gemma4_router.py
new file mode 100644
index 00000000..0572f766
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_router.py
@@ -0,0 +1,189 @@
+"""Gemma4-VL router (real subclass of Megatron's :class:`Router`).
+
+Faithfully reproduces HF ``Gemma4TextRouter`` semantics inside Megatron's MoE
+plumbing. The HF router is a self-contained module — no algebraic fusion with
+any neighbouring norm is possible (see plan §0 point 2 / §1.4 / Change Log v4
+#1) — so we reimplement it verbatim here and override the ``Router`` interface
+to return the ``(scores, routing_map)`` tuple Megatron's token dispatcher
+expects.
+
+HF reference (``transformers/models/gemma4/modeling_gemma4.py``,
+transformers==5.7.0)::
+
+    class Gemma4TextRouter(nn.Module):
+        def __init__(self, config):
+            self.norm = Gemma4RMSNorm(hidden_size, eps, with_scale=False)
+            self.proj = nn.Linear(hidden_size, num_experts, bias=False)
+            self.scale = nn.Parameter(torch.ones(hidden_size))
+            self.per_expert_scale = nn.Parameter(torch.ones(num_experts))
+
+        def forward(self, hidden_states):
+            h = self.norm(hidden_states)
+            h = h * self.scale * (hidden_size ** -0.5)
+            scores = self.proj(h)
+            probs = F.softmax(scores, dim=-1)
+            top_w, top_i = torch.topk(probs, k=top_k_experts, dim=-1)
+            top_w = top_w / top_w.sum(dim=-1, keepdim=True)
+            top_w = top_w * self.per_expert_scale[top_i]
+            return probs, top_w, top_i
+
+Notes on the Megatron mapping:
+
+* Megatron's :class:`Router` base class allocates ``self.weight`` of shape
+  ``[num_experts, hidden_size]``. We reuse that buffer as HF's
+  ``proj.weight`` (same shape, same role) so the conversion script can copy
+  it verbatim and we avoid duplicating a Parameter.
+* ``self.scale`` and ``self.per_expert_scale`` are extra Parameters added on
+  top of the base class.
+* ``self.norm`` is the HF ``Gemma4RMSNorm(with_scale=False)`` — a parameter-
+  free RMSNorm with the eps **inside** the sqrt (``1/sqrt(mean(x^2)+eps)``),
+  which differs from Megatron's default ``MegatronRMSNorm``
+  (``1/sqrt(mean(x^2))+eps``). We implement it inline as a private nn.Module
+  so the numerics match HF byte-for-byte and there is no learnable weight to
+  convert.
+* The :meth:`forward` override returns ``(scores, routing_map)`` where:
+
+  - ``scores`` is a dense ``[num_tokens, num_experts]`` tensor whose
+    non-zero entries equal HF's ``top_w`` (so the dispatcher's
+    ``out = scores * dispatched_input`` reproduces HF's gating exactly);
+  - ``routing_map`` is a dense ``[num_tokens, num_experts]`` bool mask
+    with ``top_k_experts`` True per row.
+
+  The base class's ``gating()`` / ``routing()`` plumbing is bypassed entirely
+  because (a) we need a custom pre-projection norm + scalar scale, and
+  (b) HF's per-expert-scale post-multiplication has no equivalent in the
+  ``moe_router_topk_scaling_factor`` scalar slot.
+
+Conversion (P2): direct per-Parameter copy ``HF -> MG``::
+
+    proj.weight       -> router.weight              [num_experts, hidden]
+    scale             -> router.scale               [hidden]
+    per_expert_scale  -> router.per_expert_scale    [num_experts]
+"""
+
+from __future__ import annotations
+
+import torch
+from megatron.core.transformer.moe.router import Router
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class _Gemma4ParameterFreeRMSNorm(torch.nn.Module):
+    """HF ``Gemma4RMSNorm(with_scale=False)`` — parameter-free, eps-in-sqrt.
+
+    Computes ``x * pow(mean(x^2) + eps, -0.5)`` in float32 and casts back to
+    the input dtype. Matches HF byte-for-byte; intentionally different from
+    Megatron's ``MegatronRMSNorm`` (which puts eps outside the sqrt).
+    """
+
+    def __init__(self, eps: float):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        in_dtype = hidden_states.dtype
+        x = hidden_states.float()
+        mean_sq = x.pow(2).mean(-1, keepdim=True) + self.eps
+        out = x * torch.pow(mean_sq, -0.5)
+        return out.to(in_dtype)
+
+
+class Gemma4Router(Router):
+    """Self-contained router matching HF ``Gemma4TextRouter`` semantics.
+
+    Inherits Megatron's :class:`Router` so token dispatchers can consume the
+    standard ``(scores, routing_map)`` tuple, but completely overrides the
+    base ``forward`` (and does not use the base ``gating`` / ``routing``
+    helpers).
+    """
+
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__(config=config)
+        # Base class already created self.weight: [num_experts, hidden_size]
+        # — we reuse it as HF's proj.weight. Sanity-check the shape.
+        assert self.weight.shape == (config.num_moe_experts, config.hidden_size), (
+            f"unexpected router weight shape {tuple(self.weight.shape)}; "
+            f"expected ({config.num_moe_experts}, {config.hidden_size})"
+        )
+
+        self.hidden_size = config.hidden_size
+        self.scalar_root_size = self.hidden_size ** -0.5
+        self.topk = config.moe_router_topk
+
+        # HF: Gemma4RMSNorm(hidden_size, eps=rms_norm_eps, with_scale=False)
+        # `layernorm_epsilon` is Megatron's RMSNorm eps field name.
+        self.norm = _Gemma4ParameterFreeRMSNorm(eps=config.layernorm_epsilon)
+
+        # HF: scale = nn.Parameter(torch.ones(hidden_size))
+        self.scale = torch.nn.Parameter(
+            torch.ones(self.hidden_size, dtype=config.params_dtype)
+        )
+        setattr(self.scale, "sequence_parallel", config.sequence_parallel)
+
+        # HF: per_expert_scale = nn.Parameter(torch.ones(num_experts))
+        self.per_expert_scale = torch.nn.Parameter(
+            torch.ones(self.num_experts, dtype=config.params_dtype)
+        )
+        setattr(self.per_expert_scale, "sequence_parallel", config.sequence_parallel)
+
+    # The abstract ``routing`` is required by the base class but unused
+    # (forward is fully overridden). Implement as a no-op assert so accidental
+    # callers fail loudly instead of silently misbehaving.
+    def routing(self, logits: torch.Tensor, ori_dtype=None):  # type: ignore[override]
+        raise NotImplementedError(
+            "Gemma4Router overrides forward() entirely; routing() is unused."
+        )
+
+    def forward(self, input: torch.Tensor):  # noqa: A002 - mirror base name
+        """Compute HF-equivalent gating and emit Megatron-shaped outputs.
+
+        Args:
+            input: ``[seq_len, batch, hidden_size]`` activations from the
+                post-attention residual (raw — *not* pre-FFN-normed; see plan
+                §1.4 "self-contained router" bullet).
+
+        Returns:
+            (scores, routing_map):
+                - ``scores: [num_tokens, num_experts]`` dense, non-zero only
+                  at top-k positions, equal to HF's ``top_w`` after the
+                  ``per_expert_scale`` multiplication.
+                - ``routing_map: [num_tokens, num_experts]`` bool, True at
+                  top-k positions.
+        """
+        # Move weights to GPU on first call (mirror base class behaviour).
+        if self.weight.device.type == "cpu":
+            self.weight.data = self.weight.data.to(device=torch.cuda.current_device())
+
+        # HF: norm -> *scale -> *scalar_root_size -> linear(proj.weight)
+        h = self.norm(input)
+        h = h * self.scale * self.scalar_root_size
+        # self.weight: [E, H], h: [..., H] -> logits: [..., E]
+        logits = torch.nn.functional.linear(h, self.weight)
+
+        # Flatten leading dims for the dispatcher contract.
+        logits = logits.view(-1, self.num_experts)
+
+        # HF: softmax then topk (post-softmax topk).
+        probs = torch.nn.functional.softmax(logits, dim=-1, dtype=torch.float32)
+        top_w, top_i = torch.topk(probs, k=self.topk, dim=-1)
+        # Re-normalize so top-k weights sum to 1 per token.
+        top_w = top_w / top_w.sum(dim=-1, keepdim=True)
+        # Apply per-expert scale (cast index gather to float32 to match probs).
+        top_w = top_w * self.per_expert_scale.float()[top_i]
+        top_w = top_w.to(input.dtype)
+
+        # Build dense (scores, routing_map) of shape [num_tokens, num_experts].
+        num_tokens = logits.shape[0]
+        scores = torch.zeros(
+            num_tokens, self.num_experts, dtype=top_w.dtype, device=logits.device
+        )
+        scores.scatter_(1, top_i, top_w)
+        routing_map = torch.zeros(
+            num_tokens, self.num_experts, dtype=torch.bool, device=logits.device
+        )
+        routing_map.scatter_(1, top_i, True)
+
+        return scores, routing_map
+
+
+__all__ = ["Gemma4Router"]
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_transformer_layer.py b/aiak_training_llm/models/gemma4_vl/gemma4_transformer_layer.py
new file mode 100644
index 00000000..6953e886
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_transformer_layer.py
@@ -0,0 +1,127 @@
+"""Gemma4-specific transformer layer subclass.
+
+Wraps Megatron's :class:`TransformerLayer` to attach two Gemma4-only mechanisms
+that have no slot in the upstream layer:
+
+1. ``post_feedforward_layernorm`` (with weight) — applied to the merged
+   ``dense + experts`` output **before** the residual add. In HF
+   (``Gemma4TextDecoderLayer.forward``) this lives between the merged sum
+   and the residual add::
+
+       hidden_states = self.post_feedforward_layernorm(hidden_states)
+       hidden_states = residual + hidden_states
+
+   Megatron's stock layer goes ``mlp -> mlp_bda(mlp_out, residual)`` with no
+   hookable position in between, so we override ``forward`` and inject the
+   norm by wrapping ``self.mlp(...)`` instead of patching ``mlp_bda``.
+
+2. ``layer_scalar`` — non-trainable buffer of shape ``(1,)`` initialised to
+   ``1.0`` and multiplied into the layer output at the very end. Loaded
+   from the HF checkpoint (per-layer distinct scalars after pretraining).
+   Gated by ``config.use_layer_scalar`` so the buffer is not registered when
+   the flag is off.
+
+Both mechanisms are no-ops when their gates are off, which is required so
+this subclass can be used by any Gemma4-VL spec without hard-coding Gemma4
+into upstream Megatron.
+"""
+
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_layer import (
+    TransformerLayer,
+    TransformerLayerSubmodules,
+)
+
+
+@dataclass
+class Gemma4TransformerLayerSubmodules(TransformerLayerSubmodules):
+    """Adds the ``post_feedforward_layernorm`` slot used by Gemma4.
+
+    Default ``None`` means "no post-FFN norm"; ``Gemma4TransformerLayer`` only
+    wraps ``self.mlp`` when this slot is populated, so other consumers of
+    :class:`TransformerLayerSubmodules` are unaffected.
+    """
+
+    post_feedforward_layernorm: Union[ModuleSpec, type, None] = None
+
+
+class Gemma4TransformerLayer(TransformerLayer):
+    """``TransformerLayer`` plus Gemma4's post-merge norm and ``layer_scalar``.
+
+    The base ``TransformerLayer.forward`` calls ``self.mlp(...)`` then
+    ``mlp_bda(mlp_out, residual)`` (see ``transformer_layer.py:550-562``).
+    To inject ``post_feedforward_layernorm`` between them without forking the
+    full forward, we monkey-patch ``self.mlp`` at construction time so that
+    its ``__call__`` returns ``(post_ffn_norm(mlp_out), bias)``. The base
+    forward then sees the post-normed tensor and feeds it to ``mlp_bda``,
+    which performs the residual add — exactly the HF order.
+
+    ``layer_scalar`` is post-multiplied into the layer's final output and
+    handles both the single-tensor and ``(output, context)`` return shapes
+    of the base forward.
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules,
+        layer_number: int = 1,
+        hidden_dropout: float | None = None,
+    ):
+        super().__init__(
+            config=config,
+            submodules=submodules,
+            layer_number=layer_number,
+            hidden_dropout=hidden_dropout,
+        )
+
+        post_ffn_spec = getattr(submodules, "post_feedforward_layernorm", None)
+        if post_ffn_spec is not None:
+            self.post_feedforward_layernorm = build_module(
+                post_ffn_spec,
+                config=config,
+                hidden_size=config.hidden_size,
+                eps=config.layernorm_epsilon,
+            )
+            self._wrap_mlp_with_post_ffn_norm()
+
+        if getattr(self.config, "use_layer_scalar", False):
+            self.register_buffer("layer_scalar", torch.ones(1))
+
+    def _wrap_mlp_with_post_ffn_norm(self):
+        original_mlp = self.mlp
+        post_ffn = self.post_feedforward_layernorm
+
+        class _MlpWithPostFfnNorm(torch.nn.Module):
+            def __init__(inner_self):
+                super().__init__()
+                inner_self._inner_mlp = original_mlp
+                inner_self._post_ffn = post_ffn
+
+            def __getattr__(inner_self, name):
+                if name in ("_inner_mlp", "_post_ffn"):
+                    return super().__getattr__(name)
+                try:
+                    return super().__getattr__(name)
+                except AttributeError:
+                    return getattr(inner_self._inner_mlp, name)
+
+            def forward(inner_self, hidden_states, *args, **kwargs):
+                output, bias = inner_self._inner_mlp(hidden_states, *args, **kwargs)
+                output = inner_self._post_ffn(output)
+                return output, bias
+
+        self.mlp = _MlpWithPostFfnNorm()
+
+    def forward(self, *args, **kwargs):
+        out = super().forward(*args, **kwargs)
+        if not hasattr(self, "layer_scalar"):
+            return out
+        if isinstance(out, tuple):
+            return (out[0] * self.layer_scalar, *out[1:])
+        return out * self.layer_scalar
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_attention.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_attention.py
new file mode 100644
index 00000000..32044911
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_attention.py
@@ -0,0 +1,150 @@
+"""Gemma4 vision multi-head attention.
+
+Verbatim port of HF ``Gemma4VisionAttention`` and ``eager_attention_forward``
+(modeling_gemma4 lines 779-936).
+
+Three non-obvious design choices that MUST be preserved bit-for-bit with HF:
+
+1. ``self.scaling = 1.0`` (NOT ``head_dim ** -0.5``). The standard
+   1/sqrt(d) scale is *omitted* because ``q_norm`` and ``k_norm`` already
+   normalize Q and K to unit-RMS. Adding the scale here would double-scale
+   and break HF parity.
+
+2. ``v_norm`` is a parameterless RMSNorm (``with_scale=False``) applied to
+   value states *after* the V projection but *before* head transpose. This
+   is unusual (most attention impls don't normalize V); do not remove it.
+
+3. RoPE is applied via ``apply_multidimensional_rope`` to Q and K *before*
+   the head transpose, with ``unsqueeze_dim=2`` to match the
+   ``[B, N, H, D]`` layout. After RoPE, Q/K/V are transposed to
+   ``[B, H, N, D]`` for attention.
+
+The attention math is fp32 softmax (upcast inside ``eager_attention_forward``)
+and uses standard GQA via ``repeat_kv``. We deliberately do NOT route this
+through Megatron's fused TE attention because (a) vision sequences are short
+(<=2240 patches) so the kernel speedup is marginal, and (b) TE's softmax
+fp16/bf16 path diverges from HF's fp32 path at the 4th decimal which would
+fail downstream consistency checks against the HF reference.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .gemma4_vision_common import Gemma4ClippableLinear, Gemma4RMSNorm
+from .gemma4_vision_rotary import apply_multidimensional_rope
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if scaling is None:
+        scaling = module.head_dim**-0.5
+
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, attn_weights
+
+
+class Gemma4VisionAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        rms_norm_eps: float = 1e-6,
+        attention_dropout: float = 0.0,
+        use_clipped_linears: bool = False,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.num_key_value_groups = num_attention_heads // num_key_value_heads
+        self.scaling = 1.0
+        self.attention_dropout = attention_dropout
+        self.is_causal = False
+
+        self.q_proj = Gemma4ClippableLinear(
+            hidden_size, num_attention_heads * head_dim, use_clipped_linears=use_clipped_linears
+        )
+        self.k_proj = Gemma4ClippableLinear(
+            hidden_size, num_key_value_heads * head_dim, use_clipped_linears=use_clipped_linears
+        )
+        self.v_proj = Gemma4ClippableLinear(
+            hidden_size, num_key_value_heads * head_dim, use_clipped_linears=use_clipped_linears
+        )
+        self.o_proj = Gemma4ClippableLinear(
+            num_attention_heads * head_dim, hidden_size, use_clipped_linears=use_clipped_linears
+        )
+
+        self.q_norm = Gemma4RMSNorm(dim=head_dim, eps=rms_norm_eps)
+        self.k_norm = Gemma4RMSNorm(dim=head_dim, eps=rms_norm_eps)
+        self.v_norm = Gemma4RMSNorm(dim=head_dim, eps=rms_norm_eps, with_scale=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        position_ids: torch.LongTensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        cos, sin = position_embeddings
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape)
+        query_states = self.q_norm(query_states)
+        query_states = apply_multidimensional_rope(query_states, cos, sin, position_ids)
+        query_states = query_states.transpose(1, 2)
+
+        key_states = self.k_proj(hidden_states).view(hidden_shape)
+        key_states = self.k_norm(key_states)
+        key_states = apply_multidimensional_rope(key_states, cos, sin, position_ids)
+        key_states = key_states.transpose(1, 2)
+
+        value_states = self.v_proj(hidden_states).view(hidden_shape)
+        value_states = self.v_norm(value_states)
+        value_states = value_states.transpose(1, 2)
+
+        attn_output, attn_weights = eager_attention_forward(
+            module=self,
+            query=query_states,
+            key=key_states,
+            value=value_states,
+            attention_mask=attention_mask,
+            dropout=self.attention_dropout if self.training else 0.0,
+            scaling=self.scaling,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_common.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_common.py
new file mode 100644
index 00000000..27fc6b3f
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_common.py
@@ -0,0 +1,79 @@
+"""Shared Gemma4 vision/adapter primitives: RMSNorm and ClippableLinear.
+
+Verbatim port of HF ``Gemma4RMSNorm`` (lines 168-187) and ``Gemma4ClippableLinear``
+(lines 139-167) from transformers ``modeling_gemma4``.
+
+Two non-obvious facts that determine HF checkpoint compatibility and MUST NOT
+change:
+
+1. ``Gemma4ClippableLinear`` always wraps the real ``nn.Linear`` under attribute
+   name ``.linear``. HF checkpoints store keys as ``*.<projection>.linear.weight``
+   even when ``use_clipped_linears=False`` and the clamp branch is dead code.
+   Renaming this attribute (e.g. inlining the linear) breaks every HF→Megatron
+   converter mapping.
+
+2. ``Gemma4RMSNorm.with_scale=False`` produces a parameterless module (no
+   ``weight`` attribute at all). Both ``v_norm`` in vision attention and the
+   adapter's RMSNorm rely on this; do NOT add a default ``weight`` for "API
+   uniformity" because state_dict loading from HF will then fail with an
+   unexpected key.
+
+The HF clamp branch (``use_clipped_linears=True``) registers four buffers
+(input/output min/max). Gemma4-26B-A4B-it ships with ``use_clipped_linears=False``
+so we keep the branch for fidelity but it is exercised only by post-training
+quantization configs.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class Gemma4RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, with_scale: bool = True) -> None:
+        super().__init__()
+        self.eps = eps
+        self.with_scale = with_scale
+        if self.with_scale:
+            self.weight = nn.Parameter(torch.ones(dim), requires_grad=True)
+
+    def _norm(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        mean_squared = hidden_states.pow(2).mean(-1, keepdim=True) + self.eps
+        return hidden_states * torch.rsqrt(mean_squared)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        output = self._norm(hidden_states.float())
+        if self.with_scale:
+            output = output * self.weight.float()
+        return output.type_as(hidden_states)
+
+    def extra_repr(self) -> str:
+        weight_shape = tuple(self.weight.shape) if self.with_scale else None
+        return f"weight={weight_shape}, eps={self.eps}, with_scale={self.with_scale}"
+
+
+class Gemma4ClippableLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        use_clipped_linears: bool = False,
+    ) -> None:
+        super().__init__()
+        self.use_clipped_linears = use_clipped_linears
+        self.linear = nn.Linear(in_features, out_features, bias=False)
+
+        if self.use_clipped_linears:
+            self.register_buffer("input_min", torch.tensor(-float("inf")))
+            self.register_buffer("input_max", torch.tensor(float("inf")))
+            self.register_buffer("output_min", torch.tensor(-float("inf")))
+            self.register_buffer("output_max", torch.tensor(float("inf")))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.use_clipped_linears:
+            hidden_states = torch.clamp(hidden_states, self.input_min, self.input_max)
+        hidden_states = self.linear(hidden_states)
+        if self.use_clipped_linears:
+            hidden_states = torch.clamp(hidden_states, self.output_min, self.output_max)
+        return hidden_states
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_layer.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_layer.py
new file mode 100644
index 00000000..b73a8605
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_layer.py
@@ -0,0 +1,101 @@
+"""Gemma4 vision encoder layer.
+
+Verbatim port of HF ``Gemma4VisionEncoderLayer`` (modeling_gemma4 lines 939-980).
+
+Sandwich normalization layout — this is the source of the layer's 4 RMSNorms
+and the order is NOT the conventional pre-norm or post-norm transformer:
+
+    residual = h
+    h = input_layernorm(h)             # pre-attn norm
+    h, _ = self_attn(h, ...)
+    h = post_attention_layernorm(h)    # post-attn norm, BEFORE residual add
+    h = residual + h
+
+    residual = h
+    h = pre_feedforward_layernorm(h)   # pre-mlp norm
+    h = mlp(h)
+    h = post_feedforward_layernorm(h)  # post-mlp norm, BEFORE residual add
+    h = residual + h
+
+The two "post_*_layernorm" calls happen *between* the sublayer output and the
+residual sum, not after. Reordering to standard pre-norm
+(``residual + sublayer(norm(h))``) breaks HF parity because the sandwich
+construction normalizes the sublayer output before the skip-add, which changes
+the variance arithmetic of every subsequent layer. This pattern is shared with
+the LLM hybrid layers (see ``gemma4_transformer_layer.py``); see plan v5 for
+why we do not collapse the four norms into two.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .gemma4_vision_attention import Gemma4VisionAttention
+from .gemma4_vision_common import Gemma4RMSNorm
+from .gemma4_vision_mlp import Gemma4VisionMLP
+
+
+class Gemma4VisionEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        intermediate_size: int,
+        rms_norm_eps: float = 1e-6,
+        attention_dropout: float = 0.0,
+        hidden_activation: str = "gelu_pytorch_tanh",
+        use_clipped_linears: bool = False,
+        layer_idx: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.layer_idx = layer_idx
+
+        self.self_attn = Gemma4VisionAttention(
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            rms_norm_eps=rms_norm_eps,
+            attention_dropout=attention_dropout,
+            use_clipped_linears=use_clipped_linears,
+        )
+        self.mlp = Gemma4VisionMLP(
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            hidden_activation=hidden_activation,
+            use_clipped_linears=use_clipped_linears,
+        )
+        self.input_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.pre_feedforward_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_feedforward_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        position_ids: torch.LongTensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+        )
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.pre_feedforward_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.post_feedforward_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_mlp.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_mlp.py
new file mode 100644
index 00000000..acc6023e
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_mlp.py
@@ -0,0 +1,67 @@
+"""Gemma4 vision gated MLP.
+
+Verbatim port of HF ``Gemma4VisionMLP`` (modeling_gemma4 lines 643-657).
+
+The MLP shape is the SwiGLU-style ``down(act(gate) * up)`` (NOT
+``down(gate * act(up))``); the activation wraps the gate projection only.
+This matches Gemma's LLM MLP convention and is the order that HF checkpoints
+were trained with — swapping the wrap target produces silently wrong outputs
+because ``gate_proj`` and ``up_proj`` share the same in/out shapes.
+
+``hidden_activation`` is ``gelu_pytorch_tanh`` for Gemma4 vision (config
+default), which maps to ``torch.nn.functional.gelu(x, approximate="tanh")``.
+The standard exact GELU diverges from HF outputs at the 5th decimal place,
+which compounds across 27 layers and breaks consistency checks; do NOT
+substitute the exact form.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .gemma4_vision_common import Gemma4ClippableLinear
+
+
+def _gelu_pytorch_tanh(x: torch.Tensor) -> torch.Tensor:
+    return F.gelu(x, approximate="tanh")
+
+
+_ACTIVATIONS = {
+    "gelu_pytorch_tanh": _gelu_pytorch_tanh,
+    "gelu": F.gelu,
+    "silu": F.silu,
+    "relu": F.relu,
+}
+
+
+class Gemma4VisionMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_activation: str = "gelu_pytorch_tanh",
+        use_clipped_linears: bool = False,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = Gemma4ClippableLinear(
+            hidden_size, intermediate_size, use_clipped_linears=use_clipped_linears
+        )
+        self.up_proj = Gemma4ClippableLinear(
+            hidden_size, intermediate_size, use_clipped_linears=use_clipped_linears
+        )
+        self.down_proj = Gemma4ClippableLinear(
+            intermediate_size, hidden_size, use_clipped_linears=use_clipped_linears
+        )
+        if hidden_activation not in _ACTIVATIONS:
+            raise ValueError(
+                f"Unknown hidden_activation={hidden_activation!r}; "
+                f"expected one of {sorted(_ACTIVATIONS)}."
+            )
+        self.act_fn = _ACTIVATIONS[hidden_activation]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_patch_embed.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_patch_embed.py
new file mode 100644
index 00000000..7712596c
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_patch_embed.py
@@ -0,0 +1,65 @@
+"""Gemma4 vision patch embedder.
+
+Verbatim port of HF ``Gemma4VisionPatchEmbedder`` (modeling_gemma4 lines 550-582).
+
+Key non-obvious facts that callers MUST honor:
+- Pre-projection scaling is ``2 * (pixel_values - 0.5)``, NOT a learned standardize
+  pair (``std_bias``/``std_scale``). The standardize buffers live on
+  ``Gemma4VisionTower`` and are applied *after* pooling, not before patch embed.
+- ``position_embedding_table`` has shape ``[2, position_embedding_size, hidden_size]``:
+  the leading ``2`` indexes the spatial axis (x first, y second). Two PE lookups
+  per patch are summed; padding patches (position == -1) get zeroed out.
+- Position embedding uses a one-hot @ matmul rather than ``nn.Embedding`` lookup
+  so it stays within the autograd graph for fused backward when patches are
+  shared across multiple soft-token outputs.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Gemma4VisionPatchEmbedder(nn.Module):
+    def __init__(
+        self,
+        patch_size: int,
+        hidden_size: int,
+        position_embedding_size: int,
+        in_channels: int = 3,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.hidden_size = hidden_size
+        self.position_embedding_size = position_embedding_size
+        self.in_channels = in_channels
+
+        self.input_proj = nn.Linear(in_channels * patch_size * patch_size, hidden_size, bias=False)
+        self.position_embedding_table = nn.Parameter(
+            torch.ones(2, position_embedding_size, hidden_size)
+        )
+
+    def _position_embeddings(
+        self,
+        pixel_position_ids: torch.Tensor,
+        padding_positions: torch.Tensor,
+    ) -> torch.Tensor:
+        clamped_positions = pixel_position_ids.clamp(min=0)
+        one_hot = F.one_hot(clamped_positions, num_classes=self.position_embedding_size)
+        one_hot = one_hot.permute(0, 2, 1, 3).to(self.position_embedding_table)
+        position_embeddings = one_hot @ self.position_embedding_table
+        position_embeddings = position_embeddings.sum(dim=1)
+        position_embeddings = torch.where(padding_positions.unsqueeze(-1), 0.0, position_embeddings)
+        return position_embeddings
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        pixel_position_ids: torch.Tensor,
+        padding_positions: torch.Tensor,
+    ) -> torch.Tensor:
+        pixel_values = 2 * (pixel_values - 0.5)
+        hidden_states = self.input_proj(pixel_values.to(self.input_proj.weight.dtype))
+        position_embeddings = self._position_embeddings(pixel_position_ids, padding_positions)
+        return hidden_states + position_embeddings
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_pooler.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_pooler.py
new file mode 100644
index 00000000..eabc2168
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_pooler.py
@@ -0,0 +1,94 @@
+"""Gemma4 vision spatial pooler.
+
+Verbatim port of HF ``Gemma4VisionPooler`` (modeling_gemma4 lines 584-640).
+
+Three subtleties that determine numerical parity with HF:
+
+1. The pooling kernel size ``k`` is derived dynamically from the *runtime*
+   ratio between input patch count and target soft-token count:
+   ``k = int(sqrt(input_seq_len // output_length))``. There is no
+   ``pooling_kernel_size`` field used at runtime — the config field of that
+   name is only consulted at processor / config-validation time to compute
+   ``num_soft_tokens``. So this module must accept ``output_length`` per call
+   and infer ``k`` from shapes; do NOT hard-code ``k=3`` even though that is
+   the value used by Gemma4 vision under default config.
+
+2. Pooling is done as ``(one_hot(kernel_idx) / k^2)^T @ hidden_states``: this
+   is mathematically equivalent to a 2D avg-pool by spatial-bucket index but
+   stays in the autograd graph as plain matmul, which is what the HF
+   reference uses to keep the operation TPU/XLA-friendly. Replacing it with
+   ``F.avg_pool2d`` is *not* equivalent because patches arrive in raster /
+   packed order, not as a dense [B, C, H, W] grid.
+
+3. The final ``hidden_states *= sqrt(hidden_size)`` scaling is applied
+   ALWAYS, even on the no-pool branch (when input length already equals
+   output length). This compensates for the post-LN variance reduction that
+   the soft-token consumer (Gemma4MultimodalEmbedder) was trained against.
+   Removing or moving this scale silently changes the LM input statistics
+   and breaks downstream consistency.
+
+The output ``mask`` is a *new* padding mask produced from the pooling weight
+matrix: a pooled bucket is considered padding iff every input patch that
+maps to it was already padding. The caller MUST use this returned mask, not
+the input ``padding_positions``.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Gemma4VisionPooler(nn.Module):
+    def __init__(self, hidden_size: int) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.root_hidden_size = hidden_size**0.5
+
+    def _avg_pool_by_positions(
+        self,
+        hidden_states: torch.Tensor,
+        pixel_position_ids: torch.Tensor,
+        length: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_seq_len = hidden_states.shape[1]
+        k = int((input_seq_len // length) ** 0.5)
+        k_squared = k * k
+        if k_squared * length != input_seq_len:
+            raise ValueError(
+                f"Cannot pool {tuple(hidden_states.shape)} to {length}: "
+                f"k={k}^2 times length={length} must equal input_seq_len={input_seq_len}."
+            )
+
+        clamped_positions = pixel_position_ids.clamp(min=0)
+        max_x = clamped_positions[..., 0].max(dim=-1, keepdim=True)[0] + 1
+        kernel_idxs = torch.div(clamped_positions, k, rounding_mode="floor")
+        kernel_idxs = kernel_idxs[..., 0] + (max_x // k) * kernel_idxs[..., 1]
+        weights = F.one_hot(kernel_idxs.long(), length).float() / k_squared
+        output = weights.transpose(1, 2) @ hidden_states.float()
+        mask = torch.logical_not((weights == 0).all(dim=1))
+        return output.to(hidden_states.dtype), mask
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        pixel_position_ids: torch.Tensor,
+        padding_positions: torch.Tensor,
+        output_length: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if output_length > hidden_states.shape[1]:
+            raise ValueError(
+                f"Cannot output more soft tokens (requested {output_length}) than there are "
+                f"patches ({hidden_states.shape[1]}). Change `num_soft_tokens` upstream."
+            )
+
+        hidden_states = hidden_states.masked_fill(padding_positions.unsqueeze(-1), 0.0)
+
+        if hidden_states.shape[1] != output_length:
+            hidden_states, padding_positions = self._avg_pool_by_positions(
+                hidden_states, pixel_position_ids, output_length
+            )
+
+        hidden_states = hidden_states * self.root_hidden_size
+        return hidden_states, padding_positions
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_rotary.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_rotary.py
new file mode 100644
index 00000000..90581eb6
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_rotary.py
@@ -0,0 +1,181 @@
+"""Gemma4 vision 2D rotary position embedding.
+
+Verbatim port of HF ``Gemma4VisionRotaryEmbedding`` and ``apply_multidimensional_rope``
+from transformers ``modeling_gemma4`` (lines 659-866). Kept as a standalone module so
+that the Megatron port stays bit-identical to HF; do NOT route this through the
+shared Megatron RoPE plumbing because Megatron's RoPE assumes a single 1-D position
+sequence whereas Gemma4 vision splits the head dimension across two spatial axes
+(x, y) and applies independent rope to each half.
+
+Math contract:
+- ``head_dim`` is even and divisible by ``2 * ndim`` (ndim=2 here).
+- ``spatial_dim = head_dim // ndim`` (=36 for head_dim=72): each axis gets its own
+  ``inv_freq`` table of length ``spatial_dim // 2`` (=18 sin/cos pairs per axis).
+- ``forward(x, position_ids)`` returns ``(cos, sin)`` each of shape
+  ``[B, N, head_dim]`` where the last dim is the concatenation of the per-axis
+  cos/sin tables (axis-x first ``spatial_dim`` channels, axis-y last ``spatial_dim``).
+- ``apply_multidimensional_rope`` splits ``x``, ``cos``, ``sin`` along the last axis
+  into ``ndim`` equal parts and applies standard rope independently to each, then
+  concatenates back. This matches the HF reference and is the only way to reproduce
+  ``model.embed_vision.embedding_projection`` numerics.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim: int = 1,
+) -> torch.Tensor:
+    """Apply standard 1-D rotary position embedding to ``x``.
+
+    ``cos`` / ``sin`` have shape ``[B, N, D]`` and are unsqueezed at
+    ``unsqueeze_dim`` to broadcast against ``x`` whose layout is either
+    ``[B, H, N, D]`` (unsqueeze_dim=1, post-transpose) or
+    ``[B, N, H, D]`` (unsqueeze_dim=2, pre-transpose, used by Gemma4 vision).
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+def apply_multidimensional_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    position_ids: torch.Tensor,
+    unsqueeze_dim: int = 2,
+) -> torch.Tensor:
+    """Apply multi-dimensional RoPE to ``x``.
+
+    Splits ``x`` along the last (channel) dimension into ``ndim`` parts
+    (where ``ndim = position_ids.shape[-1]``), applies standard rotary embedding
+    independently to each part using its own slice of ``cos`` / ``sin``, then
+    concatenates the parts back. For Gemma4 vision ``ndim=2`` (x and y axes).
+
+    Args:
+        x: Tensor of shape ``[B, N, H, D]`` (queries / keys before transpose).
+        cos: Tensor of shape ``[B, N, D]`` produced by ``Gemma4VisionRotaryEmbedding``.
+        sin: Tensor of shape ``[B, N, D]`` produced by ``Gemma4VisionRotaryEmbedding``.
+        position_ids: Tensor of shape ``[B, N, ndim]`` with per-axis patch coordinates.
+        unsqueeze_dim: Broadcast axis when applying rope (defaults to 2 to match
+            Gemma4VisionAttention call site, where ``x`` is pre-transpose ``[B,N,H,D]``).
+
+    Returns:
+        Tensor of shape ``[B, N, H, D]`` with rope applied.
+    """
+    ndim = position_ids.shape[-1]
+    num_input_channels = x.shape[-1]
+    num_rotated_channels_per_dim = 2 * (num_input_channels // (2 * ndim))
+
+    if num_rotated_channels_per_dim <= 0:
+        raise ValueError(
+            "Invalid configuration: num_rotated_channels_per_dim must be > 0, got "
+            f"{num_rotated_channels_per_dim} (num_input_channels={num_input_channels}, "
+            f"ndim={ndim})"
+        )
+
+    split_sizes = [num_rotated_channels_per_dim] * ndim
+    x_parts = torch.split(x, split_sizes, dim=-1)
+    cos_parts = torch.split(cos, split_sizes, dim=-1)
+    sin_parts = torch.split(sin, split_sizes, dim=-1)
+    y_parts = [
+        apply_rotary_pos_emb(
+            x=x_parts[k],
+            cos=cos_parts[k],
+            sin=sin_parts[k],
+            unsqueeze_dim=unsqueeze_dim,
+        )
+        for k in range(ndim)
+    ]
+    return torch.cat(y_parts, dim=-1)
+
+
+class Gemma4VisionRotaryEmbedding(nn.Module):
+    """Gemma4 vision 2-D rotary embedding.
+
+    Computes per-axis ``cos`` / ``sin`` tables for x and y patch coordinates and
+    concatenates them along the channel dim. The result is a single
+    ``[B, N, head_dim]`` cos/sin pair that ``apply_multidimensional_rope`` knows
+    how to split back into per-axis pieces.
+
+    The frequency table is computed once per axis with
+    ``spatial_dim = head_dim // ndim`` and ``base = rope_theta``; ``inv_freq`` is
+    a non-persistent buffer so it is re-derived from config on load and never
+    stored in checkpoints.
+    """
+
+    inv_freq: torch.Tensor
+
+    def __init__(
+        self,
+        head_dim: int,
+        rope_theta: float = 100.0,
+        ndim: int = 2,
+        device: torch.device | None = None,
+    ) -> None:
+        super().__init__()
+        if head_dim % (2 * ndim) != 0:
+            raise ValueError(
+                f"head_dim={head_dim} must be divisible by 2*ndim={2 * ndim} "
+                "to support multidimensional rotary embedding."
+            )
+        self.head_dim = head_dim
+        self.ndim = ndim
+        self.rope_theta = rope_theta
+        self.attention_scaling = 1.0
+
+        spatial_dim = head_dim // ndim
+        inv_freq = 1.0 / (
+            rope_theta
+            ** (
+                torch.arange(0, spatial_dim, 2, dtype=torch.int64).to(
+                    device=device, dtype=torch.float
+                )
+                / spatial_dim
+            )
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        )
+        device_type = (
+            x.device.type
+            if isinstance(x.device.type, str) and x.device.type != "mps"
+            else "cpu"
+        )
+
+        all_cos: list[torch.Tensor] = []
+        all_sin: list[torch.Tensor] = []
+        for i in range(self.ndim):
+            dim_position_ids = position_ids[:, :, i]
+            dim_position_ids_expanded = dim_position_ids[:, None, :].float()
+
+            with torch.autocast(device_type=device_type, enabled=False):
+                freqs = (
+                    inv_freq_expanded.float() @ dim_position_ids_expanded.float()
+                ).transpose(1, 2)
+                emb = torch.cat((freqs, freqs), dim=-1)
+                cos = emb.cos() * self.attention_scaling
+                sin = emb.sin() * self.attention_scaling
+            all_cos.append(cos)
+            all_sin.append(sin)
+
+        cos = torch.cat(all_cos, dim=-1).to(dtype=x.dtype)
+        sin = torch.cat(all_sin, dim=-1).to(dtype=x.dtype)
+        return cos, sin
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_tower.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_tower.py
new file mode 100644
index 00000000..86b6fe64
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_tower.py
@@ -0,0 +1,164 @@
+"""Gemma4 vision tower.
+
+Verbatim port of HF ``Gemma4VisionEncoder`` + ``Gemma4VisionModel``
+(modeling_gemma4 lines 983-2020) bundled into a single Megatron-side
+``VisionModule`` subclass.
+
+Forward signature matches HF exactly:
+    forward(pixel_values, pixel_position_ids) -> hidden_states
+where:
+    pixel_values: [B, num_patches, in_channels * patch_size**2] (flattened patches)
+    pixel_position_ids: [B, num_patches, 2] with (-1, -1) marking padding patches.
+
+The output is a *flat* tensor ``[total_valid_soft_tokens, hidden_size]`` where
+padding has been stripped via ``hidden_states[pooler_mask]``. This matches HF
+and is what ``Gemma4Adapter`` (= ``Gemma4MultimodalEmbedder``) expects. The
+LLM-side scatter logic in Gemma4VL.forward must use the per-image patch counts
+to slice this flat tensor back into per-sample chunks before injecting into
+the LLM input embeddings.
+
+Two design points worth flagging for the converter (Phase C):
+
+1. ``std_bias`` and ``std_scale`` are registered ONLY when
+   ``config.standardize=True``. Gemma4-26B-A4B-it has ``standardize=true``,
+   so both buffers exist as real ckpt entries. The Megatron checkpoint MUST
+   include them as buffer tensors (non-persistent=False), and the converter
+   must round-trip them. They live on the VisionModel itself, NOT inside
+   ``patch_embedder`` (a common misreading from earlier OV2 ports).
+
+2. The encoder runs a non-causal additive attention mask derived from
+   ``~padding_positions``. We build it here as a ``[B, 1, 1, N]`` broadcast
+   mask (valid=0, pad=finfo.min), avoiding the full ``[B, 1, N, N]`` materialize
+   that ``create_bidirectional_mask`` produces in HF. The arithmetic is
+   identical because the Gemma vision encoder applies the mask on the
+   key dimension only (no cross-sample masking, no per-q masking).
+"""
+
+from __future__ import annotations
+
+import torch
+from megatron.core.models.common.vision_module.vision_module import VisionModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from torch import nn
+
+from .gemma4_vision_layer import Gemma4VisionEncoderLayer
+from .gemma4_vision_patch_embed import Gemma4VisionPatchEmbedder
+from .gemma4_vision_pooler import Gemma4VisionPooler
+from .gemma4_vision_rotary import Gemma4VisionRotaryEmbedding
+
+
+class Gemma4VisionTower(VisionModule):
+    def __init__(
+        self,
+        transformer_config: TransformerConfig,
+        hidden_size: int,
+        num_hidden_layers: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        intermediate_size: int,
+        patch_size: int,
+        in_channels: int,
+        position_embedding_size: int,
+        pooling_kernel_size: int,
+        rope_theta: float = 100.0,
+        rms_norm_eps: float = 1e-6,
+        attention_dropout: float = 0.0,
+        hidden_activation: str = "gelu_pytorch_tanh",
+        use_clipped_linears: bool = False,
+        standardize: bool = True,
+    ) -> None:
+        super().__init__(config=transformer_config)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.patch_size = patch_size
+        self.pooling_kernel_size = pooling_kernel_size
+        self.standardize = standardize
+
+        self.patch_embedder = Gemma4VisionPatchEmbedder(
+            patch_size=patch_size,
+            hidden_size=hidden_size,
+            position_embedding_size=position_embedding_size,
+            in_channels=in_channels,
+        )
+        self.rotary_emb = Gemma4VisionRotaryEmbedding(
+            head_dim=head_dim,
+            rope_theta=rope_theta,
+            ndim=2,
+        )
+        self.layers = nn.ModuleList(
+            [
+                Gemma4VisionEncoderLayer(
+                    hidden_size=hidden_size,
+                    num_attention_heads=num_attention_heads,
+                    num_key_value_heads=num_key_value_heads,
+                    head_dim=head_dim,
+                    intermediate_size=intermediate_size,
+                    rms_norm_eps=rms_norm_eps,
+                    attention_dropout=attention_dropout,
+                    hidden_activation=hidden_activation,
+                    use_clipped_linears=use_clipped_linears,
+                    layer_idx=i,
+                )
+                for i in range(num_hidden_layers)
+            ]
+        )
+        self.pooler = Gemma4VisionPooler(hidden_size=hidden_size)
+
+        if self.standardize:
+            self.register_buffer(
+                "std_bias", torch.empty(hidden_size), persistent=True
+            )
+            self.register_buffer(
+                "std_scale", torch.empty(hidden_size), persistent=True
+            )
+
+    @staticmethod
+    def _build_additive_mask(
+        valid_mask: torch.Tensor, dtype: torch.dtype
+    ) -> torch.Tensor:
+        neg_inf = torch.finfo(dtype).min
+        additive = torch.zeros_like(valid_mask, dtype=dtype)
+        additive = additive.masked_fill(~valid_mask, neg_inf)
+        return additive[:, None, None, :]
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        pixel_position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        output_length = pixel_values.shape[-2] // (
+            self.pooling_kernel_size * self.pooling_kernel_size
+        )
+
+        padding_positions = (pixel_position_ids == -1).all(dim=-1)
+        valid_positions = ~padding_positions
+
+        inputs_embeds = self.patch_embedder(
+            pixel_values, pixel_position_ids, padding_positions
+        )
+        attention_mask = self._build_additive_mask(valid_positions, inputs_embeds.dtype)
+        position_embeddings = self.rotary_emb(inputs_embeds, pixel_position_ids)
+
+        hidden_states = inputs_embeds
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                position_embeddings=position_embeddings,
+                position_ids=pixel_position_ids,
+                attention_mask=attention_mask,
+            )
+
+        hidden_states, pooler_mask = self.pooler(
+            hidden_states=hidden_states,
+            pixel_position_ids=pixel_position_ids,
+            padding_positions=padding_positions,
+            output_length=output_length,
+        )
+
+        hidden_states = hidden_states[pooler_mask]
+
+        if self.standardize:
+            hidden_states = (hidden_states - self.std_bias) * self.std_scale
+
+        return hidden_states
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_attention_mask.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_attention_mask.py
new file mode 100644
index 00000000..71d3c3e5
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_attention_mask.py
@@ -0,0 +1,126 @@
+"""Gemma4-VL multimodal attention mask with bidirectional vision-span overlay.
+
+This module implements HuggingFace ``create_causal_mask_mapping`` semantics for
+Megatron's ``TransformerBlock.forward(attention_mask=...)`` contract:
+
+  1. Base causal mask (lower-triangular True-is-VISIBLE inside the helper, then
+     inverted to Megatron True-is-MASKED at return).
+
+  2. Vision overlay: tokens belonging to the same vision span (contiguous run of
+     ``mm_token_type_ids in {1, 2}``) attend bidirectionally to one another.
+     Mathematically ``mask = causal | (q_group == kv_group) & (q_group >= 0)``,
+     mirroring HF ``token_type_ids_mask_function`` and ``or_masks`` composition
+     in ``transformers.models.gemma4.modeling_gemma4``.
+
+The helper returns a SINGLE BoolTensor ``[B, 1, S, S]`` (True = MASKED), shared
+by every ``TransformerLayer``. Sliding layers further restrict this base mask at
+runtime via TE's ``window_size=(sliding_window-1, 0)`` parameter, which is wired
+per-layer by ``gemma4_attention.py:_resolve_per_layer_overrides``. Because each
+Gemma4 vision span is bounded by the ViT grid (16x16 = 256 patches) and the
+sliding window is 1024, the OR overlay is always strictly contained inside the
+sliding window — TE's hard window cutoff therefore composes losslessly with the
+overlay, recovering exact HF semantics WITHOUT a per-layer mask dict.
+
+A runtime guard (``_assert_vision_spans_within_window``) enforces the
+``max_vision_span <= sliding_window`` invariant and fails fast if a future video
+packing scheme produces a span that would break the equivalence.
+
+External callers: ``Gemma4VLModel.forward`` (P5.5) — invoked when the batch
+contains vision tokens. Pure-text batches skip this helper and pass the stock
+causal mask produced by the data path.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+VISION_TOKEN_TYPE_IDS: tuple[int, ...] = (1, 2)
+
+
+def _compute_vision_group_ids(mm_token_type_ids: torch.Tensor) -> torch.Tensor:
+    """Assign each token a vision-span group id (text tokens get -1).
+
+    Mirrors HF ``get_block_sequence_ids_for_mask`` (modeling_gemma4.py:2078-2087):
+    a new group starts whenever a vision token follows a non-vision token.
+    """
+    is_vision = torch.zeros_like(mm_token_type_ids, dtype=torch.bool)
+    for token_id in VISION_TOKEN_TYPE_IDS:
+        is_vision = is_vision | (mm_token_type_ids == token_id)
+
+    is_prev_vision = torch.roll(is_vision, shifts=1, dims=-1)
+    is_prev_vision[..., 0] = False
+    new_starts = is_vision & ~is_prev_vision
+    vision_group_ids = torch.cumsum(new_starts.to(torch.long), dim=-1) - 1
+    return torch.where(is_vision, vision_group_ids, torch.full_like(vision_group_ids, -1))
+
+
+def _assert_vision_spans_within_window(
+    vision_group_ids: torch.Tensor,
+    sliding_window: int,
+) -> None:
+    """Fail fast if any vision span exceeds the sliding window.
+
+    Gemma4 single-image spans are physically bounded by ViT grid (<=256 patches),
+    so this guard only triggers on hypothetical future video packing. Without it,
+    TE's hard sliding cutoff would silently truncate the OR overlay, producing
+    NaN-equivalent silent divergence from HF rather than a fail-fast error.
+    """
+    valid_mask = vision_group_ids >= 0
+    if not valid_mask.any():
+        return
+    flat_groups = vision_group_ids[valid_mask]
+    span_lengths = torch.bincount(flat_groups)
+    max_span = int(span_lengths.max().item())
+    assert max_span <= sliding_window, (
+        f"Gemma4-VL vision span length {max_span} exceeds sliding_window {sliding_window}; "
+        "the OR-overlay no longer composes losslessly with TE per-layer window cutoff. "
+        "Either disable sliding window for vision batches or shorten the span."
+    )
+
+
+def build_gemma4_mm_attention_mask(
+    input_ids: torch.Tensor,
+    mm_token_type_ids: torch.Tensor,
+    sliding_window: int,
+) -> torch.Tensor:
+    """Build a Megatron-format multimodal attention mask for Gemma4-VL.
+
+    Args:
+        input_ids: ``[B, S]`` token ids (only shape and device are used).
+        mm_token_type_ids: ``[B, S]`` int tensor with values in
+            ``{0=text, 1=image, 2=video}``.
+        sliding_window: Sliding window size from ``Gemma4VLConfig.sliding_window``
+            (1024 for 26B-A4B). Used only by the runtime guard; sliding layers
+            apply the cutoff themselves via TE ``window_size``.
+
+    Returns:
+        BoolTensor ``[B, 1, S, S]`` with ``True = MASKED`` (Megatron convention).
+        Layers select between the full and sliding behavior via TE's per-layer
+        ``window_size``; this single mask serves both.
+    """
+    if input_ids.shape != mm_token_type_ids.shape:
+        raise ValueError(
+            f"input_ids shape {tuple(input_ids.shape)} does not match "
+            f"mm_token_type_ids shape {tuple(mm_token_type_ids.shape)}"
+        )
+    if input_ids.dim() != 2:
+        raise ValueError(f"input_ids must be 2D [B, S]; got shape {tuple(input_ids.shape)}")
+
+    batch_size, seq_len = input_ids.shape
+    device = input_ids.device
+
+    causal_visible = torch.tril(
+        torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
+    ).unsqueeze(0).expand(batch_size, -1, -1)
+
+    vision_group_ids = _compute_vision_group_ids(mm_token_type_ids.to(device))
+    _assert_vision_spans_within_window(vision_group_ids, sliding_window)
+
+    q_groups = vision_group_ids.unsqueeze(-1)
+    kv_groups = vision_group_ids.unsqueeze(-2)
+    vision_overlay_visible = (q_groups == kv_groups) & (q_groups >= 0)
+
+    visible = causal_visible | vision_overlay_visible
+    masked = ~visible
+    return masked.unsqueeze(1)
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_config.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_config.py
new file mode 100644
index 00000000..bb084cbc
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_config.py
@@ -0,0 +1,242 @@
+"""Configuration dataclasses and registrations for Gemma4-VL family.
+
+Mirrors the structure of ``aiak_training_llm.models.llava_onevision2.llava_onevision2_config``
+and adds the Gemma4-specific knobs:
+
+- ``final_logit_softcapping`` (final logits tanh-cap, gated)
+- ``use_layer_scalar`` (per-layer scalar buffer multiplied at end of layer forward)
+- ``per_layer_kv_channels`` / ``layer_pattern`` (sliding vs global hybrid attention)
+- ``attention_k_eq_v`` (Gemma4 K=V tying flag — combined with ``layer_pattern[i]=='global'``
+  this drives the no-V-projection runtime aliasing on global layers)
+- ``kv_tied_layers`` (precomputed list of global layer indices for fast runtime lookup)
+- ``scale_emb_by_sqrt_hidden`` (Gemma family: x = embed(ids) * sqrt(hidden))
+- ``sliding_window`` (sliding-attention window size)
+
+Note: this checkpoint has dual RoPE — sliding layers use ``rope_type='default'``
+with ``rope_theta=10000`` and full rotation; global layers use HF
+``rope_type='proportional'`` with ``rope_theta=1_000_000``,
+``partial_rotary_factor=0.25``, and the proportional zero-pad scheme over
+``global_head_dim=512`` (rotate the first 64 angle pairs of 256, leave the
+remaining 192 as identity). The sliding/global split lives in
+``layer_pattern``; ``Gemma4Model`` constructs two ``RotaryEmbedding`` instances
+and ships them as a dict keyed by layer-type to ``Gemma4SelfAttention``.
+
+These fields land on ``TransformerConfig`` via the matching aiak_megatron patch
+(P1.13A) so ``core_transformer_config_from_args`` can pump them straight from
+the parsed args.
+"""
+
+from dataclasses import dataclass, field
+
+import torch
+from torch.nn.functional import gelu
+
+from aiak_training_llm.models.factory import register_model_config
+from aiak_training_llm.utils.constants import VisionLanguageModelFamilies
+
+
+@dataclass
+class AdapterConfig:
+    """Configuration for the Gemma4-VL projector.
+
+    The fields need to be consistent with the definitions in args.
+    """
+
+    normalization: str
+    activation_func: torch.nn.Module = gelu
+    add_bias_linear: bool = False
+    layernorm_epsilon: float = 1e-06
+
+
+@dataclass
+class Gemma4VLConfig:
+    num_layers: int
+    hidden_size: int
+    ffn_hidden_size: int
+    num_attention_heads: int
+    group_query_attention: bool = True
+    num_query_groups: int = 8
+    position_embedding_type: str = "rope"
+    add_position_embedding: bool = False
+    rotary_interleaved: bool = False
+    normalization: str = "RMSNorm"
+    swiglu: bool = True
+    attention_dropout: float = 0
+    hidden_dropout: float = 0
+    add_bias_linear: bool = False
+    add_qkv_bias: bool = False
+    qk_layernorm: bool = True
+    untie_embeddings_and_output_weights: bool = False
+    vocab_size_in_config_file: int = 262144
+    make_vocab_size_divisible_by: int = 128
+    norm_epsilon: float = 1e-06
+    rotary_base: int = 1000000
+    kv_channels: int = 256
+    num_experts: int = None
+    moe_ffn_hidden_size: int = None
+    # ---- Gemma4-specific extensions (mirrored on TransformerConfig via P1.13A) ----
+    final_logit_softcapping: float = None
+    use_layer_scalar: bool = False
+    sliding_window: int = None
+    rotary_base_sliding: int = None
+    partial_rotary_factor: float = 1.0
+    # Per-layer hybrid attention overrides. ``layer_pattern`` is a list of
+    # ``"sliding"`` / ``"global"`` strings of length ``num_layers``;
+    # ``per_layer_kv_channels`` and ``per_layer_num_query_groups`` are dicts
+    # keyed by layer-type ("sliding"/"global") -> int.
+    layer_pattern: list = field(default_factory=list)
+    per_layer_kv_channels: dict = field(default_factory=dict)
+    per_layer_num_query_groups: dict = field(default_factory=dict)
+    # K=V tying configuration. ``attention_k_eq_v=True`` enables HF's "no V projection,
+    # alias V from K" mechanism on layers listed in ``kv_tied_layers`` (= global layers
+    # under HF's ``use_alternative_attention = attention_k_eq_v AND not is_sliding``).
+    attention_k_eq_v: bool = False
+    kv_tied_layers: list = field(default_factory=list)
+    scale_emb_by_sqrt_hidden: bool = False
+
+
+@register_model_config(
+    model_family=VisionLanguageModelFamilies.GEMMA4_VL,
+    model_arch="gemma4-26b-a4b-vl",
+)
+def gemma4_26b_a4b_vl():
+    """Gemma4-26B-A4B-it (instruction-tuned MoE VLM).
+
+    Architecture (verified from /ov2/pretrain_models/google/gemma-4-26B-A4B-it/config.json):
+      - 30 layers with repeating ``[s,s,s,s,s,g]`` pattern (5 sliding + 1 global)
+      - hidden=2816, num_attention_heads=16
+      - dense intermediate=2112, moe intermediate=704 (per routed expert)
+      - 128 routed experts top-k=8; dense MLP runs as a parallel branch from the
+        same post-attention residual (NOT a Megatron "shared expert" — see
+        Gemma4ParallelDenseMoE / plan v5 §531 R5)
+      - sliding layers: head_dim=256, num_kv=8,  RoPE theta=10000,  window=1024
+      - global  layers: head_dim=512, num_kv=2,  RoPE theta=1000000, K=V tied (no V proj)
+      - full RoPE (no partial), qk_layernorm, parameter-free V LayerNorm
+      - final logit softcap 30.0, per-layer scalar buffer multiplied at layer end
+      - vocab=262144, ctx_len=262144
+      - tie_word_embeddings=True
+    """
+    layer_pattern = []
+    for i in range(30):
+        layer_pattern.append("global" if (i + 1) % 6 == 0 else "sliding")
+    kv_tied_layers = [i for i, t in enumerate(layer_pattern) if t == "global"]
+    return Gemma4VLConfig(
+        num_layers=30,
+        hidden_size=2816,
+        ffn_hidden_size=2112,
+        num_attention_heads=16,
+        group_query_attention=True,
+        num_query_groups=8,
+        vocab_size_in_config_file=262144,
+        make_vocab_size_divisible_by=128,
+        qk_layernorm=True,
+        kv_channels=256,
+        add_qkv_bias=False,
+        rotary_base=1000000,
+        rotary_base_sliding=10000,
+        partial_rotary_factor=0.25,
+        sliding_window=1024,
+        num_experts=128,
+        moe_ffn_hidden_size=704,
+        final_logit_softcapping=30.0,
+        use_layer_scalar=True,
+        scale_emb_by_sqrt_hidden=True,
+        untie_embeddings_and_output_weights=False,
+        layer_pattern=layer_pattern,
+        per_layer_kv_channels={"sliding": 256, "global": 512},
+        per_layer_num_query_groups={"sliding": 8, "global": 2},
+        attention_k_eq_v=True,
+        kv_tied_layers=kv_tied_layers,
+    )
+
+
+@dataclass
+class VisionConfig:
+    """Configuration for the Gemma4-VL vision tower.
+
+    Verbatim mirror of HF ``Gemma4VisionConfig`` for ``gemma-4-26B-A4B-it``
+    (see ``/ov2/pretrain_models/google/gemma-4-26B-A4B-it/config.json`` →
+    ``vision_config``). This is NOT a SigLIP tower — Gemma4 uses its own ViT
+    with: learned 2-D position embeddings via one-hot + matmul against a
+    ``[2, position_embedding_size, hidden_size]`` table; sandwich RMSNorm per
+    layer; v_proj followed by a parameterless RMSNorm (``with_scale=False``);
+    multidimensional RoPE applied per-axis on the head_dim/2 split; pooling via
+    a dynamic-kernel 2-D average pool whose kernel is derived from
+    ``pooling_kernel_size`` and the input sequence length; and an optional
+    output standardization stage gated by ``standardize``.
+
+    Field semantics that DIFFER from OneVision-Encoder / SigLIP and MUST NOT
+    be silently aliased to the OV2 defaults:
+
+    - ``patch_size=16`` (NOT 14) and ``in_channels=3`` give patch_embed input
+      ``in_features = 3 * 16 * 16 = 768``.
+    - ``num_attention_heads = num_key_value_heads = 16`` → MHA (not GQA);
+      ``head_dim=72`` → ``Q/K/V`` projection out_features ``= 16 * 72 = 1152``.
+    - ``hidden_activation='gelu_pytorch_tanh'`` is the *only* activation that
+      reproduces HF — exact ``F.gelu`` (erf form) silently diverges.
+    - ``use_clipped_linears=False`` for this checkpoint; the buffer branch in
+      :class:`Gemma4ClippableLinear` stays inactive.
+    - ``standardize=True`` → :class:`Gemma4VisionTower` registers persistent
+      ``std_bias``/``std_scale`` buffers and applies them after pooling.
+    - ``rope_theta=100.0`` (not 10000) for the vision multidimensional RoPE.
+    - ``position_embedding_size=10240`` is the per-axis position table length;
+      ``pooling_kernel_size=3`` controls the *output* sequence length via
+      ``pixel_values.shape[-2] // (pooling_kernel_size ** 2)``.
+    """
+
+    num_hidden_layers: int
+    hidden_size: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    patch_size: int
+    pooling_kernel_size: int
+    position_embedding_size: int
+    rms_norm_eps: float
+    rope_theta: float
+    standardize: bool
+    use_clipped_linears: bool
+    hidden_activation: str
+    in_channels: int = 3
+    initializer_range: float = 0.02
+
+
+def get_vision_config(model_family, model_name):
+    del model_family, model_name
+    return VisionConfig(
+        num_hidden_layers=27,
+        hidden_size=1152,
+        intermediate_size=4304,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        head_dim=72,
+        patch_size=16,
+        pooling_kernel_size=3,
+        position_embedding_size=10240,
+        rms_norm_eps=1e-06,
+        rope_theta=100.0,
+        standardize=True,
+        use_clipped_linears=False,
+        hidden_activation="gelu_pytorch_tanh",
+        in_channels=3,
+        initializer_range=0.02,
+    )
+
+
+def get_adapeter_config(model_family, model_name=None):
+    """Adapter config for Gemma4-VL: parameterless RMSNorm + bias-free Linear.
+
+    The Gemma4 adapter is intentionally minimal — a single trainable tensor
+    (``embedding_projection.weight``) plus a parameterless input RMSNorm
+    (``with_scale=False``, no ``weight``). Mirrors HF ``Gemma4MultimodalEmbedder``
+    (lines 2023-2047 of ``modeling_gemma4``); the LLM-side hidden_size used for
+    the projection out_features is derived at construction time from the LLM
+    config rather than carried here.
+    """
+    del model_family, model_name
+    return AdapterConfig(
+        normalization="RMSNorm",
+        add_bias_linear=False,
+        layernorm_epsilon=1e-06,
+    )
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_layer_spec.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_layer_spec.py
new file mode 100644
index 00000000..4b7949a3
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_layer_spec.py
@@ -0,0 +1,170 @@
+"""Gemma4-VL layer specs.
+
+Three exports:
+  - get_vision_layer_with_spec: returns ``None`` — Gemma4 vision tower is a
+    plain ``nn.Module`` stack (HF-style verbatim port), not a Megatron
+    ``ModuleSpec``-driven ``TransformerBlock``. Kept as a no-op so
+    ``gemma4_vl_provider`` can call it uniformly with the LlavaOnevision2
+    provider; the returned ``None`` is forwarded to ``Gemma4VL.__init__`` and
+    ignored there.
+  - get_adapter_layer_with_spec: same story — Gemma4 adapter is a plain
+    ``nn.Module`` (RMSNorm + Linear, see ``gemma4_adapter.py``).
+  - get_gemma4_layer_with_te_spec: Gemma4 LLM layer wiring
+        Gemma4TransformerLayer + Gemma4SelfAttention + Gemma4ParallelDenseMoE
+"""
+
+from megatron.core.transformer.attention import SelfAttentionSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from aiak_training_llm.models.dispatch import multiacc_modules
+from aiak_training_llm.models.gemma4_vl.gemma4_attention import Gemma4SelfAttention
+from aiak_training_llm.models.gemma4_vl.gemma4_mlp import (
+    Gemma4ParallelDenseMoE,
+    Gemma4ParallelDenseMoESubmodules,
+)
+from aiak_training_llm.models.gemma4_vl.gemma4_moe_layer import (
+    Gemma4MoELayer,
+    Gemma4MoESubmodules,
+)
+from aiak_training_llm.models.gemma4_vl.gemma4_transformer_layer import (
+    Gemma4TransformerLayer,
+    Gemma4TransformerLayerSubmodules,
+)
+from aiak_training_llm.utils import is_te_min_version
+
+
+__all__ = [
+    "get_vision_layer_with_spec",
+    "get_adapter_layer_with_spec",
+    "get_gemma4_layer_with_te_spec",
+]
+
+
+def get_vision_layer_with_spec():
+    """Gemma4 vision tower is not Megatron-spec-driven; return ``None``.
+
+    The provider calls this for API parity with the LlavaOnevision2 provider
+    and forwards the result to ``Gemma4VL.__init__``, which builds the vision
+    tower via ``Gemma4VisionTower(...)`` directly and ignores this argument.
+    """
+    return None
+
+
+def get_adapter_layer_with_spec():
+    """Gemma4 adapter is not Megatron-spec-driven; return ``None``.
+
+    Same rationale as ``get_vision_layer_with_spec``. The Gemma4 adapter is
+    constructed as a plain ``nn.Module`` (see ``Gemma4Adapter``) inside
+    ``Gemma4VL.__init__``.
+    """
+    return None
+
+
+def _gemma4_norm_module() -> type:
+    return multiacc_modules.TENorm if is_te_min_version("1.9.0") else multiacc_modules.LocalNorm
+
+
+def _gemma4_dense_mlp_spec() -> ModuleSpec:
+    return ModuleSpec(
+        module=MLP,
+        submodules=MLPSubmodules(
+            linear_fc1=multiacc_modules.TEColumnParallelLinear,
+            linear_fc2=multiacc_modules.TERowParallelLinear,
+            bias_activation_func_impl=multiacc_modules.bias_activation_func_impl,
+        ),
+    )
+
+
+def _gemma4_moe_spec(moe_grouped_gemm: bool) -> ModuleSpec:
+    if moe_grouped_gemm:
+        assert multiacc_modules.TEColumnParallelGroupedLinear is not None
+        expert_module = TEGroupedMLP
+        linear_fc1 = multiacc_modules.TEColumnParallelGroupedLinear
+        linear_fc2 = multiacc_modules.TERowParallelGroupedLinear
+    else:
+        expert_module = SequentialMLP
+        linear_fc1 = multiacc_modules.TEColumnParallelLinear
+        linear_fc2 = multiacc_modules.TERowParallelLinear
+
+    experts_spec = ModuleSpec(
+        module=expert_module,
+        submodules=MLPSubmodules(
+            linear_fc1=linear_fc1,
+            linear_fc2=linear_fc2,
+            bias_activation_func_impl=multiacc_modules.bias_activation_func_impl,
+        ),
+    )
+    # HF Gemma4 MoE has NO shared experts; the dense MLP is implemented as a
+    # peer branch in Gemma4ParallelDenseMoE (see plan v5 §531 R5). We must
+    # leave moe_shared_expert_intermediate_size unset (None) in the config —
+    # any other value, including 0, makes MoELayer's ``use_shared_expert``
+    # flag True (it checks ``is not None``) and triggers a build of the
+    # shared-experts branch which we did not wire (would crash).
+    return ModuleSpec(
+        module=Gemma4MoELayer,
+        submodules=Gemma4MoESubmodules(
+            experts=experts_spec,
+            pre_feedforward_layernorm_2=_gemma4_norm_module(),
+        ),
+    )
+
+
+def _gemma4_parallel_dense_moe_spec(moe_grouped_gemm: bool) -> ModuleSpec:
+    norm = _gemma4_norm_module()
+    return ModuleSpec(
+        module=Gemma4ParallelDenseMoE,
+        submodules=Gemma4ParallelDenseMoESubmodules(
+            dense=_gemma4_dense_mlp_spec(),
+            moe=_gemma4_moe_spec(moe_grouped_gemm=moe_grouped_gemm),
+            pre_feedforward_layernorm=norm,
+            post_feedforward_layernorm_1=norm,
+            post_feedforward_layernorm_2=norm,
+        ),
+    )
+
+
+def _gemma4_self_attention_submodules(qk_norm) -> SelfAttentionSubmodules:
+    return SelfAttentionSubmodules(
+        linear_qkv=multiacc_modules.TELayerNormColumnParallelLinear,
+        core_attention=multiacc_modules.DotProductAttention,
+        linear_proj=multiacc_modules.TERowParallelLinear,
+        q_layernorm=qk_norm,
+        k_layernorm=qk_norm,
+        apply_rotary_fn=multiacc_modules.apply_rotary_pos_emb,
+    )
+
+
+def get_gemma4_layer_with_te_spec(config: TransformerConfig) -> ModuleSpec:
+    assert not config.multi_latent_attention, (
+        "Gemma4-VL does not use multi-latent attention; got "
+        "config.multi_latent_attention=True"
+    )
+
+    qk_norm = _gemma4_norm_module()
+    mlp = _gemma4_parallel_dense_moe_spec(moe_grouped_gemm=config.moe_grouped_gemm)
+
+    # ``pre_mlp_layernorm = IdentityOp`` because Gemma4ParallelDenseMoE +
+    # Gemma4MoELayer between them own all four pre-merge RMSNorms internally
+    # and must receive the raw post-attention residual (router math depends
+    # on this — see gemma4_mlp.py module docstring).
+    return ModuleSpec(
+        module=Gemma4TransformerLayer,
+        submodules=Gemma4TransformerLayerSubmodules(
+            input_layernorm=IdentityOp,
+            self_attention=ModuleSpec(
+                module=Gemma4SelfAttention,
+                params={"attn_mask_type": AttnMaskType.causal},
+                submodules=_gemma4_self_attention_submodules(qk_norm=qk_norm),
+            ),
+            self_attn_bda=multiacc_modules.get_bias_dropout_add,
+            pre_mlp_layernorm=IdentityOp,
+            mlp=mlp,
+            mlp_bda=multiacc_modules.get_bias_dropout_add,
+            post_feedforward_layernorm=_gemma4_norm_module(),
+        ),
+    )
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_model.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_model.py
new file mode 100644
index 00000000..4a1c0223
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_model.py
@@ -0,0 +1,331 @@
+"""Top-level Gemma4-VL VLM (Gemma4 vision tower + Gemma4 adapter + Gemma4 LLM).
+
+Subclasses :class:`LlavaOnevision2` so the provider/freeze/set_input_tensor
+plumbing is reused, but overrides ``__init__`` and ``forward`` to swap in
+the HF-verbatim Gemma4 vision tower / adapter and the Gemma4 MoE LLM.
+
+``vision_layer_spec`` and ``adapter_layer_spec`` arguments are accepted for
+provider-call uniformity but are IGNORED — the Gemma4 vision tower / adapter
+are plain ``nn.Module`` stacks (see ``gemma4_vision_tower.py`` /
+``gemma4_adapter.py``), not Megatron ``ModuleSpec``-driven blocks.
+
+The vision data path (``images is not None``) currently raises
+``NotImplementedError``: the production Gemma4 image processor / dataloader
+that produces ``pixel_values`` + ``pixel_position_ids`` in the HF layout the
+tower expects has not yet been wired into the OV2 dataloader (P3.x). LM-only
+forward (``images is None``, P1 smoke + pure-text training) is fully
+functional.
+"""
+
+import logging
+from functools import partial
+from typing import Optional
+
+import torch
+from megatron.core import InferenceParams, parallel_state, tensor_parallel
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.spec_utils import ModuleSpec
+
+from aiak_training_llm.models.gemma4_vl.gemma4_adapter import Gemma4Adapter
+from aiak_training_llm.models.gemma4_vl.gemma4_model import Gemma4Model
+from aiak_training_llm.models.gemma4_vl.gemma4_vision_tower import Gemma4VisionTower
+from aiak_training_llm.models.llava_onevision2.llava_onevision2_model import (
+    LlavaOnevision2,
+    _load_state_dict_hook_ignore_param_names,
+)
+
+
+class Gemma4VL(LlavaOnevision2):
+    def __init__(
+        self,
+        language_config,
+        vision_config,
+        adapter_config,
+        language_layer_spec: ModuleSpec,
+        vision_layer_spec: ModuleSpec,
+        adapter_layer_spec: ModuleSpec,
+        language_vocab_size: int,
+        language_max_sequence_length: int,
+        allow_missing_adapter_checkpoint: bool = False,
+        parallel_output: bool = True,
+        language_position_embedding_type: str = "rope",
+        language_rotary_percent: float = 1.0,
+        pre_process: bool = True,
+        post_process: bool = True,
+        add_encoder: bool = True,
+        add_decoder: bool = True,
+        language_rotary_base: int = 10000,
+        language_rotary_base_sliding: Optional[int] = None,
+        fp16_lm_cross_entropy: bool = False,
+        share_embeddings_and_output_weights: bool = True,
+        seq_len_interpolation_factor: Optional[float] = None,
+    ) -> None:
+        del vision_layer_spec, adapter_layer_spec
+
+        from megatron.core.transformer import MegatronModule
+
+        MegatronModule.__init__(self, config=language_config)
+
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.add_encoder = add_encoder
+        self.add_decoder = add_decoder
+
+        self.encoder_hidden_state = None
+        self.vision_model = None
+        self.adapter = None
+        self.language_model = None
+
+        if self.add_encoder:
+            self.vision_model = Gemma4VisionTower(
+                transformer_config=vision_config,
+                hidden_size=vision_config.hidden_size,
+                num_hidden_layers=vision_config.num_hidden_layers,
+                num_attention_heads=vision_config.num_attention_heads,
+                num_key_value_heads=vision_config.num_key_value_heads,
+                head_dim=vision_config.head_dim,
+                intermediate_size=vision_config.intermediate_size,
+                patch_size=vision_config.patch_size,
+                in_channels=vision_config.in_channels,
+                position_embedding_size=vision_config.position_embedding_size,
+                pooling_kernel_size=vision_config.pooling_kernel_size,
+                rope_theta=vision_config.rope_theta,
+                rms_norm_eps=vision_config.rms_norm_eps,
+                hidden_activation=vision_config.hidden_activation,
+                use_clipped_linears=vision_config.use_clipped_linears,
+                standardize=vision_config.standardize,
+            )
+            self.adapter = Gemma4Adapter(
+                vision_hidden_size=vision_config.hidden_size,
+                text_hidden_size=language_config.hidden_size,
+                rms_norm_eps=adapter_config.layernorm_epsilon,
+            )
+            if allow_missing_adapter_checkpoint:
+                adapter_param_names = [
+                    f"adapter.{name}" for name in self.adapter.state_dict().keys()
+                ]
+                self.adapter.register_load_state_dict_post_hook(
+                    partial(_load_state_dict_hook_ignore_param_names, adapter_param_names)
+                )
+
+        if self.add_decoder:
+            self.language_model = Gemma4Model(
+                config=language_config,
+                transformer_layer_spec=language_layer_spec,
+                vocab_size=language_vocab_size,
+                max_sequence_length=language_max_sequence_length,
+                parallel_output=parallel_output,
+                position_embedding_type=language_position_embedding_type,
+                rotary_percent=language_rotary_percent,
+                pre_process=self.pre_process,
+                post_process=self.post_process,
+                rotary_base=language_rotary_base,
+                rotary_base_sliding=language_rotary_base_sliding,
+                share_embeddings_and_output_weights=share_embeddings_and_output_weights,
+                scatter_embedding_sequence_parallel=False,
+                fp16_lm_cross_entropy=fp16_lm_cross_entropy,
+                seq_len_interpolation_factor=seq_len_interpolation_factor,
+            )
+            self.share_embeddings_and_output_weights = (
+                self.language_model.share_embeddings_and_output_weights
+            )
+
+    def set_input_tensor(self, input_tensor) -> None:
+        if not isinstance(input_tensor, list):
+            input_tensor = [input_tensor]
+        assert len(input_tensor) == 1, (
+            "input_tensor should only be length 1 for Gemma4-VL"
+        )
+
+        if self.add_encoder and self.add_decoder:
+            raise NotImplementedError(
+                "Gemma4-VL encoder-stage set_input_tensor requires Gemma4VisionTower "
+                "to expose a Megatron-style set_input_tensor; deferred to P3.x along "
+                "with the vision data path."
+            )
+        if self.add_encoder:
+            raise NotImplementedError(
+                "Gemma4-VL encoder-only pipeline stage requires Gemma4VisionTower "
+                "to expose a Megatron-style set_input_tensor; deferred to P3.x."
+            )
+        if self.pre_process:
+            self.encoder_hidden_state = input_tensor[0]
+        else:
+            self.language_model.set_input_tensor(input_tensor[0])
+
+    def forward_debug(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Gemma4-VL forward_debug is deferred to P3.x; the inherited "
+            "LlavaOnevision2.forward_debug calls Gemma4VisionTower.forward_debug "
+            "which does not exist (the Gemma4 vision tower is a plain nn.Module, "
+            "not a Megatron block). The HF↔Megatron consistency check in Phase C "
+            "must call Gemma4VisionTower.forward / Gemma4Adapter.forward directly."
+        )
+
+    @staticmethod
+    def _prepare_vision_inputs(
+        images: torch.Tensor,
+        image_grid_thw: torch.Tensor,
+        patch_positions: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        patch_counts = image_grid_thw.prod(dim=-1).to(torch.long)
+        if int(patch_counts.sum().item()) != images.shape[0]:
+            raise ValueError(
+                f"Gemma4-VL image patch count mismatch: images has {images.shape[0]} "
+                f"patches but image_grid_thw sums to {int(patch_counts.sum().item())}."
+            )
+
+        max_patches = int(patch_counts.max().item())
+        batch_size = int(patch_counts.numel())
+        pixel_values = images.new_zeros((batch_size, max_patches, images.shape[-1]))
+        pixel_position_ids = torch.full(
+            (batch_size, max_patches, 2),
+            -1,
+            dtype=torch.long,
+            device=images.device,
+        )
+
+        offset = 0
+        for batch_idx, patch_count_tensor in enumerate(patch_counts):
+            patch_count = int(patch_count_tensor.item())
+            next_offset = offset + patch_count
+            pixel_values[batch_idx, :patch_count] = images[offset:next_offset]
+
+            if patch_positions is None:
+                _t, h, w = image_grid_thw[batch_idx].tolist()
+                h_coords = torch.arange(h, dtype=torch.long, device=images.device).repeat_interleave(w)
+                w_coords = torch.arange(w, dtype=torch.long, device=images.device).repeat(h)
+                coords = torch.stack((w_coords, h_coords), dim=-1)
+                if coords.shape[0] != patch_count:
+                    raise ValueError(
+                        f"Default Gemma4-VL pixel positions only support single-frame images; "
+                        f"got image_grid_thw={image_grid_thw[batch_idx].tolist()}."
+                    )
+            else:
+                coords = patch_positions[offset:next_offset, -2:].to(device=images.device, dtype=torch.long)
+                coords = torch.stack((coords[:, 1], coords[:, 0]), dim=-1)
+
+            pixel_position_ids[batch_idx, :patch_count] = coords
+            offset = next_offset
+
+        return pixel_values, pixel_position_ids
+
+    def forward(
+        self,
+        images: torch.Tensor,
+        image_grid_thw: torch.Tensor,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        attn_mask_type: AttnMaskType | None = None,
+        labels: torch.Tensor = None,
+        packed_seq_params: PackedSeqParams = None,
+        inference_params: InferenceParams = None,
+        pixel_values_videos: torch.Tensor = None,
+        video_grid_thw: torch.Tensor = None,
+        patch_positions: list[torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        del position_ids
+
+        if pixel_values_videos is not None or video_grid_thw is not None:
+            raise NotImplementedError(
+                "Gemma4-VL video path not implemented; pixel_values_videos / "
+                "video_grid_thw must be None."
+            )
+
+        use_inference_kv_cache = (
+            inference_params is not None
+            and "image_tokens_count" in inference_params.key_value_memory_dict
+        )
+
+        if images is not None and self.add_encoder and not use_inference_kv_cache:
+            if image_grid_thw is None:
+                raise ValueError("Gemma4-VL image_grid_thw is required when images are provided.")
+            pixel_values, pixel_position_ids = self._prepare_vision_inputs(
+                images, image_grid_thw, patch_positions
+            )
+            image_embeddings = self.vision_model(pixel_values, pixel_position_ids)
+            image_embeddings = self.adapter(image_embeddings)
+
+            n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
+            n_image_features = image_embeddings.shape[0]
+            if n_image_features > n_image_tokens:
+                logging.getLogger(__name__).warning(
+                    "Trimming %d extra Gemma4 image embedding(s) "
+                    "(n_image_features=%d, n_image_tokens=%d).",
+                    n_image_features - n_image_tokens,
+                    n_image_features,
+                    n_image_tokens,
+                )
+                image_embeddings = image_embeddings[:n_image_tokens]
+            elif n_image_features < n_image_tokens:
+                raise ValueError(
+                    f"Gemma4 image features {n_image_features} < image tokens {n_image_tokens}"
+                )
+
+            if inference_params is not None:
+                inference_params.key_value_memory_dict["image_tokens_count"] = image_embeddings.shape[0]
+        else:
+            image_embeddings = None
+
+        if not self.add_decoder:
+            raise NotImplementedError(
+                "Gemma4-VL encoder-only pipeline stage requires the vision data "
+                "path which is not yet implemented; see images-not-None branch."
+            )
+
+        if self.pre_process:
+            language_embeddings = self.language_model.embedding(
+                input_ids=input_ids, position_ids=None
+            )
+            if use_inference_kv_cache or images is None:
+                combined_embeddings = language_embeddings
+            elif (input_ids == self.config.image_token_id).any().item():
+                images_mask = (
+                    (input_ids == self.config.image_token_id)
+                    .transpose(0, 1)
+                    .unsqueeze(-1)
+                    .expand_as(language_embeddings)
+                    .to(language_embeddings.device)
+                )
+                image_embeddings = image_embeddings.to(
+                    language_embeddings.device, language_embeddings.dtype
+                )
+                combined_embeddings = language_embeddings.masked_scatter(
+                    images_mask, image_embeddings
+                )
+            else:
+                combined_embeddings = language_embeddings
+
+            if self.config.sequence_parallel:
+                seq_len = combined_embeddings.size(0)
+                tp_world_size = parallel_state.get_tensor_model_parallel_world_size()
+                remainder = seq_len % tp_world_size
+                if remainder != 0:
+                    pad = tp_world_size - remainder
+                    pad_shape = (pad,) + combined_embeddings.shape[1:]
+                    pad_tensor = combined_embeddings.new_zeros(pad_shape)
+                    combined_embeddings = torch.cat(
+                        (combined_embeddings, pad_tensor), dim=0
+                    )
+                combined_embeddings = (
+                    tensor_parallel.scatter_to_sequence_parallel_region(
+                        combined_embeddings
+                    )
+                )
+        else:
+            combined_embeddings = None
+
+        return self.language_model(
+            input_ids=None,
+            position_ids=None,
+            attention_mask=attention_mask,
+            attn_mask_type=attn_mask_type,
+            decoder_input=combined_embeddings,
+            labels=labels,
+            rotary_pos_emb=None,
+            inference_params=inference_params,
+            packed_seq_params=packed_seq_params,
+            extra_block_kwargs={},
+        )
diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_provider.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_provider.py
new file mode 100644
index 00000000..266bbe7b
--- /dev/null
+++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_provider.py
@@ -0,0 +1,141 @@
+"""Provider for the Gemma4-VL model family.
+
+Mirrors :mod:`llava_onevision2_provider` but switches:
+  - ``register_model_provider`` family to ``VisionLanguageModelFamilies.GEMMA4_VL``
+  - ``language_layer_spec`` builder to :func:`get_gemma4_layer_with_te_spec`
+  - top-level VLM class to :class:`Gemma4VL`
+
+All other knobs (TP/PP/SP gating, freeze gating, encoder PP) match the
+LlavaOnevision2 provider verbatim so any base-provider improvements port
+directly to this provider.
+"""
+
+from copy import deepcopy
+from dataclasses import asdict
+
+from megatron.core import mpu
+from megatron.core.transformer.spec_utils import import_module
+
+from aiak_training_llm.models.factory import register_model_provider
+from aiak_training_llm.models.gemma4_vl.gemma4_vl_config import (
+    get_adapeter_config,
+    get_vision_config,
+)
+from aiak_training_llm.models.gemma4_vl.gemma4_vl_layer_spec import (
+    get_adapter_layer_with_spec,
+    get_gemma4_layer_with_te_spec,
+    get_vision_layer_with_spec,
+)
+from aiak_training_llm.models.gemma4_vl.gemma4_vl_model import Gemma4VL
+from aiak_training_llm.utils import build_transformer_config, get_args, print_rank_0
+from aiak_training_llm.utils.constants import VisionLanguageModelFamilies
+
+
+@register_model_provider(model_family=[VisionLanguageModelFamilies.GEMMA4_VL])
+def gemma4_vl_model_provider(
+    pre_process: bool = True,
+    post_process: bool = True,
+    add_encoder: bool = True,
+    add_decoder: bool = True,
+    parallel_output: bool = True,
+) -> Gemma4VL:
+    """Build a Gemma4-VL model instance.
+
+    Args:
+        pre_process: include the embedding layer (pipeline-first stage).
+        post_process: include the output layer / loss (pipeline-last stage).
+        add_encoder / add_decoder: pipeline-encoder/decoder gating.
+        parallel_output: keep outputs split across TP ranks.
+    """
+    args = get_args()
+
+    print_rank_0(f"building {args.model_name} model ...")
+
+    config = build_transformer_config(args)
+
+    language_config = deepcopy(config)
+    vision_config = deepcopy(config)
+    adapter_config = deepcopy(config)
+
+    from aiak_training_llm.models import get_model_family
+
+    model_family = get_model_family(args.model_name)
+    for k, v in asdict(get_vision_config(model_family, args.model_name)).items():
+        setattr(vision_config, k, v)
+    for k, v in asdict(get_adapeter_config(model_family, args.model_name)).items():
+        setattr(adapter_config, k, v)
+
+    setattr(language_config, "image_token_id", 258880)
+    setattr(language_config, "video_token_id", 258884)
+
+    if args.encoder_pipeline_model_parallel_size in [0, None]:
+        vision_config.pipeline_model_parallel_size = 1
+        vision_config.first_pipeline_num_layers = (
+            vision_config.last_pipeline_num_layers
+        ) = None
+        vision_config.tp_comm_overlap = False
+        vision_config.context_parallel_size = 1
+        vision_config.context_parallel_ulysses_degree = 1
+        add_encoder = mpu.is_pipeline_first_stage()
+        add_decoder = True
+    else:
+        assert args.encoder_pipeline_model_parallel_size == 1, (
+            "vision model and projection can only live on 1 pipeline stage."
+        )
+        vision_config.pipeline_model_parallel_size = (
+            args.encoder_pipeline_model_parallel_size
+        )
+        if args.encoder_tensor_model_parallel_size > 0:
+            vision_config.tensor_model_parallel_size = (
+                args.encoder_tensor_model_parallel_size
+            )
+        vision_config.first_pipeline_num_layers = (
+            vision_config.last_pipeline_num_layers
+        ) = None
+        vision_config.context_parallel_size = 1
+        vision_config.tp_comm_overlap = False
+
+    if args.use_legacy_models:
+        raise ValueError("Classic Megatron-LM models are not supported.")
+
+    if args.spec is not None:
+        language_layer_spec = import_module(args.spec)
+    else:
+        adapter_layer_spec = get_adapter_layer_with_spec()
+        vision_layer_spec = get_vision_layer_with_spec()
+        language_layer_spec = get_gemma4_layer_with_te_spec(language_config)
+
+    model = Gemma4VL(
+        language_config=language_config,
+        vision_config=vision_config,
+        adapter_config=adapter_config,
+        language_layer_spec=language_layer_spec,
+        vision_layer_spec=vision_layer_spec,
+        adapter_layer_spec=adapter_layer_spec,
+        language_vocab_size=args.padded_vocab_size,
+        language_max_sequence_length=args.max_position_embeddings,
+        pre_process=pre_process,
+        post_process=post_process,
+        add_encoder=add_encoder,
+        add_decoder=add_decoder,
+        fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+        parallel_output=parallel_output,
+        share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+        language_position_embedding_type=args.position_embedding_type,
+        language_rotary_percent=args.rotary_percent,
+        language_rotary_base=args.rotary_base,
+        language_rotary_base_sliding=getattr(args, "rotary_base_sliding", None),
+        seq_len_interpolation_factor=args.rotary_seq_len_interpolation_factor,
+    )
+
+    if args.trainable_modules != ["all"]:
+        train_language_model = "language_model" in args.trainable_modules
+        train_vision_model = "vision_model" in args.trainable_modules
+        train_adapter = "adapter" in args.trainable_modules
+        model.freeze(
+            freeze_language_model=not train_language_model,
+            freeze_vision_model=not train_vision_model,
+            freeze_adapter=not train_adapter,
+        )
+
+    return model
diff --git a/aiak_training_llm/train/__init__.py b/aiak_training_llm/train/__init__.py
index f7fe5b2a..7d940744 100644
--- a/aiak_training_llm/train/__init__.py
+++ b/aiak_training_llm/train/__init__.py
@@ -1,8 +1,8 @@
 """aiak train module"""
 
 from .arguments import parse_train_args
-from .pretrain import pretrain_llava_onevision2, pretrain_llm, pretrain_qwen2_vl
-from .sft import sft_llava_onevision2, sft_llava_onevision1_5, sft_llm, sft_qwen2_vl
+from .pretrain import pretrain_gemma4_vl, pretrain_llava_onevision2, pretrain_llm, pretrain_qwen2_vl
+from .sft import sft_gemma4_vl, sft_llava_onevision1_5, sft_llava_onevision2, sft_llm, sft_qwen2_vl
 from .trainer_builder import build_model_trainer
 
 
diff --git a/aiak_training_llm/train/arguments.py b/aiak_training_llm/train/arguments.py
index 6d0aff9a..4930b8f7 100644
--- a/aiak_training_llm/train/arguments.py
+++ b/aiak_training_llm/train/arguments.py
@@ -126,6 +126,33 @@ def _add_extra_model_args(parser: argparse.ArgumentParser):
                        "this option, the head dimensions will be aligned by padding, so that fa can be used."
                        "Deprecated: use --attention-backend=flash")
 
+    # Gemma4 architecture fields. These are normally populated from the
+    # registered model config in _validate_extra_model_args; defining them here
+    # keeps fixed-architecture validation and TransformerConfig construction in
+    # sync with the Gemma4VLConfig dataclass.
+    group.add_argument('--final-logit-softcapping', type=float, default=None,
+                       help='Gemma4 final logits softcap. Usually set by --model-name.')
+    group.add_argument('--use-layer-scalar', action='store_true', default=False,
+                       help='Enable Gemma4 per-layer scalar buffers. Usually set by --model-name.')
+    group.add_argument('--sliding-window', type=int, default=None,
+                       help='Gemma4 sliding attention window. Usually set by --model-name.')
+    group.add_argument('--rotary-base-sliding', type=int, default=None,
+                       help='Gemma4 sliding-layer RoPE base. Usually set by --model-name.')
+    group.add_argument('--partial-rotary-factor', type=float, default=1.0,
+                       help='Gemma4 proportional RoPE factor. Usually set by --model-name.')
+    group.add_argument('--layer-pattern', default=[],
+                       help='Gemma4 per-layer attention pattern. Usually set by --model-name.')
+    group.add_argument('--per-layer-kv-channels', default={},
+                       help='Gemma4 per-layer-type head dim overrides. Usually set by --model-name.')
+    group.add_argument('--per-layer-num-query-groups', default={},
+                       help='Gemma4 per-layer-type KV head overrides. Usually set by --model-name.')
+    group.add_argument('--attention-k-eq-v', action='store_true', default=False,
+                       help='Enable Gemma4 K=V tying. Usually set by --model-name.')
+    group.add_argument('--kv-tied-layers', default=[],
+                       help='Gemma4 K=V tied layer indices. Usually set by --model-name.')
+    group.add_argument('--scale-emb-by-sqrt-hidden', action='store_true', default=False,
+                       help='Enable Gemma-style embedding scaling. Usually set by --model-name.')
+
     return parser
 
 
diff --git a/aiak_training_llm/train/pretrain/pretrain_gemma4_vl.py b/aiak_training_llm/train/pretrain/pretrain_gemma4_vl.py
new file mode 100644
index 00000000..2ef59bef
--- /dev/null
+++ b/aiak_training_llm/train/pretrain/pretrain_gemma4_vl.py
@@ -0,0 +1,406 @@
+"""default pretrain for generative models like GPTS"""
+
+import os
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from megatron.core import mpu, tensor_parallel
+from megatron.core.enums import ModelType
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.utils import StragglerDetector
+from megatron.training import get_timers
+
+from aiak_training_llm.data.multimodal.dataloader_provider import get_train_dataset, get_train_loader
+from aiak_training_llm.data.multimodal.gemma4_vl_task_encoder import Gemma4VLTaskEncoder
+from aiak_training_llm.models import get_model_family, get_model_provider
+from aiak_training_llm.models.gemma4_vl.gemma4_vl_attention_mask import (
+    build_gemma4_mm_attention_mask,
+)
+from aiak_training_llm.models.qwen_vl.utils import get_inputs_on_this_cp_rank
+from aiak_training_llm.train.megatron_trainer import MegatronTrainer
+from aiak_training_llm.train.sft.utils import build_sft_data_collator
+from aiak_training_llm.train.trainer_builder import register_model_trainer
+from aiak_training_llm.utils import constants, get_args
+from transformers import DataCollatorForSeq2Seq
+
+
+stimer = StragglerDetector()
+
+# Gemma4-26B-A4B special token ids (from HF config.json).
+image_token_id = 258880
+video_token_id = 258884
+vision_start_token_id = 255999
+
+
+def gemma4_vl_embedding_ranks(pp_ranks):
+    """gemma4_vl's embedding ranks consist of the decoder's first and last ranks (ie, the ViT has no embeddings).
+    Args:
+        pp_ranks: A list of global ranks that constitute a pipeline group.
+    """
+    args = get_args()
+
+    # encoder size is also the index to the first rank of the decoder.
+    epp = args.encoder_pipeline_model_parallel_size or 0
+
+    last_rank = pp_ranks[-1]
+    if len(pp_ranks) == 1 or pp_ranks[epp] == last_rank:
+        return [last_rank]
+    else:
+        return [pp_ranks[epp], last_rank]
+
+
+def gemma4_vl_position_embedding_ranks(pp_ranks):
+    """gemma4_vl's embedding ranks consist of the singular rank of the model or the decoder's first rank.
+    Args:
+        pp_ranks: A list of global ranks that constitute a pipeline group.
+    """
+    args = get_args()
+
+    # encoder size is also the index to the first rank of the decoder.
+    epp = args.encoder_pipeline_model_parallel_size or 0
+
+    last_rank = pp_ranks[-1]
+    if len(pp_ranks) == 1:
+        return [last_rank]
+    else:
+        return [pp_ranks[epp]]
+
+
+def model_provider(pre_process=True, post_process=True, add_encoder=True, add_decoder=True):
+    """Builds the model.
+
+    Args:
+        pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.
+        post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.
+
+    Returns:
+        MCoreModel: The returned model
+    """
+    args = get_args()
+    model_family = get_model_family(args.model_name)
+    model_provider = get_model_provider(model_family)
+    assert model_provider is not None, f"model provider for {args.model_name} not found"
+    return model_provider(pre_process, post_process, add_encoder, add_decoder)
+
+
+def get_batch(data_iterator):
+    """Generate a batch"""
+    args = get_args()
+    if data_iterator is not None and mpu.get_tensor_model_parallel_rank() == 0:
+        data = next(data_iterator)
+        if isinstance(data.get("tokens"), torch.Tensor):
+            orig_dtype = data["tokens"].dtype
+            if data["tokens"].dtype != torch.long:
+                print(
+                    f"[WARN] tokens dtype {orig_dtype} -> force cast to torch.long; shape={tuple(data['tokens'].shape)}"
+                )
+                data["tokens"] = data["tokens"].to(torch.long)
+        if isinstance(data.get("labels"), torch.Tensor) and data["labels"].dtype != torch.long:
+            data["labels"] = data["labels"].to(torch.long)
+
+        assert isinstance(data["tokens"], torch.Tensor) and data["tokens"].dtype == torch.long, (
+            f"Expected tokens torch.int64 but got {type(data['tokens'])} {getattr(data['tokens'], 'dtype', None)}"
+        )
+        assert isinstance(data["labels"], torch.Tensor) and data["labels"].dtype == torch.long, (
+            f"Expected labels torch.int64 but got {type(data['labels'])} {getattr(data['labels'], 'dtype', None)}"
+        )
+    else:
+        data = None
+
+    tokens = tensor_parallel.broadcast_data(["tokens"], data, torch.int64)["tokens"]
+    labels = tensor_parallel.broadcast_data(["labels"], data, torch.int64)["labels"]
+    attn_mask = tensor_parallel.broadcast_data(["attn_mask"], data, torch.bool)["attn_mask"]
+    cu_lengths = tensor_parallel.broadcast_data(["cu_lengths"], data, torch.int32)["cu_lengths"]
+    max_lengths = tensor_parallel.broadcast_data(["max_lengths"], data, torch.int32)["max_lengths"]
+
+    has_video = video_token_id in tokens
+    has_image = image_token_id in tokens
+    thw = None
+    video_grid_thw = None
+    imgs = None
+    pixel_values_videos = None
+    patch_positions = None
+    if has_image:
+        imgs = tensor_parallel.broadcast_data(["imgs"], data, torch.float32)["imgs"]
+        thw = tensor_parallel.broadcast_data(["image_grid_thw"], data, torch.int32)["image_grid_thw"]
+        # Synchronize whether patch_positions is available across all TP ranks.
+        # Only rank 0 of each TP group has `data`; other ranks have `data = None`.
+        # Using a collective (all_reduce MAX) ensures every rank agrees before
+        # deciding whether to call broadcast_data, preventing a collective mismatch
+        # that would otherwise corrupt the communication state and cause NaN loss.
+        has_pp = int(data is not None and "patch_positions" in data and data["patch_positions"] is not None)
+        has_pp_tensor = torch.tensor([has_pp], dtype=torch.int32, device=torch.cuda.current_device())
+        torch.distributed.all_reduce(
+            has_pp_tensor,
+            op=torch.distributed.ReduceOp.MAX,
+            group=mpu.get_tensor_model_parallel_group(),
+        )
+        if has_pp_tensor.item():
+            patch_positions = tensor_parallel.broadcast_data(["patch_positions"], data, torch.int64)["patch_positions"]
+    if has_video:
+        pixel_values_videos = tensor_parallel.broadcast_data(["pixel_values_videos"], data, torch.float32)[
+            "pixel_values_videos"
+        ]
+        video_grid_thw = tensor_parallel.broadcast_data(["video_grid_thw"], data, torch.int32)["video_grid_thw"]
+
+    packed_seq_params = None
+    is_video = video_token_id in tokens
+
+    attn_mask_type = AttnMaskType.padding_causal if attn_mask.any() else AttnMaskType.causal
+
+    labels = torch.roll(labels, shifts=-1, dims=1)
+    loss_mask = (labels != -100).long()
+
+    if cu_lengths.shape == torch.Size([1, 1]):
+        for i in range(attn_mask.shape[0]):
+            loss_mask[i, (attn_mask[i] == False).sum() - 1] = 0
+    else:
+        assert cu_lengths.shape[0] == 1, "micro-batch-size must be 1 for packing"
+        # for i in range(cu_lengths.shape[0]):
+        #     for j in range(1, cu_lengths[i].shape[0]):
+        #         loss_mask[i, cu_lengths[i][j] - 1] = 0
+
+        attn_mask = None
+        packed_seq_params = PackedSeqParams(
+            qkv_format="thd",
+            cu_seqlens_q=cu_lengths[0],
+            cu_seqlens_kv=cu_lengths[0],
+            max_seqlen_q=max_lengths[0].item(),
+            max_seqlen_kv=max_lengths[0].item(),
+        )
+
+    # Pad sequence for Sequence Parallelism (SP) divisibility.
+    # SP requires sequence length to be divisible by TP size;
+    # when combined with CP, it must be divisible by TP * CP * 2.
+    # This padding is required for both non-packed and packed sequences because
+    # scatter_to_sequence_parallel_region asserts divisibility by TP size.
+    if args.sequence_parallel:
+        tp_size = args.tensor_model_parallel_size
+        cp_size = args.context_parallel_size
+        if cp_size > 1:
+            padding_factor = tp_size * cp_size * 2
+        else:
+            padding_factor = tp_size
+        seq_len = tokens.shape[1]
+        if seq_len % padding_factor != 0:
+            padded_len = (seq_len + padding_factor - 1) // padding_factor * padding_factor
+            pad_size = padded_len - seq_len
+            pad_token_id = getattr(args, "pad_token_id", 0)
+            tokens = F.pad(tokens, (0, pad_size), value=pad_token_id)
+            labels = F.pad(labels, (0, pad_size), value=-100)
+            loss_mask = F.pad(loss_mask, (0, pad_size))
+            if packed_seq_params is not None:
+                # Append a dummy sequence entry so the packed-sequence attention
+                # kernel covers every element in the padded tensor.
+                # cu_seqlens ends with the original total token count; adding a
+                # new entry at (original_total + pad_size) represents the padding
+                # tokens as one extra sequence.
+                new_end_q = packed_seq_params.cu_seqlens_q[-1:] + pad_size
+                new_end_kv = packed_seq_params.cu_seqlens_kv[-1:] + pad_size
+                max_seqlen_q = max(packed_seq_params.max_seqlen_q, pad_size)
+                max_seqlen_kv = max(packed_seq_params.max_seqlen_kv, pad_size)
+                packed_seq_params = PackedSeqParams(
+                    qkv_format=packed_seq_params.qkv_format,
+                    cu_seqlens_q=torch.cat([packed_seq_params.cu_seqlens_q, new_end_q]),
+                    cu_seqlens_kv=torch.cat([packed_seq_params.cu_seqlens_kv, new_end_kv]),
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_kv=max_seqlen_kv,
+                )
+
+    if args.context_parallel_size > 1:
+        labels = get_inputs_on_this_cp_rank(labels.transpose(0, 1)).transpose(0, 1)
+        loss_mask = get_inputs_on_this_cp_rank(loss_mask.transpose(0, 1)).transpose(0, 1)
+
+    if packed_seq_params is None and (has_image or has_video):
+        mm_token_type_ids = torch.zeros_like(tokens, dtype=torch.int32)
+        mm_token_type_ids[tokens == image_token_id] = 1
+        mm_token_type_ids[tokens == video_token_id] = 2
+        sliding_window = int(getattr(args, "sliding_window", 1024) or 1024)
+        attn_mask = build_gemma4_mm_attention_mask(
+            tokens, mm_token_type_ids, sliding_window=sliding_window
+        )
+        attn_mask_type = AttnMaskType.causal
+    else:
+        attn_mask_type = AttnMaskType.causal
+        attn_mask = None
+    position_ids = None
+    return (
+        imgs,
+        thw,
+        pixel_values_videos,
+        video_grid_thw,
+        tokens,
+        position_ids,
+        attn_mask,
+        labels,
+        loss_mask,
+        attn_mask_type,
+        packed_seq_params,
+        patch_positions,
+    )
+
+
+def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
+    """Loss function.
+
+    Args:
+        loss_mask (torch.Tensor): Used to mask out some portions of the loss
+        output_tensor (torch.Tensor): The tensor with the losses
+
+    Returns:
+        the loss scalar for this micro-batch
+        the number of non-padded tokens in this microbatch
+        a dict containing reporting metrics on the loss and number of tokens across the data parallel ranks
+    """
+    args = get_args()
+
+    losses = output_tensor.float()
+    loss_mask = loss_mask.view(-1).float()
+
+    total_tokens = loss_mask.sum()
+    loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), total_tokens.view(1)])
+
+    if args.context_parallel_size > 1:
+        torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group())
+
+    # Check individual rank losses are not NaN prior to DP all-reduce.
+    if args.check_for_nan_in_loss_and_grad:
+        global_rank = torch.distributed.get_rank()
+        assert not loss[0].isnan(), (
+            f"Rank {global_rank}: found NaN in local forward loss calculation. "
+            f"Device: {torch.cuda.current_device()}, node: {os.uname()[1]}"
+        )
+
+    # Reduce loss for logging.
+    reporting_loss = loss.clone().detach()
+    torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group())
+
+    local_num_tokens = loss[1].clone().detach().to(torch.int)
+
+    loss_reduced_dict = {"lm loss": (reporting_loss[0], reporting_loss[1])}
+
+    if args.variable_seq_lengths:
+        # for variable seq length, we need to calculate the number of tokens on fly
+        # model output tensor shape is [B, S, H]
+        num_input_tokens = output_tensor.shape[0] * output_tensor.shape[1]
+        input_tokens = torch.tensor(num_input_tokens, dtype=torch.int, device=output_tensor.device)
+        # sum across all dp ranks
+        torch.distributed.all_reduce(input_tokens, group=mpu.get_data_parallel_group())
+        loss_reduced_dict["total_inputs"] = input_tokens.item() * args.context_parallel_size
+
+    return (loss[0] * args.context_parallel_size, local_num_tokens, loss_reduced_dict)
+
+
+def forward_step(data_iterator, model):
+    """Forward training step.
+
+    Args:
+        data_iterator : Input data iterator
+        model: Megatron Model
+    """
+    timers = get_timers()
+
+    # Get the batch.
+    timers("batch-generator", log_level=2).start()
+
+    global stimer
+    with stimer(bdata=True):
+        (
+            images,
+            image_grid_thw,
+            pixel_values_videos,
+            video_grid_thw,
+            input_ids,
+            position_ids,
+            attention_mask,
+            labels,
+            loss_mask,
+            attn_mask_type,
+            packed_seq_params,
+            patch_positions,
+        ) = get_batch(data_iterator)
+
+    timers("batch-generator").stop()
+
+    with stimer:
+        output_tensor = model(
+            images,
+            image_grid_thw,
+            input_ids,
+            position_ids,
+            attention_mask,
+            attn_mask_type,
+            labels,
+            packed_seq_params,
+            pixel_values_videos=pixel_values_videos,
+            video_grid_thw=video_grid_thw,
+            patch_positions=patch_positions,
+        )
+
+    # HOTFIX (text-only-sample backward): when a packed sample contains no
+    # images/video AND the LLM is frozen (stage-1: --trainable-modules adapter
+    # vision_model), the vision_model+adapter are never invoked, so output_tensor
+    # has no grad_fn and custom_backward (schedules.py:161) crashes — desyncing
+    # NCCL collectives. Fix: tie a zero-magnitude contribution from the first
+    # trainable param so backward succeeds with zero gradients and DP all-reduce
+    # stays aligned. Looks like a no-op — DO NOT REMOVE.
+    if torch.is_tensor(output_tensor) and not output_tensor.requires_grad:
+        dummy_param = None
+        for _p in model.parameters():
+            if _p.requires_grad:
+                dummy_param = _p
+                break
+        if dummy_param is not None:
+            _rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+            _host = os.uname()[1]
+            print(
+                f"[hotfix-zero-grad] rank={_rank} host={_host} "
+                f"output_tensor.requires_grad=False -> tying dummy zero loss to "
+                f"trainable param shape={tuple(dummy_param.shape)} dtype={dummy_param.dtype}",
+                flush=True,
+            )
+            output_tensor = output_tensor + (dummy_param.sum() * 0.0).to(output_tensor.dtype)
+        else:
+            print(
+                "[hotfix-zero-grad] WARNING: output_tensor has no grad_fn AND no "
+                "trainable parameters found in model — backward will fail.",
+                flush=True,
+            )
+
+    return output_tensor, partial(loss_func, loss_mask)
+
+
+def train_valid_test_dataset_provider(train_val_test_num_samples):
+    """Provides the datasets used by the trainer"""
+
+    args = get_args()
+    task_encoder = Gemma4VLTaskEncoder(args)
+    train_dataset = get_train_dataset(task_encoder)
+    collator = build_sft_data_collator(DataCollatorForSeq2Seq)
+    train_dataloader = get_train_loader(train_dataset, collator)
+    return train_dataloader, None, None
+
+
+@register_model_trainer(
+    model_family=[constants.VisionLanguageModelFamilies.GEMMA4_VL],
+    training_phase=constants.TrainingPhase.PRETRAIN,
+)
+def default_pretrain_trainer(train_args):
+    """build trainer"""
+    if train_args.encoder_pipeline_model_parallel_size in [None, 0]:
+        model_type = ModelType.encoder_or_decoder
+    else:
+        model_type = ModelType.encoder_and_decoder
+    trainer = MegatronTrainer(
+        train_args=train_args,
+        train_valid_test_dataset_provider=train_valid_test_dataset_provider,
+        model_provider=model_provider,
+        model_type=model_type,
+        forward_step_func=forward_step,
+        get_embedding_ranks=gemma4_vl_embedding_ranks,
+        get_position_embedding_ranks=gemma4_vl_position_embedding_ranks,
+    )
+
+    return trainer
diff --git a/aiak_training_llm/train/sft/sft_gemma4_vl.py b/aiak_training_llm/train/sft/sft_gemma4_vl.py
new file mode 100644
index 00000000..0f96654f
--- /dev/null
+++ b/aiak_training_llm/train/sft/sft_gemma4_vl.py
@@ -0,0 +1,25 @@
+from megatron.core.enums import ModelType
+
+from aiak_training_llm.train.megatron_trainer import MegatronTrainer
+from aiak_training_llm.train.trainer_builder import register_model_trainer
+from aiak_training_llm.utils.constants import TrainingPhase, VisionLanguageModelFamilies
+
+
+@register_model_trainer(model_family=[VisionLanguageModelFamilies.GEMMA4_VL], training_phase=TrainingPhase.SFT)
+def default_pretrain_trainer(train_args):
+    """build trainer"""
+    from aiak_training_llm.train.pretrain import pretrain_gemma4_vl
+
+    if train_args.encoder_pipeline_model_parallel_size in [0, None]:
+        model_type = ModelType.encoder_or_decoder
+    else:
+        model_type = ModelType.encoder_and_decoder
+    trainer = MegatronTrainer(
+        train_args=train_args,
+        train_valid_test_dataset_provider=pretrain_gemma4_vl.train_valid_test_dataset_provider,
+        model_provider=pretrain_gemma4_vl.model_provider,
+        model_type=model_type,
+        forward_step_func=pretrain_gemma4_vl.forward_step,
+    )
+
+    return trainer
diff --git a/aiak_training_llm/utils/constants.py b/aiak_training_llm/utils/constants.py
index eefdfe64..a822e581 100644
--- a/aiak_training_llm/utils/constants.py
+++ b/aiak_training_llm/utils/constants.py
@@ -94,3 +94,4 @@ class VisionLanguageModelFamilies(_BaseFamilies):
     QWEN2_5_VL = "qwen2_5_vl"
     LLAVA_ONEVISION1_5 = "llava_onevision1_5"
     LLAVA_ONEVISION2 = "llava_onevision2"
+    GEMMA4_VL = "gemma4_vl"
diff --git a/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_hf_to_mcore.sh b/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_hf_to_mcore.sh
new file mode 100644
index 00000000..9bc7af1a
--- /dev/null
+++ b/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_hf_to_mcore.sh
@@ -0,0 +1,138 @@
+# =============================================================================
+# Gemma4-VL 26B-A4B – Convert HuggingFace checkpoint to Megatron-Core
+# =============================================================================
+#
+# Usage:
+#   bash convert_gemma4_26b_a4b_hf_to_mcore.sh <LOAD> <SAVE> <PP> <EP>
+#   bash convert_gemma4_26b_a4b_hf_to_mcore.sh <LOAD> <SAVE> <PP> <EP> <CUSTOM_PIPELINE_LAYERS>
+#   bash convert_gemma4_26b_a4b_hf_to_mcore.sh <LOAD> <SAVE> <TP> <PP> <EP>
+#   bash convert_gemma4_26b_a4b_hf_to_mcore.sh <LOAD> <SAVE> <TP> <PP> <EP> <CUSTOM_PIPELINE_LAYERS>
+#
+# Arguments:
+#   LOAD  Path to the source HuggingFace checkpoint
+#   SAVE  Path to save the Megatron-Core checkpoint
+#   TP    Tensor parallel size (optional, defaults to 1)
+#   PP    Pipeline parallel size (recommended 1 for this script)
+#   EP    Expert parallel size
+#   CUSTOM_PIPELINE_LAYERS  (optional) Comma-separated layer counts per PP stage,
+#                           used for custom PP layer layouts.
+#
+# Notes:
+#   * Gemma4-VL uses standalone Python converters under
+#     tools/convert_checkpoint/custom/gemma4_vl/ for both LLM and ViT
+#     (no JSON config for those two; ViT/LLM read num_layers/num_experts
+#     directly from the HuggingFace config.json).
+#   * patch + adapter still use small JSON configs under
+#     tools/convert_checkpoint/config/gemma4-26b-a4b/.
+#   * P2.5 ships with TP=PP=EP=1 verified end to end (16680 keys).
+#   * P2.7 adds TP/EP sharding support: LLM emits 2D mp_rank_TT_EEE when
+#     EP>1 (Qwen3 convention), and ViT/adapter/patch always emit 1D
+#     mp_rank_TT (broadcast across the EP axis at merge time).
+#   * P2.8 adds PP sharding support: LLM emits 3D mp_rank_TT_PPP_EEE
+#     when PP>1 (any EP, including EP=1, Qwen3 convention), with layers
+#     locally renumbered (0..N_local-1) per PP rank. ViT/adapter/patch
+#     still emit 1D mp_rank_TT and the merger broadcasts them across
+#     EP and into PP stage 0 only.
+# =============================================================================
+
+AIAK_TRAINING_PATH="${AIAK_TRAINING_PATH:-/workspace/LLaVA-OneVision-2}"
+AIAK_MAGATRON_PATH="${AIAK_MAGATRON_PATH:-${AIAK_TRAINING_PATH%/}/aiak_megatron}"
+CONVERT_CHECKPOINT_PATH="$AIAK_TRAINING_PATH/tools/convert_checkpoint"
+
+LOAD=$1
+SAVE=$2
+CUSTOM_PIPELINE_LAYERS=
+
+if [[ $# -eq 4 ]]; then
+    TP=1
+    PP=$3
+    EP=$4
+elif [[ $# -eq 5 ]]; then
+    if [[ "$5" == *","* ]]; then
+        TP=1
+        PP=$3
+        EP=$4
+        CUSTOM_PIPELINE_LAYERS=$5
+    else
+        TP=${3:-1}
+        PP=${4:-1}
+        EP=${5:-1}
+    fi
+elif [[ $# -ge 6 ]]; then
+    TP=${3:-1}
+    PP=${4:-1}
+    EP=${5:-1}
+    CUSTOM_PIPELINE_LAYERS=$6
+else
+    TP=${3:-1}
+    PP=${4:-1}
+    EP=${5:-1}
+fi
+
+mkdir -p ./tmp/
+SAVE_LANGUAGE_MODEL=./tmp/language-mcore
+SAVE_VISION_MODEL=./tmp/vision-model-mcore
+SAVE_ADAPTER=./tmp/adapter-mcore
+SAVE_PATCH=./tmp/patch-mcore
+
+
+# llm (moe) – standalone converter, no JSON config
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/llm.py \
+    --load_platform=huggingface \
+    --save_platform=mcore \
+    --tensor_model_parallel_size=$TP \
+    --pipeline_model_parallel_size=$PP \
+    --expert_parallel_size=$EP \
+    ${CUSTOM_PIPELINE_LAYERS:+--custom_pipeline_layers=$CUSTOM_PIPELINE_LAYERS} \
+    --load_ckpt_path=$LOAD \
+    --save_ckpt_path=$SAVE_LANGUAGE_MODEL \
+    --safetensors \
+    --no_save_optim \
+    --no_load_optim
+
+# vit – standalone converter, no JSON config
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/vit.py \
+    --load_platform=huggingface \
+    --save_platform=mcore \
+    --tensor_model_parallel_size=$TP \
+    --load_ckpt_path=$LOAD \
+    --save_ckpt_path=$SAVE_VISION_MODEL \
+    --safetensors \
+    --no_save_optim \
+    --no_load_optim
+
+# adapter
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/adapter.py \
+    --load_platform=huggingface \
+    --save_platform=mcore \
+    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/gemma4-26b-a4b/adapter.json \
+    --tensor_model_parallel_size=$TP \
+    --load_ckpt_path=$LOAD \
+    --save_ckpt_path=$SAVE_ADAPTER
+
+# vision patch in vit
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/vision_patch.py \
+    --load_platform=huggingface \
+    --save_platform=mcore \
+    --tensor_model_parallel_size=$TP \
+    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/gemma4-26b-a4b/vision-patch.json \
+    --load_ckpt_path=$LOAD \
+    --save_ckpt_path=$SAVE_PATCH
+
+# merge (1D mp_rank, no TP/EP sharding at P2.5)
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/merge_megatron.py \
+    --megatron_path $AIAK_MAGATRON_PATH \
+    --language_model_path $SAVE_LANGUAGE_MODEL/release \
+    --vision_model_path $SAVE_VISION_MODEL/release \
+    --vision_patch $SAVE_PATCH/release \
+    --adapter_path $SAVE_ADAPTER/release \
+    --save_ckpt_path $SAVE/release \
+    --tensor_model_parallel_size=$TP \
+    --pipeline_model_parallel_size=$PP \
+    --expert_parallel_size=$EP
+
+echo release > $SAVE/latest_checkpointed_iteration.txt
+rm -rf $SAVE_LANGUAGE_MODEL
+rm -rf $SAVE_VISION_MODEL
+rm -rf $SAVE_ADAPTER
+rm -rf $SAVE_PATCH
diff --git a/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_mcore_to_hf.sh b/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_mcore_to_hf.sh
new file mode 100644
index 00000000..ba715c80
--- /dev/null
+++ b/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_mcore_to_hf.sh
@@ -0,0 +1,186 @@
+# =============================================================================
+# Gemma4-VL 26B-A4B – Convert Megatron-Core checkpoint to HuggingFace
+# =============================================================================
+#
+# Usage:
+#   HF_REF=/path/to/source/hf bash convert_gemma4_26b_a4b_mcore_to_hf.sh <LOAD> <SAVE> <PP> <EP>
+#   HF_REF=/path/to/source/hf bash convert_gemma4_26b_a4b_mcore_to_hf.sh <LOAD> <SAVE> <PP> <EP> <CUSTOM_PIPELINE_LAYERS>
+#   HF_REF=/path/to/source/hf bash convert_gemma4_26b_a4b_mcore_to_hf.sh <LOAD> <SAVE> <TP> <PP> <EP>
+#   HF_REF=/path/to/source/hf bash convert_gemma4_26b_a4b_mcore_to_hf.sh <LOAD> <SAVE> <TP> <PP> <EP> <CUSTOM_PIPELINE_LAYERS>
+#
+# Arguments:
+#   LOAD  Path to the source Megatron-Core checkpoint (parent of `release/`,
+#         the path produced by convert_gemma4_26b_a4b_hf_to_mcore.sh).
+#   SAVE  Path to save the HuggingFace checkpoint.
+#   TP    Tensor parallel size (optional, defaults to 1).
+#   PP    Pipeline parallel size.
+#   EP    Expert parallel size.
+#   CUSTOM_PIPELINE_LAYERS  (optional) Comma-separated layer counts per PP stage.
+#
+# Required environment:
+#   HF_REF  Path to the original HuggingFace checkpoint directory. Needed
+#           because the standalone llm.py / vit.py read num_layers,
+#           num_experts, layer_types and other architecture metadata directly
+#           from the source `config.json` (data-driven, no JSON template).
+#           The shell wrapper seeds this `config.json` into the per-component
+#           save dirs before running each converter.
+#
+# Notes:
+#   * Gemma4-VL LLM round-trip is bit-equivalent (657/657 tensors verified at
+#     P2.5). ViT/adapter/patch MG->HF conversion paths are wired through the
+#     same standalone scripts.
+#   * The final HF dict is assembled by reusing the OV2 generic
+#     custom/llava_onevision2/merge_huggingface.py which simply merges 4
+#     state-dicts by key (model-agnostic).
+#   * P2.8: when PP>1, the LLM ckpt uses 3D mp_rank_TT_PPP_EEE naming; vit /
+#     adapter / vision_patch live solely on PP stage 0, so we materialise a
+#     single 1D-named mp_rank_TT view of stage-0 EP-rank-0 shards and feed it
+#     to all three of those standalone converters.
+# =============================================================================
+
+if [[ -z "$HF_REF" ]]; then
+    echo "ERROR: HF_REF environment variable must be set to the source HuggingFace dir" >&2
+    echo "  Example: HF_REF=/ov2/pretrain_models/google/gemma-4-26B-A4B-it bash $0 ..." >&2
+    exit 1
+fi
+if [[ ! -f "$HF_REF/config.json" ]]; then
+    echo "ERROR: HF_REF=$HF_REF does not contain config.json" >&2
+    exit 1
+fi
+
+AIAK_TRAINING_PATH="${AIAK_TRAINING_PATH:-/workspace/LLaVA-OneVision-2}"
+AIAK_MAGATRON_PATH="${AIAK_MAGATRON_PATH:-${AIAK_TRAINING_PATH%/}/aiak_megatron}"
+CONVERT_CHECKPOINT_PATH="$AIAK_TRAINING_PATH/tools/convert_checkpoint"
+
+LOAD=$1
+SAVE=$2
+CUSTOM_PIPELINE_LAYERS=
+
+if [[ $# -eq 4 ]]; then
+    TP=1
+    PP=$3
+    EP=$4
+elif [[ $# -eq 5 ]]; then
+    if [[ "$5" == *","* ]]; then
+        TP=1
+        PP=$3
+        EP=$4
+        CUSTOM_PIPELINE_LAYERS=$5
+    else
+        TP=${3:-1}
+        PP=${4:-1}
+        EP=${5:-1}
+    fi
+elif [[ $# -ge 6 ]]; then
+    TP=${3:-1}
+    PP=${4:-1}
+    EP=${5:-1}
+    CUSTOM_PIPELINE_LAYERS=$6
+else
+    TP=${3:-1}
+    PP=${4:-1}
+    EP=${5:-1}
+fi
+
+mkdir -p ./tmp/
+SAVE_LANGUAGE_MODEL=./tmp/language-hf
+SAVE_VISION_MODEL=./tmp/vision-model-hf
+SAVE_ADAPTER=./tmp/adapter-hf
+SAVE_PATCH=./tmp/patch-hf
+
+mkdir -p $SAVE_LANGUAGE_MODEL $SAVE_VISION_MODEL
+cp $HF_REF/config.json $SAVE_LANGUAGE_MODEL/config.json
+cp $HF_REF/config.json $SAVE_VISION_MODEL/config.json
+
+if [[ -d "$LOAD/release" ]]; then
+    MCORE_LOAD=$LOAD/release
+else
+    MCORE_LOAD=$LOAD
+fi
+
+
+if [[ $PP -eq 1 ]]; then
+    NONLLM_LOAD=$MCORE_LOAD
+    NONLLM_EP=$EP
+else
+    NONLLM_LOAD=$MCORE_LOAD/tmp_nonllm_view/
+    NONLLM_EP=1
+    rm -rf $NONLLM_LOAD
+    mkdir -p $NONLLM_LOAD
+    for ((i=0;i<$TP;i++)); do
+        from=$(printf "mp_rank_%02d_000_000" $i)
+        to=$(printf "mp_rank_%02d" $i)
+        cp -r $MCORE_LOAD/$from $NONLLM_LOAD/$to
+    done
+fi
+
+
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/llm.py \
+    --load_platform=mcore \
+    --save_platform=huggingface \
+    --megatron_path $AIAK_MAGATRON_PATH \
+    --tensor_model_parallel_size=$TP \
+    --pipeline_model_parallel_size=$PP \
+    --expert_parallel_size=$EP \
+    ${CUSTOM_PIPELINE_LAYERS:+--custom_pipeline_layers=$CUSTOM_PIPELINE_LAYERS} \
+    --load_ckpt_path=$MCORE_LOAD \
+    --save_ckpt_path=$SAVE_LANGUAGE_MODEL \
+    --safetensors \
+    --no_save_optim \
+    --no_load_optim
+
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/vit.py \
+    --load_platform=mcore \
+    --save_platform=huggingface \
+    --megatron_path $AIAK_MAGATRON_PATH \
+    --tensor_model_parallel_size=$TP \
+    --pipeline_model_parallel_size=1 \
+    --expert_parallel_size=$NONLLM_EP \
+    --load_ckpt_path=$NONLLM_LOAD \
+    --save_ckpt_path=$SAVE_VISION_MODEL \
+    --safetensors \
+    --no_save_optim \
+    --no_load_optim
+
+
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/adapter.py \
+    --load_platform=mcore \
+    --save_platform=huggingface \
+    --megatron_path $AIAK_MAGATRON_PATH \
+    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/gemma4-26b-a4b/adapter.json \
+    --tensor_model_parallel_size=$TP \
+    --pipeline_model_parallel_size=1 \
+    --expert_parallel_size=$NONLLM_EP \
+    --load_ckpt_path=$NONLLM_LOAD \
+    --save_ckpt_path=$SAVE_ADAPTER
+
+python $CONVERT_CHECKPOINT_PATH/custom/gemma4_vl/vision_patch.py \
+    --load_platform=mcore \
+    --save_platform=huggingface \
+    --megatron_path $AIAK_MAGATRON_PATH \
+    --tensor_model_parallel_size=$TP \
+    --pipeline_model_parallel_size=1 \
+    --expert_parallel_size=$NONLLM_EP \
+    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/gemma4-26b-a4b/vision-patch.json \
+    --load_ckpt_path=$NONLLM_LOAD \
+    --save_ckpt_path=$SAVE_PATCH
+
+if [[ $MCORE_LOAD != $NONLLM_LOAD ]]; then
+    rm -rf $NONLLM_LOAD
+fi
+
+# merge – reuse OV2 generic 4-dict merger (model-agnostic)
+python $CONVERT_CHECKPOINT_PATH/custom/llava_onevision2/merge_huggingface.py \
+    --megatron_path $AIAK_MAGATRON_PATH \
+    --language_model_path $SAVE_LANGUAGE_MODEL \
+    --vision_model_path $SAVE_VISION_MODEL \
+    --vision_patch $SAVE_PATCH \
+    --adapter_path $SAVE_ADAPTER \
+    --save_ckpt_path $SAVE
+
+cp $HF_REF/config.json $SAVE/config.json
+
+rm -rf $SAVE_LANGUAGE_MODEL
+rm -rf $SAVE_VISION_MODEL
+rm -rf $SAVE_ADAPTER
+rm -rf $SAVE_PATCH
diff --git a/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_mcore_to_release.sh b/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_mcore_to_release.sh
new file mode 100644
index 00000000..d95d16e2
--- /dev/null
+++ b/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_mcore_to_release.sh
@@ -0,0 +1,74 @@
+# =============================================================================
+# Gemma4-VL 26B-A4B – Re-shard Megatron-Core checkpoint (mcore -> mcore)
+# =============================================================================
+#
+# Usage (target layout via positional args; source layout via env vars):
+#   HF_REF=/path/to/source/hf bash convert_gemma4_26b_a4b_mcore_to_release.sh <LOAD> <SAVE> <PP> <EP>
+#   HF_REF=... SRC_TP=1 SRC_PP=1 SRC_EP=1 bash ... <LOAD> <SAVE> <TP> <PP> <EP>
+#   HF_REF=... SRC_TP=1 SRC_PP=1 SRC_EP=1 bash ... <LOAD> <SAVE> <TP> <PP> <EP> <CUSTOM_PIPELINE_LAYERS>
+#
+# Required environment:
+#   HF_REF   Path to the original HuggingFace checkpoint directory (forwarded
+#            to mcore_to_hf.sh; see that script for why it is required).
+#
+# Optional environment (source mcore layout, defaults TP=PP=EP=1):
+#   SRC_TP   Source tensor parallel size      (default 1)
+#   SRC_PP   Source pipeline parallel size    (default 1)
+#   SRC_EP   Source expert parallel size      (default 1)
+#
+# Implementation: round-trips through HF as an intermediate format, reusing the
+# two sibling scripts (mcore_to_hf -> hf_to_mcore). The intermediate HF
+# directory tmp_hf is removed at the end. The mcore_to_hf step is invoked
+# with the SOURCE layout (so the standalone converters select the correct
+# loader); the hf_to_mcore step is invoked with the TARGET layout.
+# =============================================================================
+
+if [[ -z "$HF_REF" ]]; then
+    echo "ERROR: HF_REF environment variable must be set to the source HuggingFace dir" >&2
+    echo "  Example: HF_REF=/ov2/pretrain_models/google/gemma-4-26B-A4B-it bash $0 ..." >&2
+    exit 1
+fi
+
+AIAK_TRAINING_PATH="${AIAK_TRAINING_PATH:-/workspace/LLaVA-OneVision-2}"
+
+LOAD=$1
+SAVE=$2
+CUSTOM_PIPELINE_LAYERS=
+
+if [[ $# -eq 4 ]]; then
+    TP=1
+    PP=$3
+    EP=$4
+elif [[ $# -eq 5 ]]; then
+    if [[ "$5" == *","* ]]; then
+        TP=1
+        PP=$3
+        EP=$4
+        CUSTOM_PIPELINE_LAYERS=$5
+    else
+        TP=${3:-1}
+        PP=${4:-1}
+        EP=${5:-1}
+    fi
+elif [[ $# -ge 6 ]]; then
+    TP=${3:-1}
+    PP=${4:-1}
+    EP=${5:-1}
+    CUSTOM_PIPELINE_LAYERS=$6
+else
+    TP=${3:-1}
+    PP=${4:-1}
+    EP=${5:-1}
+fi
+
+SRC_TP="${SRC_TP:-1}"
+SRC_PP="${SRC_PP:-1}"
+SRC_EP="${SRC_EP:-1}"
+
+bash $AIAK_TRAINING_PATH/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_mcore_to_hf.sh \
+    $LOAD tmp_hf $SRC_TP $SRC_PP $SRC_EP $CUSTOM_PIPELINE_LAYERS
+
+bash $AIAK_TRAINING_PATH/examples/gemma4_vl/convert/convert_gemma4_26b_a4b_hf_to_mcore.sh \
+    tmp_hf $SAVE $TP $PP $EP $CUSTOM_PIPELINE_LAYERS
+
+rm -rf tmp_hf
diff --git a/examples/gemma4_vl/quick_start_26b_a4b/stage1_projector.sh b/examples/gemma4_vl/quick_start_26b_a4b/stage1_projector.sh
new file mode 100644
index 00000000..f490e38b
--- /dev/null
+++ b/examples/gemma4_vl/quick_start_26b_a4b/stage1_projector.sh
@@ -0,0 +1,193 @@
+# =============================================================================
+# Gemma4-VL 26B-A4B - Stage-1 Projector (Adapter-only) Training
+# =============================================================================
+# Plan v5 default recipe:
+#   TP=4, PP=2, EP=4, EPP=1, SEQ_LEN=8192, MBS=1, GBS=128
+#   --trainable-modules adapter (vision_model + LLM frozen)
+#   --moe-router-topk 8, --num-experts 128, --moe-aux-loss-coeff 1e-3
+#
+# Smoke override (single 8x80GB node):
+#   TP=2 PP=1 EP=4 SEQ_LEN=4096 MBS=1 GBS=8 NSTEP=100
+#
+# Usage:
+#   bash stage1_projector.sh [TP] [PP] [EP] [SEQ_LEN] [MBS] [GBS] [NSTEP]
+# =============================================================================
+
+TP="${1:-4}"
+PP="${2:-2}"
+EP="${3:-4}"
+SEQ_LEN="${4:-8192}"
+MBS="${5:-1}"
+GBS="${6:-128}"
+NSTEP="${7:-500}"
+
+AIAK_TRAINING_PATH="${AIAK_TRAINING_PATH:-/workspace/LLaVA-OneVision-2}"
+AIAK_MAGATRON_PATH="${AIAK_MAGATRON_PATH:-${AIAK_TRAINING_PATH%/}/aiak_megatron}"
+
+OUTPUT_DIR="${OUTPUT_DIR:-/ov2/xiangan/ckpts_gemma4_26b_a4b}"
+DATA_PATH=${DATA_PATH:-"/workspace/dataset/LLaVA-558K-Webdataset"}
+TOKENIZER_PATH=${TOKENIZER_PATH:-"/ov2/pretrain_models/google/gemma-4-26B-A4B-it"}
+CHECKPOINT_PATH=${CHECKPOINT_PATH:-"/workspace/LLaVA-OneVision-2/stage_0_gemma4_26b_a4b_release"}
+
+#! /bin/bash
+# The script needs to be run on at least 1 nodes.
+
+# --- Multi-node configuration ---
+declare -a list_ip=(
+    127.0.0.1
+)
+
+CURRENT_IP=$(hostname -I | awk '{print $1}')
+if [ -z "$CURRENT_IP" ]; then
+    CURRENT_IP=$(hostname -i 2>/dev/null | awk '{print $1}')
+fi
+
+SINGLE_NODE=0
+if [[ ${#list_ip[@]} -eq 1 && ( "${list_ip[0]}" == "localhost" || "${list_ip[0]}" == "127.0.0.1" ) ]]; then
+    SINGLE_NODE=1
+fi
+
+NNODES=${#list_ip[@]}
+MASTER_ADDR=${list_ip[0]}
+
+if [[ $SINGLE_NODE -eq 1 ]]; then
+    NNODES=1
+    MASTER_ADDR=127.0.0.1
+    NODE_RANK=0
+else
+    NODE_RANK=-1
+    for i in "${!list_ip[@]}"; do
+        if [[ "${list_ip[$i]}" == "${CURRENT_IP}" ]]; then
+            NODE_RANK=$i
+            break
+        fi
+    done
+    if [ "$NODE_RANK" -eq -1 ]; then
+        echo "Error: Current IP ($CURRENT_IP) not found in list_ip."
+        exit 1
+    fi
+fi
+# --- End of Multi-node configuration ---
+
+SAVE_CKPT_PATH=$OUTPUT_DIR/$(basename "$0" .sh)
+TENSORBOARD_PATH="${SAVE_CKPT_PATH}/tensorboard"
+
+mkdir -p "$SAVE_CKPT_PATH"
+mkdir -p "$TENSORBOARD_PATH"
+mkdir -p "$SAVE_CKPT_PATH/dataloader"
+GPUS_PER_NODE=8
+
+MASTER_ADDR=${MASTER_ADDR:-"${list_ip[0]}"}
+MASTER_PORT=${MASTER_PORT:-"26000"}
+
+if [[ $SINGLE_NODE -eq 1 ]]; then
+    DISTRIBUTED_ARGS=(
+        --nproc_per_node "$GPUS_PER_NODE"
+    )
+else
+    DISTRIBUTED_ARGS=(
+        --nproc_per_node "$GPUS_PER_NODE"
+        --nnodes "$NNODES"
+        --node_rank "$NODE_RANK"
+        --master_addr "$MASTER_ADDR"
+        --master_port "$MASTER_PORT"
+    )
+fi
+
+MODEL_ARGS=(
+    --model-name gemma4-26b-a4b-vl
+)
+
+DATA_ARGS=(
+    --tokenizer-type HFTokenizer
+    --hf-tokenizer-path "$TOKENIZER_PATH"
+    --data-path "$DATA_PATH"
+    --dataloader-type external
+    --split 100,0,0
+    --num-workers 4
+)
+
+TRAINING_ARGS=(
+    --training-phase sft
+    --chat-template gemma4
+    --trainable-modules adapter
+    --seq-length "${SEQ_LEN}"
+    --max-position-embeddings "${SEQ_LEN}"
+    --init-method-std 0.02
+    --micro-batch-size "${MBS}"
+    --global-batch-size "${GBS}"
+    --lr 1.0e-4
+    --min-lr 1.0e-6
+    --clip-grad 1.0
+    --weight-decay 0
+    --optimizer adam
+    --adam-beta1 0.9
+    --adam-beta2 0.99
+    --adam-eps 1e-05
+    --norm-epsilon 1e-6
+    --train-iters "$NSTEP"
+    --lr-decay-iters "$NSTEP"
+    --lr-decay-style cosine
+    --lr-warmup-fraction 0.002
+    --initial-loss-scale 65536
+    --bf16
+    --load "$CHECKPOINT_PATH"
+    --save "$SAVE_CKPT_PATH"
+    --save-interval 100
+    --ckpt-format torch
+    --dataloader-save "${SAVE_CKPT_PATH}/dataloader"
+
+    --ckpt-fully-parallel-load
+    --recompute-granularity full
+    --recompute-method uniform
+    --recompute-num-layers 4
+)
+
+MODEL_PARALLEL_ARGS=(
+    --attention-backend flash
+    --moe-token-dispatcher-type alltoall
+    --moe-per-layer-logging
+    --pipeline-model-parallel-size "${PP}"
+    --tensor-model-parallel-size "${TP}"
+    --expert-model-parallel-size "${EP}"
+    --encoder-pipeline-model-parallel-size 1
+    --num-experts 128
+    --use-distributed-optimizer
+    --distributed-backend nccl
+)
+
+MOE_ARGS=(
+    --moe-router-topk 8
+    --moe-aux-loss-coeff 1e-3
+    --moe-router-dtype fp32
+)
+
+LOGGING_ARGS=(
+    --log-interval 1
+    --tensorboard-dir "${TENSORBOARD_PATH}"
+    --log-timers-to-tensorboard
+)
+
+if [ -n "${WANDB_API_KEY}" ]; then
+    LOGGING_ARGS+=(
+        --wandb-project "${WANDB_PROJECT}"
+        --wandb-exp-name "${WANDB_NAME}"
+    )
+fi
+
+TM=$(date "+%Y-%m-%d_%H:%M:%S")
+logfile="${SAVE_CKPT_PATH}/run_${TM}_tp${TP}_pp${PP}_ep${EP}_seqlen${SEQ_LEN}_mbs${MBS}_gbs${GBS}_${NSTEP}steps.log"
+
+export OFFLINE_PACKING_BMR=1
+
+PYTHONPATH="$AIAK_MAGATRON_PATH:$AIAK_TRAINING_PATH:$PYTHONPATH" \
+    torchrun "${DISTRIBUTED_ARGS[@]}" \
+    "$AIAK_TRAINING_PATH/aiak_training_llm/train.py" \
+    "${MODEL_ARGS[@]}" \
+    "${DATA_ARGS[@]}" \
+    ${IMG_ARGS:+${IMG_ARGS[@]}} \
+    "${TRAINING_ARGS[@]}" \
+    "${MODEL_PARALLEL_ARGS[@]}" \
+    "${MOE_ARGS[@]}" \
+    "${LOGGING_ARGS[@]}" \
+    2>&1 | tee "$logfile"
diff --git a/examples/gemma4_vl/quick_start_26b_a4b/stage2_full.sh b/examples/gemma4_vl/quick_start_26b_a4b/stage2_full.sh
new file mode 100644
index 00000000..1a851236
--- /dev/null
+++ b/examples/gemma4_vl/quick_start_26b_a4b/stage2_full.sh
@@ -0,0 +1,194 @@
+# =============================================================================
+# Gemma4-VL 26B-A4B - Stage-2 Full (Vision + Adapter + LLM) Training
+# =============================================================================
+# Plan v5 default recipe:
+#   TP=4, PP=2, EP=4, EPP=1, SEQ_LEN=16384, MBS=1, GBS=128
+#   --trainable-modules all (vision_model + adapter + LLM all trainable)
+#   --moe-router-topk 8, --num-experts 128, --moe-aux-loss-coeff 1e-3
+#   --lr 1e-5 (lower than stage-1 since LLM is unfrozen)
+#
+# Smoke override (single 8x80GB node):
+#   TP=2 PP=1 EP=4 SEQ_LEN=4096 MBS=1 GBS=8 NSTEP=100
+#
+# Usage:
+#   bash stage2_full.sh [TP] [PP] [EP] [SEQ_LEN] [MBS] [GBS] [NSTEP]
+# =============================================================================
+
+TP="${1:-4}"
+PP="${2:-2}"
+EP="${3:-4}"
+SEQ_LEN="${4:-16384}"
+MBS="${5:-1}"
+GBS="${6:-128}"
+NSTEP="${7:-500}"
+
+AIAK_TRAINING_PATH="${AIAK_TRAINING_PATH:-/workspace/LLaVA-OneVision-2}"
+AIAK_MAGATRON_PATH="${AIAK_MAGATRON_PATH:-${AIAK_TRAINING_PATH%/}/aiak_megatron}"
+
+OUTPUT_DIR="${OUTPUT_DIR:-/ov2/xiangan/ckpts_gemma4_26b_a4b}"
+DATA_PATH=${DATA_PATH:-"/workspace/dataset/LLaVA-558K-Webdataset"}
+TOKENIZER_PATH=${TOKENIZER_PATH:-"/ov2/pretrain_models/google/gemma-4-26B-A4B-it"}
+CHECKPOINT_PATH=${CHECKPOINT_PATH:-"/workspace/LLaVA-OneVision-2/stage_0_gemma4_26b_a4b_release"}
+
+#! /bin/bash
+# The script needs to be run on at least 1 nodes.
+
+# --- Multi-node configuration ---
+declare -a list_ip=(
+    127.0.0.1
+)
+
+CURRENT_IP=$(hostname -I | awk '{print $1}')
+if [ -z "$CURRENT_IP" ]; then
+    CURRENT_IP=$(hostname -i 2>/dev/null | awk '{print $1}')
+fi
+
+SINGLE_NODE=0
+if [[ ${#list_ip[@]} -eq 1 && ( "${list_ip[0]}" == "localhost" || "${list_ip[0]}" == "127.0.0.1" ) ]]; then
+    SINGLE_NODE=1
+fi
+
+NNODES=${#list_ip[@]}
+MASTER_ADDR=${list_ip[0]}
+
+if [[ $SINGLE_NODE -eq 1 ]]; then
+    NNODES=1
+    MASTER_ADDR=127.0.0.1
+    NODE_RANK=0
+else
+    NODE_RANK=-1
+    for i in "${!list_ip[@]}"; do
+        if [[ "${list_ip[$i]}" == "${CURRENT_IP}" ]]; then
+            NODE_RANK=$i
+            break
+        fi
+    done
+    if [ "$NODE_RANK" -eq -1 ]; then
+        echo "Error: Current IP ($CURRENT_IP) not found in list_ip."
+        exit 1
+    fi
+fi
+# --- End of Multi-node configuration ---
+
+SAVE_CKPT_PATH=$OUTPUT_DIR/$(basename "$0" .sh)
+TENSORBOARD_PATH="${SAVE_CKPT_PATH}/tensorboard"
+
+mkdir -p "$SAVE_CKPT_PATH"
+mkdir -p "$TENSORBOARD_PATH"
+mkdir -p "$SAVE_CKPT_PATH/dataloader"
+GPUS_PER_NODE=8
+
+MASTER_ADDR=${MASTER_ADDR:-"${list_ip[0]}"}
+MASTER_PORT=${MASTER_PORT:-"26000"}
+
+if [[ $SINGLE_NODE -eq 1 ]]; then
+    DISTRIBUTED_ARGS=(
+        --nproc_per_node "$GPUS_PER_NODE"
+    )
+else
+    DISTRIBUTED_ARGS=(
+        --nproc_per_node "$GPUS_PER_NODE"
+        --nnodes "$NNODES"
+        --node_rank "$NODE_RANK"
+        --master_addr "$MASTER_ADDR"
+        --master_port "$MASTER_PORT"
+    )
+fi
+
+MODEL_ARGS=(
+    --model-name gemma4-26b-a4b-vl
+)
+
+DATA_ARGS=(
+    --tokenizer-type HFTokenizer
+    --hf-tokenizer-path "$TOKENIZER_PATH"
+    --data-path "$DATA_PATH"
+    --dataloader-type external
+    --split 100,0,0
+    --num-workers 4
+)
+
+TRAINING_ARGS=(
+    --training-phase sft
+    --chat-template gemma4
+    --trainable-modules all
+    --seq-length "${SEQ_LEN}"
+    --max-position-embeddings "${SEQ_LEN}"
+    --init-method-std 0.02
+    --micro-batch-size "${MBS}"
+    --global-batch-size "${GBS}"
+    --lr 1.0e-5
+    --min-lr 1.0e-6
+    --clip-grad 1.0
+    --weight-decay 0
+    --optimizer adam
+    --adam-beta1 0.9
+    --adam-beta2 0.99
+    --adam-eps 1e-05
+    --norm-epsilon 1e-6
+    --train-iters "$NSTEP"
+    --lr-decay-iters "$NSTEP"
+    --lr-decay-style cosine
+    --lr-warmup-fraction 0.002
+    --initial-loss-scale 65536
+    --bf16
+    --load "$CHECKPOINT_PATH"
+    --save "$SAVE_CKPT_PATH"
+    --save-interval 100
+    --ckpt-format torch
+    --dataloader-save "${SAVE_CKPT_PATH}/dataloader"
+
+    --ckpt-fully-parallel-load
+    --recompute-granularity full
+    --recompute-method uniform
+    --recompute-num-layers 4
+)
+
+MODEL_PARALLEL_ARGS=(
+    --attention-backend flash
+    --moe-token-dispatcher-type alltoall
+    --moe-per-layer-logging
+    --pipeline-model-parallel-size "${PP}"
+    --tensor-model-parallel-size "${TP}"
+    --expert-model-parallel-size "${EP}"
+    --encoder-pipeline-model-parallel-size 1
+    --num-experts 128
+    --use-distributed-optimizer
+    --distributed-backend nccl
+)
+
+MOE_ARGS=(
+    --moe-router-topk 8
+    --moe-aux-loss-coeff 1e-3
+    --moe-router-dtype fp32
+)
+
+LOGGING_ARGS=(
+    --log-interval 1
+    --tensorboard-dir "${TENSORBOARD_PATH}"
+    --log-timers-to-tensorboard
+)
+
+if [ -n "${WANDB_API_KEY}" ]; then
+    LOGGING_ARGS+=(
+        --wandb-project "${WANDB_PROJECT}"
+        --wandb-exp-name "${WANDB_NAME}"
+    )
+fi
+
+TM=$(date "+%Y-%m-%d_%H:%M:%S")
+logfile="${SAVE_CKPT_PATH}/run_${TM}_tp${TP}_pp${PP}_ep${EP}_seqlen${SEQ_LEN}_mbs${MBS}_gbs${GBS}_${NSTEP}steps.log"
+
+export OFFLINE_PACKING_BMR=1
+
+PYTHONPATH="$AIAK_MAGATRON_PATH:$AIAK_TRAINING_PATH:$PYTHONPATH" \
+    torchrun "${DISTRIBUTED_ARGS[@]}" \
+    "$AIAK_TRAINING_PATH/aiak_training_llm/train.py" \
+    "${MODEL_ARGS[@]}" \
+    "${DATA_ARGS[@]}" \
+    ${IMG_ARGS:+${IMG_ARGS[@]}} \
+    "${TRAINING_ARGS[@]}" \
+    "${MODEL_PARALLEL_ARGS[@]}" \
+    "${MOE_ARGS[@]}" \
+    "${LOGGING_ARGS[@]}" \
+    2>&1 | tee "$logfile"
diff --git a/tests/_shared/test_gemma4_per_layer_window_size.py b/tests/_shared/test_gemma4_per_layer_window_size.py
new file mode 100644
index 00000000..3f756bc1
--- /dev/null
+++ b/tests/_shared/test_gemma4_per_layer_window_size.py
@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import torch
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+from aiak_training_llm.models.gemma4_vl.gemma4_attention import (
+    _resolve_per_layer_overrides,
+)
+from aiak_training_llm.models.gemma4_vl.gemma4_vl_config import gemma4_26b_a4b_vl
+
+
+def _build_tc() -> TransformerConfig:
+    cfg = gemma4_26b_a4b_vl()
+    tc_fields = {f.name for f in TransformerConfig.__dataclass_fields__.values()}
+    tc_kwargs = {k: v for k, v in cfg.__dict__.items() if k in tc_fields}
+    tc_kwargs.update(
+        {
+            "bf16": True,
+            "params_dtype": torch.bfloat16,
+            "pipeline_dtype": torch.bfloat16,
+            "tensor_model_parallel_size": 1,
+            "pipeline_model_parallel_size": 1,
+            "expert_model_parallel_size": 1,
+            "context_parallel_size": 1,
+            "sequence_parallel": False,
+            "add_bias_linear": False,
+            "gated_linear_unit": True,
+            "activation_func": torch.nn.functional.gelu,
+        }
+    )
+    for k in (
+        "layer_pattern",
+        "per_layer_kv_channels",
+        "per_layer_num_query_groups",
+        "kv_tied_layers",
+        "sliding_window",
+        "rotary_base_sliding",
+    ):
+        if hasattr(cfg, k):
+            tc_kwargs[k] = getattr(cfg, k)
+    return TransformerConfig(**tc_kwargs)
+
+
+def test_per_layer_window_size_matches_layer_pattern() -> None:
+    tc = _build_tc()
+    sliding_w = tc.sliding_window
+    assert sliding_w == 1024
+
+    expected_te_window = (sliding_w - 1, 0)
+
+    for layer_number in range(1, len(tc.layer_pattern) + 1):
+        eff = _resolve_per_layer_overrides(tc, layer_number)
+        layer_type = tc.layer_pattern[layer_number - 1]
+        if layer_type == "sliding":
+            assert eff.window_size == expected_te_window, (
+                f"layer {layer_number} (sliding): window_size={eff.window_size}, expected {expected_te_window}"
+            )
+            assert eff.kv_channels == 256
+            assert eff.num_query_groups == 8
+        elif layer_type == "global":
+            assert eff.window_size is None, (
+                f"layer {layer_number} (global): window_size={eff.window_size}, expected None"
+            )
+            assert eff.kv_channels == 512
+            assert eff.num_query_groups == 2
+        else:
+            raise AssertionError(f"unknown layer_type: {layer_type}")
+
+
+def test_global_layers_match_hf_full_attention_positions() -> None:
+    tc = _build_tc()
+    actual_global_1based = [i + 1 for i, t in enumerate(tc.layer_pattern) if t == "global"]
+    assert actual_global_1based == [6, 12, 18, 24, 30]
+
+
+def test_resolve_returns_unchanged_for_empty_layer_pattern() -> None:
+    tc = _build_tc()
+    tc.layer_pattern = []
+    out = _resolve_per_layer_overrides(tc, 1)
+    assert out is tc
+
+
+def test_resolve_clones_config_when_pattern_present() -> None:
+    tc = _build_tc()
+    out = _resolve_per_layer_overrides(tc, 1)
+    assert out is not tc
+    assert out.window_size == (tc.sliding_window - 1, 0)
+    assert tc.window_size != out.window_size or tc.window_size is None
diff --git a/tests/_shared/test_gemma4_rope_inv_freq.py b/tests/_shared/test_gemma4_rope_inv_freq.py
new file mode 100644
index 00000000..c3876641
--- /dev/null
+++ b/tests/_shared/test_gemma4_rope_inv_freq.py
@@ -0,0 +1,110 @@
+"""Numerical equivalence tests for Gemma4 dual-RoPE inv_freq buffers.
+
+Guards the two RoPE construction paths in :class:`Gemma4Model` against silent
+drift from the HF reference. Run on CPU; no Megatron initialisation required.
+
+  - Sliding layers use stock :class:`RotaryEmbedding` with ``head_dim=256``,
+    ``rotary_base=1e4``, full rotation. HF formula:
+        inv_freq[k] = 1 / 1e4 ** (2k / 256)   for k in [0, 128)
+
+  - Global layers use :class:`Gemma4ProportionalRotaryEmbedding` with
+    ``head_dim=512``, ``partial_rotary_factor=0.25``, ``rotary_base=1e6``.
+    HF proportional formula:
+        rope_angles = int(0.25 * 512 // 2)             # = 64
+        inv_freq[k] = 1 / 1e6 ** (2k / 512)            for k in [0, 64)
+        inv_freq[k] = 0                                for k in [64, 256)
+
+A regression in either path silently corrupts attention logits across the
+LLM. These tests catch the corruption at construction time, before any forward
+pass or checkpoint conversion.
+"""
+
+from __future__ import annotations
+
+import torch
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+
+from aiak_training_llm.models.gemma4_vl.gemma4_proportional_rotary import (
+    Gemma4ProportionalRotaryEmbedding,
+)
+
+
+def _expected_default_inv_freq(head_dim: int, base: float) -> torch.Tensor:
+    return 1.0 / (
+        base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+    )
+
+
+def _expected_proportional_inv_freq(
+    head_dim: int, partial_rotary_factor: float, base: float
+) -> torch.Tensor:
+    rope_angles = int(partial_rotary_factor * head_dim // 2)
+    inv_freq = torch.zeros(head_dim // 2, dtype=torch.float32)
+    inv_freq[:rope_angles] = 1.0 / (
+        base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float32) / head_dim)
+    )
+    return inv_freq
+
+
+def test_sliding_inv_freq_matches_hf_default() -> None:
+    rope = RotaryEmbedding(
+        kv_channels=256,
+        rotary_percent=1.0,
+        rotary_interleaved=False,
+        rotary_base=10000,
+        use_cpu_initialization=True,
+    )
+    expected = _expected_default_inv_freq(head_dim=256, base=10000.0)
+    assert rope.inv_freq.shape == expected.shape
+    assert torch.allclose(rope.inv_freq, expected, rtol=0.0, atol=0.0), (
+        f"Sliding RoPE inv_freq mismatch: max abs diff "
+        f"{(rope.inv_freq - expected).abs().max().item():.3e}"
+    )
+
+
+def test_global_inv_freq_matches_hf_proportional() -> None:
+    rope = Gemma4ProportionalRotaryEmbedding(
+        head_dim=512,
+        partial_rotary_factor=0.25,
+        rotary_base=1_000_000,
+        rotary_interleaved=False,
+        use_cpu_initialization=True,
+    )
+    expected = _expected_proportional_inv_freq(
+        head_dim=512, partial_rotary_factor=0.25, base=1_000_000.0
+    )
+
+    assert rope.inv_freq.shape == expected.shape, (
+        f"Global RoPE inv_freq shape {tuple(rope.inv_freq.shape)} != expected "
+        f"{tuple(expected.shape)}"
+    )
+
+    rope_angles = int(0.25 * 512 // 2)
+    assert torch.allclose(
+        rope.inv_freq[:rope_angles], expected[:rope_angles], rtol=0.0, atol=0.0
+    ), (
+        f"Global RoPE rotated head ({rope_angles} angles) mismatch: max abs "
+        f"diff {(rope.inv_freq[:rope_angles] - expected[:rope_angles]).abs().max().item():.3e}"
+    )
+
+    assert torch.equal(
+        rope.inv_freq[rope_angles:], torch.zeros(512 // 2 - rope_angles)
+    ), (
+        f"Global RoPE non-rotated tail must be exactly zero (HF zero-pad); got "
+        f"max abs {rope.inv_freq[rope_angles:].abs().max().item():.3e}"
+    )
+
+
+def test_global_forward_emits_head_dim_long_emb() -> None:
+    rope = Gemma4ProportionalRotaryEmbedding(
+        head_dim=512,
+        partial_rotary_factor=0.25,
+        rotary_base=1_000_000,
+        rotary_interleaved=False,
+        use_cpu_initialization=True,
+    )
+    seq_len = 16
+    emb = rope(seq_len)
+    assert emb.shape[-1] == 512, (
+        f"Global RoPE emb last dim must equal head_dim=512; got {emb.shape[-1]}"
+    )
diff --git a/tests/_shared/test_gemma4_vision_state_dict.py b/tests/_shared/test_gemma4_vision_state_dict.py
new file mode 100644
index 00000000..4648ac03
--- /dev/null
+++ b/tests/_shared/test_gemma4_vision_state_dict.py
@@ -0,0 +1,98 @@
+"""Narrow unit tests for Gemma4 vision tower + adapter state-dict fidelity.
+
+Verifies the HF-verbatim port preserves exactly the state-dict keys that the
+Phase C HF↔Megatron checkpoint converter will round-trip:
+
+- ``Gemma4Adapter`` exposes exactly one trainable tensor:
+  ``embedding_projection.weight`` of shape ``[text_hidden, vision_hidden]``.
+  The ``embedding_pre_projection_norm`` (RMSNorm with ``with_scale=False``)
+  contributes NO key — HF stores no ``embedding_pre_projection_norm.weight``
+  for this checkpoint, so registering one would break ckpt loading.
+
+- ``Gemma4VisionTower`` registers ``std_bias`` and ``std_scale`` ONLY when
+  ``standardize=True`` (persistent buffers visible in state_dict). With
+  ``standardize=False`` they must NOT appear.
+
+These tests are CPU-only and do not require Megatron initialisation; they
+guard the converter contract independently of the GPU smoke test.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from aiak_training_llm.models.gemma4_vl.gemma4_adapter import Gemma4Adapter
+from aiak_training_llm.models.gemma4_vl.gemma4_vision_tower import Gemma4VisionTower
+
+
+class _DummyTransformerConfig:
+    def __init__(self) -> None:
+        self.tensor_model_parallel_size = 1
+        self.pipeline_model_parallel_size = 1
+        self.context_parallel_size = 1
+        self.expert_model_parallel_size = 1
+        self.sequence_parallel = False
+        self.params_dtype = torch.float32
+        self.timers = None
+        self.perform_initialization = False
+        self.deterministic_mode = False
+        self.use_cpu_initialization = True
+
+
+def _build_tower(*, standardize: bool) -> Gemma4VisionTower:
+    return Gemma4VisionTower(
+        transformer_config=_DummyTransformerConfig(),
+        hidden_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=16,
+        intermediate_size=64,
+        patch_size=4,
+        in_channels=3,
+        position_embedding_size=8,
+        pooling_kernel_size=2,
+        rope_theta=100.0,
+        rms_norm_eps=1e-6,
+        hidden_activation="gelu_pytorch_tanh",
+        use_clipped_linears=False,
+        standardize=standardize,
+    )
+
+
+def test_adapter_state_dict_has_only_embedding_projection_weight():
+    adapter = Gemma4Adapter(
+        vision_hidden_size=1152,
+        text_hidden_size=2816,
+        rms_norm_eps=1e-6,
+    )
+    keys = sorted(adapter.state_dict().keys())
+    assert keys == ["embedding_projection.weight"], (
+        f"Gemma4Adapter must expose exactly one ckpt key "
+        f"(embedding_projection.weight); got {keys}"
+    )
+    w = adapter.state_dict()["embedding_projection.weight"]
+    assert w.shape == (2816, 1152), f"projection weight shape wrong: {w.shape}"
+
+    pre_norm = adapter.embedding_pre_projection_norm
+    assert not any(True for _ in pre_norm.parameters()), (
+        "embedding_pre_projection_norm must be parameter-free "
+        "(with_scale=False); registering a weight would inject an "
+        "unexpected ckpt key and break HF loading."
+    )
+
+
+def test_tower_registers_std_buffers_iff_standardize():
+    tower_with_std = _build_tower(standardize=True)
+    keys = set(tower_with_std.state_dict().keys())
+    assert "std_bias" in keys, "std_bias must be in state_dict when standardize=True"
+    assert "std_scale" in keys, "std_scale must be in state_dict when standardize=True"
+
+    tower_no_std = _build_tower(standardize=False)
+    keys = set(tower_no_std.state_dict().keys())
+    assert "std_bias" not in keys, (
+        "std_bias must NOT be in state_dict when standardize=False"
+    )
+    assert "std_scale" not in keys, (
+        "std_scale must NOT be in state_dict when standardize=False"
+    )
diff --git a/tests/_shared/test_gemma4_vl_attention_mask.py b/tests/_shared/test_gemma4_vl_attention_mask.py
new file mode 100644
index 00000000..3fde1773
--- /dev/null
+++ b/tests/_shared/test_gemma4_vl_attention_mask.py
@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from aiak_training_llm.models.gemma4_vl.gemma4_vl_attention_mask import (
+    _compute_vision_group_ids,
+    build_gemma4_mm_attention_mask,
+)
+
+
+def _make_inputs(mtti_row: list[int]) -> tuple[torch.Tensor, torch.Tensor]:
+    mtti = torch.tensor([mtti_row], dtype=torch.long)
+    input_ids = torch.zeros_like(mtti)
+    return input_ids, mtti
+
+
+def test_pure_text_returns_strict_causal_mask():
+    input_ids, mtti = _make_inputs([0, 0, 0, 0, 0])
+    mask = build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
+
+    assert mask.shape == (1, 1, 5, 5)
+    assert mask.dtype == torch.bool
+
+    expected_visible = torch.tril(torch.ones(5, 5, dtype=torch.bool))
+    expected_masked = ~expected_visible
+    assert torch.equal(mask[0, 0], expected_masked)
+
+
+def test_single_image_span_attends_bidirectionally_inside_span():
+    input_ids, mtti = _make_inputs([0, 1, 1, 1, 0])
+    mask = build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
+    visible = ~mask[0, 0]
+
+    for q in (1, 2, 3):
+        for kv in (1, 2, 3):
+            assert visible[q, kv], f"expected vision token q={q} to see kv={kv}"
+
+
+def test_vision_overlay_does_not_break_causality_outside_span():
+    input_ids, mtti = _make_inputs([0, 1, 1, 1, 0, 0])
+    mask = build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
+    visible = ~mask[0, 0]
+
+    assert not visible[0, 1]
+    assert not visible[0, 4]
+    assert not visible[4, 5]
+    assert visible[4, 1]
+    assert visible[5, 3]
+
+
+def test_two_image_spans_do_not_cross_attend():
+    input_ids, mtti = _make_inputs([1, 1, 0, 1, 1])
+    mask = build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
+    visible = ~mask[0, 0]
+
+    assert visible[0, 1] and visible[1, 0]
+    assert visible[3, 4] and visible[4, 3]
+    assert not visible[0, 3]
+    assert not visible[0, 4]
+    assert visible[3, 0] and visible[3, 1]
+    assert visible[4, 0] and visible[4, 1]
+
+
+def test_video_token_type_id_2_treated_as_vision():
+    input_ids, mtti = _make_inputs([0, 2, 2, 2, 0])
+    mask = build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
+    visible = ~mask[0, 0]
+
+    for q in (1, 2, 3):
+        for kv in (1, 2, 3):
+            assert visible[q, kv]
+
+
+def test_compute_vision_group_ids_matches_hf_block_seq_ids():
+    mtti = torch.tensor([[0, 1, 1, 0, 2, 2, 2, 0, 1]], dtype=torch.long)
+    group_ids = _compute_vision_group_ids(mtti)
+    expected = torch.tensor([[-1, 0, 0, -1, 1, 1, 1, -1, 2]], dtype=torch.long)
+    assert torch.equal(group_ids, expected)
+
+
+def test_runtime_guard_fires_when_span_exceeds_sliding_window():
+    seq_len = 10
+    mtti = torch.ones((1, seq_len), dtype=torch.long)
+    input_ids = torch.zeros_like(mtti)
+    with pytest.raises(AssertionError, match="exceeds sliding_window"):
+        build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=4)
+
+
+def test_runtime_guard_passes_when_span_equals_sliding_window():
+    seq_len = 8
+    mtti = torch.ones((1, seq_len), dtype=torch.long)
+    input_ids = torch.zeros_like(mtti)
+    mask = build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=8)
+    assert mask.shape == (1, 1, 8, 8)
+
+
+def test_batch_dimension_is_independent_per_sequence():
+    mtti = torch.tensor(
+        [
+            [0, 1, 1, 0, 0],
+            [0, 0, 1, 1, 0],
+        ],
+        dtype=torch.long,
+    )
+    input_ids = torch.zeros_like(mtti)
+    mask = build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
+
+    assert mask.shape == (2, 1, 5, 5)
+
+    assert (~mask[0, 0])[1, 2] and (~mask[0, 0])[2, 1]
+    assert not (~mask[1, 0])[1, 2]
+    assert (~mask[1, 0])[2, 3] and (~mask[1, 0])[3, 2]
+
+
+def test_shape_mismatch_raises_value_error():
+    input_ids = torch.zeros((1, 5), dtype=torch.long)
+    mtti = torch.zeros((1, 6), dtype=torch.long)
+    with pytest.raises(ValueError, match="does not match"):
+        build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
+
+
+def test_non_2d_input_raises_value_error():
+    input_ids = torch.zeros((5,), dtype=torch.long)
+    mtti = torch.zeros((5,), dtype=torch.long)
+    with pytest.raises(ValueError, match="must be 2D"):
+        build_gemma4_mm_attention_mask(input_ids, mtti, sliding_window=1024)
diff --git a/tests/_shared/test_gemma4_vl_skeleton.py b/tests/_shared/test_gemma4_vl_skeleton.py
new file mode 100644
index 00000000..9be45a5d
--- /dev/null
+++ b/tests/_shared/test_gemma4_vl_skeleton.py
@@ -0,0 +1,286 @@
+"""GPU smoke test for Gemma4-VL skeleton (P1.14).
+
+Boots Megatron in-process on a single GPU with shrunken LLM dims, instantiates
+the registered ``gemma4_vl`` provider with the Gemma4 vision tower + adapter
+monkey-patched out (their dims are real Gemma4-26B-A4B-it production sizes —
+27 layers / hidden=1152 / intermediate=4304 — and would dwarf the
+shrunken-LLM smoke test), runs one forward pass on dummy text tokens, and
+asserts no crash + correct logits shape.
+
+Stubbing the vision tower / adapter keeps this smoke test focused on the
+Gemma4 LLM stack (Gemma4Model + Gemma4TransformerLayer + Gemma4ParallelDenseMoE)
+which is the moving part in P1.
+
+Skip rules:
+- requires CUDA (skipped on CPU-only hosts)
+- requires single-GPU only (TP=1 PP=1 EP=1)
+
+Run inside Docker:
+    pytest tests/_shared/test_gemma4_vl_skeleton.py -v -s
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import pytest
+import torch
+
+
+_MICRO_ARGV = [
+    "pytest-gemma4-skeleton",
+    "--model-name", "gemma4_vl",
+    "--tokenizer-type", "NullTokenizer",
+    "--vocab-size", "1024",
+    "--micro-batch-size", "1",
+    "--global-batch-size", "1",
+    "--seq-length", "16",
+    "--max-position-embeddings", "16",
+    "--num-layers", "2",
+    "--hidden-size", "128",
+    "--num-attention-heads", "4",
+    "--num-query-groups", "2",
+    "--kv-channels", "32",
+    "--ffn-hidden-size", "64",
+    "--num-experts", "4",
+    "--moe-router-topk", "2",
+    "--bf16",
+    "--disable-bias-linear",
+    "--attention-backend", "flash",
+    "--pipeline-model-parallel-size", "1",
+    "--tensor-model-parallel-size", "1",
+    "--expert-model-parallel-size", "1",
+    "--distributed-backend", "nccl",
+    "--position-embedding-type", "rope",
+    "--no-rope-fusion",
+    "--use-mcore-models",
+    "--transformer-impl", "transformer_engine",
+]
+
+
+@pytest.fixture(scope="module")
+def megatron_micro_init():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+    from aiak_training_llm.train.arguments import (
+        aiak_extra_train_args_provider,
+        parse_arguments,
+        validate_aiak_extra_args,
+    )
+    from aiak_training_llm.utils import initialize_aiak_megatron
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "12377")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("LOCAL_RANK", "0")
+
+    saved_argv = sys.argv
+    sys.argv = list(_MICRO_ARGV)
+    try:
+        args = parse_arguments(
+            extra_args_provider=aiak_extra_train_args_provider,
+            validate_extra_args_provider=validate_aiak_extra_args,
+        )
+        initialize_aiak_megatron(args=args)
+    finally:
+        sys.argv = saved_argv
+
+    yield args
+
+
+def test_provider_lookup_resolves_gemma4_vl():
+    from aiak_training_llm.models import get_model_family, get_model_provider
+
+    fam = get_model_family("gemma4-26b-a4b-vl")
+    assert fam == "gemma4_vl", f"family lookup wrong: {fam}"
+
+    provider = get_model_provider(fam)
+    assert provider is not None, "gemma4_vl provider not registered"
+    assert provider.__name__ == "gemma4_vl_model_provider"
+    assert "gemma4_vl_provider" in provider.__module__
+
+
+def test_gemma4_vl_micro_forward(megatron_micro_init, monkeypatch):
+    args = megatron_micro_init
+    from aiak_training_llm.models.gemma4_vl import gemma4_vl_model as gvm
+    from aiak_training_llm.models.gemma4_vl.gemma4_vl_provider import (
+        gemma4_vl_model_provider,
+    )
+
+    class _StubVision(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+        def forward(self, *args, **kwargs):
+            raise RuntimeError("vision tower stubbed in skeleton smoke")
+
+    class _StubAdapter(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+        def forward(self, *args, **kwargs):
+            raise RuntimeError("adapter stubbed in skeleton smoke")
+
+    monkeypatch.setattr(gvm, "Gemma4VisionTower", _StubVision)
+    monkeypatch.setattr(gvm, "Gemma4Adapter", _StubAdapter)
+
+    model = gemma4_vl_model_provider(
+        pre_process=True,
+        post_process=True,
+        add_encoder=True,
+        add_decoder=True,
+        parallel_output=False,
+    )
+    model = model.to(device="cuda", dtype=torch.bfloat16).eval()
+    micro_image_token_id = 1
+    model.config.image_token_id = micro_image_token_id
+
+    seq_len = 8
+    input_ids = torch.randint(
+        low=0, high=args.padded_vocab_size, size=(1, seq_len),
+        dtype=torch.long, device="cuda",
+    )
+    position_ids = torch.arange(seq_len, dtype=torch.long, device="cuda").unsqueeze(0)
+    attention_mask = None
+
+    with torch.no_grad():
+        logits = model(
+            images=None,
+            image_grid_thw=None,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+        )
+
+    assert logits is not None, "model returned None"
+    assert isinstance(logits, torch.Tensor), f"expected Tensor, got {type(logits)}"
+    assert logits.shape[0] == 1, f"batch dim wrong: {logits.shape}"
+    assert args.padded_vocab_size in logits.shape, (
+        f"vocab dim {args.padded_vocab_size} not in logits shape {logits.shape}"
+    )
+    assert torch.isfinite(logits).all(), "logits contain NaN/Inf"
+
+
+def test_gemma4_vl_vision_path_requires_grid(megatron_micro_init, monkeypatch):
+    args = megatron_micro_init
+    from aiak_training_llm.models.gemma4_vl import gemma4_vl_model as gvm
+    from aiak_training_llm.models.gemma4_vl.gemma4_vl_provider import (
+        gemma4_vl_model_provider,
+    )
+
+    class _StubVision(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+        def forward(self, *args, **kwargs):
+            raise RuntimeError("should not be reached: forward must short-circuit")
+
+    class _StubAdapter(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+        def forward(self, *args, **kwargs):
+            raise RuntimeError("should not be reached: forward must short-circuit")
+
+    monkeypatch.setattr(gvm, "Gemma4VisionTower", _StubVision)
+    monkeypatch.setattr(gvm, "Gemma4Adapter", _StubAdapter)
+
+    model = gemma4_vl_model_provider(
+        pre_process=True,
+        post_process=True,
+        add_encoder=True,
+        add_decoder=True,
+        parallel_output=False,
+    )
+    model = model.to(device="cuda", dtype=torch.bfloat16).eval()
+
+    dummy_images = torch.zeros(1, 1, 768, dtype=torch.bfloat16, device="cuda")
+    input_ids = torch.randint(
+        low=0, high=args.padded_vocab_size, size=(1, 4),
+        dtype=torch.long, device="cuda",
+    )
+    position_ids = torch.arange(4, dtype=torch.long, device="cuda").unsqueeze(0)
+
+    with pytest.raises(ValueError, match="image_grid_thw is required"):
+        with torch.no_grad():
+            model(
+                images=dummy_images,
+                image_grid_thw=None,
+                input_ids=input_ids,
+                position_ids=position_ids,
+                attention_mask=None,
+            )
+
+
+def test_gemma4_vl_vision_path_micro_forward(megatron_micro_init, monkeypatch):
+    args = megatron_micro_init
+    from aiak_training_llm.models.gemma4_vl import gemma4_vl_model as gvm
+    from aiak_training_llm.models.gemma4_vl.gemma4_vl_provider import (
+        gemma4_vl_model_provider,
+    )
+
+    class _StubVision(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+        def forward(self, pixel_values, pixel_position_ids):
+            assert pixel_values.shape == (1, 4, 768)
+            assert pixel_position_ids.shape == (1, 4, 2)
+            return torch.ones(4, 32, dtype=pixel_values.dtype, device=pixel_values.device)
+
+    class _StubAdapter(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+
+        def forward(self, hidden_states):
+            return torch.ones(
+                hidden_states.shape[0],
+                args.hidden_size,
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+
+    monkeypatch.setattr(gvm, "Gemma4VisionTower", _StubVision)
+    monkeypatch.setattr(gvm, "Gemma4Adapter", _StubAdapter)
+
+    model = gemma4_vl_model_provider(
+        pre_process=True,
+        post_process=True,
+        add_encoder=True,
+        add_decoder=True,
+        parallel_output=False,
+    )
+    model = model.to(device="cuda", dtype=torch.bfloat16).eval()
+    micro_image_token_id = 1
+    model.config.image_token_id = micro_image_token_id
+
+    dummy_images = torch.zeros(4, 768, dtype=torch.bfloat16, device="cuda")
+    image_grid_thw = torch.tensor([[1, 2, 2]], dtype=torch.int32, device="cuda")
+    patch_positions = torch.tensor(
+        [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1]],
+        dtype=torch.long,
+        device="cuda",
+    )
+    input_ids = torch.randint(
+        low=0, high=args.padded_vocab_size, size=(1, 8),
+        dtype=torch.long, device="cuda",
+    )
+    input_ids[0, :4] = micro_image_token_id
+    position_ids = torch.arange(8, dtype=torch.long, device="cuda").unsqueeze(0)
+
+    with torch.no_grad():
+        logits = model(
+            images=dummy_images,
+            image_grid_thw=image_grid_thw,
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=None,
+            patch_positions=patch_positions,
+        )
+
+    assert logits is not None
+    assert isinstance(logits, torch.Tensor)
+    assert torch.isfinite(logits).all()
diff --git a/tests/consistency_gemma4/__init__.py b/tests/consistency_gemma4/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/consistency_gemma4/conftest.py b/tests/consistency_gemma4/conftest.py
new file mode 100644
index 00000000..c68028ef
--- /dev/null
+++ b/tests/consistency_gemma4/conftest.py
@@ -0,0 +1,312 @@
+"""pytest fixtures for Gemma4-VL HF↔mcore consistency tests.
+
+Mirrors ``tests/conftest.py`` (LlavaOnevision2) with the following deltas:
+
+- HF model is ``Gemma4ForConditionalGeneration`` (native ``transformers``).
+- Conversion script is
+  ``examples/gemma4_vl/convert/convert_gemma4_26b_a4b_hf_to_mcore.sh`` and
+  takes a 4th positional ``EP`` argument.
+- Megatron CLI args set ``--model-name gemma4-26b-a4b-vl`` and add
+  MoE-specific flags (``--num-experts 128 --moe-router-topk 8 ...``).
+- ``mcore_model`` fixture imports ``model_provider`` from
+  ``aiak_training_llm.train.pretrain.pretrain_gemma4_vl``.
+
+This file is intentionally a clone (not a refactor of the OV2 conftest) per
+P5(B) decision — keeps OV2 consistency tests insulated from any Gemma4 changes.
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+
+def _repo_root() -> Path:
+    return Path(__file__).resolve().parents[2]
+
+
+def _env_path(name: str, default: str | None = None, *, required: bool = False) -> str:
+    value = os.environ.get(name, default)
+    if required and not value:
+        pytest.fail(f"Missing required environment variable: {name}")
+    if value is None:
+        return ""
+    return value
+
+
+def _ensure_hf_weights(config_dir: str) -> str:
+    """Return ``config_dir`` if it contains safetensors, else materialize a
+    randomly-initialized Gemma4 checkpoint under
+    ``<repo_root>/tmp_test_gemma4_random_weights/``.
+
+    The 26B-A4B real checkpoint is ~49 GB; for unit-style smoke we may want a
+    smaller random checkpoint. In practice ``HF_MODEL_PATH`` should point at
+    the real Gemma4 directory which already contains safetensors, and this
+    function is a no-op fast path.
+    """
+    if glob.glob(os.path.join(config_dir, "*.safetensors")):
+        return config_dir
+
+    repo_root = _repo_root()
+    out_dir = repo_root / "tmp_test_gemma4_random_weights"
+
+    if out_dir.exists() and glob.glob(str(out_dir / "*.safetensors")):
+        return str(out_dir)
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    for src_file in Path(config_dir).iterdir():
+        if src_file.is_file() and not src_file.name.endswith(".safetensors"):
+            shutil.copy2(src_file, out_dir / src_file.name)
+
+    from transformers import Gemma4Config, Gemma4ForConditionalGeneration
+
+    config = Gemma4Config.from_pretrained(config_dir)
+    model = Gemma4ForConditionalGeneration(config)
+    model = model.to(dtype=torch.bfloat16)
+    model.save_pretrained(str(out_dir), safe_serialization=True)
+    del model
+    torch.cuda.empty_cache()
+
+    return str(out_dir)
+
+
+def _build_megatron_cli_args(
+    *,
+    hf_model_path: str,
+    mcore_checkpoint_path: str,
+    tp: int,
+    pp: int,
+    ep: int,
+    epp: int,
+) -> list[str]:
+    return [
+        "pytest-megatron-init",
+        "--model-name",
+        "gemma4-26b-a4b-vl",
+        "--tokenizer-type",
+        "HFTokenizer",
+        "--hf-tokenizer-path",
+        hf_model_path,
+        "--dataloader-type",
+        "external",
+        "--split",
+        "100,0,0",
+        "--num-workers",
+        "4",
+        "--chat-template",
+        "qwen2-vl",
+        "--seq-length",
+        "4096",
+        "--max-position-embeddings",
+        "4096",
+        "--micro-batch-size",
+        "1",
+        "--global-batch-size",
+        "1",
+        "--bf16",
+        "--load",
+        mcore_checkpoint_path,
+        "--ckpt-format",
+        "torch",
+        "--attention-backend",
+        "flash",
+        "--pipeline-model-parallel-size",
+        str(pp),
+        "--tensor-model-parallel-size",
+        str(tp),
+        "--expert-model-parallel-size",
+        str(ep),
+        "--encoder-pipeline-model-parallel-size",
+        str(epp),
+        "--num-experts",
+        "128",
+        "--moe-router-topk",
+        "8",
+        "--moe-token-dispatcher-type",
+        "alltoall",
+        "--moe-router-dtype",
+        "fp32",
+        "--moe-aux-loss-coeff",
+        "1e-3",
+        "--distributed-backend",
+        "nccl",
+    ]
+
+
+@pytest.fixture(scope="session")
+def hf_model_path() -> str:
+    config_path = _env_path(
+        "HF_MODEL_PATH",
+        "/ov2/pretrain_models/google/gemma-4-26B-A4B-it",
+        required=True,
+    )
+    if not Path(config_path).exists():
+        pytest.fail(f"HF model path does not exist: {config_path}")
+    return _ensure_hf_weights(config_path)
+
+
+@pytest.fixture(scope="session")
+def converted_mcore_path(hf_model_path: str) -> str:
+    provided = os.environ.get("MCORE_CHECKPOINT_PATH", "").strip()
+    if provided:
+        if not Path(provided).exists():
+            pytest.fail(f"Provided MCORE_CHECKPOINT_PATH does not exist: {provided}")
+        return provided
+
+    repo_root = _repo_root()
+    tp = int(os.environ.get("CONSISTENCY_TEST_TP", "1"))
+    pp = int(os.environ.get("CONSISTENCY_TEST_PP", "1"))
+    ep = int(os.environ.get("CONSISTENCY_TEST_EP", "1"))
+    out_dir = repo_root / f"tmp_test_gemma4_mcore_ckpt_tp{tp}_pp{pp}_ep{ep}"
+
+    env = os.environ.copy()
+    env.setdefault("AIAK_TRAINING_PATH", str(repo_root))
+
+    script = (
+        repo_root
+        / "examples"
+        / "gemma4_vl"
+        / "convert"
+        / "convert_gemma4_26b_a4b_hf_to_mcore.sh"
+    )
+    result = subprocess.run(
+        ["bash", str(script), hf_model_path, str(out_dir), str(tp), str(pp), str(ep)],
+        cwd=repo_root,
+        env=env,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        pytest.fail(f"HF->mcore conversion failed.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}")
+
+    if not out_dir.exists():
+        pytest.fail(f"Converted checkpoint path not created: {out_dir}")
+
+    os.environ["MCORE_CHECKPOINT_PATH"] = str(out_dir)
+    return str(out_dir)
+
+
+@pytest.fixture(scope="session")
+def preprocessor_path(hf_model_path: str) -> str:
+    return _env_path("PREPROCESSOR_PATH", hf_model_path)
+
+
+@pytest.fixture(scope="session")
+def test_image_path() -> str:
+    default_path = str(_repo_root() / "asset" / "performance.png")
+    path = _env_path("TEST_IMAGE_PATH", default_path, required=True)
+    if path.startswith("http://") or path.startswith("https://"):
+        pytest.fail("TEST_IMAGE_PATH must be a local file path, not remote URL")
+    if not Path(path).exists():
+        pytest.fail(f"TEST_IMAGE_PATH does not exist: {path}")
+    return path
+
+
+@pytest.fixture(scope="session")
+def megatron_init(hf_model_path: str, converted_mcore_path: str):
+    from aiak_training_llm.train.arguments import (
+        aiak_extra_train_args_provider,
+        parse_arguments,
+        validate_aiak_extra_args,
+    )
+    from aiak_training_llm.utils import initialize_aiak_megatron
+
+    tp = int(os.environ.get("CONSISTENCY_TEST_TP", "1"))
+    pp = int(os.environ.get("CONSISTENCY_TEST_PP", "1"))
+    ep = int(os.environ.get("CONSISTENCY_TEST_EP", "1"))
+    epp = int(os.environ.get("CONSISTENCY_TEST_EPP", "0"))
+
+    original_argv = sys.argv
+    sys.argv = _build_megatron_cli_args(
+        hf_model_path=hf_model_path,
+        mcore_checkpoint_path=converted_mcore_path,
+        tp=tp,
+        pp=pp,
+        ep=ep,
+        epp=epp,
+    )
+    try:
+        args = parse_arguments(
+            extra_args_provider=aiak_extra_train_args_provider,
+            validate_extra_args_provider=validate_aiak_extra_args,
+            args_defaults={},
+        )
+        initialize_aiak_megatron(args=args)
+    finally:
+        sys.argv = original_argv
+
+    return args
+
+
+@pytest.fixture(scope="session")
+def hf_config(hf_model_path: str):
+    from transformers import Gemma4Config
+
+    return Gemma4Config.from_pretrained(hf_model_path)
+
+
+@pytest.fixture(scope="session")
+def hf_vision_model(hf_model_path: str):
+    """Load HF Gemma4 full model and return ``.model.vision_tower``.
+
+    Mirrors OV2's ``hf_vision_model`` fixture which returns ``.model.visual``.
+    Gemma4 names the submodule ``vision_tower`` (HF convention).
+    """
+    from transformers import Gemma4ForConditionalGeneration
+
+    full_model = Gemma4ForConditionalGeneration.from_pretrained(
+        hf_model_path, low_cpu_mem_usage=True
+    )
+    vision_model = full_model.model.vision_tower.to(dtype=torch.bfloat16, device="cuda").eval()
+    del full_model
+    return vision_model
+
+
+@pytest.fixture(scope="session")
+def hf_cond_gen_model(hf_model_path: str):
+    from transformers import Gemma4ForConditionalGeneration
+
+    model = Gemma4ForConditionalGeneration.from_pretrained(hf_model_path, low_cpu_mem_usage=True)
+    return model.to(dtype=torch.bfloat16, device="cuda").eval()
+
+
+@pytest.fixture(scope="session")
+def mcore_model(megatron_init):
+    from megatron.core.enums import ModelType
+    from megatron.training.checkpointing import load_checkpoint
+    from megatron.training.training import get_model, unwrap_model
+
+    from aiak_training_llm.train.pretrain.pretrain_gemma4_vl import model_provider
+
+    model_type = (
+        ModelType.encoder_and_decoder
+        if megatron_init.encoder_pipeline_model_parallel_size not in [0, None]
+        else ModelType.encoder_or_decoder
+    )
+
+    # Consistency tests are inference-only. Avoid DDP wrapping here: the full
+    # Gemma4-26B-A4B model cannot fit the extra fp32 training grad buffers on
+    # 80GB GPUs. ModelType follows the selected encoder pipeline size so the
+    # checkpoint layout and runtime pipeline layout stay aligned.
+    model = get_model(
+        model_provider,
+        model_type,
+        wrap_with_ddp=False,
+    )
+    load_checkpoint(model, None, None)
+    return unwrap_model(model)[0].to("cuda").eval()
+
+
+@pytest.fixture(scope="session")
+def hf_processor(preprocessor_path: str):
+    from transformers import AutoProcessor
+
+    return AutoProcessor.from_pretrained(preprocessor_path, trust_remote_code=True)
diff --git a/tests/consistency_gemma4/test_model_consistency.py b/tests/consistency_gemma4/test_model_consistency.py
new file mode 100644
index 00000000..c7c91a3f
--- /dev/null
+++ b/tests/consistency_gemma4/test_model_consistency.py
@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import os
+
+import pytest
+
+
+def _has_full_consistency_env() -> bool:
+    has_hf = bool(os.environ.get("HF_MODEL_PATH"))
+    has_mcore = bool(os.environ.get("MCORE_CHECKPOINT_PATH"))
+    allow_convert = os.environ.get("CONSISTENCY_RUN_CONVERT", "0") == "1"
+    return has_hf and (has_mcore or allow_convert)
+
+
+def test_collection_smoke():
+    assert True
+
+
+def test_full_gemma4_consistency_fixtures_lazy(request):
+    if not _has_full_consistency_env():
+        pytest.skip(
+            "Set HF_MODEL_PATH and either MCORE_CHECKPOINT_PATH or "
+            "CONSISTENCY_RUN_CONVERT=1 to run Gemma4 full consistency."
+        )
+
+    hf_config = request.getfixturevalue("hf_config")
+    hf_vision_model = request.getfixturevalue("hf_vision_model")
+    mcore_model = request.getfixturevalue("mcore_model")
+
+    assert getattr(hf_config, "model_type", None) == "gemma4"
+    assert hasattr(hf_vision_model, "forward")
+    assert hasattr(mcore_model, "forward")
diff --git a/tools/convert_checkpoint/config/gemma4-26b-a4b/adapter.json b/tools/convert_checkpoint/config/gemma4-26b-a4b/adapter.json
new file mode 100644
index 00000000..6e4c28b3
--- /dev/null
+++ b/tools/convert_checkpoint/config/gemma4-26b-a4b/adapter.json
@@ -0,0 +1,3 @@
+{
+    "adapter.embedding_projection.weight": "model.embed_vision.embedding_projection.weight"
+}
diff --git a/tools/convert_checkpoint/config/gemma4-26b-a4b/vision-patch.json b/tools/convert_checkpoint/config/gemma4-26b-a4b/vision-patch.json
new file mode 100644
index 00000000..f9f4a144
--- /dev/null
+++ b/tools/convert_checkpoint/config/gemma4-26b-a4b/vision-patch.json
@@ -0,0 +1,6 @@
+{
+    "vision_model.patch_embedder.input_proj.weight": "model.vision_tower.patch_embedder.input_proj.weight",
+    "vision_model.patch_embedder.position_embedding_table": "model.vision_tower.patch_embedder.position_embedding_table",
+    "vision_model.std_bias": "model.vision_tower.std_bias",
+    "vision_model.std_scale": "model.vision_tower.std_scale"
+}
diff --git a/tools/convert_checkpoint/custom/gemma4_vl/__init__.py b/tools/convert_checkpoint/custom/gemma4_vl/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tools/convert_checkpoint/custom/gemma4_vl/adapter.py b/tools/convert_checkpoint/custom/gemma4_vl/adapter.py
new file mode 100644
index 00000000..3c35dd38
--- /dev/null
+++ b/tools/convert_checkpoint/custom/gemma4_vl/adapter.py
@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+# -*- coding: UTF-8 -*-
+################################################################################
+#
+# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
+#
+################################################################################
+"""Gemma4-VL adapter standalone converter.
+
+Mirrors ``custom/llava_onevision2/adapter.py`` but targets Gemma4-VL.
+Differences from the OV2 adapter:
+
+- Gemma4 ``Gemma4Adapter`` has only ONE trainable parameter:
+  ``embedding_projection.weight``. The ``embedding_pre_projection_norm`` is a
+  ``Gemma4RMSNorm`` with ``with_scale=False`` so it has no learnable weight,
+  and is therefore absent from both HF and Megatron state_dict.
+- ``embedding_projection`` is a plain ``nn.Linear``, NOT a ColumnParallelLinear,
+  so the weight is replicated across TP ranks (no sharding).
+- No ``_extra_state`` blobs are written: the Gemma4 adapter does not contain
+  TransformerEngine modules.
+
+The mapping is loaded from ``args.common_config_path``
+(Megatron-key → HF-key); the only entry is ``adapter.embedding_projection.weight
+-> model.embed_vision.embedding_projection.weight``.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from copy import deepcopy
+from os.path import dirname
+
+
+SCRIPT_DIR = dirname(os.path.abspath(__file__))
+sys.path.append(dirname(dirname(dirname(SCRIPT_DIR))))
+
+from convert_checkpoint.arguments import parse_args  # noqa: E402
+from convert_checkpoint.custom.llava_onevision2.util import (  # noqa: E402
+    load_huggingface_checkpoint,
+    load_megatron_checkpoint,
+    load_megatron_checkpoint_tp_ep,
+    save_huggingface_checkpoint,
+    save_megatron_checkpoint,
+)
+
+
+args = parse_args()
+name_map: dict[str, str] = {}
+with open(args.common_config_path, "r", encoding="utf-8") as f:
+    name_map = json.loads(f.read())
+
+
+def _get_non_ep_model_source(state_dict):
+    first = state_dict[0]
+    if isinstance(first, dict):
+        return first["model"]
+    first_rank = first[0]
+    if isinstance(first_rank, dict):
+        return first_rank["model"]
+    raise TypeError("Unsupported non-EP checkpoint structure")
+
+
+if (args.load_platform, args.save_platform) == ("mcore", "huggingface"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL adapter from Megatron Core to HuggingFace ======")
+    ep = args.expert_parallel_size
+    target = {}
+    if ep is not None and ep > 1:
+        state_dict = load_megatron_checkpoint_tp_ep(args.load_ckpt_path)
+        source = state_dict[0][0]["model"]
+    else:
+        state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+        source = _get_non_ep_model_source(state_dict)
+    for mg_key, hf_key in name_map.items():
+        target[hf_key] = source[mg_key]
+        print(f" > {mg_key} -> {hf_key}  (shape {list(target[hf_key].shape)})")
+    save_huggingface_checkpoint(target, args.save_ckpt_path)
+
+elif (args.load_platform, args.save_platform) == ("huggingface", "mcore"):
+    print(" ====== convert Gemma4-VL adapter from HuggingFace to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    source = load_huggingface_checkpoint(args.load_ckpt_path)
+    target = {}
+    for mg_key, hf_key in name_map.items():
+        target[mg_key] = source[hf_key]
+        print(f" > {hf_key} -> {mg_key}  (shape {list(target[mg_key].shape)})")
+
+    state_dict = [{"model": deepcopy(target)} for _ in range(tp)]
+    save_megatron_checkpoint(state_dict, os.path.join(args.save_ckpt_path, "release"))
+
+elif (args.load_platform, args.save_platform) == ("mcore", "mcore"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL adapter from Megatron Core to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+    source = _get_non_ep_model_source(state_dict)
+    target = deepcopy(source)
+    new_state_dict = [{"model": deepcopy(target)} for _ in range(tp)]
+    save_megatron_checkpoint(new_state_dict, os.path.join(args.save_ckpt_path, "release"))
+
+else:
+    raise NotImplementedError
diff --git a/tools/convert_checkpoint/custom/gemma4_vl/llm.py b/tools/convert_checkpoint/custom/gemma4_vl/llm.py
new file mode 100644
index 00000000..c32ee598
--- /dev/null
+++ b/tools/convert_checkpoint/custom/gemma4_vl/llm.py
@@ -0,0 +1,930 @@
+#!/usr/bin/env python3
+# -*- coding: UTF-8 -*-
+################################################################################
+#
+# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
+#
+################################################################################
+"""Gemma4-VL LLM body standalone converter (HF <-> Megatron Core).
+
+Mirrors ``custom/llava_onevision2/`` pattern but targets Gemma4-26B-A4B-it
+LLM body (30 layers, MoE+dense hybrid, dual RoPE, KV-tying global layers).
+
+Per-layer schema (data-driven from HF ``text_config.layer_types``):
+  - sliding layers (25): head_dim=256, num_kv=8, has v_proj
+      HF q/k/v_proj ([4096|2048|2048, 2816]) -> MG linear_qkv ([8192, 2816]) concat dim=0
+  - global layers (5: indices 5/11/17/23/29): head_dim=512, num_kv=2, KV-tying (no v_proj)
+      HF q/k_proj ([8192|1024, 2816]) -> MG linear_qkv ([9216, 2816]) concat dim=0
+
+Common per-layer keys (HF -> MG):
+  input_layernorm.weight                -> self_attention.linear_qkv.layer_norm_weight
+  post_attention_layernorm.weight       -> self_attention.post_attention_layernorm.weight
+  pre_feedforward_layernorm.weight      -> mlp._inner_mlp.pre_feedforward_layernorm.weight
+  pre_feedforward_layernorm_2.weight    -> mlp._inner_mlp.moe.pre_feedforward_layernorm_2.weight
+  post_feedforward_layernorm.weight     -> mlp._post_ffn.weight  AND  post_feedforward_layernorm.weight (alias)
+  post_feedforward_layernorm_1.weight   -> mlp._inner_mlp.post_feedforward_layernorm_1.weight
+  post_feedforward_layernorm_2.weight   -> mlp._inner_mlp.post_feedforward_layernorm_2.weight
+  layer_scalar                          -> layer_scalar
+  self_attn.q_norm.weight               -> self_attention.q_layernorm.weight
+  self_attn.k_norm.weight               -> self_attention.k_layernorm.weight
+  self_attn.o_proj.weight               -> self_attention.linear_proj.weight
+  mlp.gate_proj.weight + mlp.up_proj.weight -> mlp._inner_mlp.dense.linear_fc1.weight (concat dim=0, gate first)
+  mlp.down_proj.weight                  -> mlp._inner_mlp.dense.linear_fc2.weight
+  experts.gate_up_proj [128, fc1_out, 2816] -> 128 x mlp._inner_mlp.moe.experts.local_experts.{i}.linear_fc1.weight
+  experts.down_proj    [128, 2816, fc2_in]  -> 128 x mlp._inner_mlp.moe.experts.local_experts.{i}.linear_fc2.weight
+  router.proj.weight [128, 2816]        -> mlp._inner_mlp.moe.router.weight (NO transpose, verified)
+  router.scale                          -> mlp._inner_mlp.moe.router.scale
+  router.per_expert_scale               -> mlp._inner_mlp.moe.router.per_expert_scale
+
+Non-layer (HF -> MG):
+  model.language_model.embed_tokens.weight -> language_model.embedding.word_embeddings.weight
+  model.language_model.norm.weight         -> language_model.decoder.final_layernorm.weight
+
+Tied weights: tie_word_embeddings=True; MG ``output_layer`` writes _extra_state=None placeholder
+(no .weight key); HF ``lm_head.weight`` is reconstructed from embed_tokens at runtime, not in safetensors.
+
+_extra_state placeholders: every TE-managed metadata sibling is written as
+``None``. Transformer Engine 2.2 treats tensor-valued extra_state as pickled
+bytes; an empty tensor therefore raises EOFError during load.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import sys
+from copy import deepcopy
+from os.path import dirname
+
+import torch
+
+
+SCRIPT_DIR = dirname(os.path.abspath(__file__))
+sys.path.append(dirname(dirname(dirname(SCRIPT_DIR))))
+
+from convert_checkpoint.arguments import parse_args  # noqa: E402
+from convert_checkpoint.custom.llava_onevision2.util import (  # noqa: E402
+    load_huggingface_checkpoint,
+    load_megatron_checkpoint,
+    load_megatron_checkpoint_tp_ep,
+    load_megatron_checkpoint_tp_pp_ep,
+    save_huggingface_checkpoint,
+    save_megatron_checkpoint,
+    save_megatron_checkpoint_tp_ep,
+    save_megatron_checkpoint_tp_pp_ep,
+)
+
+
+HF_LLM_PREFIX = "model.language_model."
+MG_LLM_PREFIX = "language_model."
+MG_LAYER_PREFIX = "language_model.decoder.layers."
+
+# HF non-layer LLM key mapping
+NON_LAYER_HF_TO_MG = {
+    "model.language_model.embed_tokens.weight": "language_model.embedding.word_embeddings.weight",
+    "model.language_model.norm.weight": "language_model.decoder.final_layernorm.weight",
+}
+
+# Per-layer norm/layer-scalar mappings: HF tail -> MG tail
+NORM_AND_SCALAR_MAP = {
+    "input_layernorm.weight": "self_attention.linear_qkv.layer_norm_weight",
+    "post_attention_layernorm.weight": "self_attention.post_attention_layernorm.weight",
+    "pre_feedforward_layernorm.weight": "mlp._inner_mlp.pre_feedforward_layernorm.weight",
+    "pre_feedforward_layernorm_2.weight": "mlp._inner_mlp.moe.pre_feedforward_layernorm_2.weight",
+    "post_feedforward_layernorm_1.weight": "mlp._inner_mlp.post_feedforward_layernorm_1.weight",
+    "post_feedforward_layernorm_2.weight": "mlp._inner_mlp.post_feedforward_layernorm_2.weight",
+    "self_attn.q_norm.weight": "self_attention.q_layernorm.weight",
+    "self_attn.k_norm.weight": "self_attention.k_layernorm.weight",
+    "self_attn.o_proj.weight": "self_attention.linear_proj.weight",
+    "layer_scalar": "layer_scalar",
+}
+
+# Router 3-tuple: HF tail -> MG tail (NO transpose, verified [128, 2816] both sides)
+ROUTER_MAP = {
+    "router.proj.weight": "mlp._inner_mlp.moe.router.weight",
+    "router.scale": "mlp._inner_mlp.moe.router.scale",
+    "router.per_expert_scale": "mlp._inner_mlp.moe.router.per_expert_scale",
+}
+
+# Per-layer _extra_state keys to emit (one per fused TE op or normalization).
+# Discovered from dump_full_v3.txt layer 0: 9 weight ops with sibling _extra_state.
+# Plus self_attention.core_attention._extra_state (single, no weight pair).
+PER_LAYER_EXTRA_STATE_TAILS = (
+    "mlp._inner_mlp.dense.linear_fc1._extra_state",
+    "mlp._inner_mlp.dense.linear_fc2._extra_state",
+    "mlp._inner_mlp.moe.pre_feedforward_layernorm_2._extra_state",
+    "mlp._inner_mlp.post_feedforward_layernorm_1._extra_state",
+    "mlp._inner_mlp.post_feedforward_layernorm_2._extra_state",
+    "mlp._inner_mlp.pre_feedforward_layernorm._extra_state",
+    "mlp._post_ffn._extra_state",
+    "post_feedforward_layernorm._extra_state",
+    "self_attention.core_attention._extra_state",
+    "self_attention.k_layernorm._extra_state",
+    "self_attention.linear_proj._extra_state",
+    "self_attention.linear_qkv._extra_state",
+    "self_attention.post_attention_layernorm._extra_state",
+    "self_attention.q_layernorm._extra_state",
+)
+# Per-expert _extra_state (128 experts * 2 fc = 256 keys/layer)
+PER_EXPERT_EXTRA_STATE_TAILS = (
+    "linear_fc1._extra_state",
+    "linear_fc2._extra_state",
+)
+
+# Non-layer _extra_state (from dump line 590-593)
+NON_LAYER_EXTRA_STATE = {
+    "language_model.decoder.final_layernorm._extra_state": "tensor",
+    "language_model.output_layer._extra_state": "none",
+}
+
+
+def _read_hf_text_config(hf_dir: str) -> dict:
+    cfg_path = os.path.join(hf_dir, "config.json")
+    with open(cfg_path, "r", encoding="utf-8") as f:
+        cfg = json.load(f)
+    return cfg["text_config"]
+
+
+def _resolve_hf_dir(args) -> str:
+    """Return the HF directory for config inspection regardless of direction."""
+    if args.load_platform == "huggingface":
+        return args.load_ckpt_path
+    return args.save_ckpt_path
+
+
+def _resolve_layer_meta(text_config: dict) -> tuple[int, int, list[str]]:
+    """Return (num_layers, num_experts, layer_types). All data-driven from HF config."""
+    num_layers = int(text_config["num_hidden_layers"])
+    num_experts = int(text_config["num_experts"])
+    layer_types = list(text_config["layer_types"])
+    if len(layer_types) != num_layers:
+        raise ValueError(
+            f"layer_types length {len(layer_types)} != num_hidden_layers {num_layers}"
+        )
+    return num_layers, num_experts, layer_types
+
+
+def _is_global_layer(layer_types: list[str], layer_idx: int) -> bool:
+    return layer_types[layer_idx] == "full_attention"
+
+
+def _empty_extra_state() -> None:
+    """Placeholder extra_state compatible with Transformer Engine 2.2 loaders."""
+    return None
+
+
+def _qkv_pack_interleaved(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor | None,
+    num_q_heads: int,
+    num_kv_heads: int,
+    head_dim: int,
+    hidden: int,
+) -> torch.Tensor:
+    """Pack HF q/k/v projections into Megatron's per-query-group interleaved layout.
+
+    Megatron's :class:`Gemma4SelfAttention.get_query_key_value_tensors` reshapes
+    ``mixed_qkv`` as ``(num_kv_heads, (heads_per_group + extra) * head_dim)``
+    where ``extra=2`` (sliding, with V) or ``extra=1`` (global, KV-tied, no V).
+    That reshape only round-trips correctly when each query group's heads sit
+    contiguous with their K (and V) head along the row axis. A naive
+    ``torch.cat([Q, K, V], dim=0)`` produces block layout that this reshape
+    interprets as scrambled Q/K/V (see ``.cache/v6_qkv_layout_test.py``,
+    BLOCK diff vs HF reference ~14, INTERLEAVED diff = 0).
+    """
+    heads_per_group = num_q_heads // num_kv_heads
+    parts = [q.view(num_kv_heads, heads_per_group, head_dim, hidden), k.view(num_kv_heads, 1, head_dim, hidden)]
+    if v is not None:
+        parts.append(v.view(num_kv_heads, 1, head_dim, hidden))
+    return torch.cat(parts, dim=1).reshape(-1, hidden)
+
+
+def _qkv_unpack_interleaved(
+    qkv: torch.Tensor,
+    num_q_heads: int,
+    num_kv_heads: int,
+    head_dim: int,
+    hidden: int,
+    has_v: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Inverse of :func:`_qkv_pack_interleaved` — split MG QKV back into HF q/k/[v]."""
+    heads_per_group = num_q_heads // num_kv_heads
+    extra = 2 if has_v else 1
+    qkv_resh = qkv.view(num_kv_heads, heads_per_group + extra, head_dim, hidden)
+    q = qkv_resh[:, :heads_per_group, :, :].reshape(num_q_heads * head_dim, hidden).contiguous()
+    k = qkv_resh[:, heads_per_group:heads_per_group + 1, :, :].reshape(num_kv_heads * head_dim, hidden).contiguous()
+    if has_v:
+        v = qkv_resh[:, heads_per_group + 1:, :, :].reshape(num_kv_heads * head_dim, hidden).contiguous()
+        return q, k, v
+    return q, k, None
+
+
+def _hf_to_mg_layer(
+    source: dict,
+    target: dict,
+    layer_idx: int,
+    num_experts: int,
+    is_global: bool,
+) -> None:
+    """Convert one LLM layer from HF to MG, in-place into ``target``."""
+    hf_pfx = f"{HF_LLM_PREFIX}layers.{layer_idx}."
+    mg_pfx = f"{MG_LAYER_PREFIX}{layer_idx}."
+
+    # 1. Norms + layer_scalar + o_proj (1:1 rename)
+    for hf_tail, mg_tail in NORM_AND_SCALAR_MAP.items():
+        target[mg_pfx + mg_tail] = source[hf_pfx + hf_tail]
+
+    # 2. post_feedforward_layernorm: HF single key -> MG dual alias
+    pffn = source[hf_pfx + "post_feedforward_layernorm.weight"]
+    target[mg_pfx + "mlp._post_ffn.weight"] = pffn
+    target[mg_pfx + "post_feedforward_layernorm.weight"] = pffn
+
+    # 3. Attention QKV concat dim=0
+    q = source[hf_pfx + "self_attn.q_proj.weight"]
+    k = source[hf_pfx + "self_attn.k_proj.weight"]
+    if is_global:
+        v = None
+        num_kv = _GLOBAL_NUM_KV
+        head_dim = _GLOBAL_HEAD_DIM
+    else:
+        v = source[hf_pfx + "self_attn.v_proj.weight"]
+        num_kv = _SLIDING_NUM_KV
+        head_dim = _SLIDING_HEAD_DIM
+    qkv = _qkv_pack_interleaved(q, k, v, _NUM_Q_HEADS, num_kv, head_dim, _HIDDEN)
+    target[mg_pfx + "self_attention.linear_qkv.weight"] = qkv
+
+    # 4. Dense MLP fused fc1 (gate + up concat dim=0)
+    gate = source[hf_pfx + "mlp.gate_proj.weight"]
+    up = source[hf_pfx + "mlp.up_proj.weight"]
+    target[mg_pfx + "mlp._inner_mlp.dense.linear_fc1.weight"] = torch.cat([gate, up], dim=0)
+    target[mg_pfx + "mlp._inner_mlp.dense.linear_fc2.weight"] = source[
+        hf_pfx + "mlp.down_proj.weight"
+    ]
+
+    # 5. MoE experts: 3D HF tensor -> 128 per-expert 2D tensors
+    gate_up_3d = source[hf_pfx + "experts.gate_up_proj"]  # [128, fc1_out, hidden]
+    down_3d = source[hf_pfx + "experts.down_proj"]  # [128, hidden, fc2_in]
+    if gate_up_3d.shape[0] != num_experts or down_3d.shape[0] != num_experts:
+        raise ValueError(
+            f"layer {layer_idx}: expert dim mismatch "
+            f"gate_up={gate_up_3d.shape}, down={down_3d.shape}, num_experts={num_experts}"
+        )
+    expert_pfx = mg_pfx + "mlp._inner_mlp.moe.experts.local_experts."
+    for e in range(num_experts):
+        target[f"{expert_pfx}{e}.linear_fc1.weight"] = gate_up_3d[e].contiguous()
+        target[f"{expert_pfx}{e}.linear_fc2.weight"] = down_3d[e].contiguous()
+
+    # 6. Router 3-tuple (no transpose, verified)
+    for hf_tail, mg_tail in ROUTER_MAP.items():
+        target[mg_pfx + mg_tail] = source[hf_pfx + hf_tail]
+
+    # 7. _extra_state placeholders (per-layer common)
+    for tail in PER_LAYER_EXTRA_STATE_TAILS:
+        target[mg_pfx + tail] = _empty_extra_state()
+    # 7b. _extra_state per-expert
+    for e in range(num_experts):
+        for tail in PER_EXPERT_EXTRA_STATE_TAILS:
+            target[f"{expert_pfx}{e}.{tail}"] = _empty_extra_state()
+
+
+def _mg_to_hf_layer(
+    source: dict,
+    target: dict,
+    layer_idx: int,
+    num_experts: int,
+    is_global: bool,
+) -> None:
+    """Convert one LLM layer from MG to HF, in-place into ``target``."""
+    hf_pfx = f"{HF_LLM_PREFIX}layers.{layer_idx}."
+    mg_pfx = f"{MG_LAYER_PREFIX}{layer_idx}."
+
+    # 1. Norms + layer_scalar + o_proj (1:1 inverse rename)
+    for hf_tail, mg_tail in NORM_AND_SCALAR_MAP.items():
+        target[hf_pfx + hf_tail] = source[mg_pfx + mg_tail]
+
+    # 2. post_feedforward_layernorm: pick canonical alias source (mlp._post_ffn.weight)
+    target[hf_pfx + "post_feedforward_layernorm.weight"] = source[
+        mg_pfx + "mlp._post_ffn.weight"
+    ]
+
+    # 3. Attention QKV split dim=0 (sliding head_dim=256, global head_dim=512)
+    qkv = source[mg_pfx + "self_attention.linear_qkv.weight"]
+    if is_global:
+        q, k, _ = _qkv_unpack_interleaved(
+            qkv, _NUM_Q_HEADS, _GLOBAL_NUM_KV, _GLOBAL_HEAD_DIM, _HIDDEN, has_v=False,
+        )
+        target[hf_pfx + "self_attn.q_proj.weight"] = q
+        target[hf_pfx + "self_attn.k_proj.weight"] = k
+    else:
+        q, k, v = _qkv_unpack_interleaved(
+            qkv, _NUM_Q_HEADS, _SLIDING_NUM_KV, _SLIDING_HEAD_DIM, _HIDDEN, has_v=True,
+        )
+        target[hf_pfx + "self_attn.q_proj.weight"] = q
+        target[hf_pfx + "self_attn.k_proj.weight"] = k
+        target[hf_pfx + "self_attn.v_proj.weight"] = v
+
+    # 4. Dense MLP: split fused fc1 -> gate + up (dim=0, gate first)
+    fc1 = source[mg_pfx + "mlp._inner_mlp.dense.linear_fc1.weight"]
+    half = fc1.shape[0] // 2
+    if fc1.shape[0] != 2 * half:
+        raise ValueError(
+            f"layer {layer_idx}: dense linear_fc1 dim0 {fc1.shape[0]} not divisible by 2"
+        )
+    gate, up = torch.split(fc1, [half, half], dim=0)
+    target[hf_pfx + "mlp.gate_proj.weight"] = gate.contiguous()
+    target[hf_pfx + "mlp.up_proj.weight"] = up.contiguous()
+    target[hf_pfx + "mlp.down_proj.weight"] = source[
+        mg_pfx + "mlp._inner_mlp.dense.linear_fc2.weight"
+    ]
+
+    # 5. MoE experts: 128 per-expert 2D -> 3D HF tensor
+    expert_pfx = mg_pfx + "mlp._inner_mlp.moe.experts.local_experts."
+    gate_up_list = [source[f"{expert_pfx}{e}.linear_fc1.weight"] for e in range(num_experts)]
+    down_list = [source[f"{expert_pfx}{e}.linear_fc2.weight"] for e in range(num_experts)]
+    target[hf_pfx + "experts.gate_up_proj"] = torch.stack(gate_up_list, dim=0).contiguous()
+    target[hf_pfx + "experts.down_proj"] = torch.stack(down_list, dim=0).contiguous()
+
+    # 6. Router 3-tuple (inverse, no transpose)
+    for hf_tail, mg_tail in ROUTER_MAP.items():
+        target[hf_pfx + hf_tail] = source[mg_pfx + mg_tail]
+
+
+# Module-level QKV split sizes; resolved from HF text_config at startup.
+_SLIDING_Q_DIM = 0
+_SLIDING_K_DIM = 0
+_SLIDING_V_DIM = 0
+_GLOBAL_Q_DIM = 0
+_GLOBAL_K_DIM = 0
+_NUM_Q_HEADS = 0
+_HIDDEN = 0
+_SLIDING_NUM_KV = 0
+_SLIDING_HEAD_DIM = 0
+_GLOBAL_NUM_KV = 0
+_GLOBAL_HEAD_DIM = 0
+
+
+def _resolve_qkv_dims(text_config: dict) -> None:
+    """Set module-level QKV split sizes from HF text_config (data-driven)."""
+    global _SLIDING_Q_DIM, _SLIDING_K_DIM, _SLIDING_V_DIM, _GLOBAL_Q_DIM, _GLOBAL_K_DIM
+    global _NUM_Q_HEADS, _HIDDEN
+    global _SLIDING_NUM_KV, _SLIDING_HEAD_DIM, _GLOBAL_NUM_KV, _GLOBAL_HEAD_DIM
+    sliding_head_dim = int(text_config["head_dim"])
+    global_head_dim = int(text_config.get("global_head_dim", sliding_head_dim))
+    num_q_heads = int(text_config["num_attention_heads"])
+    num_kv_heads = int(text_config["num_key_value_heads"])
+    num_global_kv = int(text_config.get("num_global_key_value_heads", num_kv_heads))
+    hidden = int(text_config["hidden_size"])
+    _SLIDING_Q_DIM = num_q_heads * sliding_head_dim
+    _SLIDING_K_DIM = num_kv_heads * sliding_head_dim
+    _SLIDING_V_DIM = num_kv_heads * sliding_head_dim
+    _GLOBAL_Q_DIM = num_q_heads * global_head_dim
+    _GLOBAL_K_DIM = num_global_kv * global_head_dim
+    _NUM_Q_HEADS = num_q_heads
+    _HIDDEN = hidden
+    _SLIDING_NUM_KV = num_kv_heads
+    _SLIDING_HEAD_DIM = sliding_head_dim
+    _GLOBAL_NUM_KV = num_global_kv
+    _GLOBAL_HEAD_DIM = global_head_dim
+
+
+def _convert_hf_to_mg(args, text_config: dict) -> dict:
+    """Build the full MG state_dict (single rank, TP=PP=EP=1) from HF safetensors."""
+    num_layers, num_experts, layer_types = _resolve_layer_meta(text_config)
+    print(
+        f" > LLM body: num_layers={num_layers}, num_experts={num_experts}, "
+        f"global_layers={[i for i, t in enumerate(layer_types) if t == 'full_attention']}"
+    )
+
+    source = load_huggingface_checkpoint(args.load_ckpt_path)
+    target: dict = {}
+
+    # 1. Non-layer keys (embed, final norm)
+    for hf_key, mg_key in NON_LAYER_HF_TO_MG.items():
+        target[mg_key] = source[hf_key]
+
+    # 2. Per-layer body
+    for i in range(num_layers):
+        is_global = _is_global_layer(layer_types, i)
+        _hf_to_mg_layer(source, target, i, num_experts, is_global)
+
+    # 3. Non-layer _extra_state placeholders (final_layernorm tensor, output_layer None)
+    target["language_model.decoder.final_layernorm._extra_state"] = _empty_extra_state()
+    target["language_model.output_layer._extra_state"] = None
+
+    print(f" > MG LLM keys produced: {len(target)}")
+    return target
+
+
+def _convert_mg_to_hf(args, text_config: dict, source: dict) -> dict:
+    """Build the full HF state_dict from a single-rank MG ``model`` dict."""
+    num_layers, num_experts, layer_types = _resolve_layer_meta(text_config)
+
+    target: dict = {}
+    # 1. Non-layer keys
+    for hf_key, mg_key in NON_LAYER_HF_TO_MG.items():
+        target[hf_key] = source[mg_key]
+
+    # 2. Per-layer body
+    for i in range(num_layers):
+        is_global = _is_global_layer(layer_types, i)
+        _mg_to_hf_layer(source, target, i, num_experts, is_global)
+
+    print(f" > HF LLM keys produced: {len(target)}")
+    return target
+
+
+def _get_non_ep_model_source(state_dict):
+    first = state_dict[0]
+    if isinstance(first, dict):
+        return first["model"]
+    first_rank = first[0]
+    if isinstance(first_rank, dict):
+        return first_rank["model"]
+    raise TypeError("Unsupported non-EP checkpoint structure")
+
+
+# =====================================================================
+# P2.8 PP partitioning + per-stage view construction.
+#
+# Layer numbering contract (verified against Megatron Core
+# ``transformer_block.py:691-696`` and ``transformer_layer.py:372``):
+#   - TransformerLayer.layer_number is GLOBAL (1-based) and equals
+#     ``local + get_transformer_layer_offset(config)``.
+#   - When state_dict keys are constructed, TransformerBlock subtracts the
+#     stage offset, so on-disk keys use LOCAL indices: each PP rank's shard
+#     stores ``language_model.decoder.layers.{0..local_num-1}.*`` regardless
+#     of where those layers sit in the global model.
+#
+# Non-layer routing across PP stages:
+#   - ``language_model.embedding.word_embeddings.weight``     -> PP stage 0 only
+#   - ``language_model.decoder.final_layernorm.weight``       -> last PP stage only
+#   - ``language_model.decoder.final_layernorm._extra_state`` -> last PP stage only
+#   - ``language_model.output_layer._extra_state`` (= None)   -> last PP stage only
+#     (Gemma4 has tie_word_embeddings=True, so output_layer carries no .weight,
+#      only the None placeholder.)
+#   - vit / adapter / vision_patch live solely on PP stage 0; this file does
+#     NOT touch them — they are merged in by ``merge_megatron.py``.
+# =====================================================================
+
+
+def _partition_layers_for_pp(
+    num_layers: int, pp: int, custom: list[int] | None
+) -> list[tuple[int, int]]:
+    """Return the (start, end) global layer range for each PP stage.
+
+    ``end`` is exclusive; ``end - start`` equals the stage's local layer count.
+    When ``custom`` is None the layers are split evenly (Megatron's default
+    even-split policy); otherwise ``custom`` must list per-stage counts that
+    sum to ``num_layers``.
+    """
+    if pp <= 0:
+        raise ValueError(f"pp must be positive, got {pp}")
+    if custom is not None:
+        if len(custom) != pp:
+            raise ValueError(
+                f"custom_pipeline_layers length {len(custom)} != pp {pp}"
+            )
+        if sum(custom) != num_layers:
+            raise ValueError(
+                f"custom_pipeline_layers sum {sum(custom)} != num_layers {num_layers}"
+            )
+        sizes = list(custom)
+    else:
+        if num_layers % pp != 0:
+            raise ValueError(
+                f"num_layers ({num_layers}) must be divisible by pp ({pp}) when no "
+                "custom_pipeline_layers is supplied; pass --custom_pipeline_layers "
+                "to override (e.g. '8,8,7,7')."
+            )
+        per_stage = num_layers // pp
+        sizes = [per_stage] * pp
+
+    ranges: list[tuple[int, int]] = []
+    cursor = 0
+    for size in sizes:
+        ranges.append((cursor, cursor + size))
+        cursor += size
+    return ranges
+
+
+def _per_stage_target_view(
+    full_target: dict, stage_idx: int, pp: int, ranges: list[tuple[int, int]]
+) -> dict:
+    """Return the LOCAL-indexed slice of ``full_target`` belonging to PP stage ``stage_idx``.
+
+    Layer keys ``language_model.decoder.layers.{global}.*`` belonging to this
+    stage are renamed to ``language_model.decoder.layers.{local}.*`` where
+    ``local = global - ranges[stage_idx][0]``. Non-layer keys are routed:
+    embeddings -> stage 0 only; final_layernorm + output_layer -> last stage only.
+    """
+    start, end = ranges[stage_idx]
+    is_first = stage_idx == 0
+    is_last = stage_idx == pp - 1
+    layer_re = re.compile(
+        r"^(language_model\.decoder\.layers\.)(\d+)(\..*)$"
+    )
+
+    out: dict = {}
+    for key, value in full_target.items():
+        m = layer_re.match(key)
+        if m is not None:
+            gid = int(m.group(2))
+            if start <= gid < end:
+                local = gid - start
+                out[f"{m.group(1)}{local}{m.group(3)}"] = value
+            continue
+        if key == "language_model.embedding.word_embeddings.weight":
+            if is_first:
+                out[key] = value
+            continue
+        if key.startswith("language_model.decoder.final_layernorm.") or key.startswith(
+            "language_model.output_layer."
+        ):
+            if is_last:
+                out[key] = value
+            continue
+        # Defensive: any other top-level LLM key is treated as PP-stage-0-only
+        # (matches embedding's policy). Add explicit branches above if new
+        # non-layer keys appear in future Megatron versions.
+        if is_first:
+            out[key] = value
+    return out
+
+
+def _shard_target_to_tp_pp_ep(
+    full_target: dict,
+    tp: int,
+    pp: int,
+    ep: int,
+    num_experts: int,
+    pp_ranges: list[tuple[int, int]],
+) -> list[list[list[dict]]]:
+    """Return a 3D ``[pp][tp][ep]`` nested list of per-rank ``{"model": ...}`` dicts.
+
+    Each PP stage gets its LOCAL-indexed slice via ``_per_stage_target_view``,
+    then that slice is TP+EP-sharded by reusing ``_shard_target_to_tp_ep``.
+    """
+    state_dict: list[list[list[dict]]] = []
+    for p in range(pp):
+        per_stage_target = _per_stage_target_view(full_target, p, pp, pp_ranges)
+        tp_ep_2d = _shard_target_to_tp_ep(per_stage_target, tp, ep, num_experts)
+        state_dict.append(tp_ep_2d)
+    return state_dict
+
+
+def _gather_full_target_from_tp_pp_ep(
+    state_dict: list[list[list[dict]]],
+    num_experts: int,
+    pp_ranges: list[tuple[int, int]],
+) -> dict:
+    """Inverse of ``_shard_target_to_tp_pp_ep``.
+
+    For each PP stage, un-shards the [tp][ep] 2D view back to a per-stage
+    LOCAL-indexed target via ``_unshard_tp_ep_to_target``, then re-maps LOCAL
+    layer indices back to GLOBAL by adding ``pp_ranges[stage_idx][0]``.
+    Non-layer keys (embedding / final_layernorm / output_layer) are taken from
+    the stage where they live.
+    """
+    pp = len(state_dict)
+    layer_re = re.compile(
+        r"^(language_model\.decoder\.layers\.)(\d+)(\..*)$"
+    )
+
+    full: dict = {}
+    for p in range(pp):
+        per_stage = _unshard_tp_ep_to_target(state_dict[p], num_experts)
+        start, _ = pp_ranges[p]
+        for key, value in per_stage.items():
+            m = layer_re.match(key)
+            if m is not None:
+                local = int(m.group(2))
+                gid = local + start
+                full[f"{m.group(1)}{gid}{m.group(3)}"] = value
+            else:
+                full[key] = value
+    return full
+
+
+def _parse_custom_pipeline_layers(raw: str | None) -> list[int] | None:
+    if raw is None or raw == "":
+        return None
+    parts = [p.strip() for p in raw.split(",") if p.strip()]
+    return [int(p) for p in parts]
+
+
+# =====================================================================
+# P2.7 TP + EP sharding / un-sharding helpers.
+#
+# Sharding contract (validated against Megatron-Core layer specs and the
+# Qwen3-30B-A3B reference under custom/llava_onevision2/):
+#
+#   TP axis splits per the mapping below; EP axis splits the local-experts
+#   list (contiguous chunks of expert IDs).
+#
+#   Layout decision tree (P2.8):
+#       pp==1 ep==1  -> 1D  ``mp_rank_TT/``
+#       pp==1 ep>1   -> 2D  ``mp_rank_TT_EEE/``
+#       pp>1  any ep -> 3D  ``mp_rank_TT_PPP_EEE/`` (always 3D when PP>1,
+#                            with EP=1 collapsing to a single ep slot — this
+#                            follows the Qwen3 reference convention).
+#
+#   Replicated across BOTH TP and EP (every shard holds an identical copy):
+#       norms, layer_scalar, router.{proj.weight, scale, per_expert_scale},
+#       all *._extra_state placeholders, output_layer._extra_state (None).
+#
+#   TP-sharded, EP-replicated (every EP shard holds the same TP slice):
+#       embedding.word_embeddings.weight              -> chunk dim=0
+#       self_attention.linear_qkv.weight              -> chunk dim=0
+#                                                        (INTERLEAVED layout;
+#                                                         requires tp | num_kv_heads)
+#       self_attention.linear_proj.weight             -> chunk dim=1
+#       mlp._inner_mlp.dense.linear_fc1.weight        -> per-half (gate,up)
+#                                                        chunk dim=0 then re-cat
+#       mlp._inner_mlp.dense.linear_fc2.weight        -> chunk dim=1
+#
+#   EP-sharded (each EP shard owns 1/EP of the experts), with the per-expert
+#   tensors themselves TP-sharded:
+#       mlp._inner_mlp.moe.experts.local_experts.{i}.linear_fc1.weight
+#           -> per-half chunk dim=0 (gate+up fused)
+#       mlp._inner_mlp.moe.experts.local_experts.{i}.linear_fc2.weight
+#           -> chunk dim=1
+#       (per-expert ``*._extra_state`` placeholders also follow EP sharding,
+#        but stay TP-replicated as plain placeholders.)
+#
+#   Local expert indexing: each EP rank stores its local experts as
+#   ``local_experts.{0..exp_per_ep-1}`` (re-indexed from 0), matching
+#   Megatron's ``SequentialMLP`` convention.
+# =====================================================================
+
+
+def _expert_id_map(num_experts: int, ep: int) -> list[list[int]]:
+    if num_experts % ep != 0:
+        raise ValueError(
+            f"num_experts ({num_experts}) must be divisible by ep ({ep}); "
+            "Gemma4-VL ships 128 experts which divides evenly by EP=1/2/4/8/16/32/64/128."
+        )
+    exp_per_ep = num_experts // ep
+    return [list(range(e * exp_per_ep, (e + 1) * exp_per_ep)) for e in range(ep)]
+
+
+def _split_fused_gate_up(weight: torch.Tensor, tp: int, t: int) -> torch.Tensor:
+    """Split a fused ``cat([gate, up], dim=0)`` weight for TP rank ``t``.
+
+    Naive ``chunk(tp, dim=0)`` would assign rank 0 the entire ``gate`` half and
+    rank 1 the entire ``up`` half (for tp=2), which is wrong: each TP rank must
+    receive a slice of BOTH gate and up so its local SwiGLU sees matching column
+    ranges. The Megatron contract is ``cat([gate_chunk_t, up_chunk_t], dim=0)``.
+    """
+    fc1_out = weight.shape[0]
+    if fc1_out % 2 != 0:
+        raise ValueError(f"fused fc1 dim0 ({fc1_out}) is not even")
+    half = fc1_out // 2
+    gate, up = weight[:half], weight[half:]
+    if half % tp != 0:
+        raise ValueError(
+            f"fused fc1 half ({half}) not divisible by tp ({tp}); "
+            "intermediate_size must be divisible by tp"
+        )
+    return torch.cat([gate.chunk(tp, dim=0)[t], up.chunk(tp, dim=0)[t]], dim=0).contiguous()
+
+
+def _gather_fused_gate_up(shards: list[torch.Tensor]) -> torch.Tensor:
+    halves = [s.shape[0] // 2 for s in shards]
+    if any(s.shape[0] != 2 * h for s, h in zip(shards, halves)):
+        raise ValueError("each TP shard of fused fc1 must have an even dim0")
+    gates = [s[: halves[i]] for i, s in enumerate(shards)]
+    ups = [s[halves[i] :] for i, s in enumerate(shards)]
+    return torch.cat(gates + ups, dim=0).contiguous()
+
+
+def _shard_value_for_tp(key: str, value, tp: int, t: int):
+    if not isinstance(value, torch.Tensor) or value.numel() == 0 or tp == 1:
+        return value
+
+    if key == "language_model.embedding.word_embeddings.weight":
+        return value.chunk(tp, dim=0)[t].contiguous()
+
+    if key.endswith(".self_attention.linear_qkv.weight"):
+        # INTERLEAVED layout: dim=0 chunk valid only because tp | num_kv_heads
+        # (Gemma4 sliding=8, global=2, both divisible by tp=2).
+        return value.chunk(tp, dim=0)[t].contiguous()
+
+    if key.endswith(".self_attention.linear_proj.weight"):
+        return value.chunk(tp, dim=1)[t].contiguous()
+
+    if key.endswith(".mlp._inner_mlp.dense.linear_fc1.weight"):
+        return _split_fused_gate_up(value, tp, t)
+
+    if key.endswith(".mlp._inner_mlp.dense.linear_fc2.weight"):
+        return value.chunk(tp, dim=1)[t].contiguous()
+
+    if key.endswith(".linear_fc1.weight") and ".local_experts." in key:
+        return _split_fused_gate_up(value, tp, t)
+
+    if key.endswith(".linear_fc2.weight") and ".local_experts." in key:
+        return value.chunk(tp, dim=1)[t].contiguous()
+
+    return value
+
+
+def _shard_target_to_tp_ep(
+    target: dict, tp: int, ep: int, num_experts: int
+) -> list[list[dict]]:
+    if num_experts % ep != 0:
+        raise ValueError(f"num_experts ({num_experts}) not divisible by ep ({ep})")
+    expert_map = _expert_id_map(num_experts, ep)
+    expert_owner = [-1] * num_experts
+    expert_local = [-1] * num_experts
+    for e_idx, ids in enumerate(expert_map):
+        for local_idx, gid in enumerate(ids):
+            expert_owner[gid] = e_idx
+            expert_local[gid] = local_idx
+
+    expert_key_re = re.compile(
+        r"(.*\.mlp\._inner_mlp\.moe\.experts\.local_experts\.)(\d+)(\..*)"
+    )
+
+    state_dict: list[list[dict]] = []
+    for t in range(tp):
+        ep_list: list[dict] = []
+        for e in range(ep):
+            shard: dict = {}
+            for key, value in target.items():
+                m = expert_key_re.match(key)
+                if m is not None:
+                    gid = int(m.group(2))
+                    if expert_owner[gid] != e:
+                        continue
+                    local_id = expert_local[gid]
+                    new_key = f"{m.group(1)}{local_id}{m.group(3)}"
+                    shard[new_key] = _shard_value_for_tp(key, value, tp, t)
+                else:
+                    shard[key] = _shard_value_for_tp(key, value, tp, t)
+            ep_list.append({"model": shard})
+        state_dict.append(ep_list)
+    return state_dict
+
+
+def _gather_value_for_tp(key: str, shards: list, tp: int):
+    head = shards[0]
+    if not isinstance(head, torch.Tensor) or head.numel() == 0 or tp == 1:
+        return head
+
+    if key == "language_model.embedding.word_embeddings.weight":
+        return torch.cat(shards, dim=0).contiguous()
+
+    if key.endswith(".self_attention.linear_qkv.weight"):
+        return torch.cat(shards, dim=0).contiguous()
+
+    if key.endswith(".self_attention.linear_proj.weight"):
+        return torch.cat(shards, dim=1).contiguous()
+
+    if key.endswith(".mlp._inner_mlp.dense.linear_fc1.weight"):
+        return _gather_fused_gate_up(shards)
+
+    if key.endswith(".mlp._inner_mlp.dense.linear_fc2.weight"):
+        return torch.cat(shards, dim=1).contiguous()
+
+    if key.endswith(".linear_fc1.weight") and ".local_experts." in key:
+        return _gather_fused_gate_up(shards)
+
+    if key.endswith(".linear_fc2.weight") and ".local_experts." in key:
+        return torch.cat(shards, dim=1).contiguous()
+
+    return head
+
+
+def _unshard_tp_ep_to_target(
+    state_dict: list[list[dict]], num_experts: int
+) -> dict:
+    tp = len(state_dict)
+    ep = len(state_dict[0])
+    if num_experts % ep != 0:
+        raise ValueError(f"num_experts ({num_experts}) not divisible by ep ({ep})")
+    exp_per_ep = num_experts // ep
+
+    expert_local_re = re.compile(
+        r"(.*\.mlp\._inner_mlp\.moe\.experts\.local_experts\.)(\d+)(\..*)"
+    )
+
+    target: dict = {}
+    canonical_keys: list[str] = []
+    for k in state_dict[0][0]["model"].keys():
+        if expert_local_re.match(k) is not None:
+            continue
+        canonical_keys.append(k)
+    for key in canonical_keys:
+        shards = [state_dict[t][0]["model"][key] for t in range(tp)]
+        target[key] = _gather_value_for_tp(key, shards, tp)
+
+    sample_expert_keys = [
+        k for k in state_dict[0][0]["model"].keys() if expert_local_re.match(k) is not None
+    ]
+    expert_tails: set[tuple[str, str]] = set()
+    for k in sample_expert_keys:
+        m = expert_local_re.match(k)
+        if m is not None:
+            expert_tails.add((m.group(1), m.group(3)))
+
+    for e in range(ep):
+        for local_id in range(exp_per_ep):
+            gid = e * exp_per_ep + local_id
+            for prefix, suffix in expert_tails:
+                local_key = f"{prefix}{local_id}{suffix}"
+                global_key = f"{prefix}{gid}{suffix}"
+                shards = [state_dict[t][e]["model"][local_key] for t in range(tp)]
+                target[global_key] = _gather_value_for_tp(local_key, shards, tp)
+
+    return target
+
+
+# =====================================================================
+# Entry point: dispatch on (load_platform, save_platform).
+# =====================================================================
+args = parse_args()
+text_config = _read_hf_text_config(_resolve_hf_dir(args))
+_resolve_qkv_dims(text_config)
+
+
+if (args.load_platform, args.save_platform) == ("huggingface", "mcore"):
+    print(" ====== convert Gemma4-VL LLM body from HuggingFace to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    pp = args.pipeline_model_parallel_size or 1
+    ep = args.expert_parallel_size
+    custom = _parse_custom_pipeline_layers(getattr(args, "custom_pipeline_layers", None))
+    target = _convert_hf_to_mg(args, text_config)
+    num_experts = text_config["num_experts"]
+    num_layers = int(text_config["num_hidden_layers"])
+    if pp > 1:
+        ep_eff = ep if ep is not None and ep > 0 else 1
+        pp_ranges = _partition_layers_for_pp(num_layers, pp, custom)
+        state_dict = _shard_target_to_tp_pp_ep(target, tp, pp, ep_eff, num_experts, pp_ranges)
+        save_megatron_checkpoint_tp_pp_ep(
+            state_dict, os.path.join(args.save_ckpt_path, "release")
+        )
+    elif ep is None or ep == 1:
+        state_dict = [{"model": deepcopy(target)} for _ in range(tp)]
+        save_megatron_checkpoint(state_dict, os.path.join(args.save_ckpt_path, "release"))
+    else:
+        state_dict = _shard_target_to_tp_ep(target, tp, ep, num_experts)
+        save_megatron_checkpoint_tp_ep(state_dict, os.path.join(args.save_ckpt_path, "release"))
+
+elif (args.load_platform, args.save_platform) == ("mcore", "huggingface"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL LLM body from Megatron Core to HuggingFace ======")
+    pp = args.pipeline_model_parallel_size or 1
+    ep = args.expert_parallel_size
+    custom = _parse_custom_pipeline_layers(getattr(args, "custom_pipeline_layers", None))
+    num_experts = text_config["num_experts"]
+    num_layers = int(text_config["num_hidden_layers"])
+    if pp > 1:
+        ep_eff = ep if ep is not None and ep > 0 else 1
+        pp_ranges = _partition_layers_for_pp(num_layers, pp, custom)
+        state_dict = load_megatron_checkpoint_tp_pp_ep(args.load_ckpt_path)
+        source = _gather_full_target_from_tp_pp_ep(state_dict, num_experts, pp_ranges)
+    elif ep is not None and ep > 1:
+        state_dict = load_megatron_checkpoint_tp_ep(args.load_ckpt_path)
+        source = _unshard_tp_ep_to_target(state_dict, num_experts)
+    else:
+        state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+        source = _get_non_ep_model_source(state_dict)
+    target = _convert_mg_to_hf(args, text_config, source)
+    save_huggingface_checkpoint(target, args.save_ckpt_path)
+
+elif (args.load_platform, args.save_platform) == ("mcore", "mcore"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL LLM body from Megatron Core to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    pp = args.pipeline_model_parallel_size or 1
+    ep = args.expert_parallel_size
+    custom = _parse_custom_pipeline_layers(getattr(args, "custom_pipeline_layers", None))
+    num_experts = text_config["num_experts"]
+    num_layers = int(text_config["num_hidden_layers"])
+    state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+    source = _get_non_ep_model_source(state_dict)
+    if pp > 1:
+        ep_eff = ep if ep is not None and ep > 0 else 1
+        pp_ranges = _partition_layers_for_pp(num_layers, pp, custom)
+        new_state_dict = _shard_target_to_tp_pp_ep(
+            source, tp, pp, ep_eff, num_experts, pp_ranges
+        )
+        save_megatron_checkpoint_tp_pp_ep(
+            new_state_dict, os.path.join(args.save_ckpt_path, "release")
+        )
+    elif ep is None or ep == 1:
+        new_state_dict = [{"model": deepcopy(source)} for _ in range(tp)]
+        save_megatron_checkpoint(new_state_dict, os.path.join(args.save_ckpt_path, "release"))
+    else:
+        new_state_dict = _shard_target_to_tp_ep(source, tp, ep, num_experts)
+        save_megatron_checkpoint_tp_ep(
+            new_state_dict, os.path.join(args.save_ckpt_path, "release")
+        )
+
+else:
+    raise NotImplementedError(
+        f"Unsupported (load, save) = ({args.load_platform}, {args.save_platform})"
+    )
diff --git a/tools/convert_checkpoint/custom/gemma4_vl/merge_megatron.py b/tools/convert_checkpoint/custom/gemma4_vl/merge_megatron.py
new file mode 100644
index 00000000..3ec88b87
--- /dev/null
+++ b/tools/convert_checkpoint/custom/gemma4_vl/merge_megatron.py
@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+# -*- coding: UTF-8 -*-
+################################################################################
+#
+# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
+#
+################################################################################
+"""Gemma4-VL Megatron checkpoint merger.
+
+Merges 4 separately-converted Megatron checkpoint trees produced by
+``llm.py``, ``vit.py``, ``vision_patch.py`` and ``adapter.py`` into one
+Megatron release tree.
+
+Layout selection (mirrors the OV2 ``merge_megatron_qwen3_30b_a3b.py`` template):
+
+  - PP=1, EP=1            -> 1D ``mp_rank_TT``
+  - PP=1, EP>1            -> 2D ``mp_rank_TT_EEE`` (this file's P2.7 path)
+  - PP>1                  -> 3D ``mp_rank_TT_PPP_EEE`` (staged for P2.8)
+
+The LLM body uses the appropriate EP-aware loader; the non-LLM modules
+(ViT, adapter, vision_patch) are always saved as 1D ``mp_rank_TT`` by
+their standalone converters and broadcast across the EP axis at merge
+time, matching the OV2 Qwen3 convention.
+
+Inputs (each is a directory containing ``mp_rank_*``):
+    --language_model_path  : Gemma4 LLM body
+    --vision_model_path    : Gemma4 ViT body (1D TP only)
+    --vision_patch         : Gemma4 ViT patch+pos+std (1D TP only)
+    --adapter_path         : Gemma4 vision-language adapter (1D TP only)
+
+Output:
+    --save_ckpt_path       : Merged Megatron release tree
+
+Usage (TP=PP=EP=1, the P2.5 first-round target):
+    python merge_megatron.py \\
+        --language_model_path /tmp/gemma4_llm/release \\
+        --vision_model_path   /tmp/gemma4_vit/release \\
+        --vision_patch        /tmp/gemma4_patch/release \\
+        --adapter_path        /tmp/gemma4_adapter/release \\
+        --save_ckpt_path      /ov2/pretrain_models/megatron/gemma-4-26B-A4B-it-tp1pp1ep1/release \\
+        --tensor_model_parallel_size 1 \\
+        --pipeline_model_parallel_size 1 \\
+        --expert_parallel_size 1
+"""
+
+import argparse
+import os
+import sys
+from os.path import dirname
+
+
+SCRIPT_DIR = dirname(os.path.abspath(__file__))
+sys.path.append(dirname(dirname(dirname(SCRIPT_DIR))))
+
+from convert_checkpoint.custom.llava_onevision2.util import (  # noqa: E402
+    load_megatron_checkpoint,
+    load_megatron_checkpoint_tp_ep,
+    load_megatron_checkpoint_tp_pp_ep,
+    save_megatron_checkpoint,
+    save_megatron_checkpoint_tp_ep,
+    save_megatron_checkpoint_tp_pp_ep,
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Gemma4-VL Megatron Merger", allow_abbrev=False)
+    group = parser.add_argument_group(title="checkpoint")
+    group.add_argument("--language_model_path", type=str, required=True,
+                       help="Path to language model release dir (mp_rank_*).")
+    group.add_argument("--vision_model_path", type=str, required=True,
+                       help="Path to vision model release dir (mp_rank_*).")
+    group.add_argument("--vision_patch", type=str, required=True,
+                       help="Path to vision patch release dir (mp_rank_*).")
+    group.add_argument("--adapter_path", type=str, required=True,
+                       help="Path to adapter release dir (mp_rank_*).")
+    group.add_argument("--save_ckpt_path", type=str, required=True,
+                       help="Path to save merged checkpoint.")
+    group.add_argument("--megatron_path", type=str, default=None,
+                       help="Base directory of Megatron repository.")
+    group.add_argument("--tensor_model_parallel_size", type=int, default=1,
+                       help="Tensor parallel size.")
+    group.add_argument("--pipeline_model_parallel_size", type=int, default=1,
+                       help="Pipeline parallel size.")
+    group.add_argument("--expert_parallel_size", type=int, default=1,
+                       help="Expert parallel size (1 = disabled).")
+    return parser.parse_args()
+
+
+def merge_dict(source, destination):
+    for key, value in source.items():
+        if isinstance(value, dict):
+            node = destination.setdefault(key, {})
+            merge_dict(value, node)
+        else:
+            destination[key] = value
+
+
+def merge_module_1d_into_language_2d(language_model, module, module_name):
+    """Broadcast a 1D ``[tp]`` non-LLM module into all EP shards of each TP rank."""
+    tp_size = len(language_model)
+    assert tp_size > 0, "language_model tp dimension is empty"
+    ep_size = len(language_model[0])
+    assert ep_size > 0, "language_model ep dimension is empty"
+
+    assert isinstance(module, list), f"{module_name} should be a TP-sharded list"
+    assert len(module) == tp_size, (
+        f"{module_name} tp shards ({len(module)}) mismatch language_model tp shards ({tp_size})"
+    )
+
+    for t in range(tp_size):
+        src = module[t]
+        assert "model" in src, f"{module_name}[{t}] missing 'model' key"
+        for e in range(ep_size):
+            dst = language_model[t][e]
+            assert "model" in dst, f"language_model[tp={t}][ep={e}] missing 'model' key"
+            merge_dict(src["model"], dst["model"])
+
+
+def merge_module_1d_into_language_3d_stage0(language_model, module, module_name):
+    """Broadcast a 1D ``[tp]`` non-LLM module into PP stage 0 only of a 3D LLM tree.
+
+    ``language_model`` is shaped ``[pp][tp][ep]``. The non-LLM modules
+    (ViT / adapter / vision_patch) live solely on PP stage 0 in the runtime
+    Megatron model, so we only inject them into ``language_model[0][t][e]``;
+    later PP stages are untouched.
+    """
+    pp_size = len(language_model)
+    assert pp_size > 0, "language_model pp dimension is empty"
+    tp_size = len(language_model[0])
+    assert tp_size > 0, "language_model tp dimension is empty"
+    ep_size = len(language_model[0][0])
+    assert ep_size > 0, "language_model ep dimension is empty"
+
+    assert isinstance(module, list), f"{module_name} should be a TP-sharded list"
+    assert len(module) == tp_size, (
+        f"{module_name} tp shards ({len(module)}) mismatch language_model tp shards ({tp_size})"
+    )
+
+    for t in range(tp_size):
+        src = module[t]
+        assert "model" in src, f"{module_name}[{t}] missing 'model' key"
+        for e in range(ep_size):
+            dst = language_model[0][t][e]
+            assert "model" in dst, (
+                f"language_model[pp=0][tp={t}][ep={e}] missing 'model' key"
+            )
+            merge_dict(src["model"], dst["model"])
+
+
+args = parse_args()
+if args.megatron_path is not None:
+    sys.path.insert(0, args.megatron_path)
+
+print("===== merge gemma4-vl megatron checkpoints ======")
+
+pp_active = args.pipeline_model_parallel_size > 1
+ep_active = args.expert_parallel_size is not None and args.expert_parallel_size > 1
+
+if pp_active:
+    language_model = load_megatron_checkpoint_tp_pp_ep(args.language_model_path)
+elif ep_active:
+    language_model = load_megatron_checkpoint_tp_ep(args.language_model_path)
+else:
+    language_model = load_megatron_checkpoint(args.language_model_path)
+
+vision_model = load_megatron_checkpoint(args.vision_model_path)
+adapter = load_megatron_checkpoint(args.adapter_path)
+patch = load_megatron_checkpoint(args.vision_patch)
+
+if pp_active:
+    for module_name, module in [("vision", vision_model), ("adapter", adapter), ("patch", patch)]:
+        merge_module_1d_into_language_3d_stage0(language_model, module, module_name)
+    print(
+        f" > total LLM PP/TP/EP shards: {len(language_model)}/{len(language_model[0])}"
+        f"/{len(language_model[0][0])}"
+    )
+    save_megatron_checkpoint_tp_pp_ep(language_model, args.save_ckpt_path)
+elif ep_active:
+    for module_name, module in [("vision", vision_model), ("adapter", adapter), ("patch", patch)]:
+        merge_module_1d_into_language_2d(language_model, module, module_name)
+    print(f" > total LLM TP/EP shards: {len(language_model)}/{len(language_model[0])}")
+    save_megatron_checkpoint_tp_ep(language_model, args.save_ckpt_path)
+else:
+    for module_name, module in [("vision", vision_model), ("adapter", adapter), ("patch", patch)]:
+        assert len(module) == len(language_model), (
+            f"{module_name} TP shards ({len(module)}) != language_model TP shards ({len(language_model)})"
+        )
+        for i in range(len(module)):
+            merge_dict(module[i]["model"], language_model[i]["model"])
+    print(f" > total LLM TP shards: {len(language_model)}")
+    save_megatron_checkpoint(language_model, args.save_ckpt_path)
+
+print("===== merge complete ======")
diff --git a/tools/convert_checkpoint/custom/gemma4_vl/vision_patch.py b/tools/convert_checkpoint/custom/gemma4_vl/vision_patch.py
new file mode 100644
index 00000000..836baae7
--- /dev/null
+++ b/tools/convert_checkpoint/custom/gemma4_vl/vision_patch.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+# -*- coding: UTF-8 -*-
+################################################################################
+#
+# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
+#
+################################################################################
+"""Gemma4-VL vision-patch standalone converter.
+
+Mirrors ``custom/llava_onevision2/vision_patch.py`` but targets the Gemma4-VL
+HF↔Megatron-Core checkpoint mapping. Differences from the OV2 patch:
+
+- Gemma4 ``patch_embedder.input_proj`` is a plain ``nn.Linear`` in BOTH HF and
+  Megatron, so no Conv2d↔Linear reshape is needed.
+- Gemma4 patch embedder is NOT a ColumnParallelLinear, so no TP-rank sharding
+  of the patch weight is needed.
+- Gemma4 has ``position_embedding_table`` as an ``nn.Parameter`` with shape
+  ``[2, position_embedding_size, hidden_size]``; it is replicated across TP.
+- Gemma4 has ``std_bias`` and ``std_scale`` as VisionTower-level persistent
+  buffers (``standardize=True`` for Gemma4-26B-A4B-it).
+
+The mapping is loaded from ``args.common_config_path`` (Megatron-key →
+HF-key). All five mappings are pure 1:1 copies; the only Megatron-side
+operation is replicating across TP ranks.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from copy import deepcopy
+from os.path import dirname
+
+
+SCRIPT_DIR = dirname(os.path.abspath(__file__))
+sys.path.append(dirname(dirname(dirname(SCRIPT_DIR))))
+
+from convert_checkpoint.arguments import parse_args  # noqa: E402
+from convert_checkpoint.custom.llava_onevision2.util import (  # noqa: E402
+    load_huggingface_checkpoint,
+    load_megatron_checkpoint,
+    load_megatron_checkpoint_tp_ep,
+    save_huggingface_checkpoint,
+    save_megatron_checkpoint,
+)
+
+
+args = parse_args()
+name_map: dict[str, str] = {}
+with open(args.common_config_path, "r", encoding="utf-8") as f:
+    name_map = json.loads(f.read())
+
+
+def _get_non_ep_model_source(state_dict):
+    first = state_dict[0]
+    if isinstance(first, dict):
+        return first["model"]
+    first_rank = first[0]
+    if isinstance(first_rank, dict):
+        return first_rank["model"]
+    raise TypeError("Unsupported non-EP checkpoint structure")
+
+
+if (args.load_platform, args.save_platform) == ("mcore", "huggingface"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL vision patch from Megatron Core to HuggingFace ======")
+    ep = args.expert_parallel_size
+    target = {}
+    if ep is not None and ep > 1:
+        state_dict = load_megatron_checkpoint_tp_ep(args.load_ckpt_path)
+        source = state_dict[0][0]["model"]
+    else:
+        state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+        source = _get_non_ep_model_source(state_dict)
+
+    for mg_key, hf_key in name_map.items():
+        target[hf_key] = source[mg_key]
+        print(f" > {mg_key} -> {hf_key}  (shape {list(target[hf_key].shape)})")
+    save_huggingface_checkpoint(target, args.save_ckpt_path)
+
+elif (args.load_platform, args.save_platform) == ("huggingface", "mcore"):
+    print(" ====== convert Gemma4-VL vision patch from HuggingFace to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    source = load_huggingface_checkpoint(args.load_ckpt_path)
+
+    target = {}
+    for mg_key, hf_key in name_map.items():
+        target[mg_key] = source[hf_key]
+        print(f" > {hf_key} -> {mg_key}  (shape {list(target[mg_key].shape)})")
+
+    state_dict = [{"model": deepcopy(target)} for _ in range(tp)]
+    save_megatron_checkpoint(state_dict, os.path.join(args.save_ckpt_path, "release"))
+
+elif (args.load_platform, args.save_platform) == ("mcore", "mcore"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL vision patch from Megatron Core to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+    source = _get_non_ep_model_source(state_dict)
+
+    target = {}
+    for mg_key in name_map.keys():
+        target[mg_key] = source[mg_key]
+        print(f" > {mg_key}  (replicated across TP={tp})")
+
+    new_state_dict = [{"model": deepcopy(target)} for _ in range(tp)]
+    save_megatron_checkpoint(new_state_dict, os.path.join(args.save_ckpt_path, "release"))
+else:
+    raise NotImplementedError
diff --git a/tools/convert_checkpoint/custom/gemma4_vl/vit.py b/tools/convert_checkpoint/custom/gemma4_vl/vit.py
new file mode 100644
index 00000000..a53b0ed4
--- /dev/null
+++ b/tools/convert_checkpoint/custom/gemma4_vl/vit.py
@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+# -*- coding: UTF-8 -*-
+################################################################################
+#
+# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
+#
+################################################################################
+"""Gemma4-VL vision-tower body standalone converter.
+
+Mirrors ``custom/llava_onevision2/vision_model.py`` for OV2 but targets
+the Gemma4 27-layer ViT. Differences from generic ViT converters:
+
+- The MG ViT module is the entire ``Gemma4VisionTower`` (top-level
+  ``vision_model`` attribute on ``Gemma4VL``); per-layer keys live under
+  ``vision_model.layers.{i}.*``.
+- The HF ViT module is ``Gemma4VisionEncoder`` nested inside
+  ``Gemma4VisionModel``; per-layer keys live under
+  ``model.vision_tower.encoder.layers.{i}.*``.
+- All projection weights (q/k/v/o, gate/up/down) are wrapped in
+  ``Gemma4ClippableLinear`` whose underlying ``nn.Linear`` lives at
+  attribute ``.linear``, so HF and MG keys both end with ``.linear.weight``.
+- Every projection is a plain ``nn.Linear`` (NOT ColumnParallel/RowParallel),
+  so weights are TP-replicated, NOT sharded.
+- Each layer has 5 RMSNorm tensors (input / post_attention / pre_feedforward /
+  post_feedforward / q_norm + k_norm = 7 with attn norms), all replicated.
+- Patch embedder, std_bias, std_scale are handled by ``vision_patch.py``;
+  this file ONLY converts ``vision_model.layers.{i}.*``.
+
+Per-layer key set (13 keys, identical between HF and MG modulo prefix):
+    input_layernorm.weight
+    post_attention_layernorm.weight
+    pre_feedforward_layernorm.weight
+    post_feedforward_layernorm.weight
+    self_attn.{q,k,v,o}_proj.linear.weight   (4)
+    self_attn.{q,k}_norm.weight              (2)
+    mlp.{gate,up,down}_proj.linear.weight    (3)
+
+Total: 27 layers * 13 keys = 351 keys. Plus 4 patch keys (vision_patch.py)
+gives 355, matching the HF vision_tower.* count.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from copy import deepcopy
+from os.path import dirname
+
+
+SCRIPT_DIR = dirname(os.path.abspath(__file__))
+sys.path.append(dirname(dirname(dirname(SCRIPT_DIR))))
+
+from convert_checkpoint.arguments import parse_args  # noqa: E402
+from convert_checkpoint.custom.llava_onevision2.util import (  # noqa: E402
+    load_huggingface_checkpoint,
+    load_megatron_checkpoint,
+    load_megatron_checkpoint_tp_ep,
+    save_huggingface_checkpoint,
+    save_megatron_checkpoint,
+)
+
+
+HF_PREFIX = "model.vision_tower.encoder.layers."
+MG_PREFIX = "vision_model.layers."
+
+PER_LAYER_TAILS = (
+    "input_layernorm.weight",
+    "post_attention_layernorm.weight",
+    "pre_feedforward_layernorm.weight",
+    "post_feedforward_layernorm.weight",
+    "self_attn.q_proj.linear.weight",
+    "self_attn.k_proj.linear.weight",
+    "self_attn.v_proj.linear.weight",
+    "self_attn.o_proj.linear.weight",
+    "self_attn.q_norm.weight",
+    "self_attn.k_norm.weight",
+    "mlp.gate_proj.linear.weight",
+    "mlp.up_proj.linear.weight",
+    "mlp.down_proj.linear.weight",
+)
+
+
+args = parse_args()
+
+
+def _read_vit_num_layers_from_hf_config(hf_dir: str) -> int:
+    cfg_path = os.path.join(hf_dir, "config.json")
+    with open(cfg_path, "r", encoding="utf-8") as f:
+        cfg = json.load(f)
+    return int(cfg["vision_config"]["num_hidden_layers"])
+
+
+def _resolve_vit_num_layers() -> int:
+    if (args.load_platform, args.save_platform) == ("huggingface", "mcore"):
+        return _read_vit_num_layers_from_hf_config(args.load_ckpt_path)
+    if (args.load_platform, args.save_platform) == ("mcore", "huggingface"):
+        return _read_vit_num_layers_from_hf_config(args.save_ckpt_path)
+    return _read_vit_num_layers_from_hf_config(args.load_ckpt_path)
+
+
+def _get_non_ep_model_source(state_dict):
+    first = state_dict[0]
+    if isinstance(first, dict):
+        return first["model"]
+    first_rank = first[0]
+    if isinstance(first_rank, dict):
+        return first_rank["model"]
+    raise TypeError("Unsupported non-EP checkpoint structure")
+
+
+def _build_layer_mapping(num_layers: int) -> list[tuple[str, str]]:
+    pairs: list[tuple[str, str]] = []
+    for i in range(num_layers):
+        for tail in PER_LAYER_TAILS:
+            pairs.append((f"{MG_PREFIX}{i}.{tail}", f"{HF_PREFIX}{i}.{tail}"))
+    return pairs
+
+
+if (args.load_platform, args.save_platform) == ("mcore", "huggingface"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL ViT body from Megatron Core to HuggingFace ======")
+    ep = args.expert_parallel_size
+    if ep is not None and ep > 1:
+        state_dict = load_megatron_checkpoint_tp_ep(args.load_ckpt_path)
+        source = state_dict[0][0]["model"]
+    else:
+        state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+        source = _get_non_ep_model_source(state_dict)
+
+    num_layers = _resolve_vit_num_layers()
+    target = {}
+    for mg_key, hf_key in _build_layer_mapping(num_layers):
+        target[hf_key] = source[mg_key]
+    print(f" > converted {len(target)} ViT body keys ({num_layers} layers x {len(PER_LAYER_TAILS)} tails)")
+    save_huggingface_checkpoint(target, args.save_ckpt_path)
+
+elif (args.load_platform, args.save_platform) == ("huggingface", "mcore"):
+    print(" ====== convert Gemma4-VL ViT body from HuggingFace to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    source = load_huggingface_checkpoint(args.load_ckpt_path)
+
+    num_layers = _resolve_vit_num_layers()
+    target = {}
+    for mg_key, hf_key in _build_layer_mapping(num_layers):
+        target[mg_key] = source[hf_key]
+    print(f" > converted {len(target)} ViT body keys ({num_layers} layers x {len(PER_LAYER_TAILS)} tails)")
+
+    # ViT/adapter/patch always save 1D [tp] (mp_rank_TT) regardless of PP/EP, mirroring
+    # the OV2 Qwen3 convention; merge_megatron.py broadcasts them across PP stage 0
+    # (and across EP) at merge time, since these modules live solely on PP stage 0.
+    state_dict = [{"model": deepcopy(target)} for _ in range(tp)]
+    save_megatron_checkpoint(state_dict, os.path.join(args.save_ckpt_path, "release"))
+
+elif (args.load_platform, args.save_platform) == ("mcore", "mcore"):
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+    print(" ====== convert Gemma4-VL ViT body from Megatron Core to Megatron Core ======")
+    tp = args.tensor_model_parallel_size
+    state_dict = load_megatron_checkpoint(args.load_ckpt_path)
+    source = _get_non_ep_model_source(state_dict)
+
+    target = deepcopy(source)
+    new_state_dict = [{"model": deepcopy(target)} for _ in range(tp)]
+    save_megatron_checkpoint(new_state_dict, os.path.join(args.save_ckpt_path, "release"))
+
+else:
+    raise NotImplementedError