Merge remote-tracking branch 'origin/main' into codex/pr5479-hybrid-model

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

docs/bridge-rl-integration.md

@@ -76,24 +76,46 @@ Translate your RL framework config into Megatron Bridge's `ConfigContainer` for
 
 ```python
 import torch
+from megatron.bridge import AutoBridge
 from megatron.bridge.training.config import (
     ConfigContainer,
     TrainingConfig,
     OptimizerConfig,
     SchedulerConfig,
     DistributedDataParallelConfig,
     CheckpointConfig,
+    LoggerConfig,
     TokenizerConfig,
 )
 from nemo_rl.models.policy import PolicyConfig  # or your own policy cfg type
 
 # Example: map your RL config to Megatron config
 def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> ConfigContainer:
-    model_cfg = rl_cfg["megatron_cfg"].copy()
+    bridge = AutoBridge.from_hf_pretrained(rl_cfg["model_name"])
+    model_cfg = bridge.to_megatron_provider(load_weights=False)
+
+    # Keep framework-owned optimizer/scheduler/DDP mappings out of the model
+    # provider. Apply only provider fields, with Bridge validating each name.
+    provider_fields = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "context_parallel_size",
+        "expert_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "sequence_parallel",
+        "recompute_granularity",
+        "recompute_method",
+        "recompute_num_layers",
+    )
+    model_overrides = {name: rl_cfg["megatron_cfg"][name] for name in provider_fields if name in rl_cfg["megatron_cfg"]}
+
     # Precision
-    dtype = rl_cfg["precision"]
-    model_cfg["bf16"] = dtype == "bfloat16"
-    model_cfg["fp16"] = dtype == "float16"
+    dtype = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+    }[rl_cfg["precision"]]
+    model_cfg.apply_overrides_and_finalize(dtype=dtype, overrides=model_overrides)
 
     checkpoint = CheckpointConfig(
         save_interval=100,
@@ -130,17 +152,19 @@ def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> Con
         tokenizer_model=rl_cfg["model_name"],
     )
 
-    return ConfigContainer(
+    cfg = ConfigContainer(
         model=model_cfg,
         checkpoint=checkpoint,
-        logger=None,
+        logger=LoggerConfig(),
         train=train,
         optimizer=opt,
         ddp=ddp,
         scheduler=sch,
         dataset=None,
         tokenizer=tokenizer,
     )
+    cfg.validate()
+    return cfg
 ```
 
 Initialize Megatron-Core using a helper similar to `setup_megatron_model` from NeMo-RL:

docs/fern/versions/0.4.2/pages/bridge-rl-integration.mdx

@@ -72,24 +72,46 @@ Translate your RL framework config into Megatron Bridge's `ConfigContainer` for
 
 ```python
 import torch
+from megatron.bridge import AutoBridge
 from megatron.bridge.training.config import (
     ConfigContainer,
     TrainingConfig,
     OptimizerConfig,
     SchedulerConfig,
     DistributedDataParallelConfig,
     CheckpointConfig,
+    LoggerConfig,
     TokenizerConfig,
 )
 from nemo_rl.models.policy import PolicyConfig  # or your own policy cfg type
 
 # Example: map your RL config to Megatron config
 def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> ConfigContainer:
-    model_cfg = rl_cfg["megatron_cfg"].copy()
+    bridge = AutoBridge.from_hf_pretrained(rl_cfg["model_name"])
+    model_cfg = bridge.to_megatron_provider(load_weights=False)
+
+    # Keep framework-owned optimizer/scheduler/DDP mappings out of the model
+    # provider. Apply only provider fields, with Bridge validating each name.
+    provider_fields = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "context_parallel_size",
+        "expert_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "sequence_parallel",
+        "recompute_granularity",
+        "recompute_method",
+        "recompute_num_layers",
+    )
+    model_overrides = {name: rl_cfg["megatron_cfg"][name] for name in provider_fields if name in rl_cfg["megatron_cfg"]}
+
     # Precision
-    dtype = rl_cfg["precision"]
-    model_cfg["bf16"] = dtype == "bfloat16"
-    model_cfg["fp16"] = dtype == "float16"
+    dtype = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+    }[rl_cfg["precision"]]
+    model_cfg.apply_overrides_and_finalize(dtype=dtype, overrides=model_overrides)
 
     checkpoint = CheckpointConfig(
         save_interval=100,
@@ -126,17 +148,19 @@ def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> Con
         tokenizer_model=rl_cfg["model_name"],
     )
 
-    return ConfigContainer(
+    cfg = ConfigContainer(
         model=model_cfg,
         checkpoint=checkpoint,
-        logger=None,
+        logger=LoggerConfig(),
         train=train,
         optimizer=opt,
         ddp=ddp,
         scheduler=sch,
         dataset=None,
         tokenizer=tokenizer,
     )
+    cfg.validate()
+    return cfg
 ```
 
 Initialize Megatron-Core using a helper similar to `setup_megatron_model` from NeMo-RL:

