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

Author: kamran-nvidia

.main.commit

@@ -1 +1 @@
-b2a8ec7458e90216e3a8d147d8ddf81cfb895f7c
+38986a98aae6a0cc4c8ae7b435db3288a890b0cb

3rdparty/Megatron-LM

@@ -0,0 +1,34 @@
+# DeepSeek V4
+
+[DeepSeek-V4](https://github.com/deepseek-ai/DeepSeek-V4) is the next-generation Mixture-of-Experts language model from DeepSeek-AI. It extends the V3 design with **Hyper-Connections (mHC)** for multi-stream residual mixing, **Compressed Sparse Attention (CSA)** with a learned token-importance indexer (DSA), **hash-routed MoE layers** for the first few decoder blocks, and a refined **Multi-Token Prediction (MTP)** head with separate `e_proj` / `h_proj` projections.
+
+DeepSeek V4 models are supported via the Bridge system with auto-detected configuration and weight mapping.
+
+## Model Architecture Features
+
+- **Hybrid Attention (DSv4HybridSelfAttention)**: Per-layer mix of dense MLA and Compressed Sparse Attention selected by `compress_ratios`
+- **Compressed Sparse Attention (CSA)** with **DSA Indexer**: Top-k token selection over windowed keys; `index_n_heads`, `index_head_dim`, `index_topk` control the indexer
+- **Hyper-Connections (mHC)**: 4-stream residual mixing per layer (`hc_mult = 4`) with sinkhorn-iterated attention; per-MTP-layer `hc_head_*` learns output contraction
+- **Hash-Routed MoE**: First few decoder layers use a deterministic vocab → expert mapping (`tid2eid`) instead of softmax routing
+- **Multi-Token Prediction (MTP)**: One MTP layer with separate `e_proj` and `h_proj` projections (post-MCore #4518)
+- **YaRN RoPE**: `rotary_scaling_factor=16`, `original_max_position_embeddings=65536`; `mscale=mscale_all_dim=1.0` for V4
+- **Sigmoid Gating with Expert Bias**: `noaux_tc` load balancing, `sqrtsoftplus` scoring, expert bias enabled
+- **`o_groups` Output Projection**: `o_lora_rank` low-rank output projection split into `o_groups` parallel groups
+
+## Examples, Parallelism, and Limitations
+
+For checkpoint conversion and inference scripts, recommended parallelism settings, and current known limitations, see the [DeepSeek V4 examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/deepseek_v4).
+
+## Hugging Face Model Cards & References
+
+### Hugging Face Model Cards
+- DeepSeek-V4-Flash: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash
+- DeepSeek-V4-Flash-Base: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Base
+- DeepSeek-V4-Pro: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro
+- DeepSeek-V4-Pro-Base: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro-Base
+
+### Additional Resources
+- GitHub Repository: https://github.com/deepseek-ai/DeepSeek-V4
+
+## Related Docs
+- DeepSeek V4 examples: [examples/models/deepseek_v4/README.md](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/deepseek_v4)

docs/models/llm/deepseek-v4.md

@@ -7,6 +7,7 @@ This section documents Large Language Models supported by Megatron Bridge, with
 
 deepseek-v2.md
 deepseek-v3.md
+deepseek-v4.md
 gemma2.md
 gemma3.md
 glm45.md

docs/models/llm/index.md

@@ -91,5 +91,5 @@ A LoRA checkpoint only contains the learnable adapter weights. In order to conve
 uv run python examples/peft/merge_lora.py \
 --hf-model-path $HF_MODEL_ID \
 --lora-checkpoint /path/to/lora/ckpt/iter_xxxxxxx 
---output /path/to/merged/ckpt
+--output /path/to/merged/hf/ckpt
 ```

docs/models/llm/nemotron3-nano.md

@@ -149,10 +149,9 @@ format for downstream evaluation, it is necessary to merge the LoRA adapters bac
 uv run python examples/peft/merge_lora.py \
 --hf-model-path $HF_MODEL_PATH \
 --lora-checkpoint <trained LoRA checkpoint>/iter_N \
---output <LoRA checkpoint merged>
+--output <merged HF checkpoint>
 ```
-You can now run in-framework inference with `hf_to_megatron_generate_vlm.py` by supplying the merged LoRA checkpoint.
-You can also export the merged LoRA checkpoint to Hugging Face format.
+The output is a merged Hugging Face checkpoint for downstream evaluation.
 
 
 ## Example Datasets

docs/models/vlm/nemotron-nano-v2-vl.md

@@ -53,6 +53,7 @@
 """
 
 import argparse
+import datetime
 import os
 import sys
 
@@ -84,6 +85,20 @@ def _check_distributed():
         sys.exit(1)
 
 
+def _ensure_distributed_initialized(timeout_minutes: int | None):
+    _check_distributed()
+    if timeout_minutes is None:
+        return
+    if torch.distributed.is_initialized():
+        return
+
+    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0")))
+    torch.distributed.init_process_group(
+        "nccl",
+        timeout=datetime.timedelta(minutes=timeout_minutes),
+    )
+
+
 @torchrun_main
 def import_hf_to_megatron(
     hf_model: str,
@@ -94,9 +109,10 @@ def import_hf_to_megatron(
     etp: int = 1,
     torch_dtype: str = "bfloat16",
     trust_remote_code: bool = False,
+    distributed_timeout_minutes: int | None = None,
 ) -> None:
     """Import a HuggingFace model and save it as a distributed Megatron checkpoint."""
-    _check_distributed()
+    _ensure_distributed_initialized(distributed_timeout_minutes)
     dtype = _parse_dtype(torch_dtype)
 
     print_rank_0(f"Importing: {hf_model} -> {megatron_path}")
@@ -115,6 +131,15 @@ def import_hf_to_megatron(
     model_provider.expert_tensor_parallel_size = etp
     model_provider.pipeline_dtype = dtype
     model_provider.params_dtype = dtype
+    # Auto-generate pipeline layout for models that need it (e.g. DSv4 hash MoE)
+    if pp > 1 and hasattr(bridge._model_bridge, "generate_pipeline_layout"):
+        hf_config = bridge.hf_pretrained.config
+        num_layers = hf_config.num_hidden_layers
+        mtp = getattr(hf_config, "num_nextn_predict_layers", 0) or 0
+        model_provider.pipeline_model_parallel_layout = bridge._model_bridge.generate_pipeline_layout(
+            num_layers, pp, mtp
+        )
+        print_rank_0(f"  Auto-generated pipeline layout for PP={pp} ({num_layers} layers, {mtp} MTP)")
     model_provider.finalize()
     model_provider.initialize_model_parallel(seed=0)
 
@@ -151,9 +176,10 @@ def export_megatron_to_hf(
     show_progress: bool = True,
     distributed_save: bool = False,
     save_every_n_ranks: int = 1,
+    distributed_timeout_minutes: int | None = None,
 ) -> None:
     """Export a distributed Megatron checkpoint to HuggingFace format."""
-    _check_distributed()
+    _ensure_distributed_initialized(distributed_timeout_minutes)
     dtype = _parse_dtype(torch_dtype)
 
     print_rank_0(f"Exporting: {megatron_path} -> {hf_path}")
@@ -173,6 +199,23 @@ def export_megatron_to_hf(
     model_provider.expert_tensor_parallel_size = etp
     model_provider.pipeline_dtype = dtype
     model_provider.params_dtype = dtype
+    # For PP > 1 export, read pipeline layout from checkpoint
+    if pp > 1:
+        from pathlib import Path
+
+        import yaml
+
+        ckpt_path = Path(megatron_path)
+        for candidate in [ckpt_path, *ckpt_path.glob("iter_*")]:
+            rc = candidate / "run_config.yaml"
+            if rc.exists():
+                with open(rc) as f:
+                    cfg = yaml.safe_load(f)
+                saved_layout = cfg.get("model", {}).get("pipeline_model_parallel_layout")
+                if isinstance(saved_layout, list):
+                    model_provider.pipeline_model_parallel_layout = saved_layout
+                    print_rank_0(f"  Read pipeline layout from checkpoint ({len(saved_layout)} stages)")
+                    break
     model_provider.finalize()
     model_provider.initialize_model_parallel(seed=0)
 
@@ -218,6 +261,12 @@ def _add_common_args(parser: argparse.ArgumentParser) -> None:
         help="Model precision (default: bfloat16)",
     )
     parser.add_argument("--trust-remote-code", action="store_true", help="Allow custom model code execution")
+    parser.add_argument(
+        "--distributed-timeout-minutes",
+        type=int,
+        default=None,
+        help="Initialize the distributed process group with this timeout before model setup",
+    )
 
 
 def main():
@@ -255,7 +304,6 @@ def main():
         default=1,
         help="Only every N-th rank writes files (reduces I/O, only with --distributed-save)",
     )
-
     args = parser.parse_args()
 
     if not args.command:
@@ -272,6 +320,7 @@ def main():
             etp=args.etp,
             torch_dtype=args.torch_dtype,
             trust_remote_code=args.trust_remote_code,
+            distributed_timeout_minutes=args.distributed_timeout_minutes,
         )
     elif args.command == "export":
         export_megatron_to_hf(
@@ -288,6 +337,7 @@ def main():
             show_progress=not args.no_progress,
             distributed_save=args.distributed_save,
             save_every_n_ranks=args.save_every_n_ranks,
+            distributed_timeout_minutes=args.distributed_timeout_minutes,
         )
 
 

examples/conversion/convert_checkpoints_multi_gpu.py

@@ -128,6 +128,23 @@ def main(args) -> None:
         model_provider.expert_tensor_parallel_size = etp
         model_provider.pipeline_dtype = torch.bfloat16
 
+        # Read pipeline layout from checkpoint for PP > 1
+        if pp > 1:
+            from pathlib import Path
+
+            import yaml
+
+            ckpt_path = Path(args.megatron_model_path)
+            for candidate in [ckpt_path, *ckpt_path.glob("iter_*")]:
+                rc = candidate / "run_config.yaml"
+                if rc.exists():
+                    with open(rc) as f:
+                        cfg = yaml.safe_load(f)
+                    saved_layout = cfg.get("model", {}).get("pipeline_model_parallel_layout")
+                    if isinstance(saved_layout, list):
+                        model_provider.pipeline_model_parallel_layout = saved_layout
+                        break
+
         # Once all overrides are set, finalize the model provider to ensure the post initialization logic is run
         model_provider.finalize()
         model_provider.initialize_model_parallel(seed=0)

examples/conversion/hf_to_megatron_generate_text.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()