Add Seq Packing in NeMo / Neva2

.github/workflows/cicd-main.yml

@@ -4306,11 +4306,24 @@ jobs:
     with:
       RUNNER: self-hosted-azure
       SCRIPT: |
-        python tests/collections/vlm/neva_train.py \
+        python tests/collections/vlm/test_neva_train.py \
         --devices=1 \
         --max-steps=5 \
         --experiment-dir=/tmp/nemo2_neva_results/${{ github.run_id }}
 
+  L2_NeMo_2_NEVA_MOCK_PACKED_TRAINING:
+    needs: [cicd-test-container-setup]
+    uses: ./.github/workflows/_test_template.yml
+    if: contains(fromJSON(needs.cicd-test-container-setup.outputs.test_to_run), 'L2_NeMo_2_NEVA_MOCK_PACKED_TRAINING') || needs.cicd-test-container-setup.outputs.all == 'true'
+    with:
+      RUNNER: self-hosted-azure
+      SCRIPT: |
+        python tests/collections/vlm/test_neva_train.py \
+        --devices=1 \
+        --max-steps=5 \
+        --experiment-dir=/tmp/nemo2_neva_results/${{ github.run_id }} \
+        --use_packed_sequence
+
   L2_NeMo_2_MLLAMA_MOCK_TRAINING:
     needs: [cicd-test-container-setup]
     uses: ./.github/workflows/_test_template.yml
@@ -4319,7 +4332,7 @@ jobs:
       RUNNER: self-hosted-azure
       SCRIPT: |
         TRANSFORMERS_OFFLINE=1 \
-        python tests/collections/vlm/mllama_train.py \
+        python tests/collections/vlm/test_mllama_train.py \
         --devices=1 \
         --max-steps=5 \
         --experiment-dir=/tmp/nemo2_mllama_results/${{ github.run_id }}
@@ -4978,6 +4991,7 @@ jobs:
       - Speech_Checkpoints_tests
       - L2_Stable_Diffusion_Training
       - L2_NeMo_2_NEVA_MOCK_TRAINING
+      - L2_NeMo_2_NEVA_MOCK_PACKED_TRAINING
       - L2_NeMo_2_MLLAMA_MOCK_TRAINING
       - L2_NeMo_2_GPT_Pretraining_no_transformer_engine
       - L2_NeMo_2_GPT_DDP_Param_Parity_check

nemo/collections/multimodal/data/energon/base.py

@@ -68,6 +68,7 @@ def __init__(
         multimodal_sample_config: Optional[MultiModalSampleConfig] = MultiModalSampleConfig(),
         task_encoder: Optional[MultiModalTaskEncoder] = None,
         decoder_seq_length: Optional[int] = None,
+        packing_buffer_size: Optional[int] = None,
     ) -> None:
         """
         Initialize the EnergonMultiModalDataModule.
@@ -84,6 +85,8 @@ def __init__(
         Defaults to MultiModalSampleConfig().
         task_encoder (MultiModalTaskEncoder, optional): Encoder responsible for encoding and batching samples.
         If not provided, a default (MultimodalTaskEncoder) encoder will be created. Defaults to None.
+        decoder_seq_length (int, optional): The maximum sequence length for the decoder. Used in encoder-decoder models.
+        packing_buffer_size (int, optional): Size of the packing buffer for batched samples. Defaults to None.
         """
 
         super().__init__()
@@ -113,6 +116,7 @@ def __init__(
         )
         self.train_dataloader_object = None
         self.val_dataloader_object = None
+        self.packing_buffer_size = packing_buffer_size
 
     def io_init(self, **kwargs) -> fdl.Config[Self]:
 
@@ -146,6 +150,7 @@ def datasets_provider(self, worker_config, split: Literal['train', 'val'] = 'val
             task_encoder=self.task_encoder,
             worker_config=worker_config,
             max_samples_per_sequence=None,
+            packing_buffer_size=self.packing_buffer_size,
             shuffle_buffer_size=100,
             split_part=split,
         )

nemo/collections/multimodal/data/energon/config.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 
 from dataclasses import dataclass, field
-from typing import List
+from typing import List, Tuple, Union
+
 import torch
+from megatron.core.packed_seq_params import PackedSeqParams
+
 from nemo.collections.multimodal.data.energon.conversation import LLaVATemplateConfig
 
 
@@ -43,6 +46,15 @@
     loss_mask: torch.Tensor = field(default_factory=lambda: torch.empty(0, dtype=torch.float))
 
 
+@dataclass
+class PackedImageTextSample(ImageTextSample):
+    '''Sample type for packed image text sample'''
+
+    __restore_key__: Tuple[Union[str, int, tuple], ...] = ()
+    position_ids: torch.Tensor = field(default_factory=lambda: torch.empty(0, dtype=torch.float))
+    packed_seq_params: PackedSeqParams = field(default_factory=lambda: PackedSeqParams())
+
+
 @dataclass
 class ImageTextRawBatch:
     """Sample type for image text raw batch"""
@@ -56,6 +68,14 @@
     loss_mask: torch.Tensor = field(default_factory=lambda: torch.empty(0, dtype=torch.float))
 
 
