Fix Energon Support in Qwen3-VL

src/megatron/bridge/data/energon/base_energon_datamodule.py

@@ -167,17 +167,30 @@ def train_dataloader(self) -> Any:
             return self.train_dataloader_object
         if not parallel_state.is_initialized():
             logger.info(
-                f"Muiltimodal data loader parallel state is not initialized,"
-                f"using default worker config with no_workers {self.num_workers}"
+                f"Multimodal data loader parallel state is not initialized, "
+                f"using default worker config with num_workers {self.num_workers}"
             )
             worker_config = WorkerConfig.default_worker_config(self.num_workers)
         else:
+            # NOTE: We intentionally use the pure DP rank (with_context_parallel=False)
+            # rather than the combined DP-CP rank. With Megatron's rank ordering
+            # (default "tp-cp-ep-dp-pp"), all CP ranks within the same DP replica
+            # already share the same pure DP rank. This ensures that CP ranks
+            # processing different sequence portions of the same batch receive
+            # identical data from the dataloader.
+            # Using with_context_parallel=True would be INCORRECT here — it would
+            # assign each CP rank a unique rank, causing them to read different
+            # data shards.
             rank = parallel_state.get_data_parallel_rank()
             world_size = parallel_state.get_data_parallel_world_size()
             data_parallel_group = parallel_state.get_data_parallel_group()
