Integrate direct hybrid architectures with training

megatron/training/arguments.py

@@ -6,15 +6,15 @@
 import dataclasses
 import json
 import os
-from pathlib import Path
 import re
 import types
+from pathlib import Path
 
 import torch
 
+from megatron.core.msc_utils import MultiStorageClientFeature
 from megatron.core.rerun_state_machine import RerunStateMachine
 from megatron.core.transformer import TransformerConfig
-from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
 from megatron.core.transformer.cuda_graph_config import (
     ALLOWED_INFERENCE_SCOPES,
     get_deprecated_cuda_graph_modules_migration,
@@ -23,23 +23,24 @@
     validate_deprecated_cuda_graph_modules_migration_inputs,
 )
 from megatron.core.transformer.enums import AttnBackend, CudaGraphModule, InferenceCudaGraphScope
+from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
 from megatron.core.utils import (
     get_torch_version,
     is_flashinfer_min_version,
     is_te_min_version,
     is_torch_min_version,
 )
+from megatron.training.argument_utils import (  # noqa: F401 # pylint: disable=unused-import
+    ArgumentGroupFactory,
+    core_transformer_config_from_args,
+)
 from megatron.training.global_vars import set_global_variables
 from megatron.training.utils import (
     get_device_arch_version,
-    update_use_dist_ckpt,
     print_rank_0,
+    update_use_dist_ckpt,
     warn_rank_0,
 )
-from megatron.core.msc_utils import MultiStorageClientFeature
-
-from megatron.training.argument_utils import ArgumentGroupFactory, core_transformer_config_from_args  # noqa: F401 # pylint: disable=unused-import
-
 
 
 def add_megatron_arguments(parser: argparse.ArgumentParser):
@@ -398,8 +399,9 @@ def validate_args(args, defaults={}):
         'Currently only global and local checkpoints are supported'
     if args.non_persistent_ckpt_type == 'local':
         try:
-            from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import \
-                LocalCheckpointManager
+            from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import (
+                LocalCheckpointManager,
+            )
         except ModuleNotFoundError as e:
             raise RuntimeError('nvidia_resiliency_ext is required for local checkpointing') from e
 
@@ -740,8 +742,10 @@ def validate_args(args, defaults={}):
         )
 
     from megatron.core.models.hybrid.hybrid_layer_allocation import (
-        Symbols, parse_hybrid_pattern, get_hybrid_total_layer_count,
+        Symbols,
+        get_hybrid_total_layer_count,
         get_hybrid_total_pipeline_segment_count,
+        parse_hybrid_pattern,
     )
     sep = Symbols.MTP_SEPARATOR
 
@@ -948,20 +952,10 @@ def validate_args(args, defaults={}):
         if args.hybrid_layer_pattern is None:
             args.virtual_pipeline_model_parallel_size = None
 
-        if args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None:
-            # Divisibility check not applicable for T5 models which specify encoder_num_layers
-            # and decoder_num_layers, or for hybrid models using --hybrid-layer-pattern.
-            if args.num_layers is not None and args.hybrid_layer_pattern is None:
-                num_layers = args.num_layers
-
-                if args.account_for_embedding_in_pipeline_split:
-                    num_layers += 1
-
-                if args.account_for_loss_in_pipeline_split:
-                    num_layers += 1
-
-                assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
-                    'Number of layers should be divisible by the pipeline-model-parallel size'
+        # Defer layer-count divisibility to the selected model builder. Direct
+        # HybridModel architecture specs are Python objects constructed after
+        # CLI validation and may legally define uneven or empty explicit chunks.
+        # Conventional transformer builders retain their own divisibility checks.
 
     if args.virtual_pipeline_model_parallel_size is not None:
         if args.overlap_p2p_comm:
@@ -975,6 +969,11 @@ def validate_args(args, defaults={}):
                 'p2p sends and recvs between same 2 ranks per communication batch'
     else:
         # Overlap P2P communication is disabled if not using the interleaved schedule.
+        # Preserve the requested values because a Python HybridModel architecture
+        # can infer VPP only after this CLI-only validation pass.
+        if args.hybrid_layer_pattern is None:
+            args._overlap_p2p_comm_before_direct_vpp = args.overlap_p2p_comm
+            args._align_param_gather_before_direct_vpp = args.align_param_gather
         args.overlap_p2p_comm = False
         args.align_param_gather = False
         # Only print warning if PP size > 1.
@@ -1034,8 +1033,13 @@ def validate_args(args, defaults={}):
             '--overlap-param-gather-with-optimizer-step only supported with distributed optimizer'
         assert args.overlap_param_gather, \
             'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather'
-        assert args.virtual_pipeline_model_parallel_size is not None, \
-            '--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism'
+        if args.hybrid_layer_pattern is not None:
+            assert args.virtual_pipeline_model_parallel_size is not None, (
+                '--overlap-param-gather-with-optimizer-step only supported with '
+                'interleaved pipeline parallelism'
+            )
+        # Otherwise, the interleaved-pipeline requirement is checked in pretrain after
+        # Python model builders have had a chance to infer direct VPP splits.
         assert not args.use_dist_ckpt, \
             '--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet'
 
@@ -2714,8 +2718,7 @@ def _add_rl_args(parser):
     return parser
 
 def _add_training_args(parser):
