diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a35a3666ccf..6e6a753f0d1 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -55,6 +55,7 @@ Changelog (``layerwise.get_qdq_activations_from_prev_layer=True``). Set it to ``False`` to preserve full-precision activations for subsequent layers (the default behavior for max calibration without layerwise calibration). +- ``get_te_hybrid_stack_spec`` was removed from ``modelopt.torch.nas.plugins.megatron``; it had no use outside tests. Use ``modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp`` for the SequentialMLP layout, or ``megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec`` for grouped GEMM. - Unified HuggingFace export now fails with ``NotImplementedError`` when it meets an MoE block whose expert projection names it does not know, instead of assuming Mixtral's ``w1``/``w2``/``w3``. If you hit this, register a ``ModelSpec`` for the model under ``modelopt/torch/models/``. Every MoE architecture ModelOpt exported correctly before this change is registered, so no supported model regresses. - ``--recipe`` (and ``modelopt.recipe.load_recipe``) now resolve a recipe path **filesystem-first**: a recipe of the same relative path in the current working directory takes precedence over the shipped built-in of that name, matching how recipe ``$import`` paths already resolve. Previously the built-in won. @@ -78,6 +79,7 @@ Changelog - Fix a DDP hang in DFlash training at scale where a rank whose batch contained no valid anchor skipped the draft forward, leaving its rotary buffer list shorter than other ranks' and causing ``broadcast_buffers`` to hang. The buffer is now created during ``modify()`` before training begins. - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected. - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. +- Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. 0.47.0 (2026-09-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/nas/plugins/megatron.py b/modelopt/torch/nas/plugins/megatron.py index 2376422a998..c2d31b6f3de 100644 --- a/modelopt/torch/nas/plugins/megatron.py +++ b/modelopt/torch/nas/plugins/megatron.py @@ -15,7 +15,6 @@ """Plugin to add NAS/Pruning support for megatron-core Language models like GPT and Mamba.""" -import copy import types from abc import ABC from collections.abc import Callable, Sequence @@ -35,10 +34,6 @@ ) from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding from megatron.core.models.gpt import GPTModel -from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec -from megatron.core.models.hybrid.hybrid_layer_specs import ( - hybrid_stack_spec as _te_hybrid_stack_spec, -) from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.parallel_state import is_pipeline_first_stage, is_pipeline_last_stage from megatron.core.ssm.gated_delta_net import GatedDeltaNet @@ -56,7 +51,6 @@ from megatron.core.transformer.moe.router import TopKRouter from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.multi_latent_attention import MLASelfAttention -from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_layer import TransformerLayer from modelopt.torch.nas.modules import DynamicModuleList @@ -91,21 +85,8 @@ # Attention module types that _DynamicTransformerLayer converts. _ATTENTION_TYPES: tuple[type, ...] = (SelfAttention, MLASelfAttention, GatedDeltaNet) -__all__ = ["get_te_hybrid_stack_spec"] - - -def get_te_hybrid_stack_spec(moe_grouped_gemm: bool = False) -> ModuleSpec: - """Return the TE Hybrid stack spec.""" - if moe_grouped_gemm: - return _te_hybrid_stack_spec - - # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. - # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency). - te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec) - te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec( - use_te=True, num_experts=8, moe_grouped_gemm=False - ) - return te_hybrid_stack_spec +# This module only registers DynamicModules; it exports no public API. +__all__ = [] # Local Parallel Linear DynamicModules ########################################################################## diff --git a/modelopt/torch/utils/plugins/__init__.py b/modelopt/torch/utils/plugins/__init__.py index da40fe9e565..ff251eaf133 100644 --- a/modelopt/torch/utils/plugins/__init__.py +++ b/modelopt/torch/utils/plugins/__init__.py @@ -23,6 +23,9 @@ with import_plugin("megatron_generate"): from .megatron_generate import * +with import_plugin("megatron_layer_specs"): + from .megatron_layer_specs import * + with import_plugin("megatron_mmlu"): from .megatron_mmlu import * @@ -33,6 +36,7 @@ from .prepare_megatron_data_blend import * # NOTE: Dont pre-import megatron bridge plugin here to avoid circular dependency issues. -# We dont register anything so this isnt a problem. +# It registers an instantiate allowlist prefix on import, which only the ModelOpt entrypoints +# that import it need, so leaving it out here is still fine. # with import_plugin("megatron bridge"): # from .mbridge import * diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index 05713c7b282..7b85fb9d4fb 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -22,13 +22,17 @@ from megatron.bridge import AutoBridge from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.utils import is_safe_repo -from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.models.hybrid.hybrid_provider import ( + HybridModelProvider, + transformer_engine_hybrid_stack_spec, +) from megatron.bridge.training.checkpointing import _load_model_weights_from_checkpoint from megatron.bridge.training.post_training.checkpointing import ( _get_modelopt_checkpoint_path, has_modelopt_state, load_modelopt_state, ) +from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix from megatron.core.models.gpt import GPTModel from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.transformer.module import MegatronModule @@ -37,8 +41,16 @@ from transformers import AutoConfig, AutoTokenizer from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping -from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec from modelopt.torch.utils import print_rank_0, warn_rank_0 +from modelopt.torch.utils.plugins.megatron_layer_specs import te_hybrid_stack_spec_sequential_mlp + +# ``set_moe_expert_layout`` records ``te_hybrid_stack_spec_sequential_mlp`` in a SequentialMLP +# checkpoint's ``run_config.yaml``, and rebuilding that config resolves the target against this +# allowlist. Only a process that imports this module registers it, so such a checkpoint must be +# converted through a ModelOpt entrypoint, not stock ``scripts/conversion/convert.sh``. The prefix +# covers all of ``modelopt`` rather than one module, which assumes a ``run_config.yaml`` is trusted +# input -- it comes from a checkpoint the caller is already choosing to load. +register_allowed_target_prefix("modelopt.") __all__ = [ "get_language_model", @@ -114,10 +126,24 @@ def set_moe_expert_layout(provider, moe_grouped_gemm: bool) -> None: Set ``moe_grouped_gemm`` on the provider (the bridge's native, possibly custom/hybrid spec reads it at build time) rather than replacing the whole layer spec -- overwriting it would drop custom layers (e.g. Qwen3.5's GatedDeltaNet or Gemma3's custom spec). A hybrid provider - additionally needs its stack spec rebuilt, since the native one pins ``TEGroupedMLP``. + additionally has its stack spec set, since the native one pins ``TEGroupedMLP``: the bridge's + own factory for grouped GEMM, a ModelOpt one that swaps in SequentialMLP otherwise. + + Assign a *factory function*, never a built ``ModuleSpec``: the provider is serialized into + every checkpoint's ``run_config.yaml``, and Megatron-LM's YAML writer drops the fields of a + dataclass nested inside a ``functools.partial`` keyword, which is how a stack spec holds + ``MLPSubmodules`` / ``MoESubmodules``. Such a checkpoint cannot be reloaded or exported. The + provider calls the factory at build time, so behavior is unchanged. """ if isinstance(provider, HybridModelProvider): - provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm) + # The grouped-GEMM factory is Megatron-Bridge's, and returns Megatron-Core's + # ``hybrid_stack_spec`` unchanged -- the layer composition is identical either way. It is + # named from the bridge so stock tooling resolves the target without importing ModelOpt. + provider.hybrid_stack_spec = ( + transformer_engine_hybrid_stack_spec + if moe_grouped_gemm + else te_hybrid_stack_spec_sequential_mlp + ) provider.moe_grouped_gemm = moe_grouped_gemm elif (provider.num_moe_experts or 0) > 0: provider.moe_grouped_gemm = moe_grouped_gemm diff --git a/modelopt/torch/utils/plugins/megatron_layer_specs.py b/modelopt/torch/utils/plugins/megatron_layer_specs.py new file mode 100644 index 00000000000..d9b08b7b134 --- /dev/null +++ b/modelopt/torch/utils/plugins/megatron_layer_specs.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Megatron-Core layer specs used to build models for ModelOpt workflows.""" + +import copy + +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec +from megatron.core.models.hybrid.hybrid_layer_specs import ( + hybrid_stack_spec as _te_hybrid_stack_spec, +) +from megatron.core.transformer.spec_utils import ModuleSpec + +__all__ = ["te_hybrid_stack_spec_sequential_mlp"] + + +def te_hybrid_stack_spec_sequential_mlp() -> ModuleSpec: + """Return the TE Hybrid stack spec with SequentialMLP MoE experts. + + Named and zero-argument so a provider can store this function instead of the ModuleSpec it + builds; see ``set_moe_expert_layout`` for why a built spec cannot be serialized. + + Its module path and name are written into ``run_config.yaml`` as a ``_target_``, so moving or + renaming it breaks every SequentialMLP hybrid checkpoint already saved. + """ + # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. + # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency). + # num_experts only has to be non-zero to select the MoE branch; the real count comes from the + # model config at build time. + te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec) + te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec( + use_te=True, num_experts=8, moe_grouped_gemm=False + ) + return te_hybrid_stack_spec diff --git a/tests/_test_utils/torch/megatron/models.py b/tests/_test_utils/torch/megatron/models.py index b2bd2e90921..efdc030a9dc 100644 --- a/tests/_test_utils/torch/megatron/models.py +++ b/tests/_test_utils/torch/megatron/models.py @@ -25,6 +25,7 @@ get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec, ) +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec as te_hybrid_stack_spec from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.parallel_state import ( get_pipeline_model_parallel_rank, @@ -37,7 +38,7 @@ from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig from modelopt.torch.export.unified_export_megatron import import_mcore_gpt_from_hf -from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec +from modelopt.torch.utils.plugins.megatron_layer_specs import te_hybrid_stack_spec_sequential_mlp try: from megatron.core.extensions.transformer_engine import TENorm @@ -444,11 +445,11 @@ def get_mcore_hybrid_model( "share_embeddings_and_output_weights": False, "position_embedding_type": "none", } - spec = ( - get_te_hybrid_stack_spec(moe_grouped_gemm) - if transformer_impl == "transformer_engine" - else get_hybrid_stack_modelopt_spec(remap_te_layernorm=True) - ) + if transformer_impl == "transformer_engine": + # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. + spec = te_hybrid_stack_spec if moe_grouped_gemm else te_hybrid_stack_spec_sequential_mlp() + else: + spec = get_hybrid_stack_modelopt_spec(remap_te_layernorm=True) model = HybridModel( hybrid_stack_spec=spec, hybrid_layer_pattern=hybrid_layer_pattern, **common_kwargs ) diff --git a/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py b/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py new file mode 100644 index 00000000000..3b5e7ea78c8 --- /dev/null +++ b/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 pytest +import yaml +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.utils.instantiate_utils import instantiate +from megatron.bridge.utils.yaml_utils import dump_dataclass_to_yaml + +from modelopt.torch.utils.plugins.mbridge import set_moe_expert_layout + + +def _round_trip(value): + """Serialize through the writer used for run_config.yaml, then reload.""" + node = yaml.safe_load(dump_dataclass_to_yaml({"spec": value}))["spec"] + return node["_target_"], instantiate(node) + + +@pytest.mark.parametrize( + ("moe_grouped_gemm", "expected_experts", "expected_target"), + [ + ( + True, + "TEGroupedMLP", + "megatron.bridge.models.hybrid.hybrid_provider.transformer_engine_hybrid_stack_spec", + ), + ( + False, + "SequentialMLP", + "modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp", + ), + ], +) +def test_set_moe_expert_layout_survives_run_config_round_trip( + moe_grouped_gemm, expected_experts, expected_target +): + """A provider's stack spec must still build real submodules after a run_config round trip. + + A built ``ModuleSpec`` loses its ``MLPSubmodules`` / ``MoESubmodules`` when written to + ``run_config.yaml``, so ``set_moe_expert_layout`` stores a factory function instead. + """ + provider = HybridModelProvider(num_layers=2, hidden_size=64, num_attention_heads=4) + set_moe_expert_layout(provider, moe_grouped_gemm=moe_grouped_gemm) + assert provider.moe_grouped_gemm == moe_grouped_gemm + + assert callable(provider.hybrid_stack_spec) + + target, factory = _round_trip(provider.hybrid_stack_spec) + # The target is an on-disk contract: renaming or moving the factory breaks saved checkpoints. + assert target == expected_target + + provider.hybrid_stack_spec = factory + spec = provider._resolve_hybrid_stack_spec() + + mlp = spec.submodules.mlp_layer.submodules.mlp.keywords["submodules"] + assert mlp.linear_fc1 is not None + assert mlp.linear_fc2 is not None + + moe = spec.submodules.moe_layer.submodules.mlp.keywords["submodules"] + assert moe.experts is not None + # Experts are built through a partial for the grouped-GEMM layout. + assert getattr(moe.experts, "func", moe.experts).__name__ == expected_experts