Merge branch 'main' into rob-luke/training/mlflow-description

Author: rob-luke

examples/conversion/hf_to_megatron_generate_vlm.py

@@ -38,6 +38,8 @@
     pad_input_ids_to_tp_multiple,
     patch_kimi_vision_processor,
     process_image_inputs,
+    process_multi_image_inputs,
+    process_video_inputs,
     to_cuda,
 )
 
@@ -64,6 +66,8 @@ def __init__(
         image_grid_thw=None,
         image_sizes=None,
         mm_token_type_ids=None,
+        pixel_values_videos=None,
+        video_grid_thw=None,
         image_position_ids=None,
     ):
         self.batch = dict(
@@ -79,6 +83,10 @@ def __init__(
             self.batch["image_sizes"] = image_sizes
         if mm_token_type_ids is not None:
             self.batch["mm_token_type_ids"] = mm_token_type_ids
+        if pixel_values_videos is not None:
+            self.batch["pixel_values_videos"] = pixel_values_videos
+        if video_grid_thw is not None:
+            self.batch["video_grid_thw"] = video_grid_thw
         if image_position_ids is not None:
             self.batch["image_position_ids"] = image_position_ids
         self._yielded = False
@@ -101,7 +109,15 @@ def vlm_forward_step(data_iterator, model, **kwargs) -> torch.Tensor:
         "position_ids": batch["position_ids"],
         "attention_mask": batch.get("attention_mask"),
     }
-    for key in ("pixel_values", "image_grid_thw", "image_sizes", "mm_token_type_ids", "image_position_ids"):
+    for key in (
+        "pixel_values",
+        "image_grid_thw",
+        "image_sizes",
+        "mm_token_type_ids",
+        "pixel_values_videos",
+        "video_grid_thw",
+        "image_position_ids",
+    ):
         if key in batch:
             forward_args[key] = batch[key]
 
@@ -215,27 +231,41 @@ def _disable_mtp(m):
     # ------------------------------------------------------------------
     # Process inputs
     # ------------------------------------------------------------------
+    pixel_values = image_grid_thw = image_sizes = mm_token_type_ids = image_position_ids = None
+    pixel_values_videos = video_grid_thw = None
+
+    if args.video_path:
+        input_ids_raw, pixel_values_videos, video_grid_thw = process_video_inputs(
+            processor, args.video_path, args.prompt, fps=args.video_fps
+        )
+    elif args.image_paths:
+        input_ids_raw, pixel_values, image_grid_thw = process_multi_image_inputs(
+            processor, args.image_paths, args.prompt
+        )
+    else:
+        (
+            input_ids_raw,
+            pixel_values,
+            image_grid_thw,
+            image_sizes,
+            mm_token_type_ids,
+            image_position_ids,
+        ) = process_image_inputs(
+            processor,
+            args.image_path,
+            args.prompt,
+            is_gemma4=is_gemma4,
+            is_kimi=is_kimi,
+            image_token_id=image_token_id,
+        )
 
-    (
-        input_ids_raw,
-        pixel_values,
-        image_grid_thw,
-        image_sizes,
-        mm_token_type_ids,
-        image_position_ids,
-    ) = process_image_inputs(
-        processor,
-        args.image_path,
-        args.prompt,
-        is_gemma4=is_gemma4,
-        is_kimi=is_kimi,
-        image_token_id=image_token_id,
-    )
     input_ids_raw = input_ids_raw.cuda()
     pixel_values = to_cuda(pixel_values)
     image_grid_thw = to_cuda(image_grid_thw)
     image_sizes = to_cuda(image_sizes)
     mm_token_type_ids = to_cuda(mm_token_type_ids)
+    pixel_values_videos = to_cuda(pixel_values_videos)
+    video_grid_thw = to_cuda(video_grid_thw)
     image_position_ids = to_cuda(image_position_ids)
 
     # ------------------------------------------------------------------
@@ -270,7 +300,9 @@ def _disable_mtp(m):
                 image_grid_thw,
                 image_sizes,
                 mm_ids_padded,
-                image_position_ids,
+                pixel_values_videos=pixel_values_videos,
+                video_grid_thw=video_grid_thw,
+                image_position_ids=image_position_ids,
             )
 
             output = fwd_bwd_function(
@@ -343,7 +375,18 @@ def _disable_mtp(m):
         "--pp_layout", type=str, default=None, help="Pipeline model parallel layout (e.g. 'Et*15|t*15|t*16|t*15L')"
     )
     parser.add_argument("--megatron_model_path", type=str, default=None, help="Path to Megatron model checkpoint")