docs/fern/versions/nightly/pages/bridge-rl-integration.mdx

@@ -72,24 +72,46 @@ Translate your RL framework config into Megatron Bridge's `ConfigContainer` for
 
 ```python
 import torch
+from megatron.bridge import AutoBridge
 from megatron.bridge.training.config import (
     ConfigContainer,
     TrainingConfig,
     OptimizerConfig,
     SchedulerConfig,
     DistributedDataParallelConfig,
     CheckpointConfig,
+    LoggerConfig,
     TokenizerConfig,
 )
 from nemo_rl.models.policy import PolicyConfig  # or your own policy cfg type
 
 # Example: map your RL config to Megatron config
 def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> ConfigContainer:
-    model_cfg = rl_cfg["megatron_cfg"].copy()
+    bridge = AutoBridge.from_hf_pretrained(rl_cfg["model_name"])
+    model_cfg = bridge.to_megatron_provider(load_weights=False)
+
+    # Keep framework-owned optimizer/scheduler/DDP mappings out of the model
+    # provider. Apply only provider fields, with Bridge validating each name.
+    provider_fields = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "context_parallel_size",
+        "expert_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "sequence_parallel",
+        "recompute_granularity",
+        "recompute_method",
+        "recompute_num_layers",
+    )
+    model_overrides = {name: rl_cfg["megatron_cfg"][name] for name in provider_fields if name in rl_cfg["megatron_cfg"]}
+
     # Precision
-    dtype = rl_cfg["precision"]
-    model_cfg["bf16"] = dtype == "bfloat16"
-    model_cfg["fp16"] = dtype == "float16"
+    dtype = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+    }[rl_cfg["precision"]]
+    model_cfg.apply_overrides_and_finalize(dtype=dtype, overrides=model_overrides)
 
     checkpoint = CheckpointConfig(
         save_interval=100,
@@ -126,17 +148,19 @@ def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> Con
         tokenizer_model=rl_cfg["model_name"],
     )
 
-    return ConfigContainer(
+    cfg = ConfigContainer(
         model=model_cfg,
         checkpoint=checkpoint,
-        logger=None,
+        logger=LoggerConfig(),
         train=train,
         optimizer=opt,
         ddp=ddp,
         scheduler=sch,
         dataset=None,
         tokenizer=tokenizer,
     )
+    cfg.validate()
+    return cfg
 ```
 
 Initialize Megatron-Core using a helper similar to `setup_megatron_model` from NeMo-RL:

examples/model_verification_cards/nemotron-3-super-120b-a12b/card.yaml

@@ -4,12 +4,13 @@
 title: nemotron_3_super_120b_a12b
 summary: >
   Performance scope: pretrain_performance.H100 and pretrain_performance.GB200
-  track tuned canonical 64-GPU performance recipes. Timing and throughput from
+  and pretrain_performance.GB300 track tuned canonical 64-GPU performance
+  recipes, with NVFP4 on GB300. Timing and throughput from
   functional H100 and GB200 training are support-verification sanity checks
   rather than optimized performance results. Hardware evidence is scoped
   strictly by accelerator. Nemotron 3 Super 120B-A12B verification covers
   conversion, inference, bounded real-data training, checkpoint resume, and
-  canonical H100 and GB200 benchmarks.
+  canonical H100, GB200, and GB300 benchmarks.
 verification_index:
   model_level:
     verified:
@@ -30,6 +31,7 @@ verification_index:
   performance:
     H100: verified
     GB200: verified
+    GB300: verified
 model:
   hf_id: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16
   hf_revision: d51eab0d1f979ebc26b546e634a04f450d99158e  # pragma: allowlist secret
@@ -702,4 +704,30 @@ items:
         11000 ms and at least 500 TFLOP/s/GPU, and the resolved configuration
         is persisted. Mock data and forced routing make the loss a finiteness
         check rather than convergence evidence. These numbers are GB200-only
