Introduce refactored model builder abstractions

src/megatron/bridge/models/common/__init__.py

@@ -0,0 +1,24 @@
+# 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 megatron.bridge.models.common.base import (
+    BuildConfigT,
+    ModelBuilder,
+    ModelConfig,
+    ModelProvider,
+    ModelT,
+    Serializable,
+    compose_hooks,
+)
+from megatron.bridge.models.common.unimodal import build_virtual_pipeline_stages, unimodal_build_distributed_models

src/megatron/bridge/models/common/base.py

@@ -0,0 +1,359 @@
+# 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.
+
+import abc
+import importlib
+from dataclasses import dataclass, is_dataclass
+from dataclasses import fields as dataclass_fields
+from typing import Any, Callable, ClassVar, Generic, Protocol, TypeVar
+
+from megatron.core.distributed import DistributedDataParallelConfig
+from megatron.core.enums import ModelType
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer import MegatronModule
+from megatron.core.transformer.module import Float16Module
+
+
+try:
+    from megatron.core.fp8_utils import correct_amax_history_if_needed
+except ImportError:
+    correct_amax_history_if_needed = None
+
+
+class Serializable(Protocol):
+    """Protocol for serializable configurations."""
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to dictionary with _target_ for class identification."""
+        ...
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Serializable":
+        """Deserialize from dictionary using _target_ to identify class."""
+        ...
+
+
+@dataclass
+class ModelConfig(abc.ABC, Serializable):
+    """Abstract base class for model build configurations.
+
+    Each model type (GPT, T5, Mamba, etc.) has its own build config subclass.
+    The build config contains:
+    1. Builder path (serializable string) to link to the correct builder
+    2. Model-specific parameters not in TransformerConfig
+    3. HuggingFace metadata for checkpoint conversion
+
+    Each subclass must define `builder` as a ClassVar string pointing to
+    the appropriate ModelBuilder subclass path.
+    """
+
+    # === Builder Metadata (Serializable) ===
+    builder: ClassVar[str]
+    """Class variable with full path to builder class (e.g.,
+    'megatron.bridge.builders.GPTModelBuilder').
+    """
+
+    # === ModelOpt ===
+    restore_modelopt_state: bool = False
+    """Restore ModelOpt quantization/sparsity state."""
+
+    # === HuggingFace Metadata ===
+    hf_model_id: str | None = None
+    """HuggingFace model identifier."""
+
+    generation_config: Any | None = None
+    """Generation configuration."""
+
+    def get_builder_cls(self) -> type:
+        """Get the appropriate builder instance for this config.
+        Dynamically imports and instantiates the builder from the string path.
+        """
+        module_path, class_name = self.builder.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        builder_cls = getattr(module, class_name)
+        return builder_cls
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize config to dictionary for saving.
+
+        Includes:
+        - _target_: Full class path for deserialization
+        - _builder_: Full builder class path (serialized from ClassVar)
+        - All dataclass fields, including nested dataclasses
+        """
+
+        def _to_dict(config):
+            result = {
+                "_target_": f"{config.__class__.__module__}.{config.__class__.__qualname__}",
+            }
+            for f in dataclass_fields(config):
+                value = getattr(config, f.name)
+                # Skip non-serializable fields
+                if callable(value) or f.name.startswith("_"):
+                    continue
+
+                if is_dataclass(value):
+                    result[f.name] = _to_dict(value)  # recurse on nested dataclasses
+                else:
+                    result[f.name] = value
+
+            return result
+
+        result = _to_dict(self)
+        result["_builder_"] = self.builder  # Serialize the builder path
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ModelConfig":
+        """Deserialize config from dictionary.
+
+        Uses _target_ to determine the correct class to instantiate.
+        The builder is restored from _builder_ or from the class's ClassVar.
+
+        Args:
+            data: Dictionary with _target_ and config fields
+
+        Returns:
+            Instance of the appropriate ModelBuildConfig subclass
+        """
+
+        def _from_dict(subdata):
+            target = subdata.get("_target_")
+            if target is None:
+                raise ValueError("Cannot deserialize: missing '_target_' field")
+
+            # Import the class from the target path
+            module_path, class_name = target.rsplit(".", 1)
+            module = importlib.import_module(module_path)
+            config_cls = getattr(module, class_name)
+
+            # Filter to valid fields for this class
+            valid_fields = {f.name for f in dataclass_fields(config_cls)}
+            filtered_data = {k: v for k, v in subdata.items() if k in valid_fields and not k.startswith("_")}
+
+            # recurse on serialized nested dataclasses
+            subconfigs = {}
+            for k, v in filtered_data.items():
+                if isinstance(v, dict):
+                    subconfigs[k] = _from_dict(v)
+            filtered_data.update(subconfigs)
+
+            return config_cls(**filtered_data)
+
+        result = _from_dict(data)
+        result.builder = data["_builder_"]
+
+        return result
+
+
+ModelT = TypeVar("ModelT", bound=MegatronModule)
+BuildConfigT = TypeVar("BuildConfigT", bound=ModelConfig)
+
+
+class ModelBuilder(abc.ABC, Generic[ModelT, BuildConfigT]):
+    """Abstract base class for model builders.
+
+    A builder takes configuration(s) and produces model instances.
+
+    Each builder subclass should:
+    1. Implement build_model() for the specific model type
+    2. Be linked to its corresponding ModelBuildConfig via the builder string
+
+    Builders are factory objects, therefore any state saved in __init__ should not be modified
+    and only used to build the model.
+
+    Type Parameters:
+        ModelT: The type of model this builder produces (e.g., MCoreGPTModel)
+        BuildConfigT: The type of build config this builder accepts (e.g., GPTModelBuildConfig)
+    """
+
+    def __init__(self, model_config: ModelConfig):
+        self._model_config = model_config
+        self._pre_wrap_hooks = []
+        self._post_wrap_hooks = []
+
+    @abc.abstractmethod
+    def build_model(
+        self,
+        pg_collection: ProcessGroupCollection,
+        pre_process: bool | None = None,
+        post_process: bool | None = None,
+        vp_stage: int | None = None,
+    ) -> ModelT:
+        """Build a model from the provided configurations.
+
+        Args:
+            pg_collection: Process groups for distributed training
+            pre_process: Include embedding layer
+            post_process: Include output layer
+            vp_stage: Virtual pipeline stage
+
+        Returns:
+            The constructed model
+        """
+        ...
+
+    def build_distributed_models(
+        self,
+        pg_collection: ProcessGroupCollection,
+        ddp_config: DistributedDataParallelConfig | None = None,
+        overlap_param_gather_with_optimizer_step: bool = False,
+        use_megatron_fsdp: bool = False,
+        use_torch_fsdp2: bool = False,
+        wrap_with_ddp: bool = True,
+        data_parallel_random_init: bool = True,
+        mixed_precision_wrapper: Callable[[Any, MegatronModule], MegatronModule] | None = Float16Module,
+        model_type: ModelType = ModelType.encoder_or_decoder,
+    ) -> list[ModelT]:
+        """Build model stages and wrap for distributed training."""
+        ...
+
+    def register_pre_wrap_hook(
+        self,
+        hook: Callable[[list[MegatronModule]], list[MegatronModule]],
+        prepend: bool = False,
+    ) -> None:
+        """Registers a hook to be executed before the model is wrapped.
+
+        When the hooks are executed is left up to the implementation of child class.
+
+        The hook should be a callable that accepts a list of `MegatronModule` instances
+        and returns a (potentially modified) list of `MegatronModule` instances.
+
+        Args:
+            hook: The hook to register.
+            prepend: If True, the hook is inserted at the beginning of the execution
+                chain. Otherwise, it is appended to the end.
+        """
+        if prepend:
+            self._pre_wrap_hooks.insert(0, hook)
+        else:
+            self._pre_wrap_hooks.append(hook)
+
+    def register_post_wrap_hook(
+        self,
+        hook: Callable[[list[MegatronModule]], list[MegatronModule]],
+        prepend: bool = False,
+    ) -> None:
+        """Registers a hook to be executed after the model is wrapped.
+
+        When the hooks are executed is left up to the implementation of child class.
+
+        The hook should be a callable that accepts a list of `MegatronModule` instances
+        and returns a (potentially modified) list of `MegatronModule` instances.
+
+        Args:
+            hook: The hook to register.
+            prepend: If True, the hook is inserted at the beginning of the execution
+                chain. Otherwise, it is appended to the end.
+        """
+        if prepend:
+            self._post_wrap_hooks.insert(0, hook)
+        else:
+            self._post_wrap_hooks.append(hook)
+
+
+def compose_hooks(
+    hooks: list[Callable[[list[MegatronModule]], list[MegatronModule]]],
+) -> Callable[[list[MegatronModule]], list[MegatronModule]]:
+    """Utility to compose pre/post-wrap hooks into a single function, preserving order.
+
+    Args:
+        hooks: the list of hooks.
+
+    Returns:
+        A single function that executes all functions in `hooks`.
+    """
+
+    def composed_hook(model: list[MegatronModule]) -> list[MegatronModule]:
+        for hook in hooks:
+            model = hook(model)
+        return model
+
+    return composed_hook
+
+
+class ModelProvider(Generic[ModelT]):
+    """General provider that takes model config + build config and builds models.
+
+    This is the main entry point for model construction. It automatically
+    selects the correct builder based on the model_config's builder attribute.
+
+    Example:
+        >>> provider = ModelProvider(model_cfg)
+        >>> model = provider.provide(pg_collection)
+        >>>
+        >>> # Or for distributed training with DDP
+        >>> models = provider.provide_distributed(pg_collection, wrap_with_ddp=True)
+    """
+
+    def __init__(self, model_config: ModelConfig) -> None:
+        self.model_config = model_config
+
+    def provide(
+        self,
+        pg_collection: ProcessGroupCollection,
+        pre_process: bool | None = None,
+        post_process: bool | None = None,
+        vp_stage: int | None = None,
+    ) -> ModelT:
+        """Build and return a model.
+
+        Automatically selects the correct builder based on model_config.builder.
+
+        Args:
+            pg_collection: Process groups for distributed training
+            pre_process: Include embedding layer (default: based on PP stage)
+            post_process: Include output layer (default: based on PP stage)
+            vp_stage: Virtual pipeline stage
+
+        Returns:
+            The constructed model
+        """
+        builder_cls = self.model_config.get_builder_cls()
+        return builder_cls(self.model_config).build_model(
+            pg_collection,
+            pre_process=pre_process,
+            post_process=post_process,
+            vp_stage=vp_stage,
+        )
+
+    def provide_distributed(
+        self,
+        pg_collection: ProcessGroupCollection,
+        wrap_with_ddp: bool = True,
+        fp16: bool = False,
+        bf16: bool = False,
+    ) -> list[ModelT]:
+        """Build models wrapped for distributed training.
+
+        Handles virtual pipeline parallelism, DDP wrapping, and
+        mixed precision configuration.
+
+        Args:
+            pg_collection: Process groups for distributed training
+            wrap_with_ddp: Whether to wrap with DDP
+            fp16: Use FP16 mixed precision
+            bf16: Use BF16 mixed precision
+
+        Returns:
+            List of models (multiple for virtual pipeline parallelism)
+        """
+        builder_cls = self.model_config.get_builder_cls()
+        return builder_cls(self.model_config).build_distributed_models(
+            pg_collection,
+            wrap_with_ddp=wrap_with_ddp,
+            fp16=fp16,
+            bf16=bf16,
+        )