[ckpt] feat: Add --auto-config for self-contained Megatron-to-HF checkpoint export

- Add --auto-config flag to convert_checkpoints.py for config-free HF export
- Add filter_hf_config_fields utility for filtering reconstructed HF configs
- Add megatron_to_hf_config to model bridges for Megatron→HF config mapping
- Add auto-config round-trip test reusing Nemotron-3-Nano toy model fixture

Signed-off-by: david <davidzhaiyang@gmail.com>

Author: davidyang05

examples/conversion/convert_checkpoints.py

@@ -50,6 +50,13 @@
     --megatron-path ./checkpoints/custom_model \
     --hf-path ./exports/custom_model_hf \
     --no-progress
+
+  # Export using auto-config (synthesizes config from checkpoint + HF model reference)
+  uv run python examples/conversion/convert_checkpoints.py export \
+    --hf-model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
+    --auto-config \
+    --megatron-path ./checkpoints/nemotron_h_custom \
+    --hf-path ./exports/nemotron_h_custom_hf
 """
 
 import argparse
@@ -140,16 +147,22 @@ def export_megatron_to_hf(
     hf_model: str,
     megatron_path: str,
     hf_path: str,
+    auto_config: bool = False,
     show_progress: bool = True,
     strict: bool = True,
 ) -> None:
     """
     Export a Megatron checkpoint to HuggingFace format.
 
     Args:
+        hf_model: HuggingFace model ID or path (loads full model with weights).
         megatron_path: Directory path where the Megatron checkpoint is stored
         hf_path: Directory path where the HuggingFace model will be saved
+        auto_config: If True, synthesizes the HF config from run_config.yaml using
+            the model bridge class resolved from --hf-model, instead of loading
+            the full pretrained weights from HuggingFace.
         show_progress: Display progress bar during weight export
+        strict: Whether to perform strict validation during weight export
     """
     print(f"🔄 Starting export: {megatron_path} -> {hf_path}")
 
@@ -174,9 +187,26 @@ def export_megatron_to_hf(
 
     print(f"📋 Found configuration: {config_files[0]}")
 
-    # For demonstration, we'll create a bridge from a known config
-    # This would typically be extracted from the checkpoint metadata
-    bridge = AutoBridge.from_hf_pretrained(hf_model, trust_remote_code=True)
+    # Create bridge: auto-config (synthesize from run_config.yaml) or full pretrained
+    if auto_config:
+        from transformers import AutoConfig, PretrainedConfig
+
+        from megatron.bridge.models.conversion.utils import filter_hf_config_fields
+        from megatron.bridge.training.model_load_save import load_model_config
+
+        megatron_cfg, _ = load_model_config(str(config_files[0].parent))
+        config = AutoConfig.from_pretrained(hf_model, trust_remote_code=True)
+        bridge = AutoBridge.from_hf_config(config)
+        if "megatron_to_hf_config" not in type(bridge._model_bridge).__dict__:
+            raise ValueError(f"--auto-config does not currently support {type(bridge._model_bridge).__name__}")
+        hf_cfg_dict = bridge._model_bridge.megatron_to_hf_config(megatron_cfg)
+        hf_cfg_dict = filter_hf_config_fields(hf_cfg_dict, config.to_dict())
+        config = PretrainedConfig(**hf_cfg_dict)
+        config._bridge_original_dict = dict(hf_cfg_dict)
+        bridge = AutoBridge.from_hf_config(config)
+        bridge._auto_config_source_repo = hf_model
+    else:
+        bridge = AutoBridge.from_hf_pretrained(hf_model, trust_remote_code=True)
 
     # Export using the convenience method
     print("📤 Exporting to HuggingFace format...")
@@ -227,6 +257,11 @@ def main():
     # Export subcommand (Megatron -> HF)
     export_parser = subparsers.add_parser("export", help="Export Megatron checkpoint to HuggingFace format")
     export_parser.add_argument("--hf-model", required=True, help="HuggingFace model ID or path to model directory")
+    export_parser.add_argument(
+        "--auto-config",
+        action="store_true",
+        help="Export custom models with different architectures than the base --hf-model",
+    )
     export_parser.add_argument(
         "--megatron-path", required=True, help="Directory path where the Megatron checkpoint is stored"
     )
@@ -237,7 +272,6 @@ def main():
     export_parser.add_argument(
         "--not-strict", action="store_true", help="Allow source and target checkpoint to have different keys"
     )
-
     args = parser.parse_args()
 
     if not args.command:
@@ -258,6 +292,7 @@ def main():
             hf_model=args.hf_model,
             megatron_path=args.megatron_path,
             hf_path=args.hf_path,
+            auto_config=args.auto_config,
             show_progress=not args.no_progress,
             strict=not args.not_strict,
         )

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

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import dataclasses
+import logging
 from functools import cached_property, partial
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generic, Iterable, List, Optional, Type, TypeVar, Union
@@ -30,6 +31,7 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig
 from modelopt.torch.quantization.utils import is_quantized
+from safetensors.torch import save_file
 from transformers.configuration_utils import PretrainedConfig
 from typing_extensions import Unpack
 
@@ -47,6 +49,8 @@
 from megatron.bridge.models.model_provider import GetModelKwargs, ModelParallelKwargs, ModelProviderMixin
 
 
+logger = logging.getLogger(__name__)
+
 MegatronModelT = TypeVar("MegatronModelT", bound=MegatronModule)
 DataclassT = TypeVar("DataclassT")
 
@@ -577,29 +581,65 @@ def save_hf_pretrained(
             saves the configuration files, while weight saving is coordinated
             across all ranks.
         """
-        if not isinstance(self.hf_pretrained, PreTrainedCausalLM):
-            raise ValueError(
-                "save_hf_pretrained requires a pretrained HuggingFace model. "
-                "AutoBridge.from_hf_config() creates a config-only bridge; "
-                "use AutoBridge.from_hf_pretrained(...) instead."
-            )
+        if not isinstance(self.hf_pretrained, (PreTrainedCausalLM, PretrainedConfig)):
+            raise ValueError("save_hf_pretrained requires a pretrained HuggingFace model or config.")
+        is_config_only = isinstance(self.hf_pretrained, PretrainedConfig)
+
+        def _save_artifacts():
+            if is_config_only:
+                import json
+
+                # Config-only path: write config.json and download modeling files from Hub.
+                Path(path).mkdir(parents=True, exist_ok=True)
+                config_dict = getattr(self.hf_pretrained, "_bridge_original_dict", None)
+                if config_dict is None:
+                    config_dict = self.hf_pretrained.to_dict()
+                    # Apply dtype/model_type fixups directly here
+                    if "model_type" in self.hf_pretrained.__dict__:
+                        config_dict["model_type"] = self.hf_pretrained.__dict__["model_type"]
+                    if "dtype" in config_dict and config_dict.get("torch_dtype") is None:
+                        config_dict["torch_dtype"] = config_dict.pop("dtype")
+
+                with open(Path(path) / "config.json", "w") as _f:
+                    json.dump(config_dict, _f, indent=2, sort_keys=True, allow_nan=True)
+
+                # Download custom modeling files so the output is loadable with from_pretrained().
+                hub_repo = getattr(self, "_auto_config_source_repo", None)
+                if hub_repo:
+                    try:
+                        from huggingface_hub import hf_hub_download, list_repo_files
+
+                        repo_files = list_repo_files(hub_repo)
+                        py_files = [f for f in repo_files if f.endswith(".py")]
+                        for py_file in py_files:
+                            hf_hub_download(
+                                repo_id=hub_repo,
+                                filename=py_file,
+                                local_dir=path,
+                            )
+                    except Exception as exc:
+                        logger.error(
+                            "Could not download modeling files from %s: %s. "
+                            "The exported model may not be loadable with from_pretrained().",
+                            hub_repo,
+                            exc,
+                        )
 
-        # Get bridge-level ADDITIONAL_FILE_PATTERNS if configured
-        additional_files = None
-        if hasattr(self._model_bridge, "ADDITIONAL_FILE_PATTERNS") and self._model_bridge.ADDITIONAL_FILE_PATTERNS:
-            additional_files = self._model_bridge.ADDITIONAL_FILE_PATTERNS
-
-        if dist.is_available() and dist.is_initialized():
-            # Distributed training, only rank 0 saves artifacts
-            if dist.get_rank() == 0:
+            else:
+                # Get bridge-level ADDITIONAL_FILE_PATTERNS if configured
+                mb = self._model_bridge
+                additional_files = None
+                if hasattr(mb, "ADDITIONAL_FILE_PATTERNS") and mb.ADDITIONAL_FILE_PATTERNS:
+                    additional_files = mb.ADDITIONAL_FILE_PATTERNS
                 self.hf_pretrained.save_artifacts(
                     path, original_source_path=source_path, additional_files=additional_files
                 )
+
+        if dist.is_available() and dist.is_initialized():
+            if dist.get_rank() == 0:
+                _save_artifacts()
         else:
-            # No distributed training, save artifacts
-            self.hf_pretrained.save_artifacts(
-                path, original_source_path=source_path, additional_files=additional_files
-            )
+            _save_artifacts()
 
         self.save_hf_weights(
             model,
@@ -701,7 +741,29 @@ def _filter_quant(gen):
                 save_every_n_ranks=save_every_n_ranks,
             )
         else:
-            raise ValueError("The state source is not a SafeTensorsStateSource, cannot save in streaming mode.")
+            # Config-only path: shard and write safetensors directly
+            import json
+
+            from huggingface_hub import split_torch_state_dict_into_shards
+
+            # NOTE: Collects the full state dict into CPU memory before sharding.
+            # For very large models (>70B), this may require significant host RAM.
+            state_dict = {name: tensor.contiguous().cpu() for name, tensor in generator}
+
+            is_distributed = dist.is_available() and dist.is_initialized()
+            rank = dist.get_rank() if is_distributed else 0
+
+            if rank == 0:
+                plan = split_torch_state_dict_into_shards(state_dict)
+                safe_dir = Path(path)
+                safe_dir.mkdir(parents=True, exist_ok=True)
+                for filename, tensors in plan.filename_to_tensors.items():
+                    shard = {k: state_dict[k] for k in tensors}
+                    save_file(shard, safe_dir / filename)
+                if plan.is_sharded:
+                    index = {"metadata": plan.metadata, "weight_map": plan.tensor_to_filename}
+                    with open(safe_dir / "model.safetensors.index.json", "w") as f:
+                        json.dump(index, f, indent=2)
 
         # Save quantizer/amax sidecar after the main generator is consumed (rank 0 only).
         if quant_tensors:
@@ -899,6 +961,11 @@ def import_ckpt(
         hf_tokenizer_kwargs = None
         if hasattr(bridge._model_bridge, "get_hf_tokenizer_kwargs"):
             hf_tokenizer_kwargs = bridge._model_bridge.get_hf_tokenizer_kwargs()
+        # Forward trust_remote_code to the tokenizer (needed for repos with custom code)
+        if kwargs.get("trust_remote_code"):
+            if hf_tokenizer_kwargs is None:
+                hf_tokenizer_kwargs = {}
+            hf_tokenizer_kwargs.setdefault("trust_remote_code", True)
         bridge.save_megatron_model(
             megatron_model,
             megatron_path,
@@ -921,8 +988,6 @@ def export_ckpt(
         This is a convenience method that loads a Megatron checkpoint and
         exports it to HuggingFace format. This is useful for sharing trained
         models or deploying them with HuggingFace inference tools.
-        Requires a bridge created with `from_hf_pretrained` so the tokenizer
-        and other HuggingFace artifacts can be saved alongside the weights.
 
         Args:
             megatron_path: Directory path where the Megatron checkpoint is stored
@@ -941,7 +1006,6 @@ def export_ckpt(
             ...     "./megatron_checkpoints/my_model",
             ...     "./hf_exports/my_model"
             ... )
-
             >>> # Export with specific settings
             >>> bridge.export_ckpt(
             ...     "./megatron_checkpoints/my_model",
@@ -953,12 +1017,6 @@ def export_ckpt(
             >>> from transformers import AutoModelForCausalLM
             >>> hf_model = AutoModelForCausalLM.from_pretrained("./hf_exports/my_model")
         """
-        if not isinstance(self.hf_pretrained, PreTrainedCausalLM):
-            raise ValueError(
-                "export_ckpt requires a pretrained HuggingFace model. "
-                "AutoBridge.from_hf_config() creates a config-only bridge; "
-                "use AutoBridge.from_hf_pretrained(...) instead."
-            )
         try:
             from megatron.bridge.training.model_load_save import temporary_distributed_context
         except ImportError:
@@ -1295,7 +1353,7 @@ def mla_transformer_config(self) -> MLATransformerConfig:
 
         return self._create_config_from_provider(_model_provider, MLATransformerConfig)
 
-    @property
+    @cached_property
     def _model_bridge(self) -> "MegatronModelBridge":
         hf_config = getattr(self.hf_pretrained, "hf_config", None)
         if hf_config is None:

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

@@ -42,6 +42,7 @@
     unwrap_model,
 )
 from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
+from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_utils import PreTrainedModel
 
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
@@ -1109,6 +1110,9 @@ def stream_weights_megatron_to_hf(
                         expected_keys = hf_pretrained.state.source.get_all_keys()
                         if "lm_head.weight" in expected_keys:
                             yield HFWeightTuple("lm_head.weight", final_tensor.clone().detach())
+                    elif isinstance(hf_pretrained, PretrainedConfig):
+                        # Always emit lm_head.weight for config-only
+                        yield HFWeightTuple("lm_head.weight", final_tensor.clone().detach())
                 elif embeddings_are_tied and hf_name == "lm_head.weight":
                     # This should not happen when embeddings are tied - assert error
                     raise ValueError(
@@ -1310,12 +1314,11 @@ def build_conversion_tasks(
         populated.
         """
 
-        # Ensure hf_pretrained has the required state structure
-        if not (hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source")):
+        has_hf_state = hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source")
+        if not isinstance(hf_pretrained, PretrainedConfig) and not has_hf_state:
             raise ValueError("hf_pretrained.state.source is required for weight ordering")
 
-        hf_keys: Iterable[str] = hf_pretrained.state.source.get_all_keys()
-
+        hf_keys: Optional[Iterable[str]] = hf_pretrained.state.source.get_all_keys() if has_hf_state else None
         mapping_registry = self.mapping_registry()
         unwrapped_model = unwrap_model(megatron_model)[0]
         model_config = unwrapped_model.config
@@ -1351,8 +1354,8 @@ def build_conversion_tasks(
                     logger.warning(f"WARNING: No mapping found for megatron_param: {global_name}")
                     continue
 
-                # ensure hf weights exist
-                if not mapping.allow_hf_name_mismatch:
+                # Ensure hf weights exist (skip for config-only export where hf_keys is None)
+                if hf_keys is not None and not mapping.allow_hf_name_mismatch:
                     if isinstance(mapping.hf_param, str):
                         if mapping.hf_param not in hf_keys:
                             logger.warning(f"WARNING: Can't find {mapping.hf_param} in hf_keys")
@@ -1546,7 +1549,7 @@ def _megatron_to_hf_registered_impl(
     ) -> Iterable[HFWeightTuple]:
         bridge = bridge_class()
 
-        # allow bridge to access model config (config-only shims or raw configs lack .config)
+        # Allow bridge to access model config (config-only shims or raw configs lack .config)
         bridge.hf_config = hf_pretrained.config if hasattr(hf_pretrained, "config") else hf_pretrained
 
         return bridge.stream_weights_megatron_to_hf(

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

@@ -273,6 +273,12 @@ def get_causal_lm_class_name_via_auto_map(
     return None
 
 
+def filter_hf_config_fields(hf_cfg_dict: dict[str, object], reference_cfg: dict[str, object]) -> dict[str, object]:
+    """Filter an HF config dict to only contain keys present in a reference config."""
+    reference_cfg_keys = set(reference_cfg.keys())
+    return {key: value for (key, value) in hf_cfg_dict.items() if key in reference_cfg_keys}
+
+
 def persistent_buffers(model: torch.nn.Module) -> Iterable[Tuple[str, torch.Tensor]]:
     """Return an iterator over persistent module buffers, yielding both the name of the buffer as well as the buffer itself."""