+@dataclass
+class PackedImageTextRawBatch(ImageTextRawBatch):
+    """Sample type for image text raw batch"""
+
+    position_ids: torch.Tensor = field(default_factory=lambda: torch.empty(0, dtype=torch.float))
+    packed_seq_params: PackedSeqParams = field(default_factory=lambda: PackedSeqParams())
+
+
 @dataclass
 class MultiModalSampleConfig:
     image_token: ImageToken = field(default_factory=ImageToken)

nemo/collections/multimodal/data/energon/conversation.py

@@ -30,7 +30,7 @@ class LLaVATemplateConfig(BaseConversationTemplateConfig):
     """LLava-specific template configuration which extends the base config"""
 
     system: str = field(
-        default="A chat between a curious user and artificial assistant agent. "
+        default="A chat between a curious user and an artificial intelligence assistant. "
         "The assistant gives helpful, detailed and polite answers to user's questions."
     )
     roles: List[str] = field(default_factory=lambda: ['user', 'assistant'])

nemo/collections/multimodal/data/energon/task_encoder.py

@@ -25,14 +25,21 @@
     batch_list,
     batch_pad_stack,
 )
+from megatron.energon.task_encoder.base import stateless
 
-from nemo.collections.multimodal.data.energon.config import ImageTextRawBatch, ImageTextSample
+from nemo.collections.multimodal.data.energon.config import (
+    ImageTextRawBatch,
+    ImageTextSample,
+    PackedImageTextRawBatch,
+    PackedImageTextSample,
+)
 from nemo.collections.multimodal.data.energon.sample_encoder import (
     InterleavedSampleEncoder,
     SampleEncoder,
     SimilarityInterleavedEncoder,
     VQASampleEncoder,
 )
