diff --git a/src/megatron/bridge/models/common/__init__.py b/src/megatron/bridge/models/common/__init__.py new file mode 100644 index 0000000000..03e7396120 --- /dev/null +++ b/src/megatron/bridge/models/common/__init__.py @@ -0,0 +1,35 @@ +# 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, + ModelT, + Serializable, + compose_hooks, +) +from megatron.bridge.models.common.unimodal import build_virtual_pipeline_stages, unimodal_build_distributed_models + + +__all__ = [ + "BuildConfigT", + "ModelBuilder", + "ModelConfig", + "ModelT", + "Serializable", + "compose_hooks", + "build_virtual_pipeline_stages", + "unimodal_build_distributed_models", +] diff --git a/src/megatron/bridge/models/common/base.py b/src/megatron/bridge/models/common/base.py new file mode 100644 index 0000000000..7fbbc33cd7 --- /dev/null +++ b/src/megatron/bridge/models/common/base.py @@ -0,0 +1,278 @@ +# 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, field, 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 + + +class Serializable(Protocol): + """Protocol for serializable configurations.""" + + def as_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: + """Base class for model configurations. + + Each model type (GPT, T5, Mamba, etc.) defines a concrete subclass with its + own model-specific parameters. This class is a pure data container - all model + construction logic lives in the corresponding ``ModelBuilder`` subclass. + + Subclasses must define: + - ``builder``: a ``ClassVar[str]`` with the full import path to the + associated ``ModelBuilder`` (e.g. + ``'megatron.bridge.models.mamba.MambaModelBuilder'``). + + Subclasses may also embed nested configs (e.g. ``TransformerConfig``) and + proxy attribute access to them via ``__getattr__``/``__setattr__`` overrides. + + Serialization: + Use ``as_dict()`` to serialize to a plain dict (includes a ``_target_`` key + for class resolution and a ``_builder_`` key for builder resolution). + Use ``from_dict()`` to reconstruct an instance from such a dict. + + Builder resolution: + Call ``get_builder_cls()`` to dynamically import and return the builder + class identified by the ``builder`` ClassVar. + """ + + # === 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.""" + + # === pre-wrap and post-wrap hooks === + pre_wrap_hooks: list[Callable[[list[MegatronModule]], list[MegatronModule]]] = field(default_factory=list) + """List of functions that are executed before the model is wrapped with DDP/FSDP. + Should take the model as the only argument and return a new model as the only return value. + """ + + post_wrap_hooks: list[Callable[[list[MegatronModule]], list[MegatronModule]]] = field(default_factory=list) + """List of functions that are executed after model initialization is complete. + Should take the model as the only argument and return a new model as the only return value. + """ + + def get_builder_cls(self) -> type: + """Get the appropriate builder type for this config. + Dynamically imports 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 as_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 _as_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("_") or f.name in ["pre_wrap_hooks", "post_wrap_hooks"]: + continue + + if is_dataclass(value): + result[f.name] = _as_dict(value) # recurse on nested dataclasses + else: + result[f.name] = value + + return result + + result = _as_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 ModelConfig 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) and "_target_" in v: + 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 a ``ModelConfig`` and produces distributed model instances - + either a single pipeline stage via ``build_model()``, or a list of stages + wrapped for distributed training via ``build_distributed_models()``. + + Each builder subclass should: + 1. Implement ``build_model()`` for the specific model type + 2. Implement ``build_distributed_models()`` to handle virtual pipeline parallelism, + DDP/FSDP wrapping, and pre/post-wrap hook execution + 3. Be linked to its corresponding ``ModelConfig`` via the ``builder`` ClassVar + + 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 + + @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 + """ + ... + + @abc.abstractmethod + 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. + + Args: + pg_collection: Model communication process groups. + ddp_config: DistributedDataParallel configuration + overlap_param_gather_with_optimizer_step: Whether to overlap parameter gather with optimizer step + use_megatron_fsdp: Whether to use Megatron FSDP + use_torch_fsdp2: Whether to use Torch FSDP 2.0 + wrap_with_ddp: Set to False to skip DDP wrapper + data_parallel_random_init: Whether to use data parallel random initialization + mixed_precision_wrapper: Mixed precision wrapper, e.g. ``Float16Module`` + model_type: Deprecated flag, only used for backwards compatibility. + + Returns: + List of model stages. If the model does not support virtual pipeline parallelism, + this function should still return a single-item list. + """ + ... + + +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. + + If `hooks` is empty, the returned function is an identity operation. + + 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 diff --git a/src/megatron/bridge/models/common/unimodal.py b/src/megatron/bridge/models/common/unimodal.py new file mode 100644 index 0000000000..cc882035f1 --- /dev/null +++ b/src/megatron/bridge/models/common/unimodal.py @@ -0,0 +1,328 @@ +# 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 logging + + +logger = logging.getLogger(__name__) + +from typing import Any, Callable + +import torch +from megatron.core import tensor_parallel +from megatron.core.distributed import ( + DistributedDataParallel, + DistributedDataParallelConfig, + FullyShardedDataParallel, + TorchFullyShardedDataParallel, +) +from megatron.core.enums import ModelType +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer import MegatronModule, TransformerConfig +from megatron.core.transformer.module import Float16Module +from megatron.core.utils import get_model_config + + +try: + from megatron.core.fp8_utils import correct_amax_history_if_needed +except ImportError: + correct_amax_history_if_needed = None + + +def unimodal_build_distributed_models( + build_model_func: Callable, + transformer_config: TransformerConfig, + 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, + pre_wrap_hook: Callable[[list[MegatronModule]], list[MegatronModule]] | None = None, + model_type: ModelType = ModelType.encoder_or_decoder, +) -> list[MegatronModule]: + """Build model stages and wrap for distributed training. + + Shared helper for unimodal models (GPT, Mamba, etc.) that share the same procedure + for distributed model initialization. Performs the following steps in order: + + 1. Build virtual pipeline stages (one per VP rank, or a single stage if no VP) + 2. Apply ``pre_wrap_hook`` + 3. Set tensor model parallel attributes on all parameters + 4. Move model to GPU (unless using FSDP2 or CPU/meta-device initialization) + 5. Apply mixed precision wrapper (e.g. ``Float16Module``) + 6. Materialize meta-device tensors if ``init_model_with_meta_device`` is set + 7. Optionally wrap with DDP/FSDP + + Args: + build_model_func: Callable that builds a single model stage (e.g. ``ModelBuilder.build_model``). + transformer_config: TransformerConfig; used for VP size, precision, and device placement. + pg_collection: Model communication process groups. + ddp_config: DistributedDataParallel configuration. Required when ``wrap_with_ddp=True``. + overlap_param_gather_with_optimizer_step: Whether to overlap parameter gather with optimizer step. + use_megatron_fsdp: Whether to use Megatron FSDP. + use_torch_fsdp2: Whether to use Torch FSDP 2.0. + wrap_with_ddp: Set to False to skip the DDP/FSDP wrapper. + data_parallel_random_init: Whether to broadcast parameters from data-parallel rank 0. + mixed_precision_wrapper: Mixed precision wrapper applied per model stage, e.g. ``Float16Module``. + Pass ``None`` to skip. + pre_wrap_hook: Hook applied to the model stage list before any wrapping. + model_type: Deprecated flag, only used for backwards compatibility. + + Returns: + List of model stages, wrapped and ready for distributed training. + """ + if wrap_with_ddp and not ddp_config: + raise ValueError("ddp_config is required when wrap_with_ddp is True") + + vp_size = transformer_config.virtual_pipeline_model_parallel_size + init_model_with_meta_device = transformer_config.init_model_with_meta_device + if init_model_with_meta_device: + with torch.device("meta"): + model_list = build_virtual_pipeline_stages(build_model_func, pg_collection, vp_size, model_type) + else: + model_list = build_virtual_pipeline_stages(build_model_func, pg_collection, vp_size, model_type) + + # Apply pre wrap hooks + if pre_wrap_hook is not None: + if not callable(pre_wrap_hook): + raise TypeError("pre_wrap_hook must be a callable") + _model = pre_wrap_hook(model_list) + if _model is not None: + model_list = _model + else: + logger.warning("Final pre wrap hook returned None, skipping pre wrap hooks.") + + # Set tensor model parallel attributes if not set. + # Only parameters that are already tensor model parallel have these + # attributes set for them. We should make sure the default attributes + # are set for all params so the optimizer can use them. + for model_module in model_list: + for param in model_module.parameters(): + tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param) + + _print_num_params(model_list, pg_collection=pg_collection) + + # GPU allocation. + # For FSDP2, we don't allocate GPU memory here. We allocate GPU memory + # in the fully_shard function of FSDP2 instead. + use_cpu_initialization = transformer_config.use_cpu_initialization + if not use_torch_fsdp2 and not use_cpu_initialization and not init_model_with_meta_device: + for model_module in model_list: + model_module.cuda(torch.cuda.current_device()) + + model_list = _wrap_with_mp_wrapper(model_list, transformer_config, mixed_precision_wrapper) + + # Materialize tensors on meta device (GPU allocation) if not using FSDP2 and not using Megatron FSDP. + if init_model_with_meta_device and not use_torch_fsdp2 and not use_megatron_fsdp: + model_list = [ + to_empty_if_meta_device(model_module, device=torch.device("cuda")) for model_module in model_list + ] + + if correct_amax_history_if_needed is not None: + correct_amax_history_if_needed(model_list) + + if wrap_with_ddp: + model_list = _ddp_wrap( + model_list, + data_parallel_random_init, + ddp_config, + overlap_param_gather_with_optimizer_step, + use_megatron_fsdp=use_megatron_fsdp, + use_torch_fsdp2=use_torch_fsdp2, + pg_collection=pg_collection, + ) + + return model_list + + +def _print_num_params(model: list[MegatronModule], pg_collection: ProcessGroupCollection) -> None: + """Print the number of parameters in the model on rank 0. + + Only prints on data parallel rank 0 to avoid duplicate output. + Shows parameter count per (tensor parallel, pipeline parallel) rank. + + Args: + model: List of model modules to count parameters from + pg_collection: Model communication process groups. + """ + if (pg_collection.dp.rank() == 0) and (pg_collection.cp.rank() == 0): + print( + " > number of parameters on (tensor, pipeline) model parallel rank ({}, {}): {}".format( + pg_collection.tp.rank(), + pg_collection.pp.rank(), + sum([sum([p.nelement() for p in model_module.parameters()]) for model_module in model]), + ), + flush=True, + ) + + +def _wrap_with_mp_wrapper( + model_list: list[MegatronModule], + transformer_config: TransformerConfig, + mixed_precision_wrapper: Callable[[Any, MegatronModule], MegatronModule] | None = Float16Module, +) -> list[MegatronModule]: + fp16 = transformer_config.fp16 + bf16 = transformer_config.bf16 + if (fp16 or bf16) and mixed_precision_wrapper is not None: + model_list = [mixed_precision_wrapper(transformer_config, model_module) for model_module in model_list] + + # Maintain expert bias in float32 wrapped in Float16Module + for model_module in model_list: + for submodule in model_module.modules(): + if hasattr(submodule, "_maintain_float32_expert_bias"): + submodule._maintain_float32_expert_bias() + + return model_list + + +def _ddp_wrap( + model: list[MegatronModule], + data_parallel_random_init: bool, + ddp_config: DistributedDataParallelConfig, + overlap_param_gather_with_optimizer_step: bool, + use_megatron_fsdp: bool = False, + use_torch_fsdp2: bool = False, + *, + pg_collection: ProcessGroupCollection, +) -> list[MegatronModule]: + """Wrap model with Distributed Data Parallel (DDP) or Fully Sharded Data Parallel (FSDP). + + Args: + model: List of model modules to wrap + data_parallel_random_init: Whether to broadcast parameters from rank 0 + ddp_config: Configuration for distributed data parallel + overlap_param_gather_with_optimizer_step: Whether to disable bucketing + for overlapping parameter gathering with optimizer step + use_megatron_fsdp: Whether to use Megatron FSDP. + use_torch_fsdp2: Whether to use PyTorch FSDP v2 instead of DDP + pg_collection: Model communication process groups. + + Returns: + list[MegatronModule]: List of DDP/FSDP wrapped model modules + """ + if use_megatron_fsdp: + DP = FullyShardedDataParallel + if use_torch_fsdp2: + raise ValueError("Using use_megatron_fsdp and use_torch_fsdp2 at the same time is not supported.") + elif use_torch_fsdp2: + DP = TorchFullyShardedDataParallel + else: + DP = DistributedDataParallel + + # DDP initialization is required to be on a side-stream for the full-iteration CUDA graph. + # this side-stream may be nested if being called from within the get_model function, but it + # is here in case someone wants to use this directly outside of get_model. + ddp_stream = torch.cuda.Stream() + ddp_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(ddp_stream): + model = [ + DP( + config=get_model_config(model_chunk), + ddp_config=ddp_config, + module=model_chunk, + # Turn off bucketing for model_chunk 2 onwards, since communication for these + # model chunks is overlapped with compute anyway. + disable_bucketing=(model_chunk_idx > 0) or overlap_param_gather_with_optimizer_step, + pg_collection=pg_collection, + ) + for (model_chunk_idx, model_chunk) in enumerate(model) + ] + # Critical: ensure side-stream work completes before touching params on default stream + torch.cuda.current_stream().wait_stream(ddp_stream) + + # Broadcast params from data parallel src rank to other data parallel ranks. + if data_parallel_random_init: + for model_module in model: + model_module.broadcast_params() + + return model + + +def build_virtual_pipeline_stages( + build_model_func: Callable, + pg_collection: ProcessGroupCollection, + vp_size: int | None, + model_type: ModelType = ModelType.encoder_or_decoder, +) -> list[MegatronModule]: + """Build virtual pipeline stages if using virtual pipeline parallelism. + + Args: + build_model_func: Function from ``ModelBuilder`` that builds a single stage of the model. + pg_collection: Model communication process groups. + vp_size: Virtual pipeline parallel size. If ``None`` or PP size is 1, a single stage is built. + model_type: Deprecated flag, only used for backwards compatibility. + + Returns: + List of model stages. Contains one entry per VP rank, or a single entry if VP is not enabled. + """ + from megatron.core.pipeline_parallel.utils import ( + is_pp_first_stage, + is_pp_last_stage, + is_vp_first_stage, + is_vp_last_stage, + ) + + pp_group = pg_collection.pp + if pp_group.size() > 1 and vp_size is not None: + # Create multiple model stages for virtual pipeline + model_list = [] + for i in range(vp_size): + pre_process = is_vp_first_stage(vp_stage=i, vp_size=vp_size) and is_pp_first_stage(pp_group) + post_process = is_vp_last_stage(vp_stage=i, vp_size=vp_size) and is_pp_last_stage(pp_group) + model = build_model_func( + pg_collection, + pre_process=pre_process, + post_process=post_process, + vp_stage=i, + ) + model.model_type = model_type + model_list.append(model) + else: + # Single stage, no VP + pre_process = is_pp_first_stage(pp_group) + post_process = is_pp_last_stage(pp_group) + model = build_model_func(pg_collection, pre_process=pre_process, post_process=post_process) + model.model_type = model_type + model_list = [model] + + return model_list + + +def to_empty_if_meta_device(module: torch.nn.Module, *, device: torch.device, recurse=True): + """Move tensors to device if not meta device; otherwise materialize with empty_like(). + + Officially, torch suggests to_empty() for meta device materialization. Under the hood, + torch.empty_like() is applied to all parameters or buffers (see _apply). This may + accidently overwrite buffers with precomputed values during construction. Given the + goal is to only materialize those tensors on meta device, this function checks the + device first and only move the tensor to the destination if it is not on meta device. + + Args: + module: The target module to apply this transformation. + device: The desired device of the parameters + and buffers in this module. + recurse: Whether parameters and buffers of submodules should + be recursively moved to the specified device. + """ + + def _empty_like_if_meta(tensor: torch.Tensor, *, device: torch.device): + if tensor.device == torch.device("meta"): + return torch.empty_like(tensor, device=device) + else: + return tensor.to(device) + + return module._apply(lambda t: _empty_like_if_meta(t, device=device), recurse=recurse) diff --git a/src/megatron/bridge/models/mamba/mamba_builder.py b/src/megatron/bridge/models/mamba/mamba_builder.py new file mode 100644 index 0000000000..0f421bf0ec --- /dev/null +++ b/src/megatron/bridge/models/mamba/mamba_builder.py @@ -0,0 +1,294 @@ +# 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 logging +from dataclasses import dataclass +from typing import Any, Callable, ClassVar, Literal, override + +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.enums import ModelType +from megatron.core.models.mamba import MambaModel as MCoreMambaModel +from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec as default_mamba_stack_spec +from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage +from megatron.core.post_training.modelopt.mamba.model_specs import get_mamba_stack_modelopt_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer import MegatronModule, ModuleSpec +from megatron.core.transformer.module import Float16Module + +from megatron.bridge.models.common import ( + ModelBuilder, + ModelConfig, + compose_hooks, + unimodal_build_distributed_models, +) +from megatron.bridge.models.transformer_config import TransformerConfig +from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size + + +logger = logging.getLogger(__name__) + + +def transformer_engine_mamba_stack_spec() -> ModuleSpec: + """Return the default Mamba stack spec with Transformer Engine layers. + + This is a named function (not a lambda) to allow proper serialization + and reconstruction from checkpoints. Named functions can be imported + via their module path, unlike lambdas. + + Returns: + Default Mamba stack specification from megatron.core + """ + return default_mamba_stack_spec + + +def modelopt_mamba_stack_spec() -> ModuleSpec: + """Mamba stack specification for quantization with ModelOpt. + + Uses Norm instead of TENorm and ColumnParallelLinear/RowParallelLinear + instead of TE layers to enable proper quantizer insertion by ModelOpt. + + Returns: + ModuleSpec: Module specification for quantization-ready Mamba stack + """ + return get_mamba_stack_modelopt_spec( + local_core_attention=False, + remap_te_layernorm=False, + ) + + +def get_default_mamba_stack_spec(config: "MambaModelConfig") -> ModuleSpec: + """Determine the most appropriate Mamba stack specification based on configuration. + + Args: + config: Mamba configuration object + + Returns: + ModuleSpec: Appropriate module specification based on config + """ + if config.restore_modelopt_state: + return modelopt_mamba_stack_spec() + else: + return transformer_engine_mamba_stack_spec() + + +@dataclass(kw_only=True) +class MambaModelConfig(ModelConfig): + """Configuration for a Megatron Core Mamba (SSM) model. + + This is purely a configuration object. All model construction + logic lives in ``MambaModelBuilder``. + + Contains a ``TransformerConfig`` alongside Mamba-specific parameters. Attributes + on the embedded ``transformer`` config are accessible directly on this object + via ``__getattr__``/``__setattr__`` proxying. + + Supports hybrid SSM/attention architectures via ``hybrid_attention_ratio``, + ``hybrid_mlp_ratio``, and ``hybrid_override_pattern``. + + Note: + ``vocab_size`` must be set before passing this config to ``MambaModelBuilder``. + """ + + builder: ClassVar[str] = "megatron.bridge.models.mamba.MambaModelBuilder" + transformer: TransformerConfig + fp16_lm_cross_entropy: bool = False + parallel_output: bool = True + share_embeddings_and_output_weights: bool = False + hybrid_attention_ratio: float = 0.0 + hybrid_mlp_ratio: float = 0.0 + hybrid_override_pattern: str | None = None + seq_length: int = 8192 + # Mamba with no attention has no need for position embeddings, so none is default + position_embedding_type: Literal["learned_absolute", "rope", "none"] = "none" + rotary_percent: float = 1.0 + rotary_base: int = 10000 + seq_len_interpolation_factor: float | None = None + make_vocab_size_divisible_by: int = 128 + mamba_stack_spec: ModuleSpec | Callable[[], ModuleSpec] | Callable[["MambaModelConfig"], ModuleSpec] = ( + get_default_mamba_stack_spec + ) + vocab_size: int | None = None + should_pad_vocab: bool = False + + @override + def __getattr__(self, name: str, /) -> Any: + # __getattr__ is only called when normal attribute lookup has already failed, + # so use object.__getattribute__ to fetch `transformer` without recursing. + try: + transformer = object.__getattribute__(self, "transformer") + except AttributeError: + raise AttributeError(f"MambaModelConfig has no attribute '{name}'") + if hasattr(transformer, name): + return getattr(transformer, name) + raise AttributeError(f"Neither MambaModelConfig nor TransformerConfig has any attribute '{name}'.") + + @override + def __setattr__(self, name: str, value: Any, /) -> None: + # Use object.__getattribute__ to avoid triggering __getattr__ while + # `transformer` may not yet exist (e.g. during dataclass __init__). + try: + transformer = object.__getattribute__(self, "transformer") + except AttributeError: + # `transformer` not yet initialised; store the attribute on self. + super().__setattr__(name, value) + return + if hasattr(transformer, name): + setattr(transformer, name, value) + else: + super().__setattr__(name, value) + + +class MambaModelBuilder(ModelBuilder[MCoreMambaModel, MambaModelConfig]): + """Builder to construct Megatron Core Mamba models. + + Example: + >>> transformer_cfg = TransformerConfig(num_layers=32, hidden_size=4096, ...) + >>> model_cfg = MambaModelConfig(transformer=transformer_cfg, vocab_size=32000, seq_length=2048, ...) + >>> + >>> # Single stage (e.g. inference) + >>> model = MambaModelBuilder(model_cfg).build_model(pg_collection) + >>> + >>> # Distributed training + >>> models = MambaModelBuilder(model_cfg).build_distributed_models(pg_collection) + """ + + def __init__(self, model_config: MambaModelConfig): + super().__init__(model_config) + + def build_model( + self, + pg_collection: ProcessGroupCollection, + pre_process: bool | None = None, + post_process: bool | None = None, + vp_stage: int | None = None, + ) -> MCoreMambaModel: + """Build a single ``MCoreMambaModel`` stage. + + 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 + + Note: + Virtual pipeline model parallelism is not supported for Mamba models. + """ + mamba_stack_spec = self._model_config.mamba_stack_spec + if not isinstance(mamba_stack_spec, ModuleSpec): + # Check if the function accepts config parameter + import inspect + + if len(inspect.signature(mamba_stack_spec).parameters) > 0: + mamba_stack_spec = mamba_stack_spec(self._model_config) + else: + mamba_stack_spec = mamba_stack_spec() + + assert ( + getattr(self._model_config.transformer, "virtual_pipeline_model_parallel_size", None) is None + and vp_stage is None + ), ( + "Virtual pipeline model parallelism is temporarily unsupported in SSM/Mamba " + "models due to upstream MCore MambaModel API dependency" + ) + + assert self._model_config.vocab_size is not None, "vocab_size must be configured before calling build_model()" + if self._model_config.should_pad_vocab: + padded_vocab_size = calculate_padded_vocab_size( + self._model_config.vocab_size, + self._model_config.make_vocab_size_divisible_by, + self._model_config.transformer.tensor_model_parallel_size, + ) + else: + padded_vocab_size = self._model_config.vocab_size + + pre_process = pre_process if pre_process is not None else is_pp_first_stage(pg_collection.pp) + post_process = post_process if post_process is not None else is_pp_last_stage(pg_collection.pp) + return MCoreMambaModel( + config=self._model_config.transformer, + mamba_stack_spec=mamba_stack_spec, + vocab_size=padded_vocab_size, + max_sequence_length=self._model_config.seq_length, + hybrid_attention_ratio=self._model_config.hybrid_attention_ratio, + hybrid_mlp_ratio=self._model_config.hybrid_mlp_ratio, + hybrid_override_pattern=self._model_config.hybrid_override_pattern, + fp16_lm_cross_entropy=self._model_config.fp16_lm_cross_entropy, + parallel_output=self._model_config.parallel_output, + share_embeddings_and_output_weights=self._model_config.share_embeddings_and_output_weights, + position_embedding_type=self._model_config.position_embedding_type, + rotary_percent=self._model_config.rotary_percent, + rotary_base=self._model_config.rotary_base, + seq_len_interpolation_factor=self._model_config.seq_len_interpolation_factor, + pre_process=pre_process, + post_process=post_process, + pg_collection=pg_collection, + vp_stage=vp_stage, + ) + + 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[MCoreMambaModel]: + """Build model stages and wrap for distributed training. + + Args: + pg_collection: Model communication process groups. + ddp_config: DistributedDataParallel configuration + overlap_param_gather_with_optimizer_step: Whether to overlap parameter + gather with optimizer step. + use_megatron_fsdp: Whether to use Megatron FSDP + use_torch_fsdp2: Whether to use Torch FSDP 2.0 + wrap_with_ddp: Set to False to skip the DDP/FSDP wrapper. + data_parallel_random_init: Whether to use data parallel random initialization + mixed_precision_wrapper: Mixed precision wrapper, e.g. ``Float16Module`` + model_type: Deprecated flag, only used for backwards compatibility. + + Returns: + List of model stages. + """ + transformer_config = self._model_config.transformer + composed_pre_wrap_hook = compose_hooks(self._model_config.pre_wrap_hooks) + model_list = unimodal_build_distributed_models( + self.build_model, + transformer_config, + pg_collection, + ddp_config, + overlap_param_gather_with_optimizer_step, + use_megatron_fsdp, + use_torch_fsdp2, + wrap_with_ddp, + data_parallel_random_init, + mixed_precision_wrapper, + composed_pre_wrap_hook, + model_type, + ) + + composed_post_wrap_hook = compose_hooks(self._model_config.post_wrap_hooks) + _model = composed_post_wrap_hook(model_list) + if _model is not None: + model_list = _model + else: + logger.warning("Final post wrap hook returned None, skipping post wrap hooks.") + + return model_list diff --git a/src/megatron/bridge/models/model_provider.py b/src/megatron/bridge/models/model_provider.py index aca5097e95..f8a5e6867a 100644 --- a/src/megatron/bridge/models/model_provider.py +++ b/src/megatron/bridge/models/model_provider.py @@ -17,6 +17,8 @@ from pathlib import Path from typing import Any, Callable, Generic, TypedDict, TypeVar, Union +from megatron.bridge.models.common.unimodal import _ddp_wrap, _print_num_params + try: from typing import Unpack @@ -34,10 +36,7 @@ import torch from megatron.core import parallel_state, tensor_parallel from megatron.core.distributed import ( - DistributedDataParallel, DistributedDataParallelConfig, - FullyShardedDataParallel, - TorchFullyShardedDataParallel, ) from megatron.core.enums import ModelType from megatron.core.pipeline_parallel.utils import ( @@ -675,84 +674,3 @@ def _create_model( tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param) return model - - -def _ddp_wrap( - model: list[MegatronModule], - data_parallel_random_init: bool, - ddp_config: DistributedDataParallelConfig, - overlap_param_gather_with_optimizer_step: bool, - use_megatron_fsdp: bool = False, - use_torch_fsdp2: bool = False, - *, - pg_collection: ProcessGroupCollection, -) -> list[MegatronModule]: - """Wrap model with Distributed Data Parallel (DDP) or Fully Sharded Data Parallel (FSDP). - - Args: - model: List of model modules to wrap - use_torch_fsdp2: Whether to use PyTorch FSDP v2 instead of DDP - data_parallel_random_init: Whether to broadcast parameters from rank 0 - ddp_config: Configuration for distributed data parallel - overlap_param_gather_with_optimizer_step: Whether to disable bucketing - for overlapping parameter gathering with optimizer step - - Returns: - list[MegatronModule]: List of DDP/FSDP wrapped model modules - """ - if use_megatron_fsdp: - DP = FullyShardedDataParallel - if use_torch_fsdp2: - raise ValueError("Using use_megatron_fsdp and use_torch_fsdp2 at the same time is not supported.") - elif use_torch_fsdp2: - DP = TorchFullyShardedDataParallel - else: - DP = DistributedDataParallel - - # DDP initialization is required to be on a side-stream for the full-iteration CUDA graph. - # this side-stream may be nested if being called from within the get_model function, but it - # is here in case someone wants to use this directly outside of get_model. - ddp_stream = torch.cuda.Stream() - ddp_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(ddp_stream): - model = [ - DP( - config=get_model_config(model_chunk), - ddp_config=ddp_config, - module=model_chunk, - # Turn off bucketing for model_chunk 2 onwards, since communication for these - # model chunks is overlapped with compute anyway. - disable_bucketing=(model_chunk_idx > 0) or overlap_param_gather_with_optimizer_step, - pg_collection=pg_collection, - ) - for (model_chunk_idx, model_chunk) in enumerate(model) - ] - # Critical: ensure side-stream work completes before touching params on default stream - torch.cuda.current_stream().wait_stream(ddp_stream) - - # Broadcast params from data parallel src rank to other data parallel ranks. - if data_parallel_random_init: - for model_module in model: - model_module.broadcast_params() - - return model - - -def _print_num_params(model: list[MegatronModule], pg_collection: ProcessGroupCollection) -> None: - """Print the number of parameters in the model on rank 0. - - Only prints on data parallel rank 0 to avoid duplicate output. - Shows parameter count per (tensor parallel, pipeline parallel) rank. - - Args: - model: List of model modules to count parameters from - """ - if (pg_collection.dp.rank() == 0) and (pg_collection.cp.rank() == 0): - print( - " > number of parameters on (tensor, pipeline) model parallel rank ({}, {}): {}".format( - pg_collection.tp.rank(), - pg_collection.pp.rank(), - sum([sum([p.nelement() for p in model_module.parameters()]) for model_module in model]), - ), - flush=True, - ) diff --git a/tests/unit_tests/models/common/__init__.py b/tests/unit_tests/models/common/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit_tests/models/common/test_base.py b/tests/unit_tests/models/common/test_base.py new file mode 100644 index 0000000000..59705f8e06 --- /dev/null +++ b/tests/unit_tests/models/common/test_base.py @@ -0,0 +1,295 @@ +# 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 dataclasses import dataclass, field +from typing import Callable, ClassVar +from unittest.mock import Mock + +import pytest + +from megatron.bridge.models.common.base import ModelBuilder, ModelConfig, compose_hooks + + +# --------------------------------------------------------------------------- +# Dummy concrete implementations +# --------------------------------------------------------------------------- + + +@dataclass +class DummyModelConfig(ModelConfig): + builder: ClassVar[str] = "" # set dynamically below + value: int = 42 + name: str = "test" + + +class DummyModelBuilder(ModelBuilder): + def build_model(self, pg_collection, pre_process=None, post_process=None, vp_stage=None): + pass + + def build_distributed_models(self, pg_collection, **kwargs): + return [] + + +# Set after both classes exist so __module__ is always the actual runtime path. +# get_builder_cls() calls importlib.import_module on the module portion of this +# string, so it must match what Python sees for this test module. +DummyModelConfig.builder = f"{DummyModelBuilder.__module__}.DummyModelBuilder" + + +# A plain nested dataclass (not a ModelConfig subclass) for testing nested +# serialization without depending on any real model classes. +@dataclass +class DummySubConfig: + x: int = 1 + y: str = "sub" + + +def _dummy_callable() -> None: + """Placeholder callable used as a field default in DummyNestedModelConfig.""" + + +@dataclass +class DummyNestedModelConfig(ModelConfig): + """ModelConfig subclass with a nested dataclass field and a callable field. + + Used to test nested serialization and callable-field exclusion without + depending on MambaModelConfig or TransformerConfig. + """ + + builder: ClassVar[str] = "" # set dynamically below + sub: DummySubConfig = field(default_factory=DummySubConfig) + fn_field: Callable = _dummy_callable + extra: int = 0 + + +DummyNestedModelConfig.builder = f"{DummyModelBuilder.__module__}.DummyModelBuilder" + + +# ============================================================================= +# Section 1 — TestModelConfigDefaults +# ============================================================================= + + +class TestModelConfigDefaults: + """Base class field defaults and ClassVar on a concrete subclass.""" + + def test_base_field_defaults(self): + cfg = DummyModelConfig() + assert cfg.restore_modelopt_state is False + assert cfg.hf_model_id is None + assert cfg.generation_config is None + assert cfg.pre_wrap_hooks == [] + assert cfg.post_wrap_hooks == [] + + def test_custom_fields_stored(self): + hook1 = Mock() + hook2 = Mock() + cfg = DummyModelConfig(value=99, name="hello", pre_wrap_hooks=[hook1], post_wrap_hooks=[hook2]) + assert cfg.value == 99 + assert cfg.name == "hello" + assert cfg.pre_wrap_hooks[0] == hook1 + assert cfg.post_wrap_hooks[0] == hook2 + + def test_builder_classvar_accessible(self): + assert DummyModelConfig.builder == "tests.unit_tests.models.common.test_base.DummyModelBuilder" + + +# ============================================================================= +# Section 2 — TestModelConfigGetBuilderCls +# ============================================================================= + + +class TestModelConfigGetBuilderCls: + """get_builder_cls() dynamically imports and returns the builder class named by the ClassVar.""" + + def test_returns_correct_type(self): + cfg = DummyModelConfig() + result = cfg.get_builder_cls() + assert result is DummyModelBuilder + + def test_return_is_class_not_instance(self): + cfg = DummyModelConfig() + result = cfg.get_builder_cls() + assert isinstance(result, type) + + +# ============================================================================= +# Section 3 — TestModelConfigToDict +# ============================================================================= + + +class TestModelConfigToDict: + """as_dict() serializes all non-callable, non-private dataclass fields, including nested dataclasses.""" + + def test_target_key_present(self): + cfg = DummyModelConfig() + result = cfg.as_dict() + assert result["_target_"] == f"{DummyModelConfig.__module__}.DummyModelConfig" + + def test_builder_key_present(self): + cfg = DummyModelConfig() + result = cfg.as_dict() + assert result["_builder_"] == DummyModelConfig.builder + + def test_own_fields_serialized(self): + cfg = DummyModelConfig(value=7, name="world") + result = cfg.as_dict() + assert result["value"] == 7 + assert result["name"] == "world" + + def test_base_class_fields_serialized(self): + cfg = DummyModelConfig() + result = cfg.as_dict() + assert "restore_modelopt_state" in result + assert "hf_model_id" in result + assert "generation_config" in result + + def test_hook_lists_excluded(self): + cfg = DummyModelConfig() + result = cfg.as_dict() + assert "pre_wrap_hooks" not in result + assert "post_wrap_hooks" not in result + + def test_callable_field_excluded(self): + cfg = DummyNestedModelConfig() + result = cfg.as_dict() + assert "fn_field" not in result + + def test_nested_dataclass_serialized_recursively(self): + cfg = DummyNestedModelConfig() + result = cfg.as_dict() + assert isinstance(result["sub"], dict) + assert "_target_" in result["sub"] + assert result["sub"]["x"] == 1 + assert result["sub"]["y"] == "sub" + + +# ============================================================================= +# Section 4 — TestModelConfigFromDict +# ============================================================================= + + +class TestModelConfigFromDict: + """from_dict() reconstructs configs from serialized dicts, handles nested dataclasses, and is robust to unknown keys.""" + + def _dummy_dict(self, **overrides): + d = { + "_target_": f"{DummyModelConfig.__module__}.DummyModelConfig", + "_builder_": DummyModelConfig.builder, + "value": 10, + "name": "from_dict_test", + } + d.update(overrides) + return d + + def test_reconstructs_flat_config(self): + d = self._dummy_dict(value=55, name="reconstructed") + cfg = ModelConfig.from_dict(d) + assert isinstance(cfg, DummyModelConfig) + assert cfg.value == 55 + assert cfg.name == "reconstructed" + + def test_builder_restored(self): + d = self._dummy_dict(_builder_="fake.builder.string") + cfg = ModelConfig.from_dict(d) + assert cfg.builder == "fake.builder.string" + + def test_ignores_unknown_fields(self): + d = self._dummy_dict(unknown_field="should_be_ignored", another_unknown=123) + cfg = ModelConfig.from_dict(d) + assert isinstance(cfg, DummyModelConfig) + assert cfg.value == 10 + assert cfg.name == "from_dict_test" + assert not hasattr(cfg, "unknown_field") + assert not hasattr(cfg, "another_unknown") + + def test_raises_if_target_missing(self): + d = {"_builder_": DummyModelConfig.builder, "value": 1} + with pytest.raises(ValueError): + ModelConfig.from_dict(d) + + def test_round_trip_flat(self): + original = DummyModelConfig(value=21, name="round_trip") + cfg = ModelConfig.from_dict(original.as_dict()) + assert isinstance(cfg, DummyModelConfig) + assert cfg.value == original.value + assert cfg.name == original.name + + def test_round_trip_with_nested_dataclass(self): + original = DummyNestedModelConfig(sub=DummySubConfig(x=7, y="nested"), extra=99) + cfg = ModelConfig.from_dict(original.as_dict()) + assert isinstance(cfg, DummyNestedModelConfig) + assert cfg.extra == 99 + assert isinstance(cfg.sub, DummySubConfig) + assert cfg.sub.x == 7 + assert cfg.sub.y == "nested" + + +# ============================================================================= +# Section 6 — TestComposeHooks +# ============================================================================= + + +class TestComposeHooks: + """compose_hooks() folds a list of hook functions into a single function, threading output through each in order.""" + + def test_empty_list_is_identity(self): + composed = compose_hooks([]) + sentinel = object() + assert composed(sentinel) is sentinel + + def test_single_hook_equivalent(self): + h = Mock(return_value="result_x") + composed = compose_hooks([h]) + result = composed("x") + h.assert_called_once_with("x") + assert result == "result_x" + + def test_multiple_hooks_chained(self): + h1 = Mock(return_value="after_h1") + h2 = Mock(return_value="after_h2") + composed = compose_hooks([h1, h2]) + result = composed("start") + h1.assert_called_once_with("start") + h2.assert_called_once_with("after_h1") + assert result == "after_h2" + + def test_hooks_called_in_list_order(self): + call_log = [] + + def h1(x): + call_log.append("h1") + return x + + def h2(x): + call_log.append("h2") + return x + + def h3(x): + call_log.append("h3") + return x + + composed = compose_hooks([h1, h2, h3]) + composed("anything") + assert call_log == ["h1", "h2", "h3"] + + def test_intermediate_values_threaded(self): + def double(x): + return x * 2 + + def add_ten(x): + return x + 10 + + composed = compose_hooks([double, add_ten]) + assert composed(5) == 20 # 5 * 2 = 10, 10 + 10 = 20 diff --git a/tests/unit_tests/models/common/test_unimodal.py b/tests/unit_tests/models/common/test_unimodal.py new file mode 100644 index 0000000000..896bbd2d6f --- /dev/null +++ b/tests/unit_tests/models/common/test_unimodal.py @@ -0,0 +1,838 @@ +# 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 unittest.mock import MagicMock, Mock, patch + +import pytest +import torch +import torch.nn as nn +from megatron.core.enums import ModelType + +from megatron.bridge.models.common.unimodal import ( + _ddp_wrap, + _print_num_params, + _wrap_with_mp_wrapper, + build_virtual_pipeline_stages, + to_empty_if_meta_device, + unimodal_build_distributed_models, +) + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _make_pg(): + """Mock ProcessGroupCollection with dp, cp, tp, pp sub-groups.""" + pg = Mock() + pg.dp.rank.return_value = 0 + pg.cp.rank.return_value = 0 + pg.tp.rank.return_value = 0 + pg.pp.rank.return_value = 0 + pg.pp.size.return_value = 1 + return pg + + +def _make_transformer_config(**kwargs): + cfg = Mock() + cfg.virtual_pipeline_model_parallel_size = None + cfg.init_model_with_meta_device = False + cfg.use_cpu_initialization = False + cfg.fp16 = False + cfg.bf16 = False + for k, v in kwargs.items(): + setattr(cfg, k, v) + return cfg + + +def _make_model_module(): + m = Mock() + m.parameters.return_value = [] + m.modules.return_value = [] + return m + + +# ============================================================================= +# Section 1 — TestToEmptyIfMetaDevice +# ============================================================================= + + +class TestToEmptyIfMetaDevice: + """to_empty_if_meta_device() materialises meta-device parameters while leaving non-meta parameters unchanged.""" + + def test_meta_parameter_becomes_empty_on_target_device(self): + module = nn.Module() + module.register_parameter("weight", nn.Parameter(torch.empty(4).to("meta"))) + result = to_empty_if_meta_device(module, device=torch.device("cpu")) + assert result.weight.device == torch.device("cpu") + + def test_non_meta_parameter_moved_to_device(self): + module = nn.Module() + module.register_parameter("weight", nn.Parameter(torch.zeros(4))) + result = to_empty_if_meta_device(module, device=torch.device("cpu")) + assert result.weight.device == torch.device("cpu") + + def test_recurse_true_applies_to_submodule_parameters(self): + parent = nn.Module() + child = nn.Module() + child.register_parameter("weight", nn.Parameter(torch.empty(4).to("meta"))) + parent.add_module("child", child) + to_empty_if_meta_device(parent, device=torch.device("cpu"), recurse=True) + assert parent.child.weight.device == torch.device("cpu") + + def test_recurse_false_skips_submodule_parameters(self): + parent = nn.Module() + child = nn.Module() + child.register_parameter("weight", nn.Parameter(torch.empty(4).to("meta"))) + parent.add_module("child", child) + to_empty_if_meta_device(parent, device=torch.device("cpu"), recurse=False) + assert parent.child.weight.device == torch.device("meta") + + +# ============================================================================= +# Section 2 — TestBuildVirtualPipelineStages +# ============================================================================= + + +class TestBuildVirtualPipelineStages: + """build_virtual_pipeline_stages() builds one stage without VP or multiple stages with VP, setting model_type on each.""" + + @patch("megatron.core.pipeline_parallel.utils.is_pp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_pp_first_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_first_stage", return_value=True) + def test_single_stage_when_pp_size_one(self, *_): + pg = _make_pg() + pg.pp.size.return_value = 1 + build_fn = Mock(return_value=Mock()) + result = build_virtual_pipeline_stages(build_fn, pg, vp_size=3) + assert len(result) == 1 + build_fn.assert_called_once() + + @patch("megatron.core.pipeline_parallel.utils.is_pp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_pp_first_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_first_stage", return_value=True) + def test_single_stage_when_vp_size_none(self, *_): + pg = _make_pg() + pg.pp.size.return_value = 2 + build_fn = Mock(return_value=Mock()) + result = build_virtual_pipeline_stages(build_fn, pg, vp_size=None) + assert len(result) == 1 + build_fn.assert_called_once() + + @patch("megatron.core.pipeline_parallel.utils.is_pp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_pp_first_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_first_stage", return_value=True) + def test_vp_builds_vp_size_stages(self, *_): + pg = _make_pg() + pg.pp.size.return_value = 2 + build_fn = Mock(return_value=Mock()) + result = build_virtual_pipeline_stages(build_fn, pg, vp_size=3) + assert len(result) == 3 + assert build_fn.call_count == 3 + + @patch("megatron.core.pipeline_parallel.utils.is_pp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_pp_first_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_first_stage", return_value=True) + def test_vp_passes_correct_vp_stage_index(self, *_): + pg = _make_pg() + pg.pp.size.return_value = 2 + build_fn = Mock(return_value=Mock()) + build_virtual_pipeline_stages(build_fn, pg, vp_size=3) + vp_stage_args = [c.kwargs["vp_stage"] for c in build_fn.call_args_list] + assert vp_stage_args == [0, 1, 2] + + @patch("megatron.core.pipeline_parallel.utils.is_pp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_pp_first_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_first_stage", return_value=True) + def test_model_type_set_on_every_stage(self, *_): + pg = _make_pg() + pg.pp.size.return_value = 2 + build_fn = Mock(return_value=Mock()) + result = build_virtual_pipeline_stages(build_fn, pg, vp_size=3) + for model in result: + assert model.model_type == ModelType.encoder_or_decoder + + @patch("megatron.core.pipeline_parallel.utils.is_pp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_pp_first_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_last_stage") + @patch("megatron.core.pipeline_parallel.utils.is_vp_first_stage") + def test_vp_pre_process_uses_vp_and_pp(self, mock_vp_first, mock_vp_last, mock_pp_first, mock_pp_last): + # First VP stage (i=0): vp_first=True → pre_process = True (and pp_first=True) + # Other VP stages: vp_first=False → pre_process = False + pg = _make_pg() + pg.pp.size.return_value = 2 + + # is_vp_first_stage returns True only when vp_stage=0 + mock_vp_first.side_effect = lambda vp_stage, vp_size: vp_stage == 0 + mock_vp_last.side_effect = lambda vp_stage, vp_size: vp_stage == (vp_size - 1) + + build_fn = Mock(return_value=Mock()) + build_virtual_pipeline_stages(build_fn, pg, vp_size=3) + + # Stage 0 should have pre_process=True, stages 1 and 2 should have pre_process=False + assert build_fn.call_args_list[0].kwargs["pre_process"] is True + assert build_fn.call_args_list[1].kwargs["pre_process"] is False + assert build_fn.call_args_list[2].kwargs["pre_process"] is False + + assert build_fn.call_args_list[0].kwargs["post_process"] is False + assert build_fn.call_args_list[1].kwargs["post_process"] is False + assert build_fn.call_args_list[2].kwargs["post_process"] is True + + @patch("megatron.core.pipeline_parallel.utils.is_pp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_pp_first_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_last_stage", return_value=True) + @patch("megatron.core.pipeline_parallel.utils.is_vp_first_stage", return_value=True) + def test_returns_list_of_models(self, *_): + pg = _make_pg() + pg.pp.size.return_value = 1 + build_fn = Mock(return_value=Mock()) + result = build_virtual_pipeline_stages(build_fn, pg, vp_size=None) + assert isinstance(result, list) + + +# ============================================================================= +# Section 3 — TestPrintNumParams +# ============================================================================= + + +class TestPrintNumParams: + """_print_num_params() prints parameter counts only on data-parallel and context-parallel rank 0.""" + + def test_prints_on_dp0_cp0(self, capsys): + pg = _make_pg() # dp.rank()=0, cp.rank()=0 by default + model = [_make_model_module()] + _print_num_params(model, pg_collection=pg) + captured = capsys.readouterr() + assert len(captured.out) > 0 + + def test_silent_on_nonzero_dp_rank(self, capsys): + pg = _make_pg() + pg.dp.rank.return_value = 1 + model = [_make_model_module()] + _print_num_params(model, pg_collection=pg) + captured = capsys.readouterr() + assert captured.out == "" + + def test_silent_on_nonzero_cp_rank(self, capsys): + pg = _make_pg() + pg.cp.rank.return_value = 1 + model = [_make_model_module()] + _print_num_params(model, pg_collection=pg) + captured = capsys.readouterr() + assert captured.out == "" + + def test_param_count_is_correct(self, capsys): + pg = _make_pg() + m1 = Mock() + m1.parameters.return_value = [torch.zeros(3)] + m2 = Mock() + m2.parameters.return_value = [torch.zeros(5)] + _print_num_params([m1, m2], pg_collection=pg) + captured = capsys.readouterr() + assert "8" in captured.out + + +# ============================================================================= +# Section 4 — TestWrapWithMpWrapper +# ============================================================================= + + +class TestWrapWithMpWrapper: + """_wrap_with_mp_wrapper() applies the mixed-precision wrapper to each stage when fp16 or bf16 is set.""" + + def test_no_wrap_when_fp16_false_bf16_false(self): + cfg = _make_transformer_config(fp16=False, bf16=False) + stages = [_make_model_module(), _make_model_module()] + wrapper = Mock() + result = _wrap_with_mp_wrapper(stages, cfg, wrapper) + assert result is stages + wrapper.assert_not_called() + + def test_wraps_each_stage_when_fp16_true(self): + cfg = _make_transformer_config(fp16=True) + stage1 = _make_model_module() + stage2 = _make_model_module() + wrapped1, wrapped2 = Mock(), Mock() + wrapped1.modules.return_value = [] + wrapped2.modules.return_value = [] + wrapper = Mock(side_effect=[wrapped1, wrapped2]) + result = _wrap_with_mp_wrapper([stage1, stage2], cfg, wrapper) + assert wrapper.call_count == 2 + assert wrapper.call_args_list[0].args == (cfg, stage1) + assert wrapper.call_args_list[1].args == (cfg, stage2) + assert result == [wrapped1, wrapped2] + + def test_wraps_each_stage_when_bf16_true(self): + cfg = _make_transformer_config(bf16=True) + stage = _make_model_module() + wrapped = Mock() + wrapped.modules.return_value = [] + wrapper = Mock(return_value=wrapped) + result = _wrap_with_mp_wrapper([stage], cfg, wrapper) + wrapper.assert_called_once_with(cfg, stage) + assert result == [wrapped] + + def test_no_wrap_when_mixed_precision_wrapper_is_none(self): + cfg = _make_transformer_config(fp16=True) + stages = [_make_model_module()] + result = _wrap_with_mp_wrapper(stages, cfg, None) + assert result is stages + + def test_expert_bias_hook_called_on_qualifying_submodule(self): + cfg = _make_transformer_config(fp16=True) + stage = _make_model_module() + submodule = Mock() + submodule._maintain_float32_expert_bias = Mock() + wrapped = Mock() + wrapped.modules.return_value = [submodule] + wrapper = Mock(return_value=wrapped) + _wrap_with_mp_wrapper([stage], cfg, wrapper) + submodule._maintain_float32_expert_bias.assert_called_once() + + +# ============================================================================= +# Section 5 — TestDdpWrap +# ============================================================================= + + +class TestDdpWrap: + """_ddp_wrap() wraps each model stage with DDP, Megatron-FSDP, or Torch FSDP2, and optionally broadcasts params.""" + + def setup_method(self): + self.pg = _make_pg() + self.ddp_config = Mock() + self.model = [_make_model_module(), _make_model_module()] + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_raises_when_both_fsdp_flags_set(self, mock_stream, mock_curr, mock_ctx, *_): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + with pytest.raises(ValueError): + _ddp_wrap( + self.model, + False, + self.ddp_config, + False, + use_megatron_fsdp=True, + use_torch_fsdp2=True, + pg_collection=self.pg, + ) + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_uses_ddp_by_default( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + _ddp_wrap(self.model, False, self.ddp_config, False, pg_collection=self.pg) + assert mock_ddp.call_count == 2 + mock_fsdp.assert_not_called() + mock_torch_fsdp.assert_not_called() + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_uses_megatron_fsdp_when_flagged( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + _ddp_wrap(self.model, False, self.ddp_config, False, use_megatron_fsdp=True, pg_collection=self.pg) + assert mock_fsdp.call_count == 2 + mock_ddp.assert_not_called() + mock_torch_fsdp.assert_not_called() + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_uses_torch_fsdp2_when_flagged( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + _ddp_wrap(self.model, False, self.ddp_config, False, use_torch_fsdp2=True, pg_collection=self.pg) + assert mock_torch_fsdp.call_count == 2 + mock_ddp.assert_not_called() + mock_fsdp.assert_not_called() + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_broadcasts_params_when_data_parallel_random_init_true( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + wrapped1 = Mock() + wrapped2 = Mock() + mock_ddp.side_effect = [wrapped1, wrapped2] + _ddp_wrap(self.model, True, self.ddp_config, False, pg_collection=self.pg) + wrapped1.broadcast_params.assert_called_once() + wrapped2.broadcast_params.assert_called_once() + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_no_broadcast_when_data_parallel_random_init_false( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + wrapped1 = Mock() + wrapped2 = Mock() + mock_ddp.side_effect = [wrapped1, wrapped2] + _ddp_wrap(self.model, False, self.ddp_config, False, pg_collection=self.pg) + wrapped1.broadcast_params.assert_not_called() + wrapped2.broadcast_params.assert_not_called() + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_first_chunk_bucketing_enabled( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + # Use single-element list so only chunk 0 is created + model = [_make_model_module()] + _ddp_wrap(model, False, self.ddp_config, False, pg_collection=self.pg) + call_kwargs = mock_ddp.call_args.kwargs + assert call_kwargs["disable_bucketing"] is False + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_second_chunk_bucketing_disabled( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + _ddp_wrap(self.model, False, self.ddp_config, False, pg_collection=self.pg) + # Second call corresponds to chunk index 1 + second_call_kwargs = mock_ddp.call_args_list[1].kwargs + assert second_call_kwargs["disable_bucketing"] is True + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_overlap_param_gather_disables_bucketing_for_all( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + _ddp_wrap(self.model, False, self.ddp_config, True, pg_collection=self.pg) + for c in mock_ddp.call_args_list: + assert c.kwargs["disable_bucketing"] is True + + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.FullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_returns_list_of_wrapped_modules( + self, mock_stream, mock_curr, mock_ctx, mock_cfg, mock_ddp, mock_fsdp, mock_torch_fsdp + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + result = _ddp_wrap(self.model, False, self.ddp_config, False, pg_collection=self.pg) + assert isinstance(result, list) + assert len(result) == 2 + + +# ============================================================================= +# Section 6 — TestUnimodalBuildDistributedModels +# ============================================================================= + +_MODULE = "megatron.bridge.models.common.unimodal" + + +class TestUnimodalBuildDistributedModels: + """unimodal_build_distributed_models() orchestrates stage building, hooks, parameter setup, GPU allocation, wrapping, and DDP.""" + + def setup_method(self): + self.pg = _make_pg() + self.mock_model = _make_model_module() + self.mock_model.parameters.return_value = [] + self.transformer_config = _make_transformer_config() + + # ------------------------------------------------------------------ + # Helpers to build the standard patch stack + # ------------------------------------------------------------------ + + def _standard_patches(self): + """Returns a dict of started patches that the caller must stop.""" + patches = { + "bvps": patch(f"{_MODULE}.build_virtual_pipeline_stages", return_value=[self.mock_model]), + "mp_wrap": patch(f"{_MODULE}._wrap_with_mp_wrapper", side_effect=lambda m, *_: m), + "ddp": patch(f"{_MODULE}._ddp_wrap", side_effect=lambda m, *args, **kwargs: m), + "print": patch(f"{_MODULE}._print_num_params"), + "tp_attr": patch(f"{_MODULE}.tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes"), + "cuda_dev": patch("torch.cuda.current_device", return_value=0), + } + started = {k: p.start() for k, p in patches.items()} + # Store originals for teardown + self._patch_objs = list(patches.values()) + return started + + def _stop_patches(self): + for p in getattr(self, "_patch_objs", []): + p.stop() + self._patch_objs = [] + + def teardown_method(self): + self._stop_patches() + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + def test_raises_when_wrap_with_ddp_true_but_no_ddp_config(self): + self._standard_patches() + try: + with pytest.raises(ValueError): + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + ddp_config=None, + wrap_with_ddp=True, + ) + finally: + self._stop_patches() + + def test_builds_stages_via_build_virtual_pipeline_stages(self): + mocks = self._standard_patches() + try: + build_fn = Mock() + unimodal_build_distributed_models( + build_fn, + self.transformer_config, + self.pg, + wrap_with_ddp=False, + ) + mocks["bvps"].assert_called_once_with( + build_fn, + self.pg, + self.transformer_config.virtual_pipeline_model_parallel_size, + ModelType.encoder_or_decoder, + ) + finally: + self._stop_patches() + + def test_meta_device_context_used_when_init_with_meta_device(self): + transformer_config = _make_transformer_config(init_model_with_meta_device=True) + mocks = self._standard_patches() + try: + # Should complete without error; meta device path goes through build_virtual_pipeline_stages + unimodal_build_distributed_models( + Mock(), + transformer_config, + self.pg, + wrap_with_ddp=False, + ) + mocks["bvps"].assert_called_once() + finally: + self._stop_patches() + + def test_pre_wrap_hook_applied_and_model_list_updated(self): + mocks = self._standard_patches() + try: + new_list = [_make_model_module()] + hook = Mock(return_value=new_list) + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + pre_wrap_hook=hook, + ) + hook.assert_called_once() + # The mp_wrapper receives the new list + mp_call_arg = mocks["mp_wrap"].call_args.args[0] + assert mp_call_arg is new_list + finally: + self._stop_patches() + + def test_pre_wrap_hook_returning_none_keeps_original_list(self): + mocks = self._standard_patches() + try: + hook = Mock(return_value=None) + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + pre_wrap_hook=hook, + ) + # mp_wrapper receives the original list [self.mock_model] + mp_call_arg = mocks["mp_wrap"].call_args.args[0] + assert mp_call_arg == [self.mock_model] + finally: + self._stop_patches() + + def test_pre_wrap_hook_not_callable_raises_type_error(self): + self._standard_patches() + try: + with pytest.raises(TypeError): + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + pre_wrap_hook="not_callable", + ) + finally: + self._stop_patches() + + def test_pre_wrap_hook_none_is_skipped(self): + self._standard_patches() + try: + # Should not raise + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + pre_wrap_hook=None, + ) + finally: + self._stop_patches() + + def test_tensor_parallel_attrs_set_for_each_param(self): + param = Mock() + self.mock_model.parameters.return_value = [param] + mocks = self._standard_patches() + try: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + ) + mocks["tp_attr"].assert_called_once_with(param) + finally: + self._stop_patches() + + def test_cuda_called_when_not_fsdp2_not_cpu_not_meta(self): + self._standard_patches() + try: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + use_torch_fsdp2=False, + ) + self.mock_model.cuda.assert_called_once() + finally: + self._stop_patches() + + def test_cuda_not_called_when_use_torch_fsdp2(self): + self._standard_patches() + try: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + use_torch_fsdp2=True, + ) + self.mock_model.cuda.assert_not_called() + finally: + self._stop_patches() + + def test_cuda_not_called_when_use_cpu_initialization(self): + self.transformer_config = _make_transformer_config(use_cpu_initialization=True) + self._standard_patches() + try: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + ) + self.mock_model.cuda.assert_not_called() + finally: + self._stop_patches() + + def test_cuda_not_called_when_init_model_with_meta_device(self): + self.transformer_config = _make_transformer_config(init_model_with_meta_device=True) + self._standard_patches() + try: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + ) + self.mock_model.cuda.assert_not_called() + finally: + self._stop_patches() + + def test_meta_materialization_called_when_meta_not_fsdp(self): + self.transformer_config = _make_transformer_config(init_model_with_meta_device=True) + self._standard_patches() + try: + with patch(f"{_MODULE}.to_empty_if_meta_device", return_value=self.mock_model) as mock_toempty: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + use_torch_fsdp2=False, + use_megatron_fsdp=False, + ) + assert mock_toempty.call_count == 1 + finally: + self._stop_patches() + + def test_meta_materialization_skipped_with_torch_fsdp2(self): + self.transformer_config = _make_transformer_config(init_model_with_meta_device=True) + self._standard_patches() + try: + with patch(f"{_MODULE}.to_empty_if_meta_device", return_value=self.mock_model) as mock_toempty: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + use_torch_fsdp2=True, + ) + mock_toempty.assert_not_called() + finally: + self._stop_patches() + + def test_meta_materialization_skipped_with_megatron_fsdp(self): + self.transformer_config = _make_transformer_config(init_model_with_meta_device=True) + self._standard_patches() + try: + with patch(f"{_MODULE}.to_empty_if_meta_device", return_value=self.mock_model) as mock_toempty: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + use_megatron_fsdp=True, + ) + mock_toempty.assert_not_called() + finally: + self._stop_patches() + + def test_mp_wrapper_applied(self): + mocks = self._standard_patches() + try: + mp_wrapper = Mock() + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + mixed_precision_wrapper=mp_wrapper, + ) + mocks["mp_wrap"].assert_called_once_with([self.mock_model], self.transformer_config, mp_wrapper) + finally: + self._stop_patches() + + def test_ddp_wrap_called_when_wrap_with_ddp_true(self): + mocks = self._standard_patches() + try: + ddp_config = Mock() + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + ddp_config=ddp_config, + wrap_with_ddp=True, + ) + mocks["ddp"].assert_called_once() + finally: + self._stop_patches() + + def test_ddp_wrap_not_called_when_wrap_with_ddp_false(self): + mocks = self._standard_patches() + try: + unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + wrap_with_ddp=False, + ) + mocks["ddp"].assert_not_called() + finally: + self._stop_patches() + + def test_returns_final_model_list(self): + mocks = self._standard_patches() + try: + ddp_config = Mock() + ddp_result = [_make_model_module()] + mocks["ddp"].side_effect = lambda m, *args, **kwargs: ddp_result + result = unimodal_build_distributed_models( + Mock(), + self.transformer_config, + self.pg, + ddp_config=ddp_config, + wrap_with_ddp=True, + ) + assert result is ddp_result + finally: + self._stop_patches() diff --git a/tests/unit_tests/models/mamba/test_mamba_builder.py b/tests/unit_tests/models/mamba/test_mamba_builder.py new file mode 100644 index 0000000000..f795d44123 --- /dev/null +++ b/tests/unit_tests/models/mamba/test_mamba_builder.py @@ -0,0 +1,525 @@ +# 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 unittest.mock import Mock, call, patch + +import pytest +from megatron.core.transformer import ModuleSpec + +from megatron.bridge.models.mamba.mamba_builder import ( + MambaModelBuilder, + MambaModelConfig, + get_default_mamba_stack_spec, + modelopt_mamba_stack_spec, + transformer_engine_mamba_stack_spec, +) +from megatron.bridge.models.transformer_config import TransformerConfig + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_transformer(**kwargs): + defaults = dict(num_layers=2, hidden_size=128, num_attention_heads=1) + defaults.update(kwargs) + return TransformerConfig(**defaults) + + +def _make_mamba_config(**kwargs): + defaults = dict(transformer=_make_transformer(), vocab_size=32000) + defaults.update(kwargs) + return MambaModelConfig(**defaults) + + +# ============================================================================= +# Section 1 — Spec Functions +# ============================================================================= + + +class TestTransformerEngineMambaStackSpec: + """Tests for the transformer_engine_mamba_stack_spec() factory function.""" + + def test_returns_module_spec(self): + result = transformer_engine_mamba_stack_spec() + assert isinstance(result, ModuleSpec) + + def test_returns_expected_spec_object(self): + with patch("megatron.bridge.models.mamba.mamba_builder.default_mamba_stack_spec") as mock_spec: + result = transformer_engine_mamba_stack_spec() + assert result is mock_spec + + +class TestModeloptMambaStackSpec: + """Tests for the modelopt_mamba_stack_spec() factory function.""" + + def test_returns_module_spec(self): + mock_spec = Mock(spec=ModuleSpec) + with patch( + "megatron.bridge.models.mamba.mamba_builder.get_mamba_stack_modelopt_spec", + return_value=mock_spec, + ): + result = modelopt_mamba_stack_spec() + assert result is mock_spec + + def test_calls_modelopt_spec_with_correct_args(self): + with patch( + "megatron.bridge.models.mamba.mamba_builder.get_mamba_stack_modelopt_spec", + ) as mock_fn: + mock_fn.return_value = Mock(spec=ModuleSpec) + modelopt_mamba_stack_spec() + mock_fn.assert_called_once_with(local_core_attention=False, remap_te_layernorm=False) + + +class TestGetDefaultMambaStackSpec: + """Tests for get_default_mamba_stack_spec(), which dispatches on restore_modelopt_state.""" + + def test_returns_te_spec_when_restore_modelopt_state_is_false(self): + mock_spec = Mock(spec=ModuleSpec) + config = Mock() + config.restore_modelopt_state = False + with patch( + "megatron.bridge.models.mamba.mamba_builder.transformer_engine_mamba_stack_spec", + return_value=mock_spec, + ) as mock_fn: + result = get_default_mamba_stack_spec(config) + mock_fn.assert_called_once() + assert result is mock_spec + + def test_returns_modelopt_spec_when_restore_modelopt_state_is_true(self): + mock_spec = Mock(spec=ModuleSpec) + config = Mock() + config.restore_modelopt_state = True + with patch( + "megatron.bridge.models.mamba.mamba_builder.modelopt_mamba_stack_spec", + return_value=mock_spec, + ) as mock_fn: + result = get_default_mamba_stack_spec(config) + mock_fn.assert_called_once() + assert result is mock_spec + + +# ============================================================================= +# Section 2 — MambaModelConfig +# ============================================================================= + + +class TestMambaModelConfigInitialization: + """Tests for MambaModelConfig field defaults and custom initialization.""" + + def test_builder_classvar(self): + assert MambaModelConfig.builder == "megatron.bridge.models.mamba.MambaModelBuilder" + + def test_default_values(self): + config = MambaModelConfig(transformer=_make_transformer()) + assert config.fp16_lm_cross_entropy is False + assert config.parallel_output is True + assert config.share_embeddings_and_output_weights is False + assert config.hybrid_attention_ratio == 0.0 + assert config.hybrid_mlp_ratio == 0.0 + assert config.hybrid_override_pattern is None + assert config.seq_length == 8192 + assert config.position_embedding_type == "none" + assert config.rotary_percent == 1.0 + assert config.rotary_base == 10000 + assert config.seq_len_interpolation_factor is None + assert config.make_vocab_size_divisible_by == 128 + assert config.vocab_size is None + assert config.should_pad_vocab is False + + def test_custom_initialization(self): + config = MambaModelConfig( + transformer=_make_transformer(), + fp16_lm_cross_entropy=True, + parallel_output=False, + hybrid_attention_ratio=0.25, + hybrid_mlp_ratio=0.1, + hybrid_override_pattern="M-M*-", + seq_length=4096, + vocab_size=50000, + ) + assert config.fp16_lm_cross_entropy is True + assert config.parallel_output is False + assert config.hybrid_attention_ratio == 0.25 + assert config.hybrid_mlp_ratio == 0.1 + assert config.hybrid_override_pattern == "M-M*-" + assert config.seq_length == 4096 + assert config.vocab_size == 50000 + + def test_mamba_stack_spec_default_is_callable(self): + config = _make_mamba_config() + assert callable(config.mamba_stack_spec) + assert config.mamba_stack_spec is get_default_mamba_stack_spec + + +class TestMambaModelConfigGetAttr: + """Tests for MambaModelConfig.__getattr__ — direct access vs. TransformerConfig proxy.""" + + def setup_method(self): + self.transformer = _make_transformer(hidden_size=256, num_layers=4) + self.config = MambaModelConfig(transformer=self.transformer, vocab_size=32000) + + def test_own_attribute_not_proxied(self): + # vocab_size is defined on MambaModelConfig but not on TransformerConfig; + # it is returned directly from config.__dict__, __getattr__ is never invoked. + assert self.config.vocab_size == 32000 + + def test_proxies_transformer_attribute(self): + # hidden_size is not a field on MambaModelConfig, so __getattr__ proxies to transformer + assert self.config.hidden_size == 256 + + def test_raises_attribute_error_for_unknown(self): + with pytest.raises(AttributeError): + _ = self.config.completely_unknown_attr_xyz + + def test_raises_before_transformer_init(self): + # Simulate the "transformer not yet set" path in __getattr__ + del self.config.__dict__["transformer"] + with pytest.raises(AttributeError): + _ = self.config.hidden_size + + def test_error_message_contains_attr_name(self): + attr_name = "completely_unknown_attr_xyz" + with pytest.raises(AttributeError, match=attr_name): + getattr(self.config, attr_name) + + +class TestMambaModelConfigSetAttr: + """Tests for MambaModelConfig.__setattr__ — own-field writes vs. TransformerConfig proxy writes.""" + + def setup_method(self): + self.transformer = _make_transformer(hidden_size=256) + self.config = MambaModelConfig(transformer=self.transformer, vocab_size=32000) + + def test_sets_own_attribute_on_self(self): + self.config.vocab_size = 50000 + assert self.config.vocab_size == 50000 + assert self.config.__dict__.get("vocab_size") == 50000 + + def test_proxies_set_to_transformer_attribute(self): + self.config.hidden_size = 512 + assert self.transformer.hidden_size == 512 + + def test_set_proxied_attr_reflects_on_transformer(self): + self.config.hidden_size = 1024 + assert self.config.hidden_size == 1024 + assert self.transformer.hidden_size == 1024 + + def test_set_before_transformer_init(self): + # Delete transformer to simulate the pre-init state in __setattr__ + del self.config.__dict__["transformer"] + self.config.vocab_size = 42 + assert self.config.__dict__["vocab_size"] == 42 + + def test_set_transformer_itself_stores_on_self(self): + new_transformer = _make_transformer(hidden_size=512) + self.config.transformer = new_transformer + assert self.config.transformer is new_transformer + assert self.config.__dict__["transformer"] is new_transformer + + def test_set_own_attr_does_not_go_to_transformer(self): + # vocab_size is not on TransformerConfig, so the write should stay on self + self.config.vocab_size = 99999 + assert self.config.__dict__.get("vocab_size") == 99999 + assert not hasattr(self.config.transformer, "vocab_size") + + def test_proxied_write_does_not_shadow_on_self(self): + self.config.hidden_size = 2048 + assert "hidden_size" not in self.config.__dict__ + + +# ============================================================================= +# Section 3 — MambaModelBuilder +# ============================================================================= + + +class TestMambaModelBuilderInit: + """Tests for MambaModelBuilder.__init__ — config storage.""" + + def setup_method(self): + self.config = _make_mamba_config() + self.builder = MambaModelBuilder(self.config) + + def test_stores_model_config(self): + assert self.builder._model_config is self.config + + +class TestMambaModelBuilderBuildModel: + """Tests for MambaModelBuilder.build_model() — spec resolution, vocab padding, pp-stage inference, and MCoreMambaModel kwargs.""" + + def setup_method(self): + self.config = _make_mamba_config(vocab_size=32000) + self.builder = MambaModelBuilder(self.config) + self.pg = Mock() + self.pg.pp = Mock() + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_raises_when_vocab_size_none(self, mock_model, *_): + self.config.vocab_size = None + with pytest.raises(AssertionError, match="vocab_size"): + self.builder.build_model(self.pg) + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_spec_already_module_spec_used_directly(self, mock_model, *_): + module_spec = ModuleSpec(module=object) + self.config.__dict__["mamba_stack_spec"] = module_spec + self.builder.build_model(self.pg, pre_process=True, post_process=True) + call_kwargs = mock_model.call_args.kwargs + assert call_kwargs["mamba_stack_spec"] is module_spec + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_spec_callable_no_params_called_without_args(self, mock_model, *_): + returned_spec = ModuleSpec(module=object) + calls = [] + + def zero_param_fn(): + calls.append(True) + return returned_spec + + self.config.__dict__["mamba_stack_spec"] = zero_param_fn + self.builder.build_model(self.pg, pre_process=True, post_process=True) + assert calls, "zero_param_fn was not called" + call_kwargs = mock_model.call_args.kwargs + assert call_kwargs["mamba_stack_spec"] is returned_spec + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_spec_callable_with_config_param_called_with_config(self, mock_model, *_): + returned_spec = ModuleSpec(module=object) + received = [] + + def one_param_fn(config): + received.append(config) + return returned_spec + + self.config.__dict__["mamba_stack_spec"] = one_param_fn + self.builder.build_model(self.pg, pre_process=True, post_process=True) + assert received == [self.config], "one_param_fn not called with config" + call_kwargs = mock_model.call_args.kwargs + assert call_kwargs["mamba_stack_spec"] is returned_spec + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_no_vocab_padding_uses_vocab_size_directly(self, mock_model, mock_first, mock_last, mock_pad): + self.config.__dict__["should_pad_vocab"] = False + self.config.__dict__["vocab_size"] = 32000 + self.builder.build_model(self.pg, pre_process=True, post_process=True) + mock_pad.assert_not_called() + assert mock_model.call_args.kwargs["vocab_size"] == 32000 + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size", return_value=32128) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_vocab_padding_calls_calculate_padded_vocab_size(self, mock_model, mock_first, mock_last, mock_pad): + self.config.__dict__["should_pad_vocab"] = True + self.config.__dict__["vocab_size"] = 32000 + self.config.__dict__["make_vocab_size_divisible_by"] = 128 + self.config.transformer.tensor_model_parallel_size = 2 + self.builder.build_model(self.pg, pre_process=True, post_process=True) + mock_pad.assert_called_once_with(32000, 128, 2) + assert mock_model.call_args.kwargs["vocab_size"] == 32128 + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_explicit_pre_post_process_passed_through(self, mock_model, *_): + self.builder.build_model(self.pg, pre_process=False, post_process=True) + call_kwargs = mock_model.call_args.kwargs + assert call_kwargs["pre_process"] is False + assert call_kwargs["post_process"] is True + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=False) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_infers_pre_process_from_pg(self, mock_model, mock_first, mock_last, *_): + self.builder.build_model(self.pg) + mock_first.assert_called_once_with(self.pg.pp) + assert mock_model.call_args.kwargs["pre_process"] is False + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_infers_post_process_from_pg(self, mock_model, mock_first, mock_last, *_): + self.builder.build_model(self.pg) + mock_last.assert_called_once_with(self.pg.pp) + assert mock_model.call_args.kwargs["post_process"] is True + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_virtual_pipeline_raises(self, mock_model, *_): + with pytest.raises(AssertionError, match="Virtual pipeline"): + self.builder.build_model(self.pg, vp_stage=0) + + @patch("megatron.bridge.models.mamba.mamba_builder.calculate_padded_vocab_size") + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_last_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.is_pp_first_stage", return_value=True) + @patch("megatron.bridge.models.mamba.mamba_builder.MCoreMambaModel") + def test_config_params_passed_to_mcore(self, mock_model, *_): + config = _make_mamba_config( + vocab_size=32000, + seq_length=4096, + hybrid_attention_ratio=0.1, + hybrid_mlp_ratio=0.2, + hybrid_override_pattern="M-A-", + fp16_lm_cross_entropy=True, + parallel_output=False, + share_embeddings_and_output_weights=True, + position_embedding_type="rope", + ) + builder = MambaModelBuilder(config) + pg = Mock() + pg.pp = Mock() + + builder.build_model(pg, pre_process=True, post_process=True) + + kw = mock_model.call_args.kwargs + assert kw["config"] is config.transformer + assert kw["vocab_size"] == 32000 + assert kw["max_sequence_length"] == 4096 + assert kw["hybrid_attention_ratio"] == 0.1 + assert kw["hybrid_mlp_ratio"] == 0.2 + assert kw["hybrid_override_pattern"] == "M-A-" + assert kw["fp16_lm_cross_entropy"] is True + assert kw["parallel_output"] is False + assert kw["share_embeddings_and_output_weights"] is True + assert kw["position_embedding_type"] == "rope" + assert kw["rotary_percent"] == 1.0 + assert kw["rotary_base"] == 10000 + assert kw["seq_len_interpolation_factor"] is None + assert kw["pg_collection"] is pg + assert kw["vp_stage"] is None + + +class TestMambaModelBuilderBuildDistributedModels: + """Tests for MambaModelBuilder.build_distributed_models() — delegation to unimodal helper, hook composition, and default kwargs.""" + + def setup_method(self): + self.config = _make_mamba_config(vocab_size=32000) + self.builder = MambaModelBuilder(self.config) + self.pg = Mock() + + @patch("megatron.bridge.models.mamba.mamba_builder.compose_hooks") + @patch("megatron.bridge.models.mamba.mamba_builder.unimodal_build_distributed_models") + def test_delegates_to_unimodal_build_distributed_models(self, mock_unimodal, mock_compose): + model_list = [Mock()] + mock_unimodal.return_value = model_list + mock_compose.return_value = Mock(return_value=None) + + self.builder.build_distributed_models(self.pg) + + assert mock_unimodal.called + + @patch("megatron.bridge.models.mamba.mamba_builder.compose_hooks") + @patch("megatron.bridge.models.mamba.mamba_builder.unimodal_build_distributed_models") + def test_returns_model_list_from_unimodal(self, mock_unimodal, mock_compose): + model_list = [Mock(), Mock()] + mock_unimodal.return_value = model_list + # post_wrap hook returns None → original list is kept + mock_compose.return_value = Mock(return_value=None) + + result = self.builder.build_distributed_models(self.pg) + + assert result is model_list + + @patch("megatron.bridge.models.mamba.mamba_builder.compose_hooks") + @patch("megatron.bridge.models.mamba.mamba_builder.unimodal_build_distributed_models") + def test_pre_wrap_hooks_composed_and_passed(self, mock_unimodal, mock_compose): + model_list = [Mock()] + mock_unimodal.return_value = model_list + composed_pre = Mock() + composed_post = Mock(return_value=None) + mock_compose.side_effect = [composed_pre, composed_post] + + hook1 = Mock() + self.config.pre_wrap_hooks = [hook1] + self.builder.build_distributed_models(self.pg) + + # First compose_hooks call must be with the pre_wrap_hooks list + assert mock_compose.call_args_list[0] == call([hook1]) + # The composed pre-wrap hook is the 11th positional arg (index 10) + unimodal_args = mock_unimodal.call_args.args + assert unimodal_args[10] is composed_pre + + @patch("megatron.bridge.models.mamba.mamba_builder.compose_hooks") + @patch("megatron.bridge.models.mamba.mamba_builder.unimodal_build_distributed_models") + def test_post_wrap_hook_applied_to_results(self, mock_unimodal, mock_compose): + model_list = [Mock()] + wrapped_list = [Mock(), Mock()] + mock_unimodal.return_value = model_list + composed_pre = Mock() + composed_post = Mock(return_value=wrapped_list) + mock_compose.side_effect = [composed_pre, composed_post] + + result = self.builder.build_distributed_models(self.pg) + + composed_post.assert_called_once_with(model_list) + assert result is wrapped_list + + @patch("megatron.bridge.models.mamba.mamba_builder.compose_hooks") + @patch("megatron.bridge.models.mamba.mamba_builder.unimodal_build_distributed_models") + def test_post_wrap_hook_returning_none_keeps_original_list(self, mock_unimodal, mock_compose): + model_list = [Mock()] + mock_unimodal.return_value = model_list + mock_compose.return_value = Mock(return_value=None) + + result = self.builder.build_distributed_models(self.pg) + + assert result is model_list + + @patch("megatron.bridge.models.mamba.mamba_builder.compose_hooks") + @patch("megatron.bridge.models.mamba.mamba_builder.unimodal_build_distributed_models") + def test_default_parameters_forwarded(self, mock_unimodal, mock_compose): + from megatron.core.enums import ModelType + from megatron.core.transformer.module import Float16Module + + model_list = [Mock()] + mock_unimodal.return_value = model_list + mock_compose.return_value = Mock(return_value=None) + + self.builder.build_distributed_models(self.pg) + + # unimodal_build_distributed_models is called with all positional args: + # build_model, transformer_config, pg_collection, ddp_config, + # overlap_param_gather_with_optimizer_step, use_megatron_fsdp, use_torch_fsdp2, + # wrap_with_ddp, data_parallel_random_init, mixed_precision_wrapper, + # composed_pre_wrap_hook, model_type + args = mock_unimodal.call_args.args + assert args[0] == self.builder.build_model + assert args[1] is self.config.transformer + assert args[2] is self.pg + assert args[3] is None # ddp_config + assert args[7] is True # wrap_with_ddp + assert args[8] is True # data_parallel_random_init + assert args[9] is Float16Module # mixed_precision_wrapper + assert args[11] is ModelType.encoder_or_decoder # model_type diff --git a/tests/unit_tests/models/test_model_instantiation.py b/tests/unit_tests/models/test_model_instantiation.py index c268825ba1..f8a0f26204 100644 --- a/tests/unit_tests/models/test_model_instantiation.py +++ b/tests/unit_tests/models/test_model_instantiation.py @@ -21,11 +21,13 @@ from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.bridge.models.common.unimodal import ( + _ddp_wrap, + _print_num_params, +) from megatron.bridge.models.model_provider import ( ModelProviderMixin, _create_model, - _ddp_wrap, - _print_num_params, get_model, ) @@ -206,7 +208,7 @@ def test_create_model_sets_tensor_parallel_attributes(self, mock_tensor_parallel class TestDDPWrap: """Test cases for _ddp_wrap function.""" - @patch("megatron.bridge.models.model_provider.DistributedDataParallel") + @patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") def test_ddp_wrap_standard(self, mock_ddp): """Test wrapping models with standard DDP.""" # Setup @@ -243,7 +245,7 @@ def test_ddp_wrap_standard(self, mock_ddp): for ddp_instance in mock_ddp_instances: ddp_instance.broadcast_params.assert_called_once() - @patch("megatron.bridge.models.model_provider.TorchFullyShardedDataParallel") + @patch("megatron.bridge.models.common.unimodal.TorchFullyShardedDataParallel") def test_ddp_wrap_fsdp2(self, mock_fsdp): """Test wrapping models with FSDP2.""" # Setup @@ -271,7 +273,7 @@ def test_ddp_wrap_fsdp2(self, mock_fsdp): def test_ddp_wrap_overlap_param_gather(self): """Test DDP wrapping with overlap_param_gather_with_optimizer_step.""" - with patch("megatron.bridge.models.model_provider.DistributedDataParallel") as mock_ddp: + with patch("megatron.bridge.models.common.unimodal.DistributedDataParallel") as mock_ddp: # Setup config = create_test_config() models = [MockMegatronModule(config)]