fix(conversion): improve Megatron-LM checkpoint compatibility (#5354)

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

Author: cuichenx

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 |

scripts/conversion/README.md

@@ -13,13 +13,21 @@ Run `./scripts/conversion/convert.sh import --help`,
 `./scripts/conversion/convert.sh export --help`, or
 `./scripts/conversion/convert.sh roundtrip --help` for the complete CLI.
 
-Megatron-LM training checkpoints normally store their arguments in `common.pt`
-instead of the Megatron Bridge `run_config.yaml` required by the export path.
-Before exporting such a checkpoint, follow
-[Export Megatron-LM checkpoints without a Bridge run config](../../docs/megatron-lm-to-megatron-bridge.md#export-megatron-lm-checkpoints-without-a-bridge-run-config)
-to generate and validate the provider configuration. Do not create an empty or
-hand-written YAML file; missing hybrid, MTP, MoE, or FP8 fields can construct a
-different model while appearing to load successfully.
+### Megatron-LM checkpoint compatibility
+
+Training with Megatron-LM and later relying on Megatron Bridge for Hugging Face
+export is **not recommended**. Megatron-LM checkpoints normally store their
+arguments in `common.pt` instead of the `run_config.yaml` used by the supported
+Bridge export launcher, so the launcher cannot consume them directly. For an
+existing trusted checkpoint, the
+[best-effort compatibility guidance](../../docs/megatron-lm-to-megatron-bridge.md#best-effort-export-of-an-existing-megatron-lm-checkpoint)
+describes how to generate and validate the required provider configuration. It
+does not guarantee compatibility, and any unclassified metadata mismatch must
+stop the export.
+
+Output from `scripts/translate_mlm_to_bridge.py` is best-effort configuration
+guidance for a new Bridge run. It is not checkpoint metadata and must not be
+renamed or inserted as `run_config.yaml`.
 
 ## Local CPU conversion