+from nemo.utils import logging
 
 
 class MultiModalTaskEncoder(
@@ -54,7 +61,15 @@ class MultiModalTaskEncoder(
     for model input.
     """
 
-    def __init__(self, tokenizer, image_processor, multimodal_sample_config):
+    def __init__(
+        self,
+        tokenizer,
+        image_processor,
+        multimodal_sample_config,
+        packed_sequence=False,
+        packing_seq_length=4096,
+        num_image_embeddings_per_tile=576,
+    ):
         """
         Initialize the MultiModalTaskEncoder with specific encoders for different sample types.
 
@@ -64,6 +79,10 @@ def __init__(self, tokenizer, image_processor, multimodal_sample_config):
         multimodal_sample_config (MultiModalSampleConfig): MultiModalSampleConfig object.
         """
         self.tokenizer = tokenizer
+        self.sample_config = multimodal_sample_config
+        self.packed_sequence = packed_sequence
+        self.num_image_embeddings_per_tile = num_image_embeddings_per_tile  # only used with seq packing
+        self.packing_seq_length = packing_seq_length
         self.encoders: Dict[str, SampleEncoder] = {
             VQASample.__name__: VQASampleEncoder(
                 tokenizer=tokenizer,
@@ -92,6 +111,7 @@ def register_encoder(self, sample_type: str, encoder: SampleEncoder) -> None:
         """
         self.encoders[sample_type] = encoder
 
+    @stateless(restore_seeds=True)
     def encode_sample(
         self, sample: Union[VQASample, InterleavedSample, SimilarityInterleavedSample, CaptioningSample]
     ) -> ImageTextSample:
@@ -118,7 +138,7 @@ def encode_sample(
         encoded_sample = encoder.encode(input_sample=sample, output_sample=ImageTextSample())
         return encoded_sample
 
-    def batch(self, samples: List[ImageTextSample]) -> ImageTextRawBatch:
+    def batch(self, samples):
         """
         Batch a list of encoded samples into a single raw batch.
 
@@ -131,26 +151,40 @@ def batch(self, samples: List[ImageTextSample]) -> ImageTextRawBatch:
         ImageTextRawBatch: The batched data, including images, tokens, labels, and loss masks.
         """
 
-        keys, images, tokens, labels, loss_mask = [], [], [], [], []
-        for sample in samples:
-            keys.append(sample.__key__)
-            images.append(sample.images)
-            tokens.append(sample.tokens)
-            labels.append(sample.labels)
-            loss_mask.append(sample.loss_mask)
-
-        batch_keys = batch_list(keys)
-        batch_images = batch_pad_stack(images)
-        batch_prompt_tokens = batch_pad_stack(tokens)
-        batch_labels = batch_pad_stack(labels)
-        batch_loss_mask = batch_pad_stack(loss_mask)
-        return ImageTextRawBatch(
-            __keys__=batch_keys,
-            images=batch_images,
-            tokens=batch_prompt_tokens,
-            labels=batch_labels,
-            loss_mask=batch_loss_mask,
-        )
+        if self.packed_sequence:
+            assert len(samples) == 1, "Must set MBS=1 when using `packed_sequence`."
+            # The batching are taken care by packing.
+            sample = samples[0]
+            return PackedImageTextRawBatch(
+                __keys__=sample.__key__,
+                images=sample.images,
+                tokens=sample.tokens,
+                labels=sample.labels,
+                loss_mask=sample.loss_mask,
+                position_ids=sample.position_ids,
+                packed_seq_params=sample.packed_seq_params,
+            )
+        else:
+            keys, images, tokens, labels, loss_mask = [], [], [], [], []
+            for sample in samples:
+                keys.append(sample.__key__)
+                images.append(sample.images)
+                tokens.append(sample.tokens)
+                labels.append(sample.labels)
+                loss_mask.append(sample.loss_mask)
+
+            batch_keys = batch_list(keys)
+            batch_images = batch_pad_stack(images)
+            batch_prompt_tokens = batch_pad_stack(tokens)
+            batch_labels = batch_pad_stack(labels)
+            batch_loss_mask = batch_pad_stack(loss_mask)
+            return ImageTextRawBatch(
+                __keys__=batch_keys,
+                images=batch_images,
+                tokens=batch_prompt_tokens,
+                labels=batch_labels,
+                loss_mask=batch_loss_mask,
+            )
 
     def encode_batch(self, batch_data: ImageTextRawBatch) -> dict:
         """
@@ -165,7 +199,7 @@ def encode_batch(self, batch_data: ImageTextRawBatch) -> dict:
         Returns:
         dict: A dictionary containing the encoded batch data, ready for model input.
         """
-        batch_dict = dataclasses.asdict(batch_data)
+        batch_dict = batch_data.__dict__
         if 'images' in batch_dict:
             batch_dict['media'] = batch_dict['images']
             del batch_dict['images']
@@ -177,3 +211,66 @@ def encode_batch(self, batch_data: ImageTextRawBatch) -> dict:
         if 'attention_mask' not in batch_dict:
             batch_dict['attention_mask'] = None
         return batch_dict
+
+    def select_samples_to_pack(self, samples):
+        """Selects which samples will be packed together.
+
+        NOTE: Energon dataloader calls this method internally if packing is used.
+        Please see https://nvidia.github.io/Megatron-Energon/packing.html
+        """
+        from nemo.collections.vlm.neva.data.sequence_packing import greedy_knapsack, predict_seq_len
+
+        media_token_id = self.sample_config.image_token.token_id
+        lengths = [
+            predict_seq_len(
+                sample.tokens,
+                media_token_index=media_token_id,
+                num_image_embeddings_per_tile=self.num_image_embeddings_per_tile,
+            )
+            for sample in samples
+        ]
+        packed_samples = greedy_knapsack(lengths, samples, self.packing_seq_length)
+        avg_samples_per_bin = round(len(lengths) / len(packed_samples))
+        logging.info(
+            f"[Seq Packing Info] - Packing seq len: {self.packing_seq_length}, "
+            f"Buffered samples: {len(lengths)}, Total number of bins: {len(packed_samples)}, "
+            f"Average samples per bin: {avg_samples_per_bin}"
+        )
+        return packed_samples
+
+    @stateless
+    def pack_selected_samples(self, samples):
+        """
+        Function to pack a list of ImageTaskSample into a single ImageTaskSamplePacked.
+
+        NOTE: Energon dataloader calls this method internally if packing is used.
+        Please see https://nvidia.github.io/Megatron-Energon/packing.html
+
+        Args:
+            samples: List of ImageTaskSample instances to pack into one sample.
+
+        Returns:
+            ImageTaskSamplePacked instance.
+        """
+        from nemo.collections.vlm.neva.data.sequence_packing import convert_to_packed
+
+        packed_images = torch.stack([sample.images for sample in samples])
+        media_token_id = self.sample_config.image_token.token_id
+        packed_tokens, packed_labels, packed_position_ids, packed_loss_mask, packed_seq_params = convert_to_packed(
+            tokens=[sample.tokens for sample in samples],
+            labels=[sample.labels for sample in samples],
+            num_image_embeddings_per_tile=self.num_image_embeddings_per_tile,
+            media_token_index=media_token_id,
+            ignore_index=self.sample_config.ignore_place_holder,
+        )
+
+        return PackedImageTextSample(
+            __key__=",".join([s.__key__ for s in samples]),
+            __restore_key__=(),  # Will be set by energon based on `samples`
+            tokens=packed_tokens,
+            labels=packed_labels,
+            images=packed_images,
+            position_ids=packed_position_ids,
+            loss_mask=packed_loss_mask,
+            packed_seq_params=packed_seq_params,
+        )

nemo/collections/vlm/neva/data/config.py

@@ -31,15 +31,15 @@ class DataConfig:
 @dataclass
 class ImageDataConfig(DataConfig):
     media_type: str = "image"
-    media_token: MultiModalToken = ImageToken
+    media_token: MultiModalToken = ImageToken()
     image_folder: Optional[str] = None
     image_process_mode: str = 'pad'
 
 
 @dataclass
 class VideoDataConfig(DataConfig):
     media_type: str = "video"
-    media_token: MultiModalToken = VideoToken
+    media_token: MultiModalToken = VideoToken()
     splice_single_frame: Optional[str] = None
     # 'first', 'middle', 'last' will represent video as first / middle / last frame only, all other frames discarded.
     num_frames: int = 8  # Selects the number of frames to use from the video