Merge branch 'main' into gptoss-120b-card

Author: kamran-nvidia

docs/bridge-guide.md

@@ -292,7 +292,7 @@ Or with a helper script:
 
 Megatron-LM training checkpoints do not normally contain the Bridge
 `run_config.yaml` required by this export command. See
-[Export Megatron-LM checkpoints without a Bridge run config](megatron-lm-to-megatron-bridge.md#export-megatron-lm-checkpoints-without-a-bridge-run-config)
+[Best-effort export of an existing Megatron-LM checkpoint](megatron-lm-to-megatron-bridge.md#best-effort-export-of-an-existing-megatron-lm-checkpoint)
 for the configuration generation and validation workflow.
 
 ### Create Megatron models and run locally

docs/fern/versions/nightly/pages/bridge-guide.mdx

@@ -260,7 +260,7 @@ Or with a helper script:
 
 Megatron-LM training checkpoints do not normally contain the Bridge
 `run_config.yaml` required by this export command. See
-[Export Megatron-LM checkpoints without a Bridge run config](megatron-lm-to-megatron-bridge.md#export-megatron-lm-checkpoints-without-a-bridge-run-config)
+[Best-effort export of an existing Megatron-LM checkpoint](megatron-lm-to-megatron-bridge.md#best-effort-export-of-an-existing-megatron-lm-checkpoint)
 for the configuration generation and validation workflow.
 
 ### Create Megatron models and run locally

docs/fern/versions/nightly/pages/megatron-lm-to-megatron-bridge.mdx

@@ -6,6 +6,8 @@ Megatron Bridge is Python-first: configure models, data, and training via typed
 
 `scripts/translate_mlm_to_bridge.py` translates bidirectionally between Megatron-LM `pretrain_gpt.py` CLI arguments and Megatron Bridge `run_recipe.py` Hydra overrides. It is useful for running loss-correlation experiments between the two frameworks and for migrating existing MLM configs.
 
+This script provides best-effort configuration translation only. It does not establish semantic equivalence, convert checkpoint weights, or emit the serialized `run_config.yaml` stored in a Bridge-native checkpoint. Do not rename its output to `run_config.yaml` or place it in an MLM checkpoint: the generated overrides or Python recipe are starting points for a new Bridge run, not checkpoint metadata. Review every skipped and unknown argument before using the output.
+
 ### MLM → Bridge (default direction)
 
 ```bash
@@ -87,9 +89,20 @@ uv run python scripts/translate_mlm_to_bridge.py --reverse \
 | `--rotary-base N` | `model.rotary_base=N` | |
 | `--num-experts N` | `model.num_moe_experts=N` | |
 | `--moe-router-topk K` | `model.moe_router_topk=K` | |
+| `--hybrid-layer-pattern P` | `model.hybrid_layer_pattern=P` | Selects `HybridModelProvider` in generated recipes |
+| `--hybrid-override-pattern P` | `model.hybrid_layer_pattern=P` | Migrates the deprecated MLM field |
+| `--mamba-state-dim N` | `model.mamba_state_dim=N` | |
+| `--mamba-head-dim N` | `model.mamba_head_dim=N` | |
+| `--mamba-num-groups N` | `model.mamba_num_groups=N` | |
+| `--mamba-num-heads N` | `model.mamba_num_heads=N` | |
+| `--mtp-hybrid-override-pattern P` | `model.mtp_hybrid_override_pattern=P` | Separate MLM MTP pattern |
+| `--mtp-use-repeated-layer` | `model.mtp_use_repeated_layer=true` | One physical MTP layer reused across prediction depths |
+| `--sft-tokenizer-prompt-format P` | `tokenizer.tokenizer_prompt_format=P` | Bridge exposes the MCore SFT compatibility alias at load time |
 | `--mock-data` | `dataset.mock=true` | Use synthetic data (no files needed) |
 
-Flags not present in Bridge (e.g., `--use-mcore-models`, `--use-flash-attn`) are silently skipped with a comment. `--mock-data` translates to `dataset.mock=true`. Unknown flags are listed in a separate section so you can handle them manually.
+Flags not present in Bridge (e.g., `--use-mcore-models`, `--use-flash-attn`) are omitted and listed in a comment. `--mock-data` translates to `dataset.mock=true`. Unknown flags are listed separately so you can handle them manually.
+
+Megatron-LM `--spec` values are intentionally not copied into generated configuration as arbitrary import targets. A recognized Mamba or Hybrid spec helps the translator identify the provider family, but a standalone Hybrid recipe still requires an explicit layer pattern. Review and select any Bridge stack specification in trusted Python code.
 
 > **Activation function CLI overrides**: `model.activation_func` can now be set via Hydra CLI string override (e.g. `model.activation_func=silu`, `model.activation_func=gelu`). The string is resolved to the callable in `TransformerConfig.finalize()`. This makes `--swiglu` → `model.gated_linear_unit=true model.activation_func=silu` round-trippable from the CLI.
 
@@ -115,27 +128,37 @@ Notes:
 - Config groups are nested: `rng`, `train`, `model`, `optimizer`, `ddp`, `scheduler`, `dataset`, `logger`, `tokenizer`, `checkpoint`, `dist`, `profiling`, `peft`, `comm_overlap`, `mixed_precision`, `inprocess_restart`.
 - After overrides are applied, runtime validation computes any dependent fields (e.g., data-parallel size, scheduler steps) and checks consistency.
 
-## Export Megatron-LM checkpoints without a Bridge run config
+## Best-effort export of an existing Megatron-LM checkpoint
 
 Megatron-LM checkpoints save their training arguments in `common.pt`; they do
 not normally contain the `run_config.yaml` written by Megatron Bridge. The
 Bridge checkpoint export launcher currently needs that file to reconstruct a
 model provider before loading the distributed weights.
 
+Training with Megatron-LM and later relying on Megatron Bridge for Hugging Face
+export is **not recommended**. This procedure is a best-effort recovery path for
+an existing checkpoint, not a supported general MLM-to-Hugging-Face conversion
+workflow. Use it only with a checkpoint and immutable local Hugging Face
+snapshot that you trust. Loading MLM common checkpoint state may deserialize
+Python objects; an allowlisted package prefix or audited Hugging Face code does
+not make an untrusted checkpoint safe.
+
 Use this procedure only when all of the following are true:
 
 - the iteration directory contains `common.pt` and the distributed checkpoint
   metadata and shards;
 - the checkpoint architecture has an existing Megatron Bridge implementation;
-- `--hf-model` identifies the exact Hugging Face architecture used to create
-  the checkpoint; fine-tuned weights may differ, but model dimensions, layer
-  patterns, vocabulary configuration, and MTP structure must match; and
+- `--hf-model` identifies an immutable local snapshot of the exact Hugging Face
+  architecture used to create the checkpoint; fine-tuned weights may differ,
+  but all architecture- and behavior-bearing configuration must match; and
 - Megatron Bridge, Megatron Core, and Transformer Engine are compatible with
   the versions that wrote the checkpoint.
 
-This procedure generates only `run_config.yaml`. It does not migrate legacy
-checkpoint keys or metadata, and it does not translate a custom Megatron-LM
-model spec into a Bridge provider.
+This procedure generates only `run_config.yaml` in the checkpoint directory.
+It does not migrate legacy checkpoint keys or metadata, translate a custom
+Megatron-LM model spec into a Bridge provider, or prove that the checkpoint and
+reference are semantically equivalent. Work on a copy if the checkpoint
+directory must remain pristine.
 
 ### Generate `run_config.yaml`
 
@@ -145,7 +168,7 @@ this generation step does not download its weights.
 
 ```bash
 MB_CKPT=/workspace/checkpoints/model/iter_0001000
-MB_HF_REF=organization/model-reference
+MB_HF_REF=/workspace/hf-snapshots/model-at-immutable-revision
 
 uv run python - "$MB_CKPT" "$MB_HF_REF" --trust-remote-code <<'PY'
 from dataclasses import fields
@@ -269,8 +292,12 @@ for name in (
 PY
 ```
 
-Compare the reported architecture, hybrid/MTP layout, and precision fields
-with the original training configuration. Do not continue if they differ.
+The printed fields are only a starting diagnostic. Compare every architecture-
+and behavior-bearing field with the original training configuration, including
+position and RoPE behavior, vocabulary and tokenizer IDs, tied weights, expert
+and router behavior, attention, Mamba, and physical/repeated-MTP layout.
+Parallel degrees may differ only for a valid resharding topology. Treat any
+unclassified difference as an incompatibility and do not continue.
 
 Run the normal export after validation. Adapt the parallel topology to the
 checkpoint and available GPUs:
@@ -290,6 +317,10 @@ checkpoint and available GPUs:
 
 Keep strict conversion enabled for the first export. Do not use
 `--not-strict` to bypass missing MTP, expert, or quantized parameters.
+Verify the emitted HF config, shard index, expected tensor keys, strict
+`from_pretrained()` loading, and at least one forward pass. These are structural
+and runtime checks, not numerical-parity evidence. Compare source and exported
+logits when a source-model baseline is available.
 
 For MXFP8 parameter checkpoints, use hardware supported by the saved recipe
 (normally Blackwell). The default BF16 export dequantizes Transformer Engine

docs/megatron-lm-to-megatron-bridge.md

@@ -6,6 +6,8 @@ Megatron Bridge is Python-first: configure models, data, and training via typed
 
 `scripts/translate_mlm_to_bridge.py` translates bidirectionally between Megatron-LM `pretrain_gpt.py` CLI arguments and Megatron Bridge `run_recipe.py` Hydra overrides. It is useful for running loss-correlation experiments between the two frameworks and for migrating existing MLM configs.
 
+This script provides best-effort configuration translation only. It does not establish semantic equivalence, convert checkpoint weights, or emit the serialized `run_config.yaml` stored in a Bridge-native checkpoint. Do not rename its output to `run_config.yaml` or place it in an MLM checkpoint: the generated overrides or Python recipe are starting points for a new Bridge run, not checkpoint metadata. Review every skipped and unknown argument before using the output.
+
 ### MLM → Bridge (default direction)
 
 ```bash
@@ -87,9 +89,20 @@ uv run python scripts/translate_mlm_to_bridge.py --reverse \
 | `--rotary-base N` | `model.rotary_base=N` | |
 | `--num-experts N` | `model.num_moe_experts=N` | |
 | `--moe-router-topk K` | `model.moe_router_topk=K` | |
+| `--hybrid-layer-pattern P` | `model.hybrid_layer_pattern=P` | Selects `HybridModelProvider` in generated recipes |
+| `--hybrid-override-pattern P` | `model.hybrid_layer_pattern=P` | Migrates the deprecated MLM field |
+| `--mamba-state-dim N` | `model.mamba_state_dim=N` | |
+| `--mamba-head-dim N` | `model.mamba_head_dim=N` | |
+| `--mamba-num-groups N` | `model.mamba_num_groups=N` | |
+| `--mamba-num-heads N` | `model.mamba_num_heads=N` | |
+| `--mtp-hybrid-override-pattern P` | `model.mtp_hybrid_override_pattern=P` | Separate MLM MTP pattern |
+| `--mtp-use-repeated-layer` | `model.mtp_use_repeated_layer=true` | One physical MTP layer reused across prediction depths |
+| `--sft-tokenizer-prompt-format P` | `tokenizer.tokenizer_prompt_format=P` | Bridge exposes the MCore SFT compatibility alias at load time |
 | `--mock-data` | `dataset.mock=true` | Use synthetic data (no files needed) |
 
-Flags not present in Bridge (e.g., `--use-mcore-models`, `--use-flash-attn`) are silently skipped with a comment. `--mock-data` translates to `dataset.mock=true`. Unknown flags are listed in a separate section so you can handle them manually.
+Flags not present in Bridge (e.g., `--use-mcore-models`, `--use-flash-attn`) are omitted and listed in a comment. `--mock-data` translates to `dataset.mock=true`. Unknown flags are listed separately so you can handle them manually.
+
+Megatron-LM `--spec` values are intentionally not copied into generated configuration as arbitrary import targets. A recognized Mamba or Hybrid spec helps the translator identify the provider family, but a standalone Hybrid recipe still requires an explicit layer pattern. Review and select any Bridge stack specification in trusted Python code.
 
 > **Activation function CLI overrides**: `model.activation_func` can now be set via Hydra CLI string override (e.g. `model.activation_func=silu`, `model.activation_func=gelu`). The string is resolved to the callable in `TransformerConfig.finalize()`. This makes `--swiglu` → `model.gated_linear_unit=true model.activation_func=silu` round-trippable from the CLI.
 
@@ -115,27 +128,37 @@ Notes:
 - Config groups are nested: `rng`, `train`, `model`, `optimizer`, `ddp`, `scheduler`, `dataset`, `logger`, `tokenizer`, `checkpoint`, `dist`, `profiling`, `peft`, `comm_overlap`, `mixed_precision`, `inprocess_restart`.
 - After overrides are applied, runtime validation computes any dependent fields (e.g., data-parallel size, scheduler steps) and checks consistency.
 
-## Export Megatron-LM checkpoints without a Bridge run config
+## Best-effort export of an existing Megatron-LM checkpoint
 
 Megatron-LM checkpoints save their training arguments in `common.pt`; they do
 not normally contain the `run_config.yaml` written by Megatron Bridge. The
 Bridge checkpoint export launcher currently needs that file to reconstruct a
 model provider before loading the distributed weights.
 
+Training with Megatron-LM and later relying on Megatron Bridge for Hugging Face
+export is **not recommended**. This procedure is a best-effort recovery path for
+an existing checkpoint, not a supported general MLM-to-Hugging-Face conversion
+workflow. Use it only with a checkpoint and immutable local Hugging Face
+snapshot that you trust. Loading MLM common checkpoint state may deserialize
+Python objects; an allowlisted package prefix or audited Hugging Face code does
+not make an untrusted checkpoint safe.
+
 Use this procedure only when all of the following are true:
 
 - the iteration directory contains `common.pt` and the distributed checkpoint
   metadata and shards;
 - the checkpoint architecture has an existing Megatron Bridge implementation;
-- `--hf-model` identifies the exact Hugging Face architecture used to create
-  the checkpoint; fine-tuned weights may differ, but model dimensions, layer
-  patterns, vocabulary configuration, and MTP structure must match; and
+- `--hf-model` identifies an immutable local snapshot of the exact Hugging Face
+  architecture used to create the checkpoint; fine-tuned weights may differ,
+  but all architecture- and behavior-bearing configuration must match; and
 - Megatron Bridge, Megatron Core, and Transformer Engine are compatible with
   the versions that wrote the checkpoint.
 
-This procedure generates only `run_config.yaml`. It does not migrate legacy
-checkpoint keys or metadata, and it does not translate a custom Megatron-LM
-model spec into a Bridge provider.
+This procedure generates only `run_config.yaml` in the checkpoint directory.
+It does not migrate legacy checkpoint keys or metadata, translate a custom
+Megatron-LM model spec into a Bridge provider, or prove that the checkpoint and
+reference are semantically equivalent. Work on a copy if the checkpoint
+directory must remain pristine.
 
 ### Generate `run_config.yaml`
 
@@ -145,7 +168,7 @@ this generation step does not download its weights.
 
 ```bash
 MB_CKPT=/workspace/checkpoints/model/iter_0001000
-MB_HF_REF=organization/model-reference
+MB_HF_REF=/workspace/hf-snapshots/model-at-immutable-revision
 
 uv run python - "$MB_CKPT" "$MB_HF_REF" --trust-remote-code <<'PY'
 from dataclasses import fields
@@ -269,8 +292,12 @@ for name in (
 PY
 ```
 
-Compare the reported architecture, hybrid/MTP layout, and precision fields
-with the original training configuration. Do not continue if they differ.
+The printed fields are only a starting diagnostic. Compare every architecture-
+and behavior-bearing field with the original training configuration, including
+position and RoPE behavior, vocabulary and tokenizer IDs, tied weights, expert
+and router behavior, attention, Mamba, and physical/repeated-MTP layout.
+Parallel degrees may differ only for a valid resharding topology. Treat any
+unclassified difference as an incompatibility and do not continue.
 
 Run the normal export after validation. Adapt the parallel topology to the
 checkpoint and available GPUs:
@@ -290,6 +317,10 @@ checkpoint and available GPUs:
 
 Keep strict conversion enabled for the first export. Do not use
 `--not-strict` to bypass missing MTP, expert, or quantized parameters.
+Verify the emitted HF config, shard index, expected tensor keys, strict
+`from_pretrained()` loading, and at least one forward pass. These are structural
+and runtime checks, not numerical-parity evidence. Compare source and exported
+logits when a source-model baseline is available.
 
 For MXFP8 parameter checkpoints, use hardware supported by the saved recipe
 (normally Blackwell). The default BF16 export dequantizes Transformer Engine
@@ -300,7 +331,6 @@ native FP8 Hugging Face export currently supports blockwise FP8 parameters.
 
 Below is a concise mapping from common `megatron-lm/megatron/training/arguments.py` flags to the new dataclass fields. If a field is not listed here (e.g., highly model-specific knobs), it typically lives under `model.*`, `optimizer.*`, `dataset.*`, or `tokenizer.*` with similar names.
 
-
 ### Model topology and parallelisms
 
 | megatron-lm arguments | Megatron Bridge config | Description |

docs/training/activation-recomputation.md

@@ -210,5 +210,5 @@ overlap. Prefer the narrow boundary that satisfies the memory target.
 
 - [docs/training/cuda-graphs.md](cuda-graphs.md)
 - [docs/training/moe-optimization.md](moe-optimization.md)
-- [skills/nemo-mbridge-perf-activation-recompute/SKILL.md](../../skills/nemo-mbridge-perf-activation-recompute/SKILL.md) — architecture-specific module selection, compatibility, and measurement guidance
-- [skills/nemo-mbridge-perf-memory-tuning/SKILL.md](../../skills/nemo-mbridge-perf-memory-tuning/SKILL.md) — expandable segments, parallelism resizing, and other memory reduction strategies
+- [skills/nemo-mbridge-perf-activation-recompute/SKILL.md](../skills/nemo-mbridge-perf-activation-recompute/SKILL.md) — architecture-specific module selection, compatibility, and measurement guidance
+- [skills/nemo-mbridge-perf-memory-tuning/SKILL.md](../skills/nemo-mbridge-perf-memory-tuning/SKILL.md) — expandable segments, parallelism resizing, and other memory reduction strategies

docs/training/checkpointing.md

@@ -312,8 +312,8 @@ The load default follows the actual source because the checkpoint restored on re
 
 ### Restore Behavior and Failure Modes
 
-- If the dataloader state directory is **absent** (e.g. a checkpoint saved before this feature existed), the dataloader starts from the beginning and a message is logged.
-- If the directory **exists** but the current rank's per-DP-rank file is **missing**, restore **raises**. This almost always means the data-parallel size changed since the checkpoint was saved; resuming would silently change the data order, so it fails loudly instead.
+- If the dataloader state root or the selected `iter_N` directory is **absent** (e.g. the selected checkpoint was saved before this feature existed), the dataloader starts from the beginning and a message is logged. Other state generations under the same root do not change this behavior.
+- If the selected `iter_N` directory **exists** but the current rank's per-DP-rank file is **missing**, restore **raises**. This almost always means the data-parallel size changed since the checkpoint was saved; resuming would silently change the data order, so it fails loudly instead.
 
 ### Determinism Requirement