-    parser.add_argument("--image_path", type=str, default=None, help="Path or URL to image (optional).")
+    parser.add_argument("--image_path", type=str, default=None, help="Path or URL to a single image (optional).")
+    parser.add_argument(
+        "--image_paths",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Paths to N image files in order (multi-image; Qwen-family only).",
+    )
+    parser.add_argument("--video_path", type=str, default=None, help="Path to a video file (Qwen-family only).")
+    parser.add_argument(
+        "--video_fps", type=float, default=2.0, help="Frames per second to sample from the video (default: 2.0)."
+    )
     parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code for HF model loading")
     args = parser.parse_args()
 

examples/conversion/vlm_generate_utils.py

@@ -65,9 +65,9 @@ def _patched_from_dict(cls, config_dict, **kwargs):
         pass
 
 
-def pre_expand_image_tokens(input_ids, grid_thws, image_token_id, spatial_merge_size=2):
-    """Pre-expand single image placeholders to N placeholders matching vision
-    feature count.
+def pre_expand_vision_tokens(input_ids, grid_thws, vision_token_id, spatial_merge_size=2):
+    """Pre-expand single vision placeholders (image or video) to N placeholders
+    matching vision feature count.
 
     With PP > 1 the pipeline schedule needs to know the actual sequence length
     upfront.  Dynamic expansion inside the model changes seq_length during
