diff --git a/megatron/core/models/multimodal/utils.py b/megatron/core/models/multimodal/utils.py new file mode 100644 index 00000000000..426e33a4181 --- /dev/null +++ b/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 diff --git a/megatron/core/models/vision/clip_vit_model.py b/megatron/core/models/vision/clip_vit_model.py index 4c7c6d5b75a..567dffbe8f8 100644 --- a/megatron/core/models/vision/clip_vit_model.py +++ b/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 diff --git a/tests/unit_tests/models/test_clip_vit_model.py b/tests/unit_tests/models/test_clip_vit_model.py index c176c188d15..33a57bac759 100644 --- a/tests/unit_tests/models/test_clip_vit_model.py +++ b/tests/unit_tests/models/test_clip_vit_model.py @@ -3,7 +3,12 @@ import torch from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec -from megatron.core.models.vision.clip_vit_model import CLIPViTModel, get_num_image_embeddings +from megatron.core.models.multimodal.utils import patchify_image, unpatchify_image +from megatron.core.models.vision.clip_vit_model import ( + CLIPViTModel, + get_num_image_embeddings, + get_num_video_embeddings, +) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -73,3 +78,74 @@ def test_get_num_image_embeddings(vision_model, pixel_shuffle, tile_tags, expect ) == expected ) + + +@pytest.mark.parametrize("batched", [False, True]) +def test_patchify_unpatchify_roundtrip(batched): + image = torch.arange(3 * 4 * 6).reshape(3, 4, 6) + if batched: + image = torch.stack((image, image + image.numel())) + + patches = patchify_image(image, patch_dim=2) + + assert torch.equal(unpatchify_image(patches, img_H=4, img_W=6, patch_dim=2), image) + + +@pytest.mark.parametrize( + "pixel_shuffle,attn_pooling,expected_spatial", + [(False, False, 32), (True, False, 8), (False, True, 4), (True, True, 1)], +) +@pytest.mark.parametrize( + "disable_vision_class_token,use_tile_tags,expected_non_spatial", + [(True, False, 0), (False, False, 2), (True, True, 6), (False, True, 8)], +) +def test_get_num_image_embeddings_applies_attention_pooling( + pixel_shuffle, + attn_pooling, + expected_spatial, + disable_vision_class_token, + use_tile_tags, + expected_non_spatial, +): + # The 4x8 patch grid is reduced by pixel shuffle and then 2x4 image pooling. + # Class tokens and tile tags must be added without spatial reduction. + assert ( + get_num_image_embeddings( + img_h=8, + img_w=16, + patch_dim=2, + vision_model_type="clip", + disable_vision_class_token=disable_vision_class_token, + class_token_len=2, + pixel_shuffle=pixel_shuffle, + use_tile_tags=use_tile_tags, + max_num_tiles=1, + tokenizer_type="nemotron5", + attn_pooling=attn_pooling, + attn_pooling_img_h=2, + attn_pooling_img_w=4, + ) + == expected_spatial + expected_non_spatial + ) + + +def test_get_num_video_embeddings_applies_temporal_and_attention_pooling(): + assert ( + get_num_video_embeddings( + num_frames=4, + video_temporal_patch_size=2, + img_h=8, + img_w=8, + patch_dim=2, + vision_model_type="siglip", + disable_vision_class_token=True, + class_token_len=0, + pixel_shuffle=False, + attn_pooling=True, + attn_pooling_img_h=1, + attn_pooling_img_w=1, + attn_pooling_video_h=2, + attn_pooling_video_w=2, + ) + == 8 + )