Fix Energon Support in Qwen3-VL

examples/models/vlm/qwen3_vl/energon_test.sh

@@ -0,0 +1,79 @@
+#!/usr/bin/env bash
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Workspace directory for checkpoints and results
+WORKSPACE=${WORKSPACE:-/workspace}
+
+# Before training, make sure to set WANDB_API_KEY or disable wandb logging
+# export WANDB_API_KEY=<your_wandb_api_key>
+# export WANDB_MODE=disabled
+
+# Test Seq Packing configurations for LoRA finetuning on the dense model
+PRETRAINED_CHECKPOINT=${WORKSPACE}/models/Qwen3-VL-8B-Instruct
+MODEL_NAME=qwen3_vl_8b
+DATASET_NAME=energon
+SEQ_LENGTH=4096
+TRAIN_ITERS=50
+GLOBAL_BATCH_SIZE=32
+MICRO_BATCH_SIZE=2
+EVAL_ITERS=10
+LR=0.00005
+MIN_LR=0.000005
+LR_WARMUP_ITERS=10
+LOG_INTERVAL=1
+WANDB_PROJECT=megatron-bridge-${DATASET_NAME}
+
+SEQ_PACKING_CONFIGS=(False True)
+
+# EP/TP/PP/CP/N_PROC combinations: "EP,TP,PP,CP,N_PROC" configurations
+# N_PROC is the total number of processes (GPUs) used for training
+# N_PROC is used to control DP size, to make the loss curves comparable
+PARALLELISM_CONFIGS=("1,1,1,4,8" "1,1,1,2,4" "1,1,1,1,2")
+
+for pack_config in "${SEQ_PACKING_CONFIGS[@]}"; do
+    for par_config in "${PARALLELISM_CONFIGS[@]}"; do
+        IFS=',' read -r EP TP PP CP N_PROC <<< "$par_config"
+        echo "Running LoRA finetuning pack_sequences_in_batch=$pack_config with EP=$EP TP=$TP PP=$PP CP=$CP DP=$N_PROC"
+        python -m torch.distributed.run --nproc_per_node=$N_PROC scripts/training/run_recipe.py \
+            --recipe ${MODEL_NAME}_finetune_config \
+            --step_func qwen3_vl_step \
+            --peft_scheme lora \
+            --dataset_type energon \
+            checkpoint.pretrained_checkpoint=$PRETRAINED_CHECKPOINT \
+            model.seq_length=$SEQ_LENGTH \
+            train.train_iters=$TRAIN_ITERS \
+            train.global_batch_size=$GLOBAL_BATCH_SIZE \
+            train.micro_batch_size=$MICRO_BATCH_SIZE \
+            train.eval_iters=$EVAL_ITERS \
+            optimizer.lr=$LR \
+            optimizer.min_lr=$MIN_LR \
+            scheduler.lr_warmup_iters=$LR_WARMUP_ITERS \
+            checkpoint.save=${WORKSPACE}/results/${MODEL_NAME}_lora_seq_pack_${pack_config}_cp${CP} \
+            logger.log_interval=$LOG_INTERVAL \
+            logger.wandb_project=$WANDB_PROJECT \
+            logger.wandb_exp_name=${MODEL_NAME}_${DATASET_NAME}_lora_seq_pack_${pack_config}_cp${CP} \
+            dataset.seq_length=$SEQ_LENGTH \
+            dataset.path=/path/to/energon/dataset \
+            dataset.pack_sequences_in_batch=$pack_config \
+            model.expert_model_parallel_size=$EP \
+            model.tensor_model_parallel_size=$TP \
+            model.pipeline_model_parallel_size=$PP \
+            model.context_parallel_size=$CP \
+            model.calculate_per_token_loss=True \
+            ddp.average_in_collective=False \
+            ddp.grad_reduce_in_fp32=True
+    done
+done
+

scripts/training/run_recipe.py

@@ -132,6 +132,12 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
         default=None,
         help="Sequence length for training",
     )
+    parser.add_argument(
+        "--dataset_type",
+        type=str,
+        default=None,
+        help="Dataset type for VLM recipes (e.g., 'energon', 'mock', 'hf', 'preloaded').",
+    )
     args, cli_overrides = parser.parse_known_args()
     return args, cli_overrides
 
@@ -141,6 +147,7 @@ def load_recipe(
     peft_scheme: str | None,
     packed_sequence: bool = False,
     seq_length: int | None = None,
+    dataset_type: str | None = None,
 ) -> ConfigContainer:
     """
     Load recipe by name from megatron.bridge.recipes.
@@ -150,6 +157,7 @@ def load_recipe(
         peft_scheme: PEFT scheme to use ('lora', 'dora', or None)
         packed_sequence: Enable packed sequence training (default: False)
         seq_length: Sequence length for training (optional)
+        dataset_type: Dataset type for VLM recipes (e.g., 'energon', 'mock', 'hf', 'preloaded')
 
     Returns:
         ConfigContainer from calling the recipe
@@ -175,11 +183,13 @@ def load_recipe(
         accepts_peft = "peft" in params or has_var_keyword
         accepts_packed_sequence = "packed_sequence" in params or has_var_keyword
         accepts_seq_length = "seq_length" in params or has_var_keyword
+        accepts_dataset_type = "dataset_type" in params or has_var_keyword
     except (ValueError, TypeError):
         # If signature inspection fails, fallback conservatively
         accepts_peft = True  # peft is widely supported, try passing it
         accepts_packed_sequence = False  # new parameter, don't pass if unsure
         accepts_seq_length = False  # new parameter, don't pass if unsure
+        accepts_dataset_type = False  # VLM-specific, don't pass if unsure
 
     # Build kwargs dynamically based on what the recipe accepts
     kwargs = {}
@@ -189,6 +199,8 @@ def load_recipe(
         kwargs["packed_sequence"] = packed_sequence
     if accepts_seq_length and seq_length is not None:
         kwargs["seq_length"] = seq_length
+    if accepts_dataset_type and dataset_type is not None:
+        kwargs["dataset_type"] = dataset_type
 
     try:
         return config_builder(**kwargs)
@@ -224,6 +236,7 @@ def main() -> None:
         args.peft_scheme,
         args.packed_sequence,
         args.seq_length,
+        args.dataset_type,
     )
 
     config = process_config_with_overrides(

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/data/energon/energon_provider.py

@@ -33,6 +33,8 @@ class EnergonProvider(DatasetProvider):
     num_workers: int_repr
     dataloader_type: str = "external"
     task_encoder: Optional[Any] = None
+    # Enable batch-level online sequence packing
+    pack_sequences_in_batch: bool = False
 
     def build_datasets(self, context: DatasetBuildContext):
         dataset = EnergonMultiModalDataModule(

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

@@ -19,14 +19,15 @@
 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 +155,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[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 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 +272,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,