feat(recipe): add Qwen3.5 VL long-context GB200 SFT recipes (#6324)

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/recipes/qwen_vl/__init__.py

@@ -17,6 +17,8 @@
     qwen35_vl_27b_pretrain_16gpu_gb200_bf16_mock_config,
     qwen35_vl_35b_a3b_peft_8gpu_gb200_bf16_functional_config,
     qwen35_vl_35b_a3b_sft_8gpu_gb200_bf16_functional_config,
+    qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config,
+    qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config,
 )
 
 # Qwen3.5/Qwen3.6-VL GB300 models
@@ -97,6 +99,8 @@
     # Qwen3.5-VL SFT configs — MoE
     "qwen35_vl_35b_a3b_sft_config",
     "qwen35_vl_35b_a3b_sft_8gpu_gb200_bf16_functional_config",
+    "qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config",
+    "qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config",
     "qwen35_vl_35b_a3b_fsdp_sft_config",
     "qwen35_vl_122b_a10b_sft_config",
     "qwen35_vl_397b_a17b_sft_config",

src/megatron/bridge/recipes/qwen_vl/gb200/__init__.py

@@ -16,11 +16,15 @@
     qwen35_vl_27b_pretrain_16gpu_gb200_bf16_mock_config,
     qwen35_vl_35b_a3b_peft_8gpu_gb200_bf16_functional_config,
     qwen35_vl_35b_a3b_sft_8gpu_gb200_bf16_functional_config,
+    qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config,
+    qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config,
 )
 
 
 __all__ = [
     "qwen35_vl_27b_pretrain_16gpu_gb200_bf16_mock_config",
     "qwen35_vl_35b_a3b_peft_8gpu_gb200_bf16_functional_config",
     "qwen35_vl_35b_a3b_sft_8gpu_gb200_bf16_functional_config",
+    "qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config",
+    "qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config",
 ]

src/megatron/bridge/recipes/qwen_vl/gb200/qwen35_vl.py

@@ -16,6 +16,9 @@
 
 from __future__ import annotations
 
+from megatron.bridge import AutoBridge
+from megatron.bridge.data.builders import EnergonDatasetConfig, QwenVLEnergonTaskEncoderConfig
+from megatron.bridge.recipes.common import _sft_common_vlm
 from megatron.bridge.recipes.qwen_vl.h100.qwen35_vl import (
     qwen35_vl_27b_pretrain_16gpu_h100_bf16_mock_config,
     qwen35_vl_35b_a3b_peft_4gpu_h100_bf16_config,
@@ -24,7 +27,7 @@
 from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS
 from megatron.bridge.training.comm_overlap import CommOverlapConfig
 from megatron.bridge.training.config import ConfigContainer
-from megatron.bridge.training.mixed_precision import get_mixed_precision_config
+from megatron.bridge.training.mixed_precision import bf16_mixed, bf16_with_mxfp8_mixed, get_mixed_precision_config
 from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules
 
 
@@ -212,6 +215,132 @@ def qwen35_vl_35b_a3b_sft_8gpu_gb200_bf16_functional_config() -> ConfigContainer
     return cfg
 
 
+_QWEN35_35B_A3B_REVISION = "59d61f3ce65a6d9863b86d2e96597125219dc754"  # pragma: allowlist secret
+
+
+def _qwen35_vl_35b_a3b_long_context_common() -> ConfigContainer:
+    """Build the shared 32-GB200 long-context SFT configuration."""
+    cfg = _sft_common_vlm()
+    hf_path = "Qwen/Qwen3.5-35B-A3B"
+    cfg.model = AutoBridge.from_hf_pretrained(hf_path, revision=_QWEN35_35B_A3B_REVISION).to_megatron_provider(
+        load_weights=False
+    )
+    cfg.model.calculate_per_token_loss = True
+    cfg.ddp.average_in_collective = False
+
+    cfg.model.mtp_num_layers = 1
+    cfg.model.bias_activation_fusion = True
+    cfg.model.cuda_graph_scope = []
+    cfg.model.moe_router_fusion = True
+    cfg.model.offload_modules = None
+    cfg.model.seq_length = 131072
+    cfg.model.tensor_model_parallel_size = 2
+    cfg.model.context_parallel_size = 8
+    cfg.model.expert_model_parallel_size = 16
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.sequence_parallel = True
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.moe_flex_dispatcher_num_sms = 32
+    cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = True
+
+    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.model.recompute_granularity = "selective"
+    cfg.model.recompute_modules = ["gdn_norm_out", "moe"]
+    cfg.model.vision_recompute_granularity = "full"
+    cfg.model.vision_recompute_method = "uniform"
+    cfg.model.vision_recompute_num_layers = 1
+
+    cfg.train.micro_batch_size = 1
+    # Nemotron Image Training v3, clevr_2 subset, prepared as Energon shards.
+    # Native packing combines samples into one physical 128K microbatch.
+    cfg.dataset = EnergonDatasetConfig(
+        seq_length=cfg.model.seq_length,
+        micro_batch_size=cfg.train.micro_batch_size,
+        num_workers=1,
+        shuffle_buffer_size=2,
+        packing_buffer_size=8,
+        dataset_kwargs={"image_decode_spec": "pilrgb"},
+        task_encoder=QwenVLEnergonTaskEncoderConfig(
+            hf_processor_path=hf_path,
+            hf_processor_revision=_QWEN35_35B_A3B_REVISION,
+            max_pixels=401408,
+            max_num_images=4,
+            max_visual_tokens=2048,
+        ),
+    )
+
+    cfg.optimizer.lr = 2e-5
+    cfg.optimizer.min_lr = 2e-6
+    cfg.scheduler.lr_warmup_iters = 200
+    cfg.scheduler.lr_decay_iters = 300000
+    cfg.scheduler.lr_wsd_decay_iters = 300000
+    cfg.comm_overlap = CommOverlapConfig(
+        tp_comm_overlap=False,
+        overlap_grad_reduce=False,
+        overlap_param_gather=False,
+        overlap_param_gather_with_optimizer_step=False,
+        overlap_moe_expert_parallel_comm=False,
+        delay_wgrad_compute=False,
+    )
+    cfg.env_vars = {
+        **COMMON_RECIPE_ENV_VARS,
+        "CUDA_DEVICE_MAX_CONNECTIONS": 1,
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
+        # Both language and vision towers use automatic TE backend selection.
+        "NVTE_FUSED_ATTN": 1,
+        "NVTE_FLASH_ATTN": 1,
+        "NVTE_UNFUSED_ATTN": 1,
+        "NCCL_GRAPH_REGISTER": 0,
+        "NCCL_NVLS_ENABLE": 0,
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 16,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
+        "USE_MNNVL": 1,
+    }
+    return cfg
+
+
+def qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config() -> ConfigContainer:
+    """Return 128K BF16 Qwen3.5-VL SFT with TP2/CP8/EP16 on 32 GB200 GPUs.
+
+    Override with ``"dataset.path=$CLEVR2_ENERGON"`` to select the prepared
+    Nemotron Image Training v3 ``clevr_2`` Energon shards, and supply a
+    pretrained checkpoint.
+
+    Prepare the shards with `prepare_nemotron_image_v3.py
+    <https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/tutorials/data/energon/prepare_nemotron_image_v3.py>`_
+    using ``--subsets clevr_2 --skip-source-integrity-check`` after verifying
+    the source revision; the built-in integrity manifest covers only ``turing``.
+    Native Energon packing uses one physical microbatch per 128K pack.
+    """
+    cfg = _qwen35_vl_35b_a3b_long_context_common()
+    cfg.mixed_precision = bf16_mixed()
+    return cfg
+
+
+def qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config() -> ConfigContainer:
+    """Return packed CLEVR2 128K MXFP8 SFT with BF16 parameter communication.
+
+    Override with ``"dataset.path=$CLEVR2_ENERGON"`` to select the prepared
+    Nemotron Image Training v3 ``clevr_2`` Energon shards, and supply a
+    pretrained checkpoint.
+
+    Prepare the shards with `prepare_nemotron_image_v3.py
+    <https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/tutorials/data/energon/prepare_nemotron_image_v3.py>`_
+    using ``--subsets clevr_2 --skip-source-integrity-check`` after verifying
+    the source revision; the built-in integrity manifest covers only ``turing``.
+    Native Energon packing uses one physical microbatch per 128K pack.
+    """
+    cfg = _qwen35_vl_35b_a3b_long_context_common()
+    cfg.mixed_precision = bf16_with_mxfp8_mixed()
+    cfg.mixed_precision.fp8_param_gather = False
+    cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False
+    return cfg
+
+
 def qwen35_vl_35b_a3b_peft_8gpu_gb200_bf16_functional_config() -> ConfigContainer:
     """Return shared Qwen3.5/Qwen3.6-VL 35B-A3B LoRA for eight GB200 GPUs."""
     cfg = qwen35_vl_35b_a3b_peft_4gpu_h100_bf16_config()
@@ -223,4 +352,6 @@ def qwen35_vl_35b_a3b_peft_8gpu_gb200_bf16_functional_config() -> ConfigContaine
     "qwen35_vl_27b_pretrain_16gpu_gb200_bf16_mock_config",
     "qwen35_vl_35b_a3b_peft_8gpu_gb200_bf16_functional_config",
     "qwen35_vl_35b_a3b_sft_8gpu_gb200_bf16_functional_config",
+    "qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config",
+    "qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config",
 ]

tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py

@@ -22,6 +22,7 @@
 
 import importlib
 import inspect
+from types import SimpleNamespace
 from typing import Callable
 
 import pytest
@@ -108,6 +109,8 @@
 _QWEN35_VL_GB200_FUNCS = [
     _qwen35_vl_gb200_module.qwen35_vl_27b_pretrain_16gpu_gb200_bf16_mock_config,
     _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_8gpu_gb200_bf16_functional_config,
+    _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config,
+    _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config,
     _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_peft_8gpu_gb200_bf16_functional_config,
 ]
 
@@ -148,7 +151,7 @@ class _FakeAutoBridge:
     """Fake AutoBridge for testing."""
 
     @staticmethod
-    def from_hf_pretrained(hf_path: str):
+    def from_hf_pretrained(hf_path: str, *, revision: str | None = None):
         return _FakeAutoBridge()
 
     def to_megatron_provider(self, load_weights: bool = False):
@@ -441,7 +444,7 @@ def test_qwen35_vl_27b_peft_lora_defaults(monkeypatch: pytest.MonkeyPatch):
 
 def test_qwen35_vl_397b_a17b_pretrain_64gpu_gb300_defaults(monkeypatch: pytest.MonkeyPatch):
     """The 64-GB300 library pretrain recipe should own the measured execution policy."""
-    patch_recipe_module_global(monkeypatch, _qwen35_vl_gb300_module, "AutoBridge", _FakeAutoBridge)
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_h100_module, "AutoBridge", _FakeAutoBridge)
 
     cfg = _qwen35_vl_gb300_module.qwen35_vl_397b_a17b_pretrain_config()
 
@@ -472,7 +475,7 @@ def test_qwen35_vl_397b_a17b_pretrain_64gpu_gb300_defaults(monkeypatch: pytest.M
 
 def test_qwen35_vl_35b_a3b_pretrain_16gpu_gb300_defaults(monkeypatch: pytest.MonkeyPatch):
     """The 16-GB300 library pretrain recipe should own the measured execution policy."""
-    patch_recipe_module_global(monkeypatch, _qwen35_vl_gb300_module, "AutoBridge", _FakeAutoBridge)
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_h100_module, "AutoBridge", _FakeAutoBridge)
 
     cfg = _qwen35_vl_gb300_module.qwen35_vl_35b_a3b_pretrain_16gpu_gb300_bf16_config()
 
@@ -659,6 +662,83 @@ def test_qwen35_vl_35b_a3b_long_context_sft_defaults(monkeypatch: pytest.MonkeyP
     assert cfg.ddp.average_in_collective is False
 
 
+def test_qwen35_vl_35b_a3b_gb200_long_context_precision_pair(monkeypatch: pytest.MonkeyPatch):
+    """The GB200 BF16 and MXFP8 recipes should share one 128K execution topology."""
+    from megatron.bridge.data.builders import EnergonDatasetConfig, QwenVLEnergonTaskEncoderConfig
+    from megatron.bridge.models.qwen_vl.qwen35_vl_provider import Qwen35VLMoEModelProvider
+
+    monkeypatch.setattr(
+        _FakeAutoBridge,
+        "to_megatron_provider",
+        lambda self, load_weights=False: Qwen35VLMoEModelProvider(
+            num_layers=40, hidden_size=2048, num_attention_heads=16, bias_activation_fusion=False
+        ),
+    )
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_gb200_module, "AutoBridge", _FakeAutoBridge)
+
+    bf16_cfg = _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config()
+    fp8mx_cfg = _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config()
+
+    for cfg in (bf16_cfg, fp8mx_cfg):
+        _assert_basic_config(cfg)
+        assert cfg.model.seq_length == 131072
+        assert cfg.model.bias_activation_fusion is True
+        assert cfg.model.tensor_model_parallel_size == 2
+        assert cfg.model.pipeline_model_parallel_size == 1
+        assert cfg.model.pipeline_dtype is None
+        assert cfg.model.virtual_pipeline_model_parallel_size is None
+        assert cfg.model.context_parallel_size == 8
+        assert cfg.model.expert_model_parallel_size == 16
+        assert cfg.model.expert_tensor_parallel_size == 1
+        assert cfg.model.sequence_parallel is True
+        assert cfg.model.moe_token_dispatcher_type == "flex"
+        assert cfg.model.moe_flex_dispatcher_backend == "hybridep"
+        assert cfg.model.moe_flex_dispatcher_num_sms == 32
+        assert cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is True
+        assert cfg.model.gdn_pre_gated_delta_rule_fusion is False
+        assert cfg.model.cross_entropy_fusion_impl == "te"
+        assert cfg.model.recompute_granularity == "selective"
+        assert cfg.model.recompute_modules == ["gdn_norm_out", "moe"]
+        assert cfg.model.recompute_method is None
+        assert cfg.model.recompute_num_layers is None
+        assert cfg.model.vision_recompute_granularity == "full"
+        assert cfg.model.vision_recompute_method == "uniform"
+        assert cfg.model.vision_recompute_num_layers == 1
+        assert cfg.model.vision_recompute_modules is None
+        assert cfg.model.attention_backend.name == "auto"
+        for variable in ("NVTE_FUSED_ATTN", "NVTE_FLASH_ATTN", "NVTE_UNFUSED_ATTN"):
+            assert cfg.env_vars[variable] == 1
+        assert cfg.train.global_batch_size == 32
+        assert cfg.train.micro_batch_size == 1
+        assert cfg.dataset.seq_length == 131072
+        assert cfg.dataset.enable_in_batch_packing is False
+        assert cfg.dataset.defer_in_batch_packing_to_step is False
+        assert isinstance(cfg.dataset, EnergonDatasetConfig)
+        assert cfg.dataset.path is None  # Caller supplies prepared CLEVR2 shards.
+        assert cfg.dataset.micro_batch_size == 1
+        assert cfg.dataset.num_workers == 1
+        assert cfg.dataset.shuffle_buffer_size == 2
+        assert cfg.dataset.packing_buffer_size == 8
+        assert cfg.dataset.dataset_kwargs == {"image_decode_spec": "pilrgb"}
+        encoder = cfg.dataset.task_encoder
+        assert isinstance(encoder, QwenVLEnergonTaskEncoderConfig)
+        assert encoder.hf_processor_path == "Qwen/Qwen3.5-35B-A3B"
+        assert encoder.hf_processor_revision == _qwen35_vl_gb200_module._QWEN35_35B_A3B_REVISION
+        assert encoder.min_pixels == 200704
+        assert encoder.max_pixels == 401408
+        assert encoder.max_num_images == 4
+        assert encoder.max_visual_tokens == 2048
+        assert cfg.mixed_precision.grad_reduce_in_fp32 is True
+        assert cfg.ddp.grad_reduce_in_fp32 is True
+        assert cfg.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] == 1
+        assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 16
+
+    assert bf16_cfg.mixed_precision.fp8 is None
+    assert fp8mx_cfg.mixed_precision.fp8_recipe == "mxfp8"
+    assert fp8mx_cfg.mixed_precision.fp8_param_gather is False
+    assert fp8mx_cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False
+
+
 def test_qwen35_vl_35b_a3b_fsdp_sft_defaults(monkeypatch: pytest.MonkeyPatch):
     """35B-A3B FSDP SFT should have FSDP-specific parallelism and settings."""
     patch_recipe_module_global(monkeypatch, _qwen35_vl_module, "AutoBridge", _FakeAutoBridge)
@@ -1453,3 +1533,55 @@ def test_qwen35_vl_gdn_conv_fusion_skipped_on_older_core(monkeypatch: pytest.Mon
     cfg = _qwen35_vl_module.qwen35_vl_9b_pretrain_mock_config()
 
     assert not hasattr(cfg.model, "gdn_pre_gated_delta_rule_fusion")
+
+
+@pytest.mark.parametrize(
+    "recipe_func",
+    [
+        _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config,
+        _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_fp8mx_config,
+    ],
+)
+def test_gb200_long_context_validates_with_real_provider(monkeypatch, recipe_func):
+    from megatron.bridge.models.qwen_vl.qwen35_vl_provider import Qwen35VLMoEModelProvider
+
+    monkeypatch.setattr(
+        torch.cuda, "get_device_properties", lambda index: SimpleNamespace(major=10, name="NVIDIA GB200")
+    )
+
+    config_module = importlib.import_module("megatron.bridge.training.config")
+    monkeypatch.setattr(config_module, "get_world_size_safe", lambda: 32)
+    monkeypatch.setattr(
+        _FakeAutoBridge,
+        "to_megatron_provider",
+        lambda self, load_weights=False: Qwen35VLMoEModelProvider(
+            num_layers=40, hidden_size=2048, num_attention_heads=16
+        ),
+    )
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_gb200_module, "AutoBridge", _FakeAutoBridge)
+    config = recipe_func()
+
+    config.dataset.path = "prepared-clevr2-energon"
+    config.dataset.validate()
+    config.validate()
+    assert config.dataset.pad_to_max_length is True  # Derived for HybridEP.
+
+    assert config.train.global_batch_size % (2 * config.train.micro_batch_size) == 0
+    config.dataset.enable_in_batch_packing = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        config.dataset.validate()
+
+
+def test_gb200_long_context_pins_model_and_processor_together(monkeypatch):
+    calls = []
+
+    def from_hf_pretrained(hf_path, **kwargs):
+        calls.append((hf_path, kwargs))
+        return _FakeAutoBridge()
+
+    monkeypatch.setattr(_FakeAutoBridge, "from_hf_pretrained", from_hf_pretrained)
+    patch_recipe_module_global(monkeypatch, _qwen35_vl_gb200_module, "AutoBridge", _FakeAutoBridge)
+    recipe = _qwen35_vl_gb200_module.qwen35_vl_35b_a3b_sft_long_context_32gpu_gb200_bf16_config
+    cfg = recipe()
+    encoder = cfg.dataset.task_encoder
+    assert calls == [(encoder.hf_processor_path, {"revision": encoder.hf_processor_revision})]