Add multimodal vision token utilities

megatron/core/models/multimodal/utils.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Portions Copyright (c) 2023 OpenGVLab and licensed under the MIT license found in LICENSE.
+
+import torch
+from einops import rearrange
+
+
+def patchify_image(x: torch.Tensor, patch_dim: int) -> torch.Tensor:
+    """Patchify an image:
+    Option 1: (C,H,W) == (3,H,W) -> (L,3*patch_dim*patch_dim)
+    Option 2: (B,C,H,W) == (B,3,H,W) -> (B,L,3*patch_dim*patch_dim)
+    """
+    assert (
+        x.shape[-2] % patch_dim == 0 and x.shape[-1] % patch_dim == 0
+    ), f"H and W must be divisible by patch_dim={patch_dim}, found {x.shape[-2:]}"
+
+    if x.ndim == 3:
+        assert x.shape[0] == 3, f"Expected (C,H,W) tensor, found {x.shape}"
+        py = x.shape[-2] // patch_dim  # H
+        px = x.shape[-1] // patch_dim  # W
+        x = rearrange(
+            x, 'c (py yy) (px xx) -> (py px) (c yy xx)', py=py, yy=patch_dim, px=px, xx=patch_dim
+        )
+    elif x.ndim == 4:
+        assert x.shape[1] == 3, f"Expected (B,C,H,W) tensor, found {x.shape}"
+        py = x.shape[-2] // patch_dim  # H
+        px = x.shape[-1] // patch_dim  # W
+        x = rearrange(
+            x,
+            'b c (py yy) (px xx) -> b (py px) (c yy xx)',
+            py=py,
+            yy=patch_dim,
+            px=px,
+            xx=patch_dim,
+        )
+    else:
+        raise NotImplementedError(
+            f"Expected (B,C,H,W) or (C,H,W) input tensor, found shape {x.shape}"
+        )
+
+    return x
+
+
+def unpatchify_image(x: torch.Tensor, img_H: int, img_W: int, patch_dim: int) -> torch.Tensor:
+    """Unpatchify an image (reverse of patchify_image()):
+    Option 1: (L,C) == (L,3*patch_dim*patch_dim) -> (3,H,W)
+    Option 2: (B,L,C) == (B,L,3*patch_dim*patch_dim) -> (B,3,H,W)
+    """
+    assert img_H % patch_dim == 0 and img_W % patch_dim == 0 and patch_dim > 0, (
+        f"Expected img_H/img_W to be divisible by patch_dim={patch_dim}, "
+        f"found img_H={img_H}, img_W={img_W}"
+    )
+
+    py = img_H // patch_dim
+    px = img_W // patch_dim
+
+    if x.ndim == 2:
+        expected_shape = (py * px, 3 * patch_dim * patch_dim)
+        assert (
+            x.shape == expected_shape
+        ), f"Expected x.shape={expected_shape}, found x.shape={x.shape}"
+
+        x = rearrange(
+            x, '(py px) (c yy xx) -> c (py yy) (px xx)', py=py, yy=patch_dim, px=px, xx=patch_dim
+        )
+
+    elif x.ndim == 3:
+        expected_shape = (x.shape[0], py * px, 3 * patch_dim * patch_dim)
+        assert (
+            x.shape == expected_shape
+        ), f"Expected x.shape={expected_shape}, found x.shape={x.shape}"
+
+        x = rearrange(
+            x,
+            'b (py px) (c yy xx) -> b c (py yy) (px xx)',
+            py=py,
+            yy=patch_dim,
+            px=px,
+            xx=patch_dim,
+        )
+
+    else:
+        raise NotImplementedError(f"Expected (B,L,C) or (L,C) input tensor, found shape {x.shape}")
+
+    return x

megatron/core/models/vision/clip_vit_model.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
-from typing import Optional, Union
+from typing import List, Optional, Union
 
 import torch
 
@@ -162,7 +162,10 @@ def set_input_tensor(self, input_tensor: torch.Tensor) -> None:
         self.decoder.set_input_tensor(input_tensor)
 
     def forward(
-        self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        num_frames: Optional[Union[List[int], torch.Tensor]] = None,
     ) -> torch.Tensor:
         """Forward function of the CLIP ViT Model. This function passes the input tensors
         through the embedding layer and then the transformer.
@@ -174,6 +177,9 @@ def forward(
         Returns:
             x (torch.Tensor): output after final transformer block of shape [b, s, h].
         """
+        if num_frames is not None:
+            raise NotImplementedError("Temporal compression is not supported for CLIP ViT.")
+
         x = self.conv1(x)  # shape = [batch, hidden_size, grid, grid]
         x = x.reshape(x.shape[0], x.shape[1], -1)  # [batch, hidden_size, grid ** 2]
         x = x.permute(0, 2, 1)  # [batch, grid ** 2, hidden_size]
@@ -202,28 +208,80 @@ def forward(
         return x
 
 
-def get_num_image_embeddings(
-    img_h,
-    img_w,
-    patch_dim,
-    vision_model_type,
-    disable_vision_class_token,
-    class_token_len,
-    pixel_shuffle,
-    use_tile_tags=False,
-    max_num_tiles=0,
-    tokenizer_type=None,
+def _get_num_spatial_embeddings(
+    img_h: int,
+    img_w: int,
+    patch_dim: int,
+    pixel_shuffle: bool,
+    attn_pooling: bool,
+    attn_pooling_img_h: int,
+    attn_pooling_img_w: int,
+    attn_pooling_video_h: int,
+    attn_pooling_video_w: int,
+    is_video: bool,
+):
+    # Check if the patch grid is divisible by the required factors
+    assert (
+        img_h % patch_dim == 0 and img_w % patch_dim == 0
+    ), f"Image dimensions ({img_h}x{img_w}) must be divisible by patch_dim ({patch_dim})"
+    num_patches_per_dim_h = img_h // patch_dim
+    num_patches_per_dim_w = img_w // patch_dim
+
+    if pixel_shuffle:
+        assert num_patches_per_dim_h % 2 == 0 and num_patches_per_dim_w % 2 == 0, (
+            f"Patch grid ({num_patches_per_dim_h}x{num_patches_per_dim_w}) must be "
+            "divisible by 2 for pixel_shuffle"
+        )
+        num_patches_per_dim_h = num_patches_per_dim_h // 2
+        num_patches_per_dim_w = num_patches_per_dim_w // 2
+
+    # Attention pooling reduces tokens similarly to pixel shuffle.
+    # Select pooling size based on content type (video uses video params, falls back to img params)
+    if attn_pooling and (attn_pooling_img_h is not None or attn_pooling_img_w is not None):
+        assert (
+            attn_pooling_img_h is not None and attn_pooling_img_w is not None
+        ), "attn_pooling_img_h and attn_pooling_img_w must be provided together"
+        if is_video:
+            effective_pooling_h = (
+                attn_pooling_video_h if attn_pooling_video_h is not None else attn_pooling_img_h
+            )
+            effective_pooling_w = (
+                attn_pooling_video_w if attn_pooling_video_w is not None else attn_pooling_img_w
+            )
+        else:
+            effective_pooling_h = attn_pooling_img_h
+            effective_pooling_w = attn_pooling_img_w
+        assert num_patches_per_dim_h % effective_pooling_h == 0, (
+            f"Patch grid height ({num_patches_per_dim_h}) must be divisible by pooling_h "
+            f"({effective_pooling_h}, is_video={is_video})"
+        )
+        assert num_patches_per_dim_w % effective_pooling_w == 0, (
+            f"Patch grid width ({num_patches_per_dim_w}) must be divisible by pooling_w "
+            f"({effective_pooling_w}, is_video={is_video})"
+        )
+        num_patches_per_dim_h = num_patches_per_dim_h // effective_pooling_h
+        num_patches_per_dim_w = num_patches_per_dim_w // effective_pooling_w
+
+    return num_patches_per_dim_h * num_patches_per_dim_w
+
+
+def _get_num_non_spatial_embeddings(
+    vision_model_type: str,
+    disable_vision_class_token: bool,
+    class_token_len: int,
+    use_tile_tags: bool,
+    max_num_tiles: int,
+    tokenizer_type: str,
 ):
-    """Get the number of image embeddings per image tile."""
     if vision_model_type == "siglip":
         keep_class_token = False
     elif vision_model_type in ("clip", "internvit", "internvit300M"):
         keep_class_token = not disable_vision_class_token
-    elif vision_model_type.startswith("radio"):
-        keep_class_token = not disable_vision_class_token
     elif vision_model_type == "cradio-g":
         class_token_len = 8
         keep_class_token = not disable_vision_class_token
+    elif "radio" in vision_model_type:
+        keep_class_token = not disable_vision_class_token
     elif vision_model_type.startswith("hf://"):
         from megatron.core.models.huggingface.module import get_hf_model_type
 
@@ -236,26 +294,143 @@ def get_num_image_embeddings(
     else:
         raise NotImplementedError(f"unknown vision model type {vision_model_type}")
 
-    num_patches_per_dim_h = img_h // patch_dim
-    num_patches_per_dim_w = img_w // patch_dim
-    num_patches = num_patches_per_dim_h * num_patches_per_dim_w
-    num_image_embeddings_per_tile = num_patches + (class_token_len if keep_class_token else 0)
-
-    if pixel_shuffle:
-        num_image_embeddings_per_tile = int(num_image_embeddings_per_tile * (0.5**2))
+    num_non_spatial_embeddings = class_token_len if keep_class_token else 0
 
     if use_tile_tags:
         if tokenizer_type in ("llama3p1", "chatml", "qwen2p0", "qwen2p5"):
-            num_image_embeddings_per_tile += 5
+            num_non_spatial_embeddings += 5
         elif tokenizer_type.startswith("nemotron5"):
-            num_image_embeddings_per_tile += 6
+            num_non_spatial_embeddings += 6
         else:
             raise ValueError("tokenizer type not defined")
 
         if 10 < max_num_tiles < 100:
             if tokenizer_type.startswith("qwen"):
-                num_image_embeddings_per_tile += 1  # add padding 0
+                num_non_spatial_embeddings += 1  # add padding 0
         elif max_num_tiles > 100:
             raise ValueError(f"max number of tiles {max_num_tiles} not supported")
 
-    return num_image_embeddings_per_tile
+    return num_non_spatial_embeddings
+
+
+def get_num_image_embeddings(
+    img_h: int,
+    img_w: int,
+    patch_dim: int,
+    vision_model_type: str,
+    disable_vision_class_token: bool,
+    class_token_len: int,
+    pixel_shuffle: bool,
+    use_tile_tags: bool = False,
+    max_num_tiles: int = 0,
+    tokenizer_type: str = None,
+    attn_pooling: bool = False,
+    attn_pooling_img_h: int = None,
+    attn_pooling_img_w: int = None,
+    allow_non_spatial_embeddings: bool = True,
+):
+    """Get the number of embeddings per image tile (LLM tokens from vision).
+
+    For IMAGES ONLY. Uses image pooling params. For videos (including per-frame
+    calculations), use get_num_video_embeddings() which handles video pooling.
+    """
+    num_spatial_embeddings = _get_num_spatial_embeddings(
+        img_h=img_h,
+        img_w=img_w,
+        patch_dim=patch_dim,
+        pixel_shuffle=pixel_shuffle,
+        attn_pooling=attn_pooling,
+        attn_pooling_img_h=attn_pooling_img_h,
+        attn_pooling_img_w=attn_pooling_img_w,
+        attn_pooling_video_h=None,
+        attn_pooling_video_w=None,
+        is_video=False,
+    )
+    num_non_spatial_embeddings = _get_num_non_spatial_embeddings(
+        vision_model_type=vision_model_type,
+        disable_vision_class_token=disable_vision_class_token,
+        class_token_len=class_token_len,
+        use_tile_tags=use_tile_tags,
+        max_num_tiles=max_num_tiles,
+        tokenizer_type=tokenizer_type,
+    )
+
+    if num_non_spatial_embeddings > 0 and not allow_non_spatial_embeddings:
+        raise RuntimeError(
+            f"Found {num_non_spatial_embeddings} non-spatial embeddings when "
+            f"allow_non_spatial_embeddings=False. "
+            f"This usually indicates we're using temporal compression but calling this function on"
+            f" individual frames during pre-processing."
+        )
+
+    # Note: Temporal compression (video_temporal_patch_size > 1) must be handled either by:
+    #   1) Calling get_num_video_embeddings() to accurately count embeddings for a sequence
+    #       of frames that have been compressed by video_temporal_patch_size. This is preferred if
+    #       processing an entire video sequence, as is done during evaluation
+    #   2) Calling get_num_image_embeddings() on individual frames during pre-processing and
+    #       using the number of embeddings every video_temporal_patch_size number of frames. For an
+    #       example see examples/multimodal/data_loading/task_encoder.py,
+    #       MultiModalTaskEncoder._group_video_frame_params_into_tubelets()
+    return num_spatial_embeddings + num_non_spatial_embeddings
+
+
+def get_num_video_embeddings(
+    num_frames: int,
+    video_temporal_patch_size: int,
+    img_h: int,
+    img_w: int,
+    patch_dim: int,
+    vision_model_type: str,
+    disable_vision_class_token: bool,
+    class_token_len: int,
+    pixel_shuffle: bool,
+    use_tile_tags: bool = False,
+    max_num_tiles: int = 0,
+    tokenizer_type: str = None,
+    attn_pooling: bool = False,
+    attn_pooling_img_h: int = None,
+    attn_pooling_img_w: int = None,
+    attn_pooling_video_h: int = None,
+    attn_pooling_video_w: int = None,
+):
+    """Get the number of embeddings produced for a video."""
+    num_spatial_embeddings_per_frame = _get_num_spatial_embeddings(
+        img_h=img_h,
+        img_w=img_w,
+        patch_dim=patch_dim,
+        pixel_shuffle=pixel_shuffle,
+        attn_pooling=attn_pooling,
+        attn_pooling_img_h=attn_pooling_img_h,
+        attn_pooling_img_w=attn_pooling_img_w,
+        attn_pooling_video_h=attn_pooling_video_h,
+        attn_pooling_video_w=attn_pooling_video_w,
+        is_video=True,
+    )
+    num_non_spatial_embeddings_per_frame = _get_num_non_spatial_embeddings(
+        vision_model_type=vision_model_type,
+        disable_vision_class_token=disable_vision_class_token,
+        class_token_len=class_token_len,
+        use_tile_tags=use_tile_tags,
+        max_num_tiles=max_num_tiles,
+        tokenizer_type=tokenizer_type,
+    )
+
+    # It's unclear right now how the rest of the pipeline will handle temporal compression w.r.t.
+    #   non-spatial embeddings such as class tokens or tile tags.
+    # Raising NotImplementedError to force us to revist this when a concrete use case arises
+    if num_non_spatial_embeddings_per_frame > 0 and video_temporal_patch_size > 1:
+        raise NotImplementedError(
+            f"Found {num_non_spatial_embeddings_per_frame} non spatial embeddings per frame. This"
+            " is not currently supported when "
+            f"video_temporal_patch_size={video_temporal_patch_size} > 1"
+        )
+
+    assert num_frames % video_temporal_patch_size == 0, (
+        f"num_frames ({num_frames}) must be a multiple of "
+        f"video_temporal_patch_size ({video_temporal_patch_size})"
+    )
+    num_output_frames = num_frames // video_temporal_patch_size
+
+    num_spatial_embeddings = num_spatial_embeddings_per_frame * num_output_frames
+    num_non_spatial_embeddings = num_non_spatial_embeddings_per_frame * num_output_frames
+    return num_spatial_embeddings + num_non_spatial_embeddings