aiak-train-2121 [Task] [AIAK-Training-Omni] Skip overlong Kimi VLM samples and handle text-only multimodal inputs

Change-Id: I70f2ddcf1cbdd85027f6f375fb7c5cd47903dca5
(cherry picked from commit 0270d66)

Author: mokai

loongforge/data/kimi_k25_plugin.py

@@ -149,7 +149,9 @@ def _get_mm_inputs(
             pixel_values: tensor with shape (num_patches, patch_dim)
             grid_thws: tensor with shape (num_images, 3) for [T, H, W]
         """
-        import torch
+        mm_inputs = {}
+        if len(images) == 0 and len(videos) == 0:
+            return mm_inputs
 
         # Use Kimi's media processor
         media_processor = getattr(processor, 'media_processor', None)
@@ -159,8 +161,6 @@ def _get_mm_inputs(
         if media_processor is None:
             raise ValueError("Processor must have media_processor or image_processor")
 
-        mm_inputs = {}
-
         if len(images) != 0:
             # Regularize images first
             images = self._regularize_images(
@@ -255,12 +255,14 @@ def process_messages(
 
         num_image_tokens, num_video_tokens = 0, 0
         messages = deepcopy(messages)
+        has_images = len(images) > 0
+        has_videos = len(videos) > 0
 
         for message in messages:
             content = message["content"]
 
             # Replace image placeholders
-            while Placeholder.IMAGE in content:
+            while has_images and Placeholder.IMAGE in content:
                 if num_image_tokens >= len(image_grid_thw):
                     raise ValueError(
                         f"`len(images)` ({len(images)}) is less than the number of "
@@ -279,7 +281,7 @@ def process_messages(
                 num_image_tokens += 1
 
             # Replace video placeholders
-            while Placeholder.VIDEO in content:
+            while has_videos and Placeholder.VIDEO in content:
                 if num_video_tokens >= len(video_grid_thw):
                     raise ValueError(
                         f"`len(videos)` ({len(videos)}) is less than the number of "
@@ -301,13 +303,13 @@ def process_messages(
             message["content"] = content
 
         # Validate counts
-        if len(images) != num_image_tokens:
+        if has_images and len(images) != num_image_tokens:
             raise ValueError(
                 f"The number of images ({len(images)}) does not match "
                 f"the number of {Placeholder.IMAGE} tokens ({num_image_tokens})"
             )
 
-        if len(videos) != num_video_tokens:
+        if has_videos and len(videos) != num_video_tokens:
             raise ValueError(
                 f"The number of videos ({len(videos)}) does not match "
                 f"the number of {Placeholder.VIDEO} tokens ({num_video_tokens})"

loongforge/data/multimodal/kimi_task_encoder.py

@@ -3,6 +3,7 @@
 
 """Kimi Task Encoder."""
 
+import logging
 import torch
 from loongforge.data.multimodal.vlm_task_encoder import VLMTaskEncoder
 from typing import Dict, List, Optional, Tuple, Union
@@ -26,6 +27,8 @@
 
 
 from loongforge.utils import constants, get_chat_template
+from megatron.energon.task_encoder.base import stateless
+from loongforge.data.multimodal import MultiMixQASample, MultiVidQASample
 from .base.task_encoder import (
     BaseTaskEncoder,
     BaseTaskSample,
@@ -96,6 +99,51 @@ def __init__(self, args):
         else:
             self.merge_kernel_size = list(merge_kernel_size)
 
+    def _sample_sequence_limit(self) -> int:
+        sequence_limit = self.args.seq_length
+        packed_limit = getattr(self.args, "max_packed_tokens", None)
+        if self.is_packing_enabled and packed_limit is not None:
+            sequence_limit = min(sequence_limit, packed_limit)
+        return sequence_limit
+
+    def _should_discard_overlong(self, sample, input_ids) -> bool:
+        if not self.args.enable_discard_sample:
+            return False
+
+        sequence_limit = self._sample_sequence_limit()
+        input_length = len(input_ids)
+        if input_length <= sequence_limit:
+            return False
+
+        logging.warning(
+            "discard overlong sample %s: input length %s > sequence limit %s",
+            sample.__key__,
+            input_length,
+            sequence_limit,
+        )
+        return True
+
+    @stateless(restore_seeds=True)
+    def encode_sample(
+        self,
+        sample: Union[CaptioningSample, VQASample, MultiVidQASample, MultiMixQASample],
+    ):
+        """Return tokenised multimodal sample."""
+        if isinstance(sample, CaptioningSample):
+            encoded_sample = self.encode_captioning(sample)
+        elif isinstance(sample, VQASample):
+            encoded_sample = self.encode_vqa(sample)
+        elif isinstance(sample, MultiVidQASample):
+            encoded_sample = self.encode_multi_vid_qa(sample)
+        elif isinstance(sample, MultiMixQASample):
+            encoded_sample = self.encode_multi_mix_qa(sample)
+        else:
+            yield from super().encode_sample(sample)
+            return
+
+        if encoded_sample is not None:
+            yield encoded_sample
+
     def _get_vision_token_ids(self):
         """Get special token IDs for vision processing."""
         media_begin_id = self.tokenizer.convert_tokens_to_ids(MEDIA_BEGIN)
@@ -352,11 +400,9 @@ def encode_captioning(self, sample: CaptioningSample) -> BaseTaskSample:
         )
         num_tiles = [len(image_grid_thw)] if image_grid_thw is not None else [0]
 
-        if self.args.enable_discard_sample:
-            assert (
-                len(input_ids) <= self.args.seq_length
-            ), f"{sample.__key__} input length {len(input_ids)}"
-        elif image_grid_thw is not None:
+        if self._should_discard_overlong(sample, input_ids):
+            return None
+        if not self.args.enable_discard_sample and image_grid_thw is not None:
             assert (
                 image_grid_thw.prod() / 4 <= self.args.seq_length
             ), f"{sample.__key__} thw {image_grid_thw}"
@@ -414,11 +460,9 @@ def encode_vqa(self, sample: VQASample) -> BaseTaskSample:
 
         num_tiles = [len(image_grid_thw)] if image_grid_thw is not None else [0]
 
-        if self.args.enable_discard_sample:
-            assert (
-                len(input_ids) <= self.args.seq_length
-            ), f"{sample.__key__} input length {len(input_ids)}"
-        elif image_grid_thw is not None:
+        if self._should_discard_overlong(sample, input_ids):
+            return None
+        if not self.args.enable_discard_sample and image_grid_thw is not None:
             assert (
                 image_grid_thw.prod() / 4 <= self.args.seq_length
             ), f"{sample.__key__} grid_thw: {image_grid_thw}"
@@ -564,15 +608,21 @@ def encode_multi_mix_qa(self, sample) -> BaseTaskSample:
                 f"Unknown training phase {self.args.training_phase}"
             )
 
-        if self.args.enable_discard_sample:
-            assert (
-                len(input_ids) <= self.args.seq_length
-            ), f"{sample.__key__} input length {len(input_ids)}"
-        elif sample.video is not None and video_grid_thw is not None:
+        if self._should_discard_overlong(sample, input_ids):
+            return None
+        if (
+            not self.args.enable_discard_sample
+            and sample.video is not None
+            and video_grid_thw is not None
+        ):
             assert (
                 video_grid_thw.prod(dim=-1).sum() / 4 <= self.args.seq_length
             ), f"{sample.__key__} grid_thw: {video_grid_thw}"
-        elif sample.image is not None and image_grid_thw is not None:
+        elif (
+            not self.args.enable_discard_sample
+            and sample.image is not None
+            and image_grid_thw is not None
+        ):
             assert (
                 image_grid_thw.prod(dim=-1).sum() / 4 <= self.args.seq_length
             ), f"{sample.__key__} grid_thw: {image_grid_thw}"
@@ -626,11 +676,9 @@ def encode_multi_vid_qa(self, sample) -> BaseTaskSample:
                 f"Unknown training phase {self.args.training_phase}"
             )
 
-        if self.args.enable_discard_sample:
-            assert (
-                len(input_ids) <= self.args.seq_length
-            ), f"{sample.__key__} input length {len(input_ids)}"
-        elif video_grid_thw is not None:
+        if self._should_discard_overlong(sample, input_ids):
+            return None
+        if not self.args.enable_discard_sample and video_grid_thw is not None:
             assert (
                 video_grid_thw.prod(dim=-1).sum() / 4 <= self.args.seq_length
             ), f"{sample.__key__} grid_thw: {video_grid_thw}"