+            cp_size = parallel_state.get_context_parallel_world_size()
+            cp_rank = parallel_state.get_context_parallel_rank()
             logger.info(
-                f" Multimodal  train dataloader initializing with"
-                f"rank {rank} world_size {world_size} data_parallel_group {data_parallel_group} ****** "
+                f"Multimodal train dataloader initializing with "
+                f"dp_rank {rank} dp_world_size {world_size} "
+                f"cp_rank {cp_rank} cp_size {cp_size} "
+                f"data_parallel_group {data_parallel_group}"
             )
             worker_config = WorkerConfig(
                 rank=rank,
@@ -207,20 +220,28 @@ def val_dataloader(self):
 
         if not parallel_state.is_initialized():
             logger.info(
-                f"Muiltimodal val data loader parallel state is not initialized,"
-                f"using default worker config with no_workers {self.num_workers}"
+                f"Multimodal val data loader parallel state is not initialized, "
+                f"using default worker config with num_workers {self.num_val_workers}"
             )
             worker_config = WorkerConfig.default_worker_config(self.num_val_workers)
         else:
+            # NOTE: Pure DP rank (with_context_parallel=False) is correct here.
+            # See train_dataloader() comment for detailed explanation of CP handling.
             rank = parallel_state.get_data_parallel_rank()
             world_size = parallel_state.get_data_parallel_world_size()
             data_parallel_group = parallel_state.get_data_parallel_group()
-
-            logger.info(f"rank {rank} world_size {world_size} data_parallel_group {data_parallel_group}")
+            cp_size = parallel_state.get_context_parallel_world_size()
+            cp_rank = parallel_state.get_context_parallel_rank()
+            logger.info(
+                f"Multimodal val dataloader initializing with "
+                f"dp_rank {rank} dp_world_size {world_size} "
+                f"cp_rank {cp_rank} cp_size {cp_size} "
+                f"data_parallel_group {data_parallel_group}"
+            )
             worker_config = WorkerConfig(
                 rank=rank,
                 world_size=world_size,
-                num_workers=self.num_workers,
+                num_workers=self.num_val_workers,
                 data_parallel_group=data_parallel_group,
                 worker_debug_path=None,
                 worker_log_level=0,

src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py

@@ -13,20 +13,22 @@
 # limitations under the License.
 
 import dataclasses
+import io
 import json
 import logging
 import pickle
 import re
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import numpy as np
 import torch
 from megatron.energon import Batch, DefaultTaskEncoder
+from megatron.energon.epathlib.epath import EPath
 from megatron.energon.flavors.base_dataset import Sample
-from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
-from PIL import Image
+from megatron.energon.flavors.webdataset import DefaultDecoderWebdatasetFactory
+from webdataset.autodecode import Decoder, imagehandler
 
 from megatron.bridge.training.utils.visual_inputs import Qwen2_5_VLVisualInputs
 
@@ -154,13 +156,53 @@ def _get(token_str: str, default_id: int) -> int:
 
 @dataclass
 class ChatMLSample(Sample):
-    """Intermediate Sample Format"""
+    """multi-turn complex samples with images and videos"""
 
-    # __key__: str
-    # __subflavors__: Dict
-    imgs: List[Image.Image]
-    videos: List[torch.Tensor | list[Image.Image]]
     conversation: str  # JSON string of GPT-format conversations
+    imgs: Optional[List[torch.Tensor]] = None
+    videos: Optional[List[torch.Tensor]] = None
+
+
+class videohandler:
+    """Create an video handler."""
+
+    def __init__(self, imagespec):
+        self.extensions = ["jpgs", "mp4s"]
+        self.extensions_mapping = {"jpgs": "jpg", "mp4s": "jpg"}
+        self.image_handler = imagehandler(imagespec)
+
+    def __call__(self, key, data):
+        """Perform nested image decoding."""
+        extension = re.sub(r".*[.]", "", key)
+        if extension.lower() not in self.extensions:
+            return None
+        data = pickle.loads(data)
+        key = self.extensions_mapping[extension]
+        if extension.lower() == "jpgs":
+            data = [self.image_handler(key, d) for d in data]
+        else:
+            data = [[self.image_handler(key, d) for d in video] for video in data]
+        return data
+
+
+class ChatMLWebdataset(DefaultDecoderWebdatasetFactory[ChatMLSample]):
+    """Webdataset factory for multi-turn ChatML samples with multimodal support.
+
+    Extends DefaultDecoderWebdatasetFactory to decode webdataset shards into
+    ChatMLSample instances, using custom handlers for image, audio, and video fields.
+    """
+
+    __sample_type__ = ChatMLSample
+
+    def __init__(self, path: EPath, *, auto_decode: bool = True, **kwargs):
+        super().__init__(path, auto_decode=auto_decode, **kwargs)
+        if auto_decode:
+            self._decoder = Decoder(
+                [
+                    imagehandler(self.image_decode),
+                    videohandler(self.image_decode),
+                ]
+            )
 
 
 @dataclass
@@ -231,51 +273,9 @@ def convert_to_qwenvl_content(user_input: str, image_pattern: str = "<image>", v
     return contents
 
 
-def cook_chatml_sample(sample: dict) -> ChatMLSample:
-    """
-    Convert crude sampel to ChatMLSample.
-
-    Args:
-        sample: Crude sample in pickle serialized format
-
-    Returns:
-        sample in ChatMLSample format
-    """
-    imgs = sample.get("jpgs", None)
-    if imgs:
-        imgs = pickle.loads(imgs)
-        if isinstance(imgs, list) and len(imgs) > 0:
-            imgs = [Image.fromarray(d) for d in imgs]
-        else:
-            imgs = None
-    videos = sample.get("videos", None)
-    if videos:
-        videos = pickle.loads(videos)
-        if isinstance(videos, list) and len(videos) > 0:
-            videos = [[d for d in video] for video in videos]
-        else:
-            videos = None
-    if "<image>" in sample["json"] and imgs is None:
-        logging.warning("<image> in conversation text but no image data")
-    if "<video>" in sample["json"] and videos is None:
-        logging.warning("<video> in conversation text but no video data")
-
-    chat_sample = ChatMLSample(
-        **basic_sample_keys(sample),
-        imgs=imgs,
-        videos=videos,
-        conversation=sample["json"],
-    )
-    return chat_sample
-
-
 class QwenVLTaskEncoder(DefaultTaskEncoder[ChatMLSample, QwenVLTaskSample, QwenVLTaskBatch, dict]):
     """A simple task encoder for captioning."""
 
-    cookers = [
-        Cooker(cook_chatml_sample),
-    ]
-
     def __init__(
         self,
         tokenizer,