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

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

@@ -0,0 +1,59 @@
+# DeepSeek V4
+
+End-to-end conversion and inference scripts for the DeepSeek V4 family on Megatron Bridge.
+
+The bridge supports four published variants out of the same code path. The on-disk quantisation differs between post-trained (Flash, Pro) and pretrained-only (Flash-Base, Pro-Base) models — see [`docs/models/llm/deepseek-v4.md`](../../../docs/models/llm/deepseek-v4.md) for the per-variant scheme.
+
+## MCore Dev Branch Requirement
+
+DSv4 imports require MCore changes that are not yet on a tagged release: PR [#3430](https://github.com/NVIDIA/Megatron-LM/pull/3430), PR [#4458](https://github.com/NVIDIA/Megatron-LM/pull/4458), PR [#4481](https://github.com/NVIDIA/Megatron-LM/pull/4481), and PR [#4518](https://github.com/NVIDIA/Megatron-LM/pull/4518), and PR [#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839). Until these merge to Megatron-LM `main` and the bridge submodule pin advances, point `3rdparty/Megatron-LM` at the Megatron-LM `dev` branch:
+
+```bash
+./scripts/switch_mcore.sh dev
+uv sync
+```
+
+Use `./scripts/switch_mcore.sh main` and `uv sync --locked` to return to the pinned main-branch submodule.
+
+| Variant | HF path | Quant scheme | Validation |
+|---------|---------|--------------|------------|
+| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | FP8 attn + MXFP4 experts | Verified on GB200, last-token logit cosine 0.96-0.99 (short prompts ~0.98, long prompts >1024 tokens ~0.96-0.99) vs official inference |
+| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | uniform FP8 (F32 scales) | Verified on GB200, last-real-token logit cosine 0.9866-0.9930, mean 0.9907 vs official inference |
+| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | FP8 attn + MXFP4 experts | Import, export, inference verified on GB200 (PP=4 EP=8) and H100 (PP=16 EP=8) |
+| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | uniform FP8 (F32 scales) | Same bridge code as Pro; end-to-end untested |
+
+## Examples
+
+- `conversion.sh` imports HF weights into Megatron Bridge and exports Megatron checkpoints back to HF format.
+- `inference.sh` runs text generation against an HF or Megatron checkpoint.
+
+Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each script's header comments for the expected environment variables and `#SBATCH` directives to edit before submitting.
+
+The bridge's `maybe_modify_loaded_hf_weight` hook dispatches dequantisation by tensor dtype:
+
+- `int8` -> MXFP4 packed nibbles -> `bfloat16` via the E2M1 lookup table and per-row 16-K-tile E8M0 scales
+- `float8_e4m3fn` with companion `.scale` -> `bfloat16` via 128x128 block-scale expansion, handling both E8M0 and F32 scale dtypes
+
+No external dequantisation script is required.
+
+## Parallelism Configurations
+
+DSv4 currently requires **TP=1** because MLA tensor parallelism is not supported alongside the DSv4 hybrid attention path. Scale via expert and pipeline parallelism instead.
+
+| Model | TP | PP | EP | GPUs | GPU | Verified |
+|-------|---:|---:|---:|-----:|-----|----------|
+| DeepSeek-V4-Flash | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference |
+| DeepSeek-V4-Flash | 1 | 1 | 8 | 8 | H100 80GB | Import, export, inference |
+| DeepSeek-V4-Flash-Base | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference |
+| DeepSeek-V4-Pro | 1 | 4 | 8 | 32 | GB200 192GB | Import, export, inference |
+| DeepSeek-V4-Pro | 1 | 16 | 8 | 128 | H100 80GB | Import, export, inference |
+
+## Known Limitations
+
+- **MTP is disabled for inference** via `disable_mtp_for_inference()`. MTP weights are mapped end-to-end and loaded into the Megatron model.
+
+- **Fused mHC is not supported on H100.** Set `use_fused_mhc=False` in the bridge config when running on Hopper GPUs. Fused mHC is enabled by default and works on GB200.
+
+- **`fast_hadamard_transform` is optional.** When unavailable, DSA falls back to a PyTorch hadamard implementation. Throughput is lower but numerical behavior is unchanged.
+
+- **Logit parity is verified for Flash and Flash-Base** against the official inference stack at last-real-token logits. The remaining gap is structural, from different attention/HC kernel decompositions and accumulation precisions between MCore and official inference.

examples/models/deepseek_v4/README.md

@@ -0,0 +1,107 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DeepSeek-V4 import + export with the Bridge.
+#
+# DSv4 currently requires TP=1; scale via expert and pipeline parallelism (EP, PP).
+# The Bridge dispatches FP8 / MXFP4 dequantisation by tensor dtype, so the
+# same script works for Flash, Flash-Base, Pro, and Pro-Base.
+#
+# Override defaults by exporting environment variables before running:
+#   WORKSPACE: directory for converted Megatron checkpoints (default: /workspace)
+#   MODEL_VARIANT: one of DeepSeek-V4-Flash, DeepSeek-V4-Flash-Base,
+#                  DeepSeek-V4-Pro, DeepSeek-V4-Pro-Base
+#                  (default: DeepSeek-V4-Flash-Base)
+#   EP: expert-parallel size (default: 4 for Flash, 8 for Pro)
+#   PP: pipeline-parallel size (default: 1 for Flash, 4 for Pro)
+#
+# Defaults below are for GB200 (192 GB). For H100 (80 GB) configs, see README.md.
+
+set -xeuo pipefail
+
+WORKSPACE=${WORKSPACE:-/workspace}
+MODEL_VARIANT=${MODEL_VARIANT:-DeepSeek-V4-Flash-Base}
+HF_MODEL_ID="deepseek-ai/${MODEL_VARIANT}"
+
+if [[ -z "${EP:-}" ]]; then
+    case "${MODEL_VARIANT}" in
+        DeepSeek-V4-Pro*) EP=8 ;;
+        *)                EP=4 ;;
+    esac
+fi
+if [[ -z "${PP:-}" ]]; then
+    case "${MODEL_VARIANT}" in
+        DeepSeek-V4-Pro*) PP=4 ;;
+        *)                PP=1 ;;
+    esac
+fi
+TP=1
+
+MEGATRON_DIR="${WORKSPACE}/models/${MODEL_VARIANT}"
+EXPORT_DIR="${WORKSPACE}/models/${MODEL_VARIANT}-hf-export"
+ITER=iter_0000000
+
+# 1) Import HF -> Megatron (FP8 / MXFP4 dequantised to bfloat16 in-flight)
+uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
+    examples/conversion/convert_checkpoints_multi_gpu.py import \
+    --hf-model "${HF_MODEL_ID}" \
+    --megatron-path "${MEGATRON_DIR}" \
+    --tp ${TP} --pp ${PP} --ep ${EP} \
+    --torch-dtype bfloat16 \
+    --trust-remote-code
+
+# 2) Compare HF and Megatron logits on a short prompt
+uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
+    examples/conversion/compare_hf_and_megatron/compare.py \
+    --hf_model_path "${HF_MODEL_ID}" \
+    --megatron_model_path "${MEGATRON_DIR}" \
+    --prompt "Hello, how are you?" \
+    --tp ${TP} --pp ${PP} --ep ${EP} \
+    --trust-remote-code
+
+# 3) Export Megatron -> HF (round-trip)
+uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
+    examples/conversion/convert_checkpoints_multi_gpu.py export \
+    --hf-model "${HF_MODEL_ID}" \
+    --megatron-path "${MEGATRON_DIR}/${ITER}" \
+    --tp ${TP} --pp ${PP} --ep ${EP} \
+    --torch-dtype bfloat16 \
+    --hf-path "${EXPORT_DIR}" \
+    --distributed-save \
+    --trust-remote-code
+
+# 4) Round-trip validation (bf16 -> Megatron -> bf16)
+# DSv4 HF weights are quantized (FP8/MXFP4), so the first import dequantises
+# to bfloat16. A true lossless roundtrip re-imports the exported bf16 checkpoint
+# and compares against the first export.
+ROUNDTRIP_DIR="${WORKSPACE}/models/${MODEL_VARIANT}-roundtrip"
+uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
+    examples/conversion/convert_checkpoints_multi_gpu.py import \
+    --hf-model "${EXPORT_DIR}" \
+    --megatron-path "${ROUNDTRIP_DIR}" \
+    --tp ${TP} --pp ${PP} --ep ${EP} \
+    --torch-dtype bfloat16 \
+    --trust-remote-code
+
+ROUNDTRIP_EXPORT_DIR="${WORKSPACE}/models/${MODEL_VARIANT}-roundtrip-export"
+uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
+    examples/conversion/convert_checkpoints_multi_gpu.py export \
+    --hf-model "${EXPORT_DIR}" \
+    --megatron-path "${ROUNDTRIP_DIR}" \
+    --hf-path "${ROUNDTRIP_EXPORT_DIR}" \
+    --tp ${TP} --pp ${PP} --ep ${EP} \
+    --torch-dtype bfloat16 \
+    --distributed-save \
+    --trust-remote-code