@@ -86,8 +86,8 @@ def pre_expand_image_tokens(input_ids, grid_thws, image_token_id, spatial_merge_
     expanded = []
     feat_idx = 0
     for token_id in input_ids[0]:
-        if token_id.item() == image_token_id and feat_idx < len(feature_counts):
-            expanded.extend([image_token_id] * feature_counts[feat_idx])
+        if token_id.item() == vision_token_id and feat_idx < len(feature_counts):
+            expanded.extend([vision_token_id] * feature_counts[feat_idx])
             feat_idx += 1
         else:
             expanded.append(token_id.item())
@@ -132,6 +132,62 @@ def to_cuda(x):
     return x.cuda()
 
 
+def process_multi_image_inputs(processor, image_paths: list[str], prompt: str):
+    """Process N ordered images + prompt into model inputs (Qwen-family).
+
+    Returns:
+        (input_ids, pixel_values, image_grid_thw)
+    """
+    if not _HAS_QWEN_VL_UTILS:
+        raise ImportError("qwen-vl-utils required: pip install qwen-vl-utils")
+    pils = [load_image(p).convert("RGB") for p in image_paths]
+    messages = [
+        {
+            "role": "user",
+            "content": [{"type": "image", "image": p} for p in pils] + [{"type": "text", "text": prompt}],
+        }
+    ]
+    image_inputs, video_inputs = process_vision_info(messages)
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
+    return inputs.input_ids, inputs.get("pixel_values"), inputs.get("image_grid_thw")
+
+
+def process_video_inputs(processor, video_path: str, prompt: str, *, fps: float = 2.0):
+    """Process a video + prompt into model inputs (Qwen-family).
+
+    Frame decoding mirrors the Qwen3-VL training pipeline: fetch_video decodes at
+    ``fps``, then video_processor is called with do_sample_frames=False to use the
+    pre-decoded frames as-is.
+
+    Returns:
+        (input_ids, pixel_values_videos, video_grid_thw)
+    """
+    if not _HAS_QWEN_VL_UTILS:
+        raise ImportError("qwen-vl-utils required: pip install qwen-vl-utils")
+    from qwen_vl_utils import fetch_video
+
+    frames = fetch_video({"video": video_path, "fps": fps})
+    messages = [
+        {
+            "role": "user",
+            "content": [{"type": "video"}, {"type": "text", "text": prompt}],
+        }
+    ]
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    text_inputs = processor(text=[text], padding=True, return_tensors="pt")
+    video_proc = processor.video_processor(videos=[frames], return_tensors="pt", do_sample_frames=False)
+    # processor(text=...) without videos produces a single <|video_pad|> placeholder.
+    # Pre-expand to match actual vision feature count so PP send/recv shapes are correct.
+    video_token_id = processor.tokenizer.convert_tokens_to_ids("<|video_pad|>")
+    input_ids = pre_expand_vision_tokens(
+        text_inputs["input_ids"],
+        video_proc["video_grid_thw"],
+        vision_token_id=video_token_id,
+    )
+    return input_ids, video_proc.get("pixel_values_videos"), video_proc.get("video_grid_thw")
+
+
 def process_image_inputs(
     processor,
     image_path: Optional[str],
@@ -179,7 +235,7 @@ def _process_kimi_inputs(processor, image_path, prompt, image_token_id):
     ]
     inputs = processor(messages=messages)
     grid_thws = getattr(inputs, "grid_thws", None)
-    input_ids = pre_expand_image_tokens(inputs.input_ids, grid_thws, image_token_id)
+    input_ids = pre_expand_vision_tokens(inputs.input_ids, grid_thws, image_token_id)
     return input_ids, inputs.pixel_values, grid_thws, None, None
 
 

examples/models/qwen/qwen3_vl/README.md

@@ -113,6 +113,12 @@ See the [peft.sh](peft.sh) script for LoRA fine-tuning with sequence-packing.
 
 **Note:** LoRA/DoRA significantly reduces memory requirements, allowing for larger batch sizes and fewer GPUs.
 
+## Controlling visual tokens computation budget
+Three independent CLI-overridable controls bound a sample's GPU cost. They compose:
+- **`dataset.min_pixels` / `dataset.max_pixels`** — image/frame resolutions lower and upper bound (defaults `200704` / `1003520`). 
+- **`dataset.max_num_images` / `dataset.max_num_frames`** - limit count of images/frames (defaults `10` / `60`). Too many images → sample is dropped. Too many frames → frame list truncated.
+- **`dataset.max_visual_tokens`** — limit total visual tokens across all images and frames in a sample, computed post-rescaling as `prod(T,H,W) // merge_size²` (default `16384`; set to `None` to disable). Catches cases the other two miss (few images at high resolution, or many at low resolution). Exceeding samples are dropped.
+
 ## Finetuning with Energon Dataset
 
 Follow the instructions [here](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/multimodal#pretraining) to prepare `LLaVA-Pretrain` dataset in Energon format. Change the file `.nv-meta/dataset.yaml` to the following:
@@ -127,10 +133,65 @@ field_map:
 
 Then, update the dataset path (`dataset.path=/path/to/energon/dataset`) in [peft_energon.sh](peft_energon.sh) and run the script.
 
-
 ### Expected Training Dynamics
 We provide a [Weights & Biases report](https://api.wandb.ai/links/nvidia-nemo-fw-public/lczz4ixx) for the expected loss curves and grad norms.
 
+## Dataset with Multiple Images
+
+Below is an example for finetuning on a dataset containing multiple images in a sample, using a subset of [TIGER-Lab/Mantis-Instruct](https://huggingface.co/datasets/TIGER-Lab/Mantis-Instruct) dataset.
+
+1. Download the `llava_665k_multi` subset of TIGER-Lab/Mantis-Instruct dataset from Hugging Face and unzip the images folder (NOTE: 44GB of disk space required):
+
+    ```
+    pip install -U "huggingface_hub[cli]"
+    huggingface-cli download TIGER-Lab/Mantis-Instruct \
+        --include "llava_665k_multi/*" \
+        --repo-type dataset \
+        --local-dir /path/to/Mantis-Instruct-LLaVA    
+    ```
+
+2. Run the following script to convert the data to webdataset format:
+
+    ```
+    python examples/models/vlm/qwen3_vl/prepare_mantis_energon.py \
+        --source-dir/path/to/Mantis-Instruct-LLaVA \
+        --output-dir /path/to/Mantis-Instruct-LLaVA/wds \
+        --max-samples-per-tar 10000
+    ```
+
+3. Run the following command to convert to megatron-energon format:
+
+    ```
+    cd /path/to/Mantis-Instruct-LLaVA/wds
+    energon prepare ./
+    ```
+
+    select the following values for the presented options:
+
+    ```
+    > Please enter a desired train/val/test split like "0.5, 0.2, 0.3" or "8,1,1": 9,1,0
+    > Do you want to create a dataset.yaml interactively? [Y/n]: Y
+    > Please enter a number to choose a class: 9 (VQASample)
+    > Do you want to set a simple field_map[Y] (or write your own sample_loader [n])? [Y/n]: Y
+    > Please enter a webdataset field name for 'image' (<class 'torch.Tensor'>): jpg
+    > Please enter a webdataset field name for 'context' (<class 'str'>): json[0][value]
+    > Please enter a webdataset field name for 'answers' (typing.Optional[typing.List[str]], default: None): json[1][value]
+    > Please enter a webdataset field name for 'answer_weights' (typing.Optional[torch.Tensor], default: None):
+    ```
+
+4. Change the file `.nv-meta/dataset.yaml` to the following:
+
+    ```yaml
+    __module__: megatron.bridge.recipes.qwen_vl.data.energon.task_encoder
+    __class__: ChatMLWebdataset
+    field_map:
+      imgs: jpgs
+      conversation: json
+    ```
+
+Follow previous instruction to run the finetuning with the prepared dataset.
+
+
 ## Evaluation
 
 Coming soon.