-    from megatron.training.config import TrainingConfig
-    from megatron.training.config import ProfilingConfig
+    from megatron.training.config import ProfilingConfig, TrainingConfig
 
     prof_factory = ArgumentGroupFactory(ProfilingConfig)
     prof_group = prof_factory.build_group(parser, "profiling")
@@ -3582,7 +3585,7 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser):
     If kitchen isn't available, nothing to do here, return unchanged parser
     """
     try:
-        from megatron.core.extensions.kitchen import KitchenSpecProvider, HAVE_KITCHEN
+        from megatron.core.extensions.kitchen import HAVE_KITCHEN, KitchenSpecProvider
 
     except (ImportError, ModuleNotFoundError):
         HAVE_KITCHEN = False

megatron/training/checkpointing.py

@@ -3,6 +3,7 @@
 """Input/output checkpointing."""
 
 import contextlib
+import copy
 import inspect
 import multiprocessing
 import os
@@ -39,7 +40,10 @@
 from megatron.core.msc_utils import maybe_msc
 from megatron.core.num_microbatches_calculator import update_num_microbatches
 from megatron.core.optimizer import DistributedOptimizer
-from megatron.core.post_training.modelopt.checkpointing import save_modelopt_state, save_sharded_modelopt_state
+from megatron.core.post_training.modelopt.checkpointing import (
+    save_modelopt_state,
+    save_sharded_modelopt_state,
+)
 from megatron.core.rerun_state_machine import get_rerun_state_machine
 from megatron.core.utils import get_pg_rank, get_pg_size, unwrap_model
 from megatron.post_training.utils import print_distributed_quant_summary
@@ -1348,7 +1352,16 @@ def generate_state_dict(
 
     # Arguments, iteration, and model.
     state_dict = {}
-    state_dict['args'] = args
+    resolved_architecture = getattr(args, 'resolved_hybrid_architecture', None)
+    if getattr(resolved_architecture, "source", None) == "direct":
+        checkpoint_args = copy.copy(args)
+        # This runtime-only summary contains live ModuleSpec/config objects. Direct
+        # architecture recipes must be supplied by Python again on resume.
+        del checkpoint_args.resolved_hybrid_architecture
+        state_dict['args'] = checkpoint_args
+    else:
+        # Preserve the historical args object identity for legacy and non-hybrid checkpoints.
+        state_dict['args'] = args
     state_dict['checkpoint_version'] = 3.0
     if iteration is not None:
         state_dict['iteration'] = iteration

megatron/training/hybrid_metrics.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Helpers for reporting metrics from a resolved hybrid architecture."""
+
+from dataclasses import dataclass
+from typing import Any, Iterator
+
+_LAYER_TYPE_ALIASES = {
+    "M": "mamba",
+    "G": "gdn",
+    "*": "attention",
+    "D": "dsa",
+    "+": "mla",
+    "-": "mlp",
+    "E": "moe",
+}
+
+
+def get_resolved_hybrid_architecture(args: Any) -> Any | None:
+    """Return a model-provided direct architecture, if one was installed."""
+
+    architecture = getattr(args, "resolved_hybrid_architecture", None)
+    if architecture is None or getattr(architecture, "source", None) != "direct":
+        return None
+    return architecture
+
+
+def get_hybrid_layer_type(layer: Any) -> str:
+    """Return the stable semantic type name for a resolved hybrid layer."""
+
+    layer_type = layer.layer_type
+    if not isinstance(layer_type, str):
+        layer_type = getattr(layer_type, "value", layer_type)
+    layer_type = str(layer_type)
+    return _LAYER_TYPE_ALIASES.get(layer_type, layer_type.lower())
+
+
+def iter_resolved_hybrid_layers(architecture: Any) -> Iterator[Any]:
+    """Iterate main layers followed by each repeated MTP-depth template."""
+
+    yield from architecture.main_layers
+    for _ in range(architecture.mtp_num_layers):
+        yield from architecture.mtp_layers
+
+
+@dataclass(frozen=True)
+class HybridMoEMetricMetadata:
+    """Arguments needed to size and normalize global MoE metric tensors."""
+
+    num_layers: int
+    moe_layer_freq: list[int]
+    mtp_num_layers: int
+    num_moe_layers: int
+
+
+def get_hybrid_moe_metric_metadata(architecture: Any) -> HybridMoEMetricMetadata:
+    """Derive MoE logging metadata from per-occurrence semantic layer types.
+
+    Direct hybrid models assign metric slots over the fully expanded global
+    architecture, including every layer in every MTP depth. Consequently MTP
+    has already been incorporated into ``num_layers`` and ``moe_layer_freq``;
+    ``mtp_num_layers`` is zero to disable the tracker's legacy implicit MTP
+    expansion.
+    """
+
+    layer_types = [
+        get_hybrid_layer_type(layer) for layer in iter_resolved_hybrid_layers(architecture)
+    ]
+    moe_layer_freq = [int(layer_type == "moe") for layer_type in layer_types]
+    return HybridMoEMetricMetadata(
+        num_layers=len(layer_types),
+        moe_layer_freq=moe_layer_freq,
+        mtp_num_layers=0,
+        num_moe_layers=sum(moe_layer_freq),
+    )