-        evidence and make no GB300 performance claim.
+        evidence.
+
+    GB300:
+      status: verified
+      precision: nvfp4
+      bridge_commit: 0480586879f2513958fd8634fc529693ae13e536  # pragma: allowlist secret
+      command: >
+        ./scripts/training/train.sh --wait --nodes 16 --gpus-per-node 4
+        --recipe nemotron_3_super_pretrain_64gpu_gb300_nvfp4_config
+        --mode pretrain --max_steps 50 --seq_length 8192
+        logger.save_config_filepath=work/model-verification/nemotron-3-super-120b-a12b/gb300-performance/ConfigContainer.yaml
+      last_verified: 2026-08-18
+      metrics:
+        initial_loss: 12.17656
+        final_loss: 0.01398012
+        last_10_steps_step_time_ms_avg: 6357.940
+        last_10_steps_model_tflops_per_gpu_avg: 873.330
+        last_10_steps_tokens_per_second_per_gpu_avg: 10307.741
+      expected_result: >
+        On exactly 64 GB300s, the canonical NVFP4 mock-data recipe completes
+        exactly 50 optimizer steps at TP1/PP1/CP1/EP64/ETP1, GBS/MBS 512/1,
+        and sequence length 8192. All 50 keyed rows have finite loss with zero
+        skipped or NaN iterations. Loss moves from 12.17656 to 0.01398012; the
+        final ten steps average 6357.940 ms, 873.330 TFLOP/s/GPU, and
+        10307.741 tokens/s/GPU. The resolved configuration persists. Mock data
+        and forced routing make the loss a finiteness check rather than
+        convergence evidence.

src/megatron/bridge/data/sft_processing.py

@@ -142,18 +142,21 @@ class TokenizedPromptCompletion:
 _CONVERSATION_KEYS = ("messages", "conversation", "conversations")
 _CANONICAL_PROMPT_KEY = "prompt"
 _CANONICAL_COMPLETION_KEY = "completion"
+_PLURAL_MEDIA_KEYS = (
+    "images",
+    "image_paths",
+    "videos",
+    "video_paths",
+    "audio_paths",
+)
 _MEDIA_KEYS = (
     "image",
-    "images",
     "image_path",
-    "image_paths",
     "video",
-    "videos",
     "video_path",
-    "video_paths",
     "audio",
     "audio_path",
-    "audio_paths",
+    *_PLURAL_MEDIA_KEYS,
 )
 
 
@@ -275,7 +278,10 @@ def is_text_only_prompt_completion_example(
     preprocessing: PromptCompletionSFTPreprocessingConfig,
 ) -> bool:
     """Return whether a row is a text-only prompt-completion example."""
-    if any(example.get(key) is not None for key in _MEDIA_KEYS):
+    for key in _MEDIA_KEYS:
+        value = example.get(key)
+        if value is None or (key in _PLURAL_MEDIA_KEYS and isinstance(value, list) and not value):
+            continue
         return False
     try:
         normalize_sft_example(example, preprocessing)

src/megatron/bridge/models/conversion/param_mapping.py

@@ -815,7 +815,7 @@ def gather_from_ep_ranks(
             else:
                 weights_dict[param_name] = gathered_weights[i].unsqueeze(0)
         for param_name in weights_dict:
-            weights_dict[param_name] = weights_dict[param_name].squeeze()
+            weights_dict[param_name] = weights_dict[param_name].squeeze(0)
         return weights_dict
 
     def gather_from_ep_ranks_scale(

tests/unit_tests/data/builders/test_direct_hf_sft.py

@@ -3,6 +3,7 @@
 import importlib
 
 import pytest
+from datasets import Dataset
 from megatron.training.config.instantiate_utils import instantiate
 
 from megatron.bridge.data.base import DatasetBuildContext
@@ -576,6 +577,29 @@ def test_builder_uses_prompt_completion_without_chat_template(monkeypatch):
     assert batch["loss_mask"].sum().item() == 2
 
 
+def test_builder_treats_empty_optional_media_as_text_only(monkeypatch):
+    row = Dataset.from_list([{"prompt": "Q", "completion": "A", "images": []}])[0]
+    monkeypatch.setattr(builder_module, "load_and_adapt_hf_dataset", lambda source: [row])
+    config = DirectHFSFTDatasetConfig(
+        seq_length=16,
+        source=HFDatasetSourceConfig(path_or_dataset="org/paired"),
+        preprocessing=PromptCompletionSFTPreprocessingConfig(),
+        pad_to_multiple_of=1,
+        do_validation=False,
+        do_test=False,
+    )
+
+    train, _, _ = DirectHFSFTDatasetBuilder(config).build(DatasetBuildContext(1, 0, 0, tokenizer=_Tokenizer()))
+
+    assert train is not None
+    assert train[0]["images"] == []
+    with pytest.raises(ValueError, match="supports text-only examples"):
+        select_direct_hf_sft_collate(
+            [{"prompt": "Q", "completion": "A", "images": [object()]}],
+            config.preprocessing,
+        )
+
+
 def test_direct_hf_sft_config_resolves_canonical_builder(monkeypatch):
     from megatron.bridge.data import utils as data_utils