auto_bridge.py

# Copyright (c) 2025, 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.

from __future__ import annotations

import dataclasses
import logging
from collections.abc import Callable
from contextlib import nullcontext
from functools import cached_property, partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Generic, Iterable, List, Literal, Optional, Tuple, Type, TypeVar, Union

import torch
import torch.distributed as dist
import transformers


if TYPE_CHECKING:
    from megatron.bridge.peft.base import PEFT

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

from megatron.bridge.models.conversion import model_bridge
from megatron.bridge.models.conversion.model_bridge import (
    HFWeightTuple,
    MegatronModelBridge,
    WeightConversionTask,
)
from megatron.bridge.models.conversion.utils import get_causal_lm_class_name_via_auto_map
from megatron.bridge.models.gpt_provider import GPTModelProvider
from megatron.bridge.models.hf_pretrained.base import PreTrainedBase
from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM, _ConfigOnlyPretrainedShim
from megatron.bridge.models.hf_pretrained.safe_config_loader import safe_load_config_with_retry
from megatron.bridge.models.hf_pretrained.state import SafeTensorsStateSource
from megatron.bridge.models.model_provider import GetModelKwargs, ModelParallelKwargs, ModelProviderMixin


logger = logging.getLogger(__name__)

MegatronModelT = TypeVar("MegatronModelT", bound=MegatronModule)
DataclassT = TypeVar("DataclassT")

# Supported HuggingFace architecture suffixes for causal generation models
SUPPORTED_HF_ARCHITECTURES: tuple[str, ...] = (
    "ForCausalLM",
    "ForConditionalGeneration",
    "NemotronH_Nano_VL_V2",
    "NemotronH_Nano_Omni_Reasoning_V3",
    "Qwen2_5OmniModel",
    "NemotronLabsDiffusionModel",
)

# Mapping from non-standard HF architecture names to their actual transformers class names.
# Some HF model configs report architecture names that don't follow the standard
# 'ForCausalLM'/'ForConditionalGeneration' convention and don't directly map to a
# transformers class. This dict resolves those aliases.
HF_ARCHITECTURE_ALIASES: dict[str, str] = {
    "Qwen2_5OmniModel": "Qwen2_5OmniForConditionalGeneration",
}

MTP_CONFIG_FIELDS: tuple[str, ...] = ("num_nextn_predict_layers", "mtp_num_hidden_layers", "mtp_num_layers")
_MISSING = object()


def _get_config_field(config: Any, field: str) -> Any:
    if isinstance(config, dict):
        return config.get(field, _MISSING)

    return getattr(config, "__dict__", {}).get(field, _MISSING)


def _config_disables_mtp(config: Any) -> bool:
    """Return True when a config object or dict explicitly disables MTP layers."""
    if config is None:
        return False

    for field in MTP_CONFIG_FIELDS:
        value = _get_config_field(config, field)
        if value is _MISSING or value is None:
            continue

        return int(value) == 0

    text_config = _get_config_field(config, "text_config")
    return text_config is not _MISSING and _config_disables_mtp(text_config)


def _saved_config_disables_mtp(path: str | Path) -> bool:
    """Check the final config.json written for an HF export."""
    import json

    config_path = Path(path) / "config.json"
    if not config_path.exists():
        return False
    with open(config_path) as f:
        return _config_disables_mtp(json.load(f))


# Preformatted display string for error/help messages
SUPPORTED_HF_ARCHITECTURES_DISPLAY = " or ".join(f"'{s}'" for s in SUPPORTED_HF_ARCHITECTURES)


def _drop_readonly_config_properties(
    config_dict: dict[str, object], config_type: Type[PretrainedConfig]
) -> dict[str, object]:
    """Remove read-only property names before constructing a HuggingFace config."""
    readonly_properties = {
        name
        for cls in config_type.mro()
        for name, attr in vars(cls).items()
        if isinstance(attr, property) and attr.fset is None
    }
    if not readonly_properties:
        return config_dict
    return {key: value for key, value in config_dict.items() if key not in readonly_properties}


class AutoBridge(Generic[MegatronModelT]):
    """
    Automatically select and instantiate the appropriate bridge for a model.

    This unified bridge class combines automatic model detection with full bridge
    functionality for converting models between HuggingFace and Megatron formats.
    It handles the conversion of causal language models (e.g., GPT, Llama, Phi)
    between HuggingFace's transformers library format and Megatron-Core's distributed
    training format. It manages weight mapping, tensor parallelism distribution, and
    configuration translation.

    The bridge supports both directions of conversion:
    - HuggingFace → Megatron: For training or inference with Megatron
    - Megatron → HuggingFace: For saving trained models in HF format

    Args:
        hf_pretrained: Either a PreTrainedCausalLM instance with loaded model,
            or a PretrainedConfig for configuration-only operations

    Example:
        >>> # Load and convert a model to Megatron format
        >>> bridge = AutoBridge.from_hf_pretrained("meta-llama/Meta-Llama-3-8B")
        >>> provider = bridge.to_megatron_provider()
        >>> megatron_model = provider.provide_distributed_model(wrap_with_ddp=False)

        >>> # Export a Megatron model back to HuggingFace format
        >>> bridge.save_hf_pretrained(megatron_model, "./exported_model")

        >>> # Convert weights with custom settings
        >>> for name, weight in bridge.export_hf_weights(
        ...     megatron_model,
        ...     cpu=True
        ... ):
        ...     print(f"Exported {name}: {weight.shape}")

        >>> # Check if a model is supported before loading
        >>> if AutoBridge.can_handle("microsoft/phi-2"):
        ...     bridge = AutoBridge.from_hf_pretrained("microsoft/phi-2")

    Note:
        The bridge automatically detects the model architecture and applies
        the appropriate weight mappings. Custom architectures require implementing
        a MegatronModelBridge subclass.
    """

    def __init__(self, hf_pretrained: PreTrainedCausalLM | PretrainedConfig):
        if not isinstance(hf_pretrained, (PreTrainedCausalLM, PretrainedConfig)):
            raise ValueError("hf_pretrained must be a PreTrainedCausalLM or PretrainedConfig instance")
        self.hf_pretrained: PreTrainedCausalLM | PretrainedConfig = hf_pretrained
        # Data type for exporting weights
        self.export_weight_dtype: Literal["bf16", "fp16", "fp8"] = "bf16"
        self.hf_model_id: Optional[str] = None
        trust_remote_code = getattr(hf_pretrained, "trust_remote_code", False)
        self.trust_remote_code = trust_remote_code if isinstance(trust_remote_code, bool) else False

    @classmethod
    def list_supported_models(cls) -> list[str]:
        """
        List all model architectures currently supported by the bridge system.

        Returns:
            List of supported HuggingFace model architecture names
        """
        # Get all registered implementations from the dispatch system
        supported = []

        # Access the dispatch registry to find all registered types

        if hasattr(model_bridge.get_model_bridge, "_exact_types"):
            for arch_type in model_bridge.get_model_bridge._exact_types.keys():
                # Support both type and string registrations
                if isinstance(arch_type, str):
                    supported.append(arch_type)
                elif hasattr(arch_type, "__name__"):
                    supported.append(arch_type.__name__)

        return sorted(supported)

    @classmethod
    def supports(cls, config: Any) -> bool:
        """
        Check if this bridge supports the given model configuration.

        A model is supported if it has at least one architecture ending with one of the
        suffixes listed in SUPPORTED_HF_ARCHITECTURES.

        Args:
            config: HuggingFace model config object

        Returns:
            True if this bridge can handle the model, False otherwise
        """
        architectures = getattr(config, "architectures", [])
        if not architectures:
            return False
        return any(arch.endswith(SUPPORTED_HF_ARCHITECTURES) for arch in architectures)

    @classmethod
    def from_auto_config(cls, megatron_path: str, hf_model_id: str, trust_remote_code: bool = False) -> "AutoBridge":
        """
        Create a config-only AutoBridge by synthesizing an HF config from a Megatron checkpoint.

        This method creates a bridge instace from a Megatron checkpoint and reference hf_model_id,
        without loading any weights. This enables exporting of:
        - Custom small models of popular architectures
        - Models pruned from a larger teacher model

        Args:
            megatron_path: Directory path where the Megatron checkpoint is stored
            hf_model_id: HuggingFace model ID or path to model directory
                Examples: "meta-llama/Meta-Llama-3-8B", "./my_model"
            trust_remote_code: Whether to trust remote code when loading config.
                Defaults to False for security. Set to True only for models that
                require custom modeling code from the repository.

        Returns:
            AutoBridge: Bridge instance configured for the architecture

        Raises:
            FileNotFoundError: If run_config.yaml is not found in the Megatron path
        """
        from transformers import AutoConfig

        from megatron.bridge.models.conversion.utils import conform_config_to_reference
        from megatron.bridge.training.model_load_save import load_model_config

        checkpoint_path = Path(megatron_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Megatron checkpoint not found: {checkpoint_path}")

        # Look for configuration files to determine the model type
        run_config = checkpoint_path / "run_config.yaml"
        if not run_config.exists():
            iter_dirs = [d for d in checkpoint_path.iterdir() if d.is_dir() and d.name.startswith("iter_")]
            if iter_dirs:
                latest_iter = max(iter_dirs, key=lambda d: int(d.name.replace("iter_", "")))
                run_config = latest_iter / "run_config.yaml"

        if not run_config.exists():
            raise FileNotFoundError(
                f"Could not find run_config.yaml in {checkpoint_path}. Ensure this is a valid Megatron checkpoint."
            )

        # 1. Load config from both sides
        megatron_cfg, _ = load_model_config(str(run_config.parent))
        if trust_remote_code:
            logger.warning(
                "Loading a model with trust_remote_code=True allows arbitrary code execution "
                "from the model repository. Only use this with models you trust."
            )
        hf_cfg = AutoConfig.from_pretrained(hf_model_id, trust_remote_code=trust_remote_code)
        # 2. Translate Megatron config -> HF, conforming to reference config
        bridge = cls.from_hf_config(hf_cfg)
        megatron_hf_cfg_dict = bridge._model_bridge.megatron_to_hf_config(megatron_cfg)
        megatron_hf_cfg_dict = conform_config_to_reference(megatron_hf_cfg_dict, hf_cfg.to_dict())
        megatron_hf_cfg_dict = _drop_readonly_config_properties(megatron_hf_cfg_dict, type(hf_cfg))
        # 3. Build final bridge from the synthesized config
        synthesized_config = type(hf_cfg)(**megatron_hf_cfg_dict)
        bridge = cls.from_hf_config(synthesized_config)
        bridge.hf_model_id = hf_model_id
        bridge.trust_remote_code = trust_remote_code

        return bridge

    @classmethod
    def from_hf_config(cls, config: PretrainedConfig) -> "AutoBridge":
        """
        Create an AutoBridge from a HuggingFace configuration.

        This method creates a bridge instance from just a model configuration,
        without loading any weights. This is useful for:
        - Creating Megatron models with random initialization
        - Working with model architectures without downloading weights
        - Testing and development scenarios

        Args:
            config: HuggingFace PretrainedConfig instance containing model
                architecture information

        Returns:
            AutoBridge: Bridge instance configured for the architecture

        Raises:
            ValueError: If the configuration is not for a supported CausalLM model

        Example:
            >>> from transformers import AutoConfig
            >>>
            >>> # Load just the configuration
            >>> config = AutoConfig.from_pretrained("meta-llama/Meta-Llama-3-8B")
            >>>
            >>> # Create bridge from config (no weights)
            >>> bridge = AutoBridge.from_hf_config(config)
            >>>
            >>> # Create Megatron model with random initialization
            >>> provider = bridge.to_megatron_provider(load_weights=False)
            >>> model = provider.provide_distributed_model(wrap_with_ddp=False)

            >>> # Or use for architecture exploration
            >>> transformer_config = bridge.transformer_config
            >>> print(f"Hidden size: {transformer_config.hidden_size}")
            >>> print(f"Num layers: {transformer_config.num_layers}")

        See Also:
            from_hf_pretrained: Create bridge with loaded weights
            transformer_config: Access the Megatron TransformerConfig
        """
        cls._validate_config(config)
        return cls(config)

    @classmethod
    def from_hf_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoBridge":
        """
        Load an AutoBridge from a pretrained model, automatically detecting the model type.

        This method loads a model from HuggingFace Hub or a local directory and
        creates a bridge instance ready for conversion operations. The model
        architecture is validated to ensure compatibility.

        Args:
            path: HuggingFace model ID or path to model directory
                Examples: "meta-llama/Meta-Llama-3-8B", "./my_model"
            **kwargs: Additional arguments passed to HuggingFace from_hf_pretrained
                Common options include:
                - torch_dtype: Model precision (torch.float16, torch.bfloat16)
                - device_map: Device placement strategy ("auto", "cuda:0", etc.)
                - trust_remote_code: Allow custom model code execution
                - attn_implementation: Attention implementation ("flash_attention_2", etc.)

        Returns:
            AutoBridge: Bridge instance with loaded model

        Raises:
            ValueError: If the model architecture is not supported

        Example:
            >>> # Basic loading
            >>> bridge = AutoBridge.from_hf_pretrained("gpt2")

            >>> # Load with specific settings
            >>> bridge = AutoBridge.from_hf_pretrained(
            ...     "meta-llama/Meta-Llama-3-8B",
            ...     torch_dtype=torch.float16,
            ...     device_map="auto"
            ... )

            >>> # Works with local paths too
            >>> bridge = AutoBridge.from_hf_pretrained("/path/to/model")
        """
        # First load just the config to check architecture support
        # Use thread-safe config loading to prevent race conditions
        config_kwargs = dict(kwargs)
        trust_remote_code = bool(config_kwargs.pop("trust_remote_code", False))
        if trust_remote_code:
            logger.warning(
                "Loading a model with trust_remote_code=True allows arbitrary code execution "
                "from the model repository. Only use this with models you trust."
            )
        config = safe_load_config_with_retry(path, trust_remote_code=trust_remote_code, **config_kwargs)

        cls._validate_config(config, str(path))

        # Transformers 5.0+ changed `rope_scaling` to a property whose setter
        # does `self.rope_parameters = value`, replacing the entire dict and
        # dropping any fields (e.g. `rope_theta`) that were set during initial
        # construction.  When a `rope_scaling` override is passed as a kwarg,
        # `PretrainedConfig.from_dict` applies it via `setattr` *after* the
        # initial construction, so those fields are silently lost and Megatron
        # falls back to defaults (e.g. `rotary_base=10000`).  Pre-populate the
        # override dict with all base-config rope fields so the setter
        # preserves them.
        if "rope_scaling" in kwargs and isinstance(kwargs["rope_scaling"], dict):
            base_rope = getattr(config, "rope_scaling", None)
            if isinstance(base_rope, dict):
                for key, value in base_rope.items():
                    if key not in kwargs["rope_scaling"]:
                        kwargs["rope_scaling"][key] = value

        try:
            return cls(PreTrainedCausalLM.from_pretrained(path, **kwargs))
        except Exception as e:
            raise ValueError(f"Failed to load model with AutoBridge: {e}") from e

    @classmethod
    def can_handle(cls, path: Union[str, Path], trust_remote_code: bool = False) -> bool:
        """
        Check if the bridge can handle the model at the given path.

        This method allows you to verify model compatibility before attempting
        to load it, which can be useful for validation or UI feedback.

        Args:
            path: Path to model directory or HuggingFace model ID
                Examples: "meta-llama/Meta-Llama-3-8B", "/models/my_model"
            trust_remote_code: Whether to trust remote code when loading config.
                Set to True for models that use custom modeling code.

        Returns:
            bool: True if the bridge supports the model, False otherwise

        Example:
            >>> # Check if a model is supported
            >>> if AutoBridge.can_handle("meta-llama/Meta-Llama-3-8B"):
            ...     print("Model is supported!")
            ... else:
            ...     print("Model requires a custom bridge implementation")
        """
        try:
            config = safe_load_config_with_retry(path, trust_remote_code=trust_remote_code)
            return cls.supports(config)
        except Exception:
            return False

    def load_hf_weights(
        self,
        model: list[MegatronModelT],
        hf_path: str | Path | None = None,
        allowed_mismatched_params: list[str] | None = None,
    ) -> None:
        """
        Load HuggingFace weights into a Megatron model.

        This method handles the conversion and distribution of weights from
        HuggingFace format to Megatron's distributed format, including proper
        tensor parallel and pipeline parallel distribution.

        Args:
            model: List of Megatron model instances (one per virtual pipeline stage)
            hf_path: Optional path to load weights from. If None, uses weights
                from the bridge's hf_pretrained instance
            allowed_mismatched_params: Optional list of parameter names or patterns
                to allow mismatch (skip instead of raise error).

        Returns:
            The input model with loaded weights

        Raises:
            ValueError: If hf_path is None and bridge was created without weights

        Example:
            >>> # Load weights from bridge's pretrained model
            >>> bridge = AutoBridge.from_hf_pretrained("gpt2")
            >>> megatron_model = create_megatron_model()  # Your model creation
            >>> bridge.load_hf_weights(megatron_model)

            >>> # Load weights from a different checkpoint
            >>> bridge.load_hf_weights(megatron_model, "./finetuned_model")

            >>> # Load weights with allowed mismatched parameters
            >>> bridge.load_hf_weights(
            ...     megatron_model,
            ...     allowed_mismatched_params=["*.bias", "decoder.layers.0.*"]
            ... )
        """
        if hf_path is None:
            if not isinstance(self.hf_pretrained, PreTrainedCausalLM):
                raise ValueError("hf_path is required when hf_pretrained is not a PreTrainedCausalLM instance")
            pre_trained = self.hf_pretrained
        else:
            # Preserve trust_remote_code setting from the original bridge instance
            trust_remote_code = getattr(self.hf_pretrained, "trust_remote_code", False)
            pre_trained = PreTrainedCausalLM.from_pretrained(hf_path, trust_remote_code=trust_remote_code)
        bridge = self._model_bridge
        bridge.load_weights_hf_to_megatron(pre_trained, model, allowed_mismatched_params=allowed_mismatched_params)
        # Get unquantized_state_dict from the bridge instance that was used for optimizer reload
        self.unquantized_state_dict = getattr(bridge, "unquantized_state_dict", None)
        return model

    def export_hf_weights(
        self,
        model: list[MegatronModelT],
        cpu: bool = False,
        show_progress: bool = True,
        conversion_tasks: Optional[List[WeightConversionTask]] = None,
        merge_adapter_weights: bool = True,
    ) -> Iterable["HFWeightTuple"]:
        """
        Export Megatron model weights to HuggingFace format.

        This method yields weight tensors in HuggingFace format, handling the
        gathering of distributed tensors and format conversion. It's useful for
        streaming weight export or custom processing. All ranks get full tensors.

        If the model contains LoRA adapters, they will be automatically merged
        into the base weights before export. This ensures the exported model
        contains the full fine-tuned weights.

        Args:
            model: Megatron model instance or list of instances
            cpu: Whether to move tensors to CPU before yielding
            show_progress: Display progress bar during export
            conversion_tasks (Optional[List[WeightConversionTask]]): Pre-built conversion tasks.
                If not provided, tasks will be built automatically from the models.
                *Please note that this is an advanced feature and should be used with caution.
                The tasks needs to be built with the `get_conversion_tasks` method first and
                carefully adjust based on your needs.*
            merge_adapter_weights: Whether to gather and merge LoRA adapter weights into the base
                tensors during export (defaults to True). Set to False to export only the base tensors.


        Yields:
            HFWeightTuple: Named tuples of (param_name, weight_tensor)

        Example:
            >>> # Export and process weights
            >>> for name, weight in bridge.export_hf_weights(model):
            ...     print(f"{name}: {weight.shape}")

            >>> # Export with specific settings
            >>> weights = list(bridge.export_hf_weights(
            ...     model,
            ...     cpu=True
            ... ))
        """
        # Build conversion tasks based on export_weight_dtype configuration
        if conversion_tasks is None and self.export_weight_dtype == "fp8":
            if not isinstance(model, list):
                model = [model]
            self._validate_fp8_export_config(model)
            # Use FP8 export tasks for blockwise FP8 weights
            conversion_tasks = self._model_bridge.build_export_fp8_tasks(self.hf_pretrained, model)

        bridge = self._model_bridge
        return bridge.stream_weights_megatron_to_hf(
            model,
            self.hf_pretrained,
            cpu=cpu,
            show_progress=show_progress,
            conversion_tasks=conversion_tasks,
            merge_adapter_weights=merge_adapter_weights,
        )

    def export_hf_weights_quant(
        self,
        model: list[MegatronModelT],
        quantization_checker: Callable[[str], bool],
        quant_fn: Callable[..., Tuple[torch.Tensor, torch.Tensor]],
        quant_block_size: Optional[Tuple[int, int]] = None,
        cpu: bool = False,
        show_progress: bool = True,
        conversion_tasks: Optional[List[WeightConversionTask]] = None,
        merge_adapter_weights: bool = False,
    ) -> Iterable["HFWeightTuple"]:
        """
        Export Megatron model weights to HuggingFace format with quantization.
        """
        dispatch_instance = (self._causal_lm_architecture, self._get_model_instance(model))
        return model_bridge.stream_weights_megatron_to_hf_quant(
            dispatch_instance,
            model,
            self.hf_pretrained,
            quantization_checker,
            quant_fn,
            quant_block_size=quant_block_size,
            cpu=cpu,
            show_progress=show_progress,
            conversion_tasks=conversion_tasks,
            merge_adapter_weights=merge_adapter_weights,
        )

    def export_adapter_weights(
        self,
        model: list[MegatronModelT],
        cpu: bool = True,
        show_progress: bool = True,
    ) -> Iterable["HFWeightTuple"]:
        """
        Export only adapter weights from a Megatron model without merging them into base tensors.

        This is useful when you want to save or inspect LoRA adapters independently from the
        underlying pretrained weights.

        Args:
            model: Megatron model instance or list of instances
            cpu: Whether to move tensors to CPU before yielding
            show_progress: Display progress bar during export

        Yields:
            HFWeightTuple: Named tuples of (param_name, weight_tensor) for adapter parameters
        """
        bridge = self._model_bridge
        return bridge.stream_adapter_weights_megatron_to_hf(model, cpu=cpu, show_progress=show_progress)

    def save_hf_adapter(
        self,
        model: list[MegatronModelT],
        path: str | Path,
        peft_config: "PEFT",
        base_model_name_or_path: Optional[str] = None,
        show_progress: bool = True,
    ) -> None:
        """Save LoRA adapter weights as a HuggingFace PEFT-compatible directory.

        The output directory contains ``adapter_config.json`` and
        ``adapter_model.safetensors`` and can be loaded directly with
        ``peft.PeftModel.from_pretrained(base_model, path)``.

        Args:
            model: Megatron model instance or list of instances.
            path: Directory path where the adapter files will be saved.
            peft_config: The LoRA / DoRA config used during training (provides
                ``dim``, ``alpha``, ``dropout``, etc.).
            base_model_name_or_path: HuggingFace model identifier or local path
                of the base model this adapter was trained on.  If *None*, the
                value is inferred from ``hf_pretrained.model_name_or_path``.
            show_progress: Display progress bar during export.

        Example:
            >>> bridge.save_hf_adapter(
            ...     megatron_model,
            ...     "./my-lora-adapter",
            ...     peft_config=lora,
            ...     base_model_name_or_path="Qwen/Qwen3-4B",
            ... )
            >>> # Load with HuggingFace PEFT
            >>> from peft import PeftModel
            >>> from transformers import AutoModelForCausalLM
            >>> base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B")
            >>> model = PeftModel.from_pretrained(base, "./my-lora-adapter")

        Note:
            This method is collective -- all ranks must call it.  Only rank 0
            writes files to disk; the other ranks participate in the generator
            to gather distributed (TP/PP/EP) tensors.
        """
        import json

        from safetensors.torch import save_file

        from megatron.bridge.models.conversion.peft_bridge import (
            build_adapter_config_dict,
            convert_adapter_weights_to_peft_state,
            infer_rank_pattern_from_adapter_weights,
            infer_target_modules_from_adapter_weights,
        )

        if dist.is_initialized():
            dist.barrier()

        raw_adapter_weights = [
            HFWeightTuple(exported_weight.param_name, exported_weight.weight.clone().float())
            for exported_weight in self.export_adapter_weights(model, cpu=True, show_progress=show_progress)
        ]
        if not raw_adapter_weights:
            raise RuntimeError(
                "No adapter weights were found on the model. "
                "Ensure the model has PEFT adapters applied before calling save_hf_adapter()."
            )
        adapter_state, module_adapter_keys, target_parameters = convert_adapter_weights_to_peft_state(
            raw_adapter_weights,
        )
        rank_pattern = infer_rank_pattern_from_adapter_weights(
            raw_adapter_weights,
            default_rank=getattr(peft_config, "dim", 32),
        )

        is_rank0 = not dist.is_initialized() or dist.get_rank() == 0
        if is_rank0:
            save_dir = Path(path)
            save_dir.mkdir(parents=True, exist_ok=True)

            if base_model_name_or_path is None:
                base_model_name_or_path = str(
                    getattr(self.hf_pretrained, "model_name_or_path", "")
                    or getattr(self.hf_pretrained, "name_or_path", "")
                )

            target_modules = infer_target_modules_from_adapter_weights(module_adapter_keys)
            adapter_config = build_adapter_config_dict(
                peft_config,
                target_modules=target_modules,
                target_parameters=target_parameters,
                base_model_name_or_path=base_model_name_or_path,
                rank_pattern=rank_pattern,
            )

            config_path = save_dir / "adapter_config.json"
            with open(config_path, "w") as f:
                json.dump(adapter_config, f, indent=2)

            weights_path = save_dir / "adapter_model.safetensors"
            save_file(adapter_state, str(weights_path))

        if dist.is_initialized():
            dist.barrier()

    def save_hf_pretrained(
        self,
        model: list[MegatronModelT],
        path: str | Path,
        show_progress: bool = True,
        source_path: Optional[Union[str, Path]] = None,
        strict: bool = True,
        merge_adapter_weights: bool = True,
        distributed_save: bool = False,
        save_every_n_ranks: int = 1,
    ) -> None:
        """
        Save a Megatron model in HuggingFace format.

        This method exports the complete model including configuration, tokenizer,
        and weights to a directory that can be loaded with HuggingFace's
        from_pretrained methods.

        If the model contains LoRA adapters, they will be automatically merged
        into the base weights before saving. This ensures the saved model
        contains the full fine-tuned weights.

        If the original model was loaded with trust_remote_code=True, any custom
        modeling files (e.g., modeling_*.py, configuration_*.py) will be preserved
        to ensure the saved model can be loaded properly.

        Config-only bridges are supported when created via the auto-config
        flow in convert_checkpoints.py.

        Args:
            model: Megatron model instance or list of instances
            path: Directory path to save the model
            show_progress: Display progress bar during weight export
            source_path: Path to the directory containing custom modeling files to be preserved.
                This is useful when converting from Megatron checkpoints where the original
                HuggingFace model with custom modeling files needs to be referenced. If not specified,
                the path will be automatically determined from the HuggingFace configuration.
            strict: Whether to perform strict validation during weight export
            merge_adapter_weights: Whether to gather/merge LoRA adapter weights into base tensors during export.
            distributed_save: Whether to enable distributed saving mode where each rank saves
                part of weights independently. When False (default), only rank 0 performs
                the save operation after gathering weights from all ranks.
            save_every_n_ranks: Interval for saving weights across ranks in distributed mode.
                For example, if set to 2, only ranks 0, 2, 4, ... will save weights.
                This is useful for reducing I/O pressure when dealing with large-scale distributed
                training. Only effective when distributed_save=True. Default is 1 (all ranks save).

        Example:
            >>> # Save model after training
            >>> bridge.save_hf_pretrained(megatron_model, "./my_finetuned_model")

            >>> # Load the saved model with HuggingFace
            >>> from transformers import AutoModelForCausalLM
            >>> hf_model = AutoModelForCausalLM.from_pretrained("./my_finetuned_model")

        Note:
            This method is collective - all ranks must call it. Only rank 0
            saves the configuration files, while weight saving is coordinated
            across all ranks.
        """
        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 optionally download modeling files from Hub.
                Path(path).mkdir(parents=True, exist_ok=True)
                config_dict = self.hf_pretrained.to_dict()
                if not self.trust_remote_code:
                    config_dict.pop("auto_map", None)
                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 only when the export is meant to preserve remote code.
                hub_repo = self.hf_model_id
                if hub_repo and self.trust_remote_code:
                    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.warning(
                            "Could not download modeling files from %s: %s. "
                            "This is expected for models that use standard transformers "
                            "modeling classes and do not define custom .py files.",
                            hub_repo,
                            exc,
                        )
                    finally:
                        PreTrainedBase._cleanup_hf_local_dir_cache(Path(path))

            else:
                # Get bridge-level ADDITIONAL_FILE_PATTERNS if configured
                additional_files = getattr(self._model_bridge, "ADDITIONAL_FILE_PATTERNS", None) or None
                self.hf_pretrained.save_artifacts(
                    path, original_source_path=source_path, additional_files=additional_files
                )

        if dist.is_initialized():
            if dist.get_rank() == 0:
                _save_artifacts()
        else:
            _save_artifacts()

        self.save_hf_weights(
            model,
            path,
            show_progress,
            strict,
            merge_adapter_weights=merge_adapter_weights,
            distributed_save=distributed_save,
            save_every_n_ranks=save_every_n_ranks,
        )

    def save_hf_weights(
        self,
        model: list[MegatronModelT],
        path: str | Path,
        show_progress: bool = True,
        strict: bool = True,
        merge_adapter_weights: bool = True,
        distributed_save: bool = False,
        save_every_n_ranks: int = 1,
    ) -> None:
        """
        Save Megatron model weights in HuggingFace safetensors format.

        This method exports only the model weights (not configuration or tokenizer)
        to safetensors files compatible with HuggingFace. It uses streaming save
        to handle large models efficiently without requiring all weights in memory
        at once.

        If the model contains LoRA adapters, they will be automatically merged
        into the base weights before saving. This ensures the saved weights
        contain the full fine-tuned parameters.

        The weights are gathered from distributed ranks and saved in the standard
        HuggingFace sharded format when the model is large.

        Args:
            model: Megatron model instance or list of instances
            path: Directory path where weight files will be saved
            show_progress: Display progress bar during export
            merge_adapter_weights: Whether to gather/merge LoRA adapter weights into base tensors during export.
            distributed_save: Whether to enable distributed saving mode where each rank saves
                part of weights independently.
            save_every_n_ranks: Interval for saving weights across ranks in distributed mode.
                For example, if set to 2, only ranks 0, 2, 4, ... will save weights.

        Raises:
            ValueError: If the state source doesn't support streaming save

        Example:
            >>> # Save just the weights
            >>> bridge.save_hf_weights(megatron_model, "./model_weights")

            >>> # Save without progress bar (useful in scripts)
            >>> bridge.save_hf_weights(megatron_model, "./weights", show_progress=False)

        Note:
            - This method is collective and must be called by all ranks
            - Uses safetensors format for efficient loading and security
            - Automatically handles model sharding for large models
            - The saved weights can be loaded with HuggingFace's from_pretrained
        """
        is_distributed = dist.is_initialized()
        if is_distributed:
            dist.barrier()
        bridge = self._model_bridge
        generator = bridge.stream_weights_megatron_to_hf(
            model,
            self.hf_pretrained,
            cpu=True,
            show_progress=show_progress,
            merge_adapter_weights=merge_adapter_weights,
        )
        model_instance = self._get_model_instance(model)
        quant_tensors = None
        if is_quantized(model_instance):
            quant_tensors = {}

            def _filter_quant(gen):
                for name, tensor in gen:
                    if "_quantizer." in name:
                        quant_tensors[name] = tensor
                        continue
                    yield name, tensor

            generator = _filter_quant(generator)

        # Check if the state source is SafeTensorsStateSource for streaming save.
        if (
            hasattr(self.hf_pretrained, "state")
            and hasattr(self.hf_pretrained.state, "source")
            and isinstance(self.hf_pretrained.state.source, SafeTensorsStateSource)
        ):
            source = self.hf_pretrained.state.source
            model_config = getattr(model_instance, "config", None)
            hf_config = getattr(self.hf_pretrained, "config", self.hf_pretrained)
            mtp_disabled = _saved_config_disables_mtp(path) or any(
                _config_disables_mtp(config) for config in (hf_config, model_config)
            )
            ignored_source_key_prefixes = ("mtp.",) if mtp_disabled and source.has_glob("mtp.*") else None
            source.save_generator(
                generator,
                path,
                strict=strict,
                distributed_save=distributed_save,
                save_every_n_ranks=save_every_n_ranks,
                ignored_source_key_prefixes=ignored_source_key_prefixes,
            )
        else:
            # 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.
            rank = dist.get_rank() if is_distributed else 0

            if rank == 0:
                state_dict = {name: tensor.contiguous().cpu() for name, tensor in generator}
            else:
                for _ in generator:
                    pass
                state_dict = None

            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:
            rank = dist.get_rank() if is_distributed else 0
            if rank == 0:
                sidecar_path = Path(path) / "modelopt_weights.pt"
                sidecar_path.parent.mkdir(parents=True, exist_ok=True)
                torch.save(quant_tensors, sidecar_path)

        if is_distributed:
            dist.barrier()

    def save_megatron_model(
        self,
        model: list[MegatronModule],
        path: str | Path,
        hf_tokenizer_path: Optional[str | Path] = None,
        low_memory_save: bool = False,
        hf_tokenizer_kwargs: Optional[dict] = None,
    ) -> None:
        """
        Save a Megatron model in native Megatron checkpoint format without optimizer
        state.

        This method saves the model in Megatron's native checkpoint format, which
        can be loaded directly by Megatron for training or inference. The checkpoint
        includes the model configuration and weights, NO optimizer state or other
        artifacts.

        Args:
            model: Megatron model instance or list of instances
            path: Directory path where the checkpoint will be saved
            hf_tokenizer_path: Optional HuggingFace model ID or path for tokenizer metadata.
                If provided, the tokenizer metadata will be included in the checkpoint.
            low_memory_save: If True, uses a memory-optimized save flow that reduces
                peak memory by ~50% for models with merged weights (e.g., gate+up
                projections). The model is deleted after state dict generation and
                cannot be used afterward. Default is False, preserving the model
                for further use.
            hf_tokenizer_kwargs: Optional dictionary of kwargs to pass to the HuggingFace tokenizer.
                Common options include trust_remote_code=True for models with custom tokenizers,
                or use_fast=True for models that require the fast tokenizer.

        Example:
            >>> # Save model checkpoint after conversion
            >>> bridge.save_megatron_model(megatron_model, "./megatron_checkpoint")

            >>> # Save model checkpoint with tokenizer metadata
            >>> bridge.save_megatron_model(
            ...     megatron_model,
            ...     "./megatron_checkpoint",
            ...     hf_tokenizer_path="meta-llama/Meta-Llama-3-8B"
            ... )

            >>> # Low-memory save (destroys model after save)
            >>> bridge.save_megatron_model(
            ...     megatron_model,
            ...     "./megatron_checkpoint",
            ...     low_memory_save=True
            ... )

        Note:
            - This method is collective and must be called by all ranks
            - The saved checkpoint can be loaded with Megatron's checkpoint loading utilities
            - The checkpoint format follows Megatron's standard structure for compatibility
            - When low_memory_save=True, the model is deleted and cannot be used afterward
        """
        try:
            from megatron.bridge.training.model_load_save import save_megatron_model
        except ImportError:
            raise ImportError("megatron.bridge.training is not available.")
        save_megatron_model(
            model,
            path,
            hf_tokenizer_path=hf_tokenizer_path,
            low_memory_save=low_memory_save,
            hf_tokenizer_kwargs=hf_tokenizer_kwargs,
        )

    def load_megatron_model(
        self, path: str | Path, *, mp_overrides: ModelParallelKwargs | None = None, **kwargs: Unpack[GetModelKwargs]
    ) -> list[MegatronModelT]:
        """
        Load a Megatron model from a native Megatron checkpoint.

        This method loads a model from a Megatron checkpoint that was saved using
        the save_megatron_model method. It reads the checkpoint configuration,
        creates the appropriate model provider, and loads the weights.

        Args:
            path: Directory path where the Megatron checkpoint is stored
            mp_overrides: Optional model-parallel overrides to apply to the loaded config.
            **kwargs: Additional arguments passed to the model provider

        Returns:
            List of Megatron model instances loaded from the checkpoint

        Example:
            >>> # Load a previously saved Megatron model
            >>> bridge = AutoBridge.from_hf_config(config)
            >>> model = bridge.load_megatron_model("./megatron_checkpoint")

            >>> # Load and specify model configuration
            >>> model = bridge.load_megatron_model(
            ...     "./megatron_checkpoint",
            ...     wrap_with_ddp=False
            ... )

        Note:
            - This method is collective and must be called by all ranks
            - The checkpoint must have been saved with save_megatron_model
            - The model architecture must match the bridge configuration
        """
        try:
            from megatron.bridge.training.model_load_save import load_megatron_model
        except ImportError:
            raise ImportError("megatron.bridge.training is not available.")

        if self.trust_remote_code:
            from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix

            register_allowed_target_prefix("transformers_modules.")

        checkpoint_path = Path(path)

        # Check for iter_* folders
        iter_folders = [f for f in checkpoint_path.iterdir() if f.is_dir() and f.name.startswith("iter_")]

        if iter_folders:
            # Find the folder with the largest iteration number
            def get_iter_number(folder_name):
                try:
                    return int(folder_name.replace("iter_", ""))
                except ValueError:
                    return -1  # Invalid format, put at the end

            latest_iter = max(iter_folders, key=lambda f: get_iter_number(f.name))
            checkpoint_path = checkpoint_path / latest_iter.name
        # else: checkpoint_path remains as the input path (no iter folders found)

        skip_temp_dist_context = dist.is_initialized()
        # Load the state dict
        model = load_megatron_model(
            str(checkpoint_path),
            use_cpu_init=(skip_temp_dist_context and dist.get_backend() == "gloo"),
            skip_temp_dist_context=skip_temp_dist_context,
            mp_overrides=mp_overrides,
        )
        return model if isinstance(model, list) else [model]

    @classmethod
    def import_ckpt(
        cls,
        hf_model_id: str | Path,
        megatron_path: str | Path,
        **kwargs,
    ) -> None:
        """
        Import a HuggingFace model and save it as a Megatron checkpoint.

        This is a convenience method that combines loading a HuggingFace model,
        converting it to Megatron format, and saving it as a native Megatron
        checkpoint. This is useful for preparing models for Megatron training
        or creating Megatron checkpoints from pretrained HuggingFace models.

        Args:
            hf_model_id: HuggingFace model ID or path to model directory
                Examples: "meta-llama/Meta-Llama-3-8B", "./my_model"
            megatron_path: Directory path where the Megatron checkpoint will be saved
            **kwargs: Additional arguments passed to from_hf_pretrained
                Common options include:
                - torch_dtype: Model precision (torch.float16, torch.bfloat16)
                - device_map: Device placement strategy ("auto", "cuda:0", etc.)
                - trust_remote_code: Allow custom model code execution
                - attn_implementation: Attention implementation ("flash_attention_2", etc.)

        Example:
            >>> # Basic import
            >>> AutoBridge.import_ckpt(
            ...     "meta-llama/Meta-Llama-3-8B",
            ...     "./megatron_checkpoints/llama3_8b"
            ... )

            >>> # Import with specific settings
            >>> AutoBridge.import_ckpt(
            ...     "meta-llama/Meta-Llama-3-8B",
            ...     "./megatron_checkpoints/llama3_8b",
            ...     torch_dtype=torch.float16,
            ...     device_map="auto"
            ... )
        """
        # Load the HuggingFace model
        bridge = cls.from_hf_pretrained(hf_model_id, **kwargs)

        # Convert to Megatron model
        megatron_model = bridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True)

        # Save as Megatron checkpoint
        hf_tokenizer_kwargs = {}
        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,
            hf_tokenizer_path=hf_model_id,
            hf_tokenizer_kwargs=hf_tokenizer_kwargs,
            low_memory_save=True,
        )

    def export_ckpt(
        self,
        megatron_path: str | Path,
        hf_path: str | Path,
        show_progress: bool = True,
        strict: bool = False,
        source_path: Optional[Union[str, Path]] = None,
    ) -> None:
        """
        Export a Megatron checkpoint to HuggingFace format.

        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.

        Also supports config-only bridges created via the auto-config
        flow in convert_checkpoints.py.

        Args:
            megatron_path: Directory path where the Megatron checkpoint is stored
            hf_path: Directory path where the HuggingFace model will be saved
            show_progress: Display progress bar during weight export
            strict: Whether to perform strict validation during weight export
            source_path: Path to the directory containing custom modeling files to be preserved.
                This is useful when converting from Megatron checkpoints where the original
                HuggingFace model with custom modeling files needs to be referenced. If not specified,
                the path will be automatically determined from the HuggingFace configuration.

        Example:
            >>> # Basic export
            >>> bridge = AutoBridge.from_hf_pretrained("meta-llama/Meta-Llama-3-8B")
            >>> bridge.export_ckpt(
            ...     "./megatron_checkpoints/my_model",
            ...     "./hf_exports/my_model"
            ... )
            >>> # Export with specific settings
            >>> bridge.export_ckpt(
            ...     "./megatron_checkpoints/my_model",
            ...     "./hf_exports/my_model",
            ...     show_progress=False
            ... )

            >>> # Load the exported model with HuggingFace
            >>> from transformers import AutoModelForCausalLM
            >>> hf_model = AutoModelForCausalLM.from_pretrained("./hf_exports/my_model")
        """
        try:
            from megatron.bridge.training.model_load_save import temporary_distributed_context
        except ImportError:
            raise ImportError("megatron.bridge.training is not available.")

        # Export ckpt performs on CPU
        with temporary_distributed_context(backend="gloo"):
            # Load the Megatron model
            megatron_model = self.load_megatron_model(megatron_path, wrap_with_ddp=False)

            # Save in HuggingFace format
            self.save_hf_pretrained(
                megatron_model,
                hf_path,
                show_progress=show_progress,
                source_path=source_path,
                strict=strict,
            )

    def export_adapter_ckpt(
        self,
        peft_checkpoint: str | Path,
        output_path: str | Path,
        show_progress: bool = True,
    ) -> None:
        """Export LoRA adapter weights from a Megatron PEFT checkpoint to HuggingFace PEFT format.

        This convenience method loads a Megatron-Bridge fine-tuning checkpoint,
        reconstructs the LoRA adapter structure from ``run_config.yaml``, and
        writes a HuggingFace PEFT-compatible directory containing
        ``adapter_config.json`` and ``adapter_model.safetensors``.

        The bridge must be created with ``from_hf_pretrained`` so that
        base weights are available for the conversion mapping.

        Args:
            peft_checkpoint: Path to the Megatron-Bridge distributed checkpoint
                that contains LoRA adapter weights (produced by a PEFT
                fine-tuning run).  May point at the top-level run directory
                (containing ``iter_*`` folders) or directly at an iteration
                directory.
            output_path: Directory where the adapter files will be saved.
            show_progress: Display progress bar during export.

        Example:
            >>> bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B")
            >>> bridge.export_adapter_ckpt(
            ...     "/path/to/finetune_ckpt",
            ...     "./my_adapter",
            ... )
            >>> # Load with HuggingFace PEFT
            >>> from peft import PeftModel
            >>> from transformers import AutoModelForCausalLM
            >>> base = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
            >>> model = PeftModel.from_pretrained(base, "./my_adapter")
        """
        import logging

        from megatron.core import dist_checkpointing

        from megatron.bridge.peft.lora import LoRA, VLMLoRA
        from megatron.bridge.training.checkpointing import (
            _generate_model_state_dict,
            apply_peft_adapter_filter_to_state_dict,
        )
        from megatron.bridge.training.model_load_save import temporary_distributed_context
        from megatron.bridge.training.utils.checkpoint_utils import read_run_config

        _logger = logging.getLogger(__name__)

        ckpt_path = Path(peft_checkpoint).expanduser().resolve()
        if not ckpt_path.exists():
            raise FileNotFoundError(f"PEFT checkpoint not found: {ckpt_path}")
        output_path = Path(output_path).expanduser().resolve()

        peft_class: type = LoRA
        peft_cfg: dict = {}
        cfg_file = ckpt_path / "run_config.yaml"
        if not cfg_file.exists() and ckpt_path.parent != ckpt_path:
            cfg_file = ckpt_path.parent / "run_config.yaml"
        if cfg_file.exists():
            try:
                run_cfg_dict = read_run_config(str(cfg_file))
                peft_cfg = run_cfg_dict.get("peft", {}) or {}
                if "VLMLoRA" in peft_cfg.get("_target_", ""):
                    peft_class = VLMLoRA
                allowed_keys = {
                    "target_modules",
                    "exclude_modules",
                    "dim",
                    "alpha",
                    "dropout",
                    "dropout_position",
                    "normalize_moe_lora",
                    "share_expert_adapters",
                    "freeze_language_model",
                    "freeze_vision_model",
                    "freeze_vision_projection",
                }
                peft_cfg = {k: v for k, v in peft_cfg.items() if k in allowed_keys}
            except Exception as err:
                _logger.warning(f"Failed to read LoRA settings from {cfg_file}: {err}. Using defaults.")
        else:
            _logger.warning("run_config.yaml not found in PEFT checkpoint; using default LoRA settings.")

        lora = peft_class(**peft_cfg)

        # Materialise model with base weights + LoRA structure.
        # Use float32 so adapter weights are exported at full precision;
        # bfloat16 matmul in downstream PEFT merges causes ~1e-3 weight
        # errors that compound into large logit diffs.
        provider = self.to_megatron_provider(load_weights=True)
        provider.pipeline_dtype = torch.float32
        provider.params_dtype = torch.float32
        provider.finalize()
        provider.register_pre_wrap_hook(lambda chunks: lora(chunks, training=False))

        def _load_and_export_adapter(model):
            """Load adapter weights into a materialized model and export them as PEFT."""

            # Load adapter weights from the PEFT checkpoint
            sharded_state_dict = _generate_model_state_dict(model, {})
            sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, lora)
            loaded_sd = dist_checkpointing.load(sharded_state_dict, str(ckpt_path))
            model_key = "model" if "model" in loaded_sd else next(k for k in loaded_sd if k.startswith("model"))
            model[0].load_state_dict(loaded_sd[model_key], strict=False)

            # Export
            base_model_name = str(
                getattr(self.hf_pretrained, "model_name_or_path", "")
                or getattr(self.hf_pretrained, "name_or_path", "")
            )
            self.save_hf_adapter(
                model,
                path=output_path,
                peft_config=lora,
                base_model_name_or_path=base_model_name,
                show_progress=show_progress,
            )

        model_context = (
            nullcontext() if torch.distributed.is_initialized() else temporary_distributed_context(backend="gloo")
        )
        with model_context:
            model = provider.provide_distributed_model(
                wrap_with_ddp=False,
                use_cpu_initialization=True,
                init_model_with_meta_device=False,
            )
            _load_and_export_adapter(model)

    def push_to_hub(self, path: str | Path) -> None: ...

    def to_megatron_model(
        self,
        load_weights: bool = True,
        hf_path: str | Path | None = None,
        **kwargs: Unpack[GetModelKwargs],
    ) -> list[MegatronModelT]:
        provider = self.to_megatron_provider(load_weights, hf_path)

        # Finalize the provider before creating models
        if hasattr(provider, "finalize"):
            provider.finalize()

        return provider.provide_distributed_model(**kwargs)

    def to_megatron_provider(self, load_weights: bool = True, hf_path: str | Path | None = None) -> GPTModelProvider:
        """
        Convert to a Megatron model provider.

        This method creates a GPTModelProvider configured to match the HuggingFace
        model's architecture. The provider can then be used to instantiate
        Megatron models for training or inference.

        Args:
            load_weights: Whether to configure the provider to load weights
                from HuggingFace format. If False, creates model with random
                initialization.
            hf_path: Optional path to load weights from. If None, uses weights
                from the bridge's hf_pretrained instance. Useful for loading
                weights from a different checkpoint.

        Returns:
            GPTModelProvider: A configured model provider ready to create
                Megatron models

        Example:
            >>> # Create provider and model with loaded weights
            >>> bridge = AutoBridge.from_hf_pretrained("meta-llama/Meta-Llama-3-8B")
            >>> provider = bridge.to_megatron_provider()
            >>> model = provider.get_model()

            >>> # Create provider without loading weights (for training from scratch)
            >>> provider = bridge.to_megatron_provider(load_weights=False)
            >>> model = provider.get_model()  # Random initialization

            >>> # Load weights from a different checkpoint
            >>> bridge = AutoBridge.from_hf_config(config)  # Config only
            >>> provider = bridge.to_megatron_provider(hf_path="./finetuned_model")
            >>> model = provider.get_model()  # Loads finetuned weights

        See Also:
            GPTModelProvider: The provider class for creating models
            load_weights: Method to load weights into existing models
        """
        provider_input = self._provider_bridge_input
        provider: ModelProviderMixin = self._model_bridge.provider_bridge(provider_input)

        if load_weights:
            if hf_path is None and not isinstance(self.hf_pretrained, PreTrainedCausalLM):
                raise ValueError(
                    "AutoBridge.from_hf_config() does not include weights. "
                    "Pass load_weights=False for random initialization or provide hf_path to load weights."
                )
            # Skip weights initialization since we are going to load weights
            provider.perform_initialization = False
            if hf_path is None:
                provider.register_pre_wrap_hook(
                    partial(self._model_bridge.load_weights_hf_to_megatron, self.hf_pretrained)
                )
            else:
                # Load from specified path
                trust_remote_code = getattr(self.hf_pretrained, "trust_remote_code", False)
                pre_trained = PreTrainedCausalLM.from_pretrained(hf_path, trust_remote_code=trust_remote_code)
                provider.register_pre_wrap_hook(partial(self._model_bridge.load_weights_hf_to_megatron, pre_trained))

        hf_identifier: str | None = None
        if hf_path is not None:
            hf_identifier = str(hf_path)
        else:
            hf_name_or_path = getattr(self.hf_pretrained, "model_name_or_path", None)
            if hf_name_or_path is None and isinstance(self.hf_pretrained, PretrainedConfig):
                hf_name_or_path = getattr(self.hf_pretrained, "name_or_path", None)
            if hf_name_or_path:
                hf_identifier = str(hf_name_or_path)

        if hf_identifier:
            setattr(provider, "hf_model_id", hf_identifier)

        return provider

    @staticmethod
    def get_hf_model_id_from_checkpoint(path: str | Path) -> str | None:
        """Get the HuggingFace model identifier stored in a checkpoint.

        Args:
            path: Path to a Megatron checkpoint directory, either the root directory
                containing iteration subdirectories or a specific iteration directory.

        Returns:
            The HuggingFace model ID or path recorded in the checkpoint metadata if present,
            otherwise `None`.
        """
        from megatron.bridge.training.utils import checkpoint_utils as _checkpoint_utils

        return _checkpoint_utils.get_hf_model_id_from_checkpoint(path)

    def get_conversion_tasks(
        self,
        megatron_model: Union[MegatronModelT, List[MegatronModelT]],
        hf_path: str | Path | None = None,
    ) -> List["WeightConversionTask"]:
        """Get the conversion tasks for weight conversion between HuggingFace and Megatron formats.

        This method returns the planned conversion tasks that would be executed during
        weight conversion in either direction. Each task contains information about parameter
        mappings, source and target parameters, and the conversion logic required.

        The tasks can be used for both HF→Megatron and Megatron→HF conversions since they
        contain bidirectional mapping information.

        Args:
            megatron_model: Megatron model instance or list of instances (one per
                virtual pipeline stage) that participate in the conversion.
            hf_path: Optional path to load HF weights from. If None, uses weights
                from the bridge's hf_pretrained instance.

        Returns:
            List[WeightConversionTask]: List of conversion tasks that would be executed.
                Each task contains:
                - param_name: Megatron parameter name
                - mapping: The parameter mapping object handling the conversion
                - pp_rank: Pipeline parallel rank that owns the parameter
                - vp_stage: Virtual pipeline stage index
                - megatron_module: Reference to the Megatron module owning the parameter
                - param_weight: The actual parameter tensor

        Example:
            >>> bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B")
            >>> megatron_model = bridge.to_megatron_model(load_weights=False, wrap_with_ddp=False)
            >>> tasks = bridge.get_conversion_tasks(megatron_model)
            >>>
            >>> for task in tasks:
            ...     # For HF→Megatron direction
            ...     print(f"HF param {task.mapping.hf_param} -> Megatron param {task.param_name}")
            ...
            ...     # For Megatron→HF direction
            ...     hf_params = task.mapping.hf_param
            ...     if isinstance(hf_params, str):
            ...         print(f"Megatron param {task.param_name} -> HF param {hf_params}")
            ...     else:
            ...         print(f"Megatron param {task.param_name} -> HF params {list(hf_params.values())}")
            ...
            ...     print(f"  Mapping type: {type(task.mapping).__name__}")
            ...     print(f"  PP rank: {task.pp_rank}, VP stage: {task.vp_stage}")

        Note:
            This method is useful for:
            - Debugging weight conversion issues in both directions
            - Understanding parameter mappings between formats
            - Custom weight conversion implementations
            - Analyzing model structure differences
            - Verifying parameter alignment and shapes
        """
        if not isinstance(megatron_model, list):
            megatron_model = [megatron_model]

        if hf_path is None:
            if not isinstance(self.hf_pretrained, PreTrainedCausalLM):
                raise ValueError("hf_path is required when hf_pretrained is not a PreTrainedCausalLM instance")
            pre_trained = self.hf_pretrained
        else:
            pre_trained = PreTrainedCausalLM.from_pretrained(hf_path)

        return self._model_bridge.build_conversion_tasks(pre_trained, megatron_model)

    @property
    def transformer_config(self) -> TransformerConfig:
        _model_provider = self.to_megatron_provider(load_weights=False)

        # Finalize the provider before extracting config
        if hasattr(_model_provider, "finalize"):
            _model_provider.finalize()

        return self._create_config_from_provider(_model_provider, TransformerConfig)

    @property
    def mla_transformer_config(self) -> MLATransformerConfig:
        _model_provider = self.to_megatron_provider(load_weights=False)

        # Finalize the provider before extracting config
        if hasattr(_model_provider, "finalize"):
            _model_provider.finalize()

        return self._create_config_from_provider(_model_provider, MLATransformerConfig)

    @property
    def _model_bridge(self) -> "MegatronModelBridge":
        hf_config = getattr(self.hf_pretrained, "hf_config", None)
        if hf_config is None:
            if isinstance(self.hf_pretrained, PreTrainedCausalLM):
                hf_config = self.hf_pretrained.config
            else:
                hf_config = self.hf_pretrained

        bridge = model_bridge.get_model_bridge(self._causal_lm_architecture, hf_config=hf_config)
        bridge.export_weight_dtype = self.export_weight_dtype
        return bridge

    @property
    def _provider_bridge_input(self) -> PreTrainedCausalLM | _ConfigOnlyPretrainedShim:
        if isinstance(self.hf_pretrained, PreTrainedCausalLM):
            return self.hf_pretrained
        return self._config_only_pretrained

    @property
    def _causal_lm_architecture(self):
        """Resolve the model's CausalLM architecture for dispatch.

        Behavior:
        - If the model can be imported from transformers directly, return the actual transformers class object.
        - Otherwise, if the model uses HuggingFace auto_map, return the architecture's class name as a string (e.g.,
        "DeepseekV2ForCausalLM").

        Returns:
            str | type: The transformers class for the CausalLM architecture or the architecture's class name as a
            string for auto_map models.

        Raises:
            ValueError: If no CausalLM architecture is found or cannot be resolved.
        """
        if isinstance(self.hf_pretrained, PreTrainedCausalLM):
            config = self.hf_pretrained.config
        else:
            config = self.hf_pretrained

        architectures = getattr(config, "architectures", [])

        if not architectures:
            raise ValueError(
                "\n✗ No architectures found in model config\n\n"
                "The model configuration does not specify any architectures.\n"
                "This is required for determining the model type."
            )

        causal_lm_arch = None
        for architecture_name in architectures:
            # TODO: Can we improve this?
            if architecture_name.endswith(SUPPORTED_HF_ARCHITECTURES):
                causal_lm_arch = architecture_name
                break

        if not causal_lm_arch:
            raise ValueError(
                f"\n✗ No CausalLM architecture found\n\n"
                f"Model architectures: {architectures}\n\n"
                f"None of the architectures end with {SUPPORTED_HF_ARCHITECTURES_DISPLAY}.\n"
                f"This bridge only supports causal language models.\n"
                f"For other model types, use a different bridge class."
            )

        # Try auto_map first (returns class name string if available)
        cls_name = get_causal_lm_class_name_via_auto_map(config=config)
        if cls_name is not None:
            # For auto_map models, return the class name as a string
            return cls_name

        # Resolve non-standard architecture names via alias mapping
        resolved_arch = HF_ARCHITECTURE_ALIASES.get(causal_lm_arch, causal_lm_arch)

        try:
            return getattr(transformers, resolved_arch)
        except AttributeError:
            # Model class not in standard transformers — fall back to class-name string.
            # This handles custom models registered via AutoConfig.register / AutoModelForCausalLM.register
            # in model bridge modules (e.g. BailingMoeV2ForCausalLM).
            return resolved_arch

    @classmethod
    def _validate_config(cls, config: PretrainedConfig, path: str | None = None) -> None:
        # Check if this is a causal LM model
        if not cls.supports(config):
            architectures = getattr(config, "architectures", [])
            raise ValueError(
                f"\n✗ Model architecture not supported by AutoBridge\n\n"
                f"Model: {path}\n"
                f"Architectures: {architectures}\n\n"
                f"AutoBridge only supports models with architectures ending in {SUPPORTED_HF_ARCHITECTURES_DISPLAY}.\n"
                f"Found architectures that don't match this pattern.\n\n"
                f"If this is a different model type (e.g., Vision, Sequence-to-Sequence),\n"
                f"you may need to use a different bridge class."
            )

        # Check if we have an implementation for this specific architecture
        architecture = None
        for arch_name in config.architectures:
            if arch_name.endswith(SUPPORTED_HF_ARCHITECTURES):
                architecture = arch_name
                break

        if architecture:
            # Try auto_map first; returns a class-name string if available
            arch_name = get_causal_lm_class_name_via_auto_map(config=config)
            if arch_name is not None:
                # For auto_map models, use class-name string
                arch_key = arch_name
            else:
                # Resolve non-standard architecture names via alias mapping
                resolved_arch = HF_ARCHITECTURE_ALIASES.get(architecture, architecture)
                try:
                    arch_class = getattr(transformers, resolved_arch)
                    arch_key = arch_class
                except AttributeError:
                    # Fall back to name-based registration
                    arch_key = architecture

            # Test if we have a registered implementation (type or class-name string)
            has_implementation = False
            if hasattr(model_bridge.get_model_bridge, "_exact_types"):
                registry = model_bridge.get_model_bridge._exact_types
                if isinstance(arch_key, str):
                    has_implementation = arch_key in registry
                else:
                    has_implementation = (arch_key in registry) or (getattr(arch_key, "__name__", None) in registry)

                if not has_implementation:
                    # Get list of supported models
                    supported_models = cls.list_supported_models()

                    raise ValueError(
                        f"\n✗ Model architecture '{architecture}' is not yet supported\n\n"
                        f"Model: {path}\n"
                        f"Architecture: {architecture}\n\n"
                        f"Currently supported architectures:\n"
                        + "\n".join(f"  • {model}" for model in supported_models)
                        + f"\n\nTo add support for {architecture}, you need to:\n"
                        f"1. Create a new bridge class that inherits from MegatronModelBridge\n"
                        f"2. Implement the required methods (provider_bridge, mapping_registry)\n"
                        f"3. Register it with @MegatronModelBridge.register_bridge decorator\n\n"
                        f"Example implementation:\n"
                        f"  from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge\n"
                        f"  from transformers import {architecture}\n"
                        f"  from megatron.core.models.gpt import GPTModel\n\n"
                        f"  @MegatronModelBridge.register_bridge(source={architecture}, target=GPTModel)\n"
                        f"  class Megatron{architecture.replace('ForCausalLM', '')}Bridge(MegatronModelBridge):\n"
                        f"      def provider_bridge(self, hf_pretrained):\n"
                        f"          # Return a ModelProvider instance\n"
                        f"          ...\n\n"
                        f"      def mapping_registry(self):\n"
                        f"          # Return a MegatronMappingRegistry with weight mappings\n"
                        f"          ...\n\n"
                        f"For reference implementations, see:\n"
                        f"  • src/megatron/bridge/models/llama/llama_bridge.py\n"
                        f"  • src/megatron/bridge/models/qwen/qwen_2_causal_bridge.py"
                    ) from None

    def _get_model_instance(self, model: list[MegatronModelT]) -> MegatronModelT:
        model_instance = model[0]
        while hasattr(model_instance, "module"):
            model_instance = model_instance.module
        return model_instance

    def _create_config_from_provider(self, source_obj: Any, target_dataclass: Type[DataclassT]) -> DataclassT:
        kwargs = {}
        for field in dataclasses.fields(target_dataclass):
            if hasattr(source_obj, field.name):
                kwargs[field.name] = getattr(source_obj, field.name)
        return target_dataclass(**kwargs)

    @cached_property
    def _config_only_pretrained(self) -> _ConfigOnlyPretrainedShim:
        if not isinstance(self.hf_pretrained, PretrainedConfig):
            raise ValueError("Config-only shim accessed when hf_pretrained is not a PretrainedConfig instance.")
        return _ConfigOnlyPretrainedShim(self.hf_pretrained)

    def __repr__(self) -> str:
        class_name = self.__class__.__name__

        lines_for_build = []

        # Format hf_pretrained
        hf_repr_actual_lines = repr(self.hf_pretrained).splitlines()
        if hf_repr_actual_lines:
            # First line of hf_pretrained part
            lines_for_build.append(f"  (hf_pretrained): {hf_repr_actual_lines[0]}")
            # Subsequent lines of hf_pretrained part, indented
            for line in hf_repr_actual_lines[1:]:
                lines_for_build.append(f"  {line}")
        else:
            lines_for_build.append("  (hf_pretrained): ")  # Fallback for empty repr

        # Format model bridge
        mb_repr_actual_lines = repr(self._model_bridge).splitlines()
        if mb_repr_actual_lines:
            # First line of model bridge part
            lines_for_build.append(f"  (model_bridge): {mb_repr_actual_lines[0]}")
            # Subsequent lines of model bridge part, indented
            for line in mb_repr_actual_lines[1:]:
                lines_for_build.append(f"  {line}")
        else:
            lines_for_build.append("  (model_bridge): ")  # Fallback for empty repr

        return f"{class_name}(\n" + "\n".join(lines_for_build) + "\n)"

    def _validate_fp8_export_config(self, model: list[MegatronModelT]) -> None:
        """Validate runtime Megatron config before enabling FP8 export tasks."""
        model_instance = self._get_model_instance(model)
        model_config = getattr(model_instance, "config", None)
        fp8 = getattr(model_config, "fp8", None)
        fp8_recipe = getattr(model_config, "fp8_recipe", None)
        fp8_param = getattr(model_config, "fp8_param", None)
        if fp8 is None or fp8_recipe != "blockwise" or not fp8_param:
            raise ValueError(
                "export_weight_dtype='fp8' only supports blockwise FP8 parameter export. "
                f"Expected fp8 to be enabled, fp8_recipe='blockwise', and fp8_param=True, "
                f"but got fp8={fp8!r}, fp8_recipe={fp8_recipe!r}, fp8_param={fp8_param!r}."
            )