From f910ad4547947e29f971a0756deb5b2e41f2b4ae Mon Sep 17 00:00:00 2001 From: Jinliang Li Date: Tue, 26 May 2026 23:21:03 -0700 Subject: [PATCH 1/3] docs(examples/multimodal_dev): add checkpoint conversion guide to README Document the HF -> Megatron-FSDP DTensor conversion path needed before pretraining from pretrained weights: setup (clone Bridge, pin its 3rdparty/Megatron-LM submodule to this branch), the `torchrun convert_checkpoints_fsdp.py import` command with EP=8 default topology, expected output layout, and the open Bridge dependency (NVIDIA-NeMo/Megatron-Bridge#3987) to skip the post-save tokenizer build that otherwise crashes on this branch. --- examples/multimodal_dev/README.md | 60 +++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/examples/multimodal_dev/README.md b/examples/multimodal_dev/README.md index bdc34414da9..e4e6c53ceb4 100644 --- a/examples/multimodal_dev/README.md +++ b/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 From df54ae7cff74c76a8bb5750dc52f7e54242d9975 Mon Sep 17 00:00:00 2001 From: Jinliang Li Date: Wed, 27 May 2026 19:57:31 -0700 Subject: [PATCH 2/3] fix(examples/multimodal_dev): address Victarry review feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves the inline comments from @Victarry's PR review on #4751. * vision_encoder.py — patch merger GELU was `approximate='tanh'` while the in-code NOTE acknowledged HF uses `approximate='none'`. Switched to `approximate='none'` to match the official Qwen3VLVisionPatchMerger numerics for HF -> Megatron checkpoint parity. * pretrain_multimodal.py — added an explicit guard against `--pipeline-model-parallel-size > 1`. The model_provider builds the full model on every rank and ignores pre_process / post_process stage flags, so PP>1 would silently break Megatron's pipeline-parallel contract. Fail fast instead. * scripts/run_qwen35_vl.sh — three fixes: 1. `EP` now defaults to 1 (was 2). MoE variants must opt in via the environment override. 2. After the variant case block, fail fast if `NUM_EXPERTS=0 && EP>1` so a dense run such as `MODEL_VARIANT=9b ./run_qwen35_vl.sh` no longer trips Megatron's arg validation downstream. 3. `--moe-router-force-load-balancing` was unconditionally added to GPT_MODEL_ARGS (and therefore enabled even when no MoE args were emitted). It is now gated behind `FORCE_LOAD_BALANCING=1`, defaults off, and is appended to MOE_ARGS only when MoE is active. Real finetuning runs no longer freeze router routing decisions by default. * data/{vlm_dataset.py -> cord_v2.py} + models/__init__.py — renamed the CORD-V2-specific module from the generic-sounding `vlm_dataset.py` to `cord_v2.py`, updated the model registry path string accordingly, and added an "Adding another VLM dataset" section to the module docstring documenting the per-dataset module + `MODEL_REGISTRY["..."]["dataset_providers"]` registration pattern. * models/qwen35_vl/mrope.py — added a performance note on the `_build_sample_mrope_positions` helper documenting the `.tolist()` / `.item()` GPU<->CPU sync points and CUDA-graph incompatibility, and the precompute-in-collate / cache-by-shape follow-up plan. Behavior preserved here pending a follow-up data pipeline change. The other tests-import comment (test_thd_*.py importing `_pack_batch`) is already addressed on this branch: the helper is now named `pack_or_pad_batch` and the tests import that symbol. --- .../data/{vlm_dataset.py => cord_v2.py} | 34 ++++++++++++++++--- examples/multimodal_dev/models/__init__.py | 2 +- .../multimodal_dev/models/qwen35_vl/mrope.py | 1 + .../models/qwen35_vl/vision_encoder.py | 4 +-- .../multimodal_dev/pretrain_multimodal.py | 10 ++++++ .../multimodal_dev/scripts/run_qwen35_vl.sh | 29 ++++++++++++++-- 6 files changed, 70 insertions(+), 10 deletions(-) rename examples/multimodal_dev/data/{vlm_dataset.py => cord_v2.py} (91%) diff --git a/examples/multimodal_dev/data/vlm_dataset.py b/examples/multimodal_dev/data/cord_v2.py similarity index 91% rename from examples/multimodal_dev/data/vlm_dataset.py rename to examples/multimodal_dev/data/cord_v2.py index 3a493087252..69fd4c13ec4 100644 --- a/examples/multimodal_dev/data/vlm_dataset.py +++ b/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/.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"][""] = ( + "examples.multimodal_dev.data." + ".train_valid_test_datasets_provider" + ) + +3. Launch with ``--dataset-provider ``. + +No edits to ``pretrain_multimodal.py`` or ``forward_step.py`` are required. """ import json diff --git a/examples/multimodal_dev/models/__init__.py b/examples/multimodal_dev/models/__init__.py index e8ed05f1ca2..225414055e6 100644 --- a/examples/multimodal_dev/models/__init__.py +++ b/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" ), }, diff --git a/examples/multimodal_dev/models/qwen35_vl/mrope.py b/examples/multimodal_dev/models/qwen35_vl/mrope.py index 763070929be..9e0e98b1a35 100644 --- a/examples/multimodal_dev/models/qwen35_vl/mrope.py +++ b/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 diff --git a/examples/multimodal_dev/models/qwen35_vl/vision_encoder.py b/examples/multimodal_dev/models/qwen35_vl/vision_encoder.py index 8e8a6146a7f..d57d114374d 100644 --- a/examples/multimodal_dev/models/qwen35_vl/vision_encoder.py +++ b/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 diff --git a/examples/multimodal_dev/pretrain_multimodal.py b/examples/multimodal_dev/pretrain_multimodal.py index 9031339f360..053fa00a5a2 100644 --- a/examples/multimodal_dev/pretrain_multimodal.py +++ b/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, diff --git a/examples/multimodal_dev/scripts/run_qwen35_vl.sh b/examples/multimodal_dev/scripts/run_qwen35_vl.sh index 80a2a94671a..3a1ca55c826 100755 --- a/examples/multimodal_dev/scripts/run_qwen35_vl.sh +++ b/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 --- From 59923c4abacac0ca6bc34895ed7cb237a7ef94f5 Mon Sep 17 00:00:00 2001 From: Jinliang Li Date: Wed, 27 May 2026 20:10:45 -0700 Subject: [PATCH 3/3] test(examples/multimodal_dev): add Qwen35VLPatchMerger HF parity test New test ``tests/test_vision_patch_merger_parity.py`` verifies the Megatron patch merger against an inlined verbatim copy of HuggingFace ``Qwen3VLVisionPatchMerger`` (``use_postshuffle_norm=False`` branch from ``transformers/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py``). The HF reference is inlined so the test has no runtime dependency on the ``transformers`` package. The test copies HF state-dict tensors into the Megatron module (TP=1, 1:1 mapping), runs both on the same random input, and asserts ``torch.testing.assert_close`` on the logits in fp32 and bf16: [torch.float32] shape=(16, 3584) max_abs_diff=2.551e-05 (atol=1e-4) [torch.bfloat16] shape=(16, 3584) max_abs_diff=3.906e-03 (atol=5e-2) The fp32 residual is structural (TE LayerNorm vs nn.LayerNorm use different fused reduction orders) and the bf16 figure is at the arithmetic floor for a two-layer MLP. This pins the GELU ``approximate='none'`` fix (commit 8aace7b5e) against future regressions. Run with:: torchrun --nproc_per_node=1 \\ examples/multimodal_dev/tests/test_vision_patch_merger_parity.py --- .../tests/test_vision_patch_merger_parity.py | 187 ++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 examples/multimodal_dev/tests/test_vision_patch_merger_parity.py diff --git a/examples/multimodal_dev/tests/test_vision_patch_merger_parity.py b/examples/multimodal_dev/tests/test_vision_patch_merger_parity.py new file mode 100644 index 00000000000..62c908ea91d --- /dev/null +++ b/examples/multimodal_dev/tests/test_vision_patch_merger_parity.py @@ -0,0 +1,187 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Numerical parity for ``Qwen35VLPatchMerger`` vs HuggingFace reference. + +The HuggingFace reference (``Qwen3VLVisionPatchMerger`` in +``transformers/models/qwen3_vl/modeling_qwen3_vl.py``, branch +``use_postshuffle_norm=False``) is reproduced inline so this test does not +require ``transformers`` to be installed. The HF module is verbatim:: + + self.norm = nn.LayerNorm(hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(merge_dim, merge_dim) + self.act_fn = nn.GELU() # default approximate='none' + self.linear_fc2 = nn.Linear(merge_dim, out_hidden_size) + + x = self.norm(x) + x = x.view(-1, merge_dim) + x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + +With matching dims and weights copied across, the Megatron +implementation must agree: + + * fp32 forward: max-abs diff <= 1e-4 (TE LayerNorm vs nn.LayerNorm + have different fused reduction order; ~1e-5 absolute residual is + structural and not a real divergence) + * bf16 forward: max-abs diff <= 5e-2 (bf16 ceiling for two-layer MLP) + +Run with:: + + torchrun --nproc_per_node=1 \\ + examples/multimodal_dev/tests/test_vision_patch_merger_parity.py +""" + +import os +import sys + +import torch +import torch.distributed as dist +import torch.nn as nn + +_REPO_ROOT = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../.."), +) +if _REPO_ROOT in sys.path: + sys.path.remove(_REPO_ROOT) +sys.path.insert(0, _REPO_ROOT) + +from megatron.core import parallel_state as ps +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.transformer_config import TransformerConfig + +from examples.multimodal_dev.models.qwen35_vl.vision_encoder import Qwen35VLPatchMerger + + +# Match Qwen3.5-VL 9B / 397B-A17B vision tower dims. +HIDDEN_SIZE = 1152 +OUT_HIDDEN_SIZE = 3584 +SPATIAL_MERGE_SIZE = 2 +NUM_PATCHES = 64 # must be divisible by spatial_merge_size ** 2 + +ATOL_FP32 = 1e-4 +RTOL_FP32 = 1e-3 +ATOL_BF16 = 5e-2 +RTOL_BF16 = 5e-2 + + +class HFPatchMergerReference(nn.Module): + """Inline HF ``Qwen3VLVisionPatchMerger`` (use_postshuffle_norm=False).""" + + def __init__(self, hidden_size: int, out_hidden_size: int, spatial_merge_size: int): + super().__init__() + self.merge_dim = hidden_size * (spatial_merge_size ** 2) + self.norm = nn.LayerNorm(hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.merge_dim, self.merge_dim) + self.act_fn = nn.GELU() # approximate='none' by default + self.linear_fc2 = nn.Linear(self.merge_dim, out_hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + x = self.norm(hidden_states) + x = x.view(-1, self.merge_dim) + x = self.linear_fc1(x) + x = self.act_fn(x) + x = self.linear_fc2(x) + return x + + +def _init_distributed() -> int: + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + return local_rank + + +def _init_megatron_parallel() -> None: + ps.destroy_model_parallel() + ps.initialize_model_parallel(tensor_model_parallel_size=1) + model_parallel_cuda_manual_seed(42) + + +def _build_config(dtype: torch.dtype) -> TransformerConfig: + is_bf16 = dtype is torch.bfloat16 + return TransformerConfig( + num_layers=1, + hidden_size=HIDDEN_SIZE, + ffn_hidden_size=HIDDEN_SIZE, + num_attention_heads=8, + kv_channels=HIDDEN_SIZE // 8, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + sequence_parallel=False, + bf16=is_bf16, + params_dtype=dtype, + pipeline_dtype=dtype, + add_bias_linear=True, + gated_linear_unit=False, + normalization="LayerNorm", + layernorm_epsilon=1e-6, + attention_dropout=0.0, + hidden_dropout=0.0, + ) + + +def _copy_hf_to_megatron(hf: HFPatchMergerReference, mcore: Qwen35VLPatchMerger) -> None: + """TP=1: 1:1 parameter copy between HF nn.Module and the MCore module.""" + with torch.no_grad(): + mcore.patch_norm.weight.copy_(hf.norm.weight.to(mcore.patch_norm.weight.dtype)) + mcore.patch_norm.bias.copy_(hf.norm.bias.to(mcore.patch_norm.bias.dtype)) + mcore.linear_fc1.weight.copy_(hf.linear_fc1.weight.to(mcore.linear_fc1.weight.dtype)) + mcore.linear_fc1.bias.copy_(hf.linear_fc1.bias.to(mcore.linear_fc1.bias.dtype)) + mcore.linear_fc2.weight.copy_(hf.linear_fc2.weight.to(mcore.linear_fc2.weight.dtype)) + mcore.linear_fc2.bias.copy_(hf.linear_fc2.bias.to(mcore.linear_fc2.bias.dtype)) + + +def _run_one(dtype: torch.dtype, atol: float, rtol: float, device: torch.device, seed: int = 42) -> None: + torch.manual_seed(seed) + + hf_ref = HFPatchMergerReference( + hidden_size=HIDDEN_SIZE, + out_hidden_size=OUT_HIDDEN_SIZE, + spatial_merge_size=SPATIAL_MERGE_SIZE, + ).to(device=device, dtype=dtype).eval() + + config = _build_config(dtype) + mcore = Qwen35VLPatchMerger( + config=config, + hidden_size=HIDDEN_SIZE, + out_hidden_size=OUT_HIDDEN_SIZE, + spatial_merge_size=SPATIAL_MERGE_SIZE, + ).to(device=device, dtype=dtype).eval() + + _copy_hf_to_megatron(hf_ref, mcore) + + x = torch.randn(NUM_PATCHES, HIDDEN_SIZE, device=device, dtype=dtype) + + with torch.no_grad(): + y_hf = hf_ref(x) + y_mcore = mcore(x) + + assert y_hf.shape == y_mcore.shape, (y_hf.shape, y_mcore.shape) + diff = (y_hf - y_mcore).abs() + print( + f"[{dtype}] shape={tuple(y_hf.shape)} " + f"max_abs_diff={diff.max().item():.3e} " + f"mean_abs_diff={diff.mean().item():.3e} " + f"hf_norm={y_hf.float().norm().item():.4f} " + f"mcore_norm={y_mcore.float().norm().item():.4f}" + ) + torch.testing.assert_close(y_mcore, y_hf, atol=atol, rtol=rtol) + + +def main() -> None: + local_rank = _init_distributed() + _init_megatron_parallel() + device = torch.device(f"cuda:{local_rank}") + + _run_one(torch.float32, ATOL_FP32, RTOL_FP32, device) + _run_one(torch.bfloat16, ATOL_BF16, RTOL_BF16, device) + + if int(os.environ.get("RANK", 0)) == 0: + print( + "\nPASS: Qwen35VLPatchMerger logits match HF Qwen3VLVisionPatchMerger " + "in both fp32 and bf16." + ) + + +if __name__ == "__main__": + main()