cosmos_framework/utils/generator/monkey_patch.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

from __future__ import annotations

import importlib

import torch
import transformers
from transformers.cache_utils import Cache
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLModel
from transformers.utils.import_utils import is_torchdynamo_compiling

from cosmos_framework.utils import log

_EXPECTED_TRANSFORMERS_VERSION_PREFIX = "4.57."


def patch_qwen3_vl_forward(model):
    """Monkey-patch a ``Qwen3VLModel`` **instance's** forward:
       **Single visual forward per batch**: Under FSDP, every rank must call
       ``self.visual(...)`` the same number of times each forward step so that
       collective all-gather operations stay in sync. Image and video inputs are
       encoded together in one visual call. When a batch contains only text, a
       lightweight dummy image (16x16 zeros) is pushed through the full ViT ->
       merger -> deepstack pipeline, then outputs are sliced to ``[0:0]`` so
       they carry ``grad_fn`` but contribute no features.
    Args:
        model: The ``Qwen3VLModel`` instance (i.e. ``model.model.model`` when
            the outer model is ``HFModel``).
    """
    if not transformers.__version__.startswith(_EXPECTED_TRANSFORMERS_VERSION_PREFIX):
        raise ValueError(f"monkey patching transformers version {transformers.__version__} is not supported")

    if not isinstance(model, Qwen3VLModel):
        raise ValueError(f"Trying to monkey patch a model that is not a Qwen3VLModel instance: {type(model)}")

    # Resolve the output dataclass from the actual runtime module
    model_module = importlib.import_module(type(model).__module__)
    Qwen3VLModelOutputWithPast = getattr(model_module, "Qwen3VLModelOutputWithPast")

    # Replaces Qwen3VLModel.forward from:
    #   transformers.models.qwen3_vl.modeling_qwen3_vl  (transformers v4.57.1)
    def patched_forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        pixel_values: torch.Tensor | None = None,
        pixel_values_videos: torch.FloatTensor | None = None,
        image_grid_thw: torch.LongTensor | None = None,
        video_grid_thw: torch.LongTensor | None = None,
        cache_position: torch.LongTensor | None = None,
        **kwargs,
    ):
        r"""
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.
        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
            The temporal, height and width of feature shape of each video in LLM.
        """
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        image_mask = None
        video_mask = None
        deepstack_image_embeds = None
        deepstack_video_embeds = None

        visual_pixel_values_list: list[torch.Tensor] = []
        visual_grid_thw_list: list[torch.Tensor] = []
        image_embed_len = 0
        video_embed_len = 0
        has_image = pixel_values is not None
        has_video = pixel_values_videos is not None

        if has_image:
            if image_grid_thw is None:
                raise ValueError("image_grid_thw must be provided when pixel_values is provided")
            visual_pixel_values_list.append(pixel_values)
            visual_grid_thw_list.append(image_grid_thw)
            image_embed_len = int((image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).sum().item())

        if has_video:
            if video_grid_thw is None:
                raise ValueError("video_grid_thw must be provided when pixel_values_videos is provided")
            visual_pixel_values_list.append(pixel_values_videos)
            visual_grid_thw_list.append(video_grid_thw)
            video_embed_len = int((video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).sum().item())

        if has_image or has_video:
            visual_pixel_values = torch.cat(visual_pixel_values_list, dim=0).type(self.visual.dtype)  # [N_patch,D]
            visual_grid_thw = torch.cat(visual_grid_thw_list, dim=0)  # [N_media,3]
            visual_embeds, deepstack_visual_feature_lists = self.visual(
                visual_pixel_values, grid_thw=visual_grid_thw
            )  # visual_embeds: [N_visual,C]
            image_embeds = visual_embeds[:image_embed_len]  # [N_image_visual,C]
            video_embeds = visual_embeds[image_embed_len : image_embed_len + video_embed_len]  # [N_video_visual,C]
            deepstack_image_embeds = [
                deepstack_visual_embeds[:image_embed_len] for deepstack_visual_embeds in deepstack_visual_feature_lists
            ]  # each: [N_image_visual,C]
            deepstack_video_embeds = [
                deepstack_visual_embeds[image_embed_len : image_embed_len + video_embed_len]
                for deepstack_visual_embeds in deepstack_visual_feature_lists
            ]  # each: [N_video_visual,C]

            if has_image:
                image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)  # [N_image_visual,C]
                image_mask, _ = self.get_placeholder_mask(
                    input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
                )
                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)  # [B,N_token,C]

            if has_video:
                video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)  # [N_video_visual,C]
                _, video_mask = self.get_placeholder_mask(
                    input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
                )
                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)  # [B,N_token,C]

        # Dummy visual forward for text-only data
        else:
            dummy_h, dummy_w = 16, 16
            dummy_pixels = torch.zeros(
                dummy_h * dummy_w,
                self.visual.config.temporal_patch_size * self.visual.config.patch_size**2 * 3,
                device=inputs_embeds.device,
                dtype=self.visual.dtype,
            )  # [N_dummy_patch,D]
            dummy_thw = torch.tensor([[1, dummy_h, dummy_w]], device=inputs_embeds.device)  # [1,3]
            image_embeds, deepstack_image_embeds = self.visual(dummy_pixels, grid_thw=dummy_thw)
            image_embeds = image_embeds[0:0]  # [0,C]
            deepstack_image_embeds = [e[0:0] for e in deepstack_image_embeds]  # each: [0,C]

            # no-op to mask scatter empty embeddings into inputs to preserve computation graph
            image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)  # [0,C]
            image_mask, _ = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
            )
            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)  # [B,N_token,C]

        visual_pos_masks = None
        deepstack_visual_embeds = None
        if image_mask is not None and video_mask is not None:
            # aggregate visual_pos_masks and deepstack_visual_embeds
            image_mask = image_mask[..., 0]
            video_mask = video_mask[..., 0]
            visual_pos_masks = image_mask | video_mask
            deepstack_visual_embeds = []
            image_mask_joint = image_mask[visual_pos_masks]
            video_mask_joint = video_mask[visual_pos_masks]
            for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds):
                embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device)
                embed_joint[image_mask_joint, :] = img_embed
                embed_joint[video_mask_joint, :] = vid_embed
                deepstack_visual_embeds.append(embed_joint)
        elif image_mask is not None:
            image_mask = image_mask[..., 0]
            visual_pos_masks = image_mask
            deepstack_visual_embeds = deepstack_image_embeds
        elif video_mask is not None:
            video_mask = video_mask[..., 0]
            visual_pos_masks = video_mask
            deepstack_visual_embeds = deepstack_video_embeds

        if position_ids is None:
            attention_mask_tensor = (
                attention_mask if not isinstance(attention_mask, dict) else attention_mask["full_attention"]
            )
            if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4:
                attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2)
                # Only apply conversion for floating point tensors (inverted masks)
                if attention_mask_tensor.dtype.is_floating_point:
                    attention_mask_tensor = attention_mask_tensor / torch.finfo(attention_mask_tensor.dtype).min
                    attention_mask_tensor = (1.0 - attention_mask_tensor).int()

            # Calculate RoPE index once per generation in the pre-fill stage only.
            # When compiling, we can't check tensor values thus we check only input length
            # It is safe to assume that `length!=1` means we're in pre-fill because compiled
            # models currently cannot do asssisted decoding
            prefill_compiled_stage = is_torchdynamo_compiling() and (
                (input_ids is not None and input_ids.shape[1] != 1)
                or (inputs_embeds is not None and inputs_embeds.shape[1] != 1)
            )
            prefill_noncompiled_stage = not is_torchdynamo_compiling() and (
                (cache_position is not None and cache_position[0] == 0)
                or (past_key_values is None or past_key_values.get_seq_length() == 0)
            )
            if (prefill_compiled_stage or prefill_noncompiled_stage) or self.rope_deltas is None:
                position_ids, rope_deltas = self.get_rope_index(
                    input_ids,
                    image_grid_thw,
                    video_grid_thw,
                    attention_mask=attention_mask_tensor,
                )
                self.rope_deltas = rope_deltas
            # then use the prev pre-calculated rope-deltas to get the correct position ids
            else:
                batch_size, seq_length, _ = inputs_embeds.shape
                delta = (
                    (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) if cache_position is not None else 0
                )
                position_ids = torch.arange(seq_length, device=inputs_embeds.device)
                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
                if cache_position is not None:  # otherwise `deltas` is an int `0`
                    delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
                position_ids = position_ids.add(delta)
                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)

        outputs = self.language_model(
            input_ids=None,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            visual_pos_masks=visual_pos_masks,
            deepstack_visual_embeds=deepstack_visual_embeds,
            **kwargs,
        )

        return Qwen3VLModelOutputWithPast(
            last_hidden_state=outputs.last_hidden_state,
            past_key_values=outputs.past_key_values,
            rope_deltas=self.rope_deltas,
        )

    # Replace the forward method
    model.forward = patched_forward.__get__(model, type(model))
    log.critical(f"Patched {type(model).__name__} instance forward with one visual call per forward")