examples/multimodal_dev: docs + Victarry review fixes + vision patch merger parity test

examples/multimodal_dev/README.md

@@ -34,6 +34,66 @@ torchrun --nproc_per_node=8 multimodal_dev/pretrain_multimodal.py \
     ... # other Megatron args (--num-layers, --hidden-size, etc.)
 ```
 
+## Checkpoint Conversion (HF → Megatron-FSDP DTensor)
+
+Convert a HuggingFace release to a Megatron-FSDP DTensor checkpoint via
+[Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) before
+pretraining from pretrained weights.
+
+### Setup
+
+Clone Bridge and pin its `3rdparty/Megatron-LM` submodule to this branch:
+
+```bash
+git clone --recurse-submodules https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
+cd Megatron-Bridge/3rdparty/Megatron-LM
+git remote add wplf https://github.com/wplf/Megatron-LM.git
+git fetch wplf feat/qwen35-vl-example
+git checkout feat/qwen35-vl-example
+cd ../..
+```
+
+### Convert
+
+Single 8×GPU node, EP=8 / TP=CP=1; substitute any Qwen3.5 variant for
+`--hf-model`:
+
+```bash
+PYTHONPATH=./src:./3rdparty/Megatron-LM/ \
+  torchrun --nproc_per_node=8 \
+  examples/conversion/mfsdp/convert_checkpoints_fsdp.py import \
+  --hf-model Qwen/Qwen3.5-35B-A3B \
+  --megatron-path ${WORKSPACE}/models/Qwen/Qwen3.5-35B-A3B-fsdp \
+  --ckpt-format fsdp_dtensor \
+  --ep 8
+```
+
+HF weights are auto-fetched on first run via `huggingface_hub`. Adjust
+`--tp` / `--cp` / `--ep` to match the training topology (must satisfy
+`WORLD_SIZE % (TP*CP*EP) == 0`).
+
+### Output
+
+```
+${WORKSPACE}/models/Qwen/Qwen3.5-35B-A3B-fsdp/
+├── iter_0000000/
+│   ├── __0_0.distcp .. __7_0.distcp   # FSDP DTensor shards, one per rank (~18 GB each for 35B-A3B)
+│   ├── .metadata
+│   ├── run_config.yaml
+│   └── train_state.pt
+├── latest_checkpointed_iteration.txt
+└── latest_train_state.pt
+```
+
+### Bridge dependency
+
+Requires
+[NVIDIA-NeMo/Megatron-Bridge#3987](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/3987)
+(skip tokenizer save). Without that fix the checkpoint is still written
+correctly but the script exits non-zero after save with
+`AttributeError: 'TokenizerConfig' object has no attribute 'make_vocab_size_divisible_by'`
+against this branch's `megatron.core.tokenizers.utils.build_tokenizer`.
+
 ## Architecture
 
 `pretrain_multimodal.py` is **model-agnostic**. All model-specific logic

examples/multimodal_dev/data/vlm_dataset.py → examples/multimodal_dev/data/cord_v2.py

@@ -1,11 +1,11 @@
 # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
 
-"""Simple VLM dataset for multimodal_dev training.
+"""CORD-V2 VLM dataset for multimodal_dev training.
 
 Single-turn image-text dataset using a HuggingFace ``AutoProcessor`` for
-tokenization and image preprocessing.  Currently supports CORD-V2 (receipt
-OCR).  No multi-turn support — each sample is one image + question →
-answer pair.
+tokenization and image preprocessing.  This module is the reference
+implementation for the CORD-V2 receipt-OCR dataset.  No multi-turn support —
+each sample is one image + question → answer pair.
 
 Each image is preprocessed via ``qwen_vl_utils.process_vision_info`` and
 fed to the processor with Qwen-VL's recommended ``min_pixels`` /
@@ -20,6 +20,32 @@
         --model-arch qwen35_vl --dataset-provider cord_v2 \\
         --hf-processor-path Qwen/Qwen3.5-397B-A17B \\
         --total-seq-length 4096 --use-vanilla-collate-fn
+
+Adding another VLM dataset
+--------------------------
+
+The dataset layer mirrors the model layer's registry pattern: each dataset
+ships its own module and a ``train_valid_test_datasets_provider`` factory,
+and the model's registry entry maps a ``--dataset-provider`` name to that
+factory's dotted path.  To add a new dataset (e.g. NLVR2):
+
+1. Create ``examples/multimodal_dev/data/<name>.py`` with::
+
+       def train_valid_test_datasets_provider(train_val_test_num_samples):
+           ...  # build datasets using args from get_args()
+           return train_ds, val_ds, test_ds
+
+2. Register it under the relevant model in
+   ``examples/multimodal_dev/models/__init__.py``::
+
+       MODEL_REGISTRY["qwen35_vl"]["dataset_providers"]["<name>"] = (
+           "examples.multimodal_dev.data.<name>"
+           ".train_valid_test_datasets_provider"
+       )
+
+3. Launch with ``--dataset-provider <name>``.
+
+No edits to ``pretrain_multimodal.py`` or ``forward_step.py`` are required.
 """
 
 import json

examples/multimodal_dev/models/__init__.py

@@ -54,7 +54,7 @@
                 ".train_valid_test_datasets_provider"
             ),
             "cord_v2": (
-                "examples.multimodal_dev.data.vlm_dataset"
+                "examples.multimodal_dev.data.cord_v2"
                 ".train_valid_test_datasets_provider"
             ),
         },

examples/multimodal_dev/models/qwen35_vl/mrope.py

@@ -55,6 +55,7 @@ def _build_sample_mrope_positions(
     vision_tokens = sample_input_ids[vision_start_indices + 1]
     image_nums = int((vision_tokens == image_token_id).sum())
     video_nums = int((vision_tokens == video_token_id).sum())
+    # TODO: fuse into a kernel to drop the per-iter GPU<->CPU sync.
     input_tokens = sample_input_ids.tolist()
     llm_pos_ids_list: list = []
     st = 0

examples/multimodal_dev/models/qwen35_vl/vision_encoder.py

@@ -249,8 +249,8 @@ def forward(self, hidden_states: Tensor) -> Tensor:
         hidden_states = self.patch_norm(hidden_states)
         merged = hidden_states.view(-1, self.merge_dim)
         merged, _ = self.linear_fc1(merged)
-        # NOTE: Official HuggingFace uses default approximate='none' in Qwen3VLVisionPatchMerger.
-        merged = torch.nn.functional.gelu(merged, approximate="tanh")
+        # Match official HuggingFace Qwen3VLVisionPatchMerger (default approximate='none').
+        merged = torch.nn.functional.gelu(merged, approximate="none")
         merged, _ = self.linear_fc2(merged)
         return merged
 

examples/multimodal_dev/pretrain_multimodal.py

@@ -148,6 +148,16 @@ def datasets_provider(train_val_test_num_samples):
         extra_args_provider=add_multimodal_args,
         args_defaults={},
     )
+    # multimodal_dev's model_provider builds the full model on every rank and
+    # does not honor pre_process / post_process pipeline-stage flags. PP>1
+    # would silently violate Megatron's pipeline-parallel contract.
+    if args.pipeline_model_parallel_size > 1:
+        raise ValueError(
+            "multimodal_dev does not support pipeline_model_parallel_size > 1 "
+            f"(got {args.pipeline_model_parallel_size}). The model provider "
+            "builds the full model on every rank; pipeline-stage splitting is "
+            "not wired through. Run with --pipeline-model-parallel-size 1."
+        )
     full_config = pretrain_cfg_container_from_args(args)
     pretrain(
         full_config,

examples/multimodal_dev/scripts/run_qwen35_vl.sh

@@ -9,9 +9,12 @@
 #   MODEL_VARIANT: proxy (default), 0.8b, 2b, 4b, 9b, 27b, 35b_a3b, 122b_a10b, 397b_a17b, 35b_a3b_light
 #   CKPT_LOAD: path to a pre-converted checkpoint to load (enables --load + --finetune)
 #   CKPT_FORMAT: checkpoint format override (e.g. torch_dist); auto-detected when empty
-#   TP, EP, PP: parallelism sizes
+#   TP, EP, PP: parallelism sizes (PP must stay 1; multimodal_dev does not
+#               support pipeline parallelism)
 #   MBS, GBS: micro/global batch sizes
 #   NUM_LAYERS, NUM_EXPERTS: override for proxy testing
+#   FORCE_LOAD_BALANCING: set to 1 to enable --moe-router-force-load-balancing
+#                         (perf / mock-data only; OFF for real finetuning)
 #   LAUNCHER: torchrun (default) or python
 #   PROFILE: set to 1 to enable Nsight Systems profiling (default: 0)
 #   PROFILE_STEP_START/PROFILE_STEP_END: profiled iteration window (default: 4-5)
@@ -48,9 +51,15 @@ GBS=${GBS:-16}
 
 # Parallelism
 TP=${TP:-1}
-EP=${EP:-2}
+# EP defaults to 1; MoE variants override via the variant case block below.
+EP=${EP:-1}
 PP=${PP:-1}
 CP=${CP:-1}
+# Gate --moe-router-force-load-balancing behind an explicit opt-in. Useful for
+# perf / mock-data benchmarking (it disables the auxiliary load-balancing loss
+# coupling so router routes are perfectly uniform), but it must be off for any
+# real fine-tuning / convergence run because it freezes data-dependent routing.
+FORCE_LOAD_BALANCING=${FORCE_LOAD_BALANCING:-0}
 
 # Variant-aware architecture defaults.
 # The model provider builds configs from the variant dict in
@@ -168,6 +177,16 @@ case "$MODEL_VARIANT" in
         VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-27}
         ;;
 esac
+
+# Fail fast on inconsistent expert-parallelism configuration. Dense variants
+# (NUM_EXPERTS=0) do not emit any --num-experts / MoE args, so forwarding
+# --expert-model-parallel-size > 1 would trip Megatron's arg validation.
+if [ "${NUM_EXPERTS:-0}" -eq 0 ] && [ "$EP" -gt 1 ]; then
+    echo "ERROR: MODEL_VARIANT=$MODEL_VARIANT has NUM_EXPERTS=0 (dense) but EP=$EP." >&2
+    echo "       Set EP=1 for dense variants, or pick a MoE variant." >&2
+    exit 1
+fi
+
 SEQ_LEN=${SEQ_LEN:-4096}
 
 WANDB_PROJECT=${WANDB_PROJECT:-'qwen35-vl-0524'}
@@ -348,7 +367,6 @@ GPT_MODEL_ARGS=(
     --linear-num-key-heads 16
     --linear-num-value-heads "$LINEAR_NUM_VALUE_HEADS"
     --make-vocab-size-divisible-by 485
-    --moe-router-force-load-balancing
 )
 
 # --- Tied / untied embeddings ---
@@ -391,6 +409,11 @@ if [ "${NUM_EXPERTS:-0}" -gt 0 ]; then
         --moe-permute-fusion
         --moe-router-fusion
     )
+    # Perf / mock-data only: forces uniform router decisions; do NOT enable for
+    # real finetuning (it freezes data-dependent routing).
+    if [ "$FORCE_LOAD_BALANCING" -eq 1 ]; then
+        MOE_ARGS+=( --moe-router-force-load-balancing )
+    fi
 fi
 
 # --- Recompute ---