[port] Megatron-Core APIs needed by Megatron-Bridge main (#4452 MLA resolver, gdn alias, #5865 FSDP names)

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -484,6 +484,29 @@ def sync_rng_states_across_tp_group(self):
         _load_rng_state_dict(broadcast_list[0])
 
 
+# Name compatibility with upstream Megatron-Core (NVIDIA/Megatron-LM#5865), which split
+# Megatron-FSDP into ``FullyShardedDataParallelV1`` (this implementation) and
+# ``FullyShardedDataParallelV2`` and turned ``FullyShardedDataParallel`` into a factory.
+# Callers written against upstream, such as Megatron-Bridge's ``unwrap_model``, import both
+# names from this module. ``miles-main`` keeps ``FullyShardedDataParallel`` as the v1 class
+# (nothing here selects v2) and does not carry the v2 implementation.
+FullyShardedDataParallelV1 = FullyShardedDataParallel
+
+
+class FullyShardedDataParallelV2(_BaseDataParallel):
+    """Placeholder for Megatron-FSDP v2 (NVIDIA/Megatron-LM#5865), which is not ported.
+
+    The class exists so that upstream-style ``isinstance`` checks import and evaluate to
+    ``False``; constructing it raises.
+    """
+
+    def __init__(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Megatron-FSDP v2 (NVIDIA/Megatron-LM#5865) is not ported to miles-main; "
+            "use FullyShardedDataParallel (FullyShardedDataParallelV1)."
+        )
+
+
 def _get_hsdp_tp_mesh(outer_fsdp_dp_group, dp_cp_group, tp_group, ep_size=1):
     assert HAVE_EINOPS, "einops is not installed. Please install it with `pip install einops`."
     world_size = dist.get_world_size()

megatron/core/models/backends.py

@@ -4,7 +4,7 @@
 import warnings
 from abc import abstractmethod
 from functools import partial
-from typing import Optional, Protocol, cast
+from typing import Literal, Optional, Protocol, cast
 
 from megatron.core.extensions.transformer_engine import (
     TEColumnParallelGroupedLinear,
@@ -200,3 +200,19 @@ def grouped_mlp_modules(self, moe_use_grouped_gemm: bool) -> ExpertsBuilder:
                 activation_func=self.activation_func(),
             ),
         )
+
+
+def get_backend(
+    transformer_impl: Literal["local", "transformer_engine", "inference_optimized"]
+) -> BackendSpecProvider:
+    """Return the backend that's selected with the given `transformer_impl`."""
+    if transformer_impl == "transformer_engine":
+        from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+
+        return TESpecProvider()
+    elif transformer_impl == "inference_optimized":
+        return InferenceSpecProvider()
+    elif transformer_impl == "local":
+        return LocalSpecProvider()
+    else:
+        raise ValueError(f"unknown transformer_impl='{transformer_impl}'")

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -37,7 +37,7 @@
     TransformerBlockSubmodules,
     get_num_layers_to_build,
 )
-from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_config import TransformerConfig, is_gated_delta_net_variant
 from megatron.core.transformer.transformer_layer import (
     HyperConnectionTransformerLayer,
     MlpBuilder,
@@ -209,7 +209,7 @@ def get_experimental_attention_variant_module_spec(
     if backend is None:
         backend = _get_backend_spec_provider(config=config)
 
-    if config.experimental_attention_variant == "gated_delta_net":
+    if is_gated_delta_net_variant(config.experimental_attention_variant):
         return get_gated_delta_net_module_spec(config=config, backend=backend)
     elif config.experimental_attention_variant == "dsa":
         return get_dsa_module_spec_for_backend(config=config, backend=backend)
@@ -411,8 +411,8 @@ def get_transformer_block_with_experimental_attention_variant_spec(
 
 def is_linear_attention_variant(experimental_attention_variant: Optional[str]) -> bool:
     """Check if the experimental attention variant is a linear attention variant."""
-    linear_attention_variants = ["gated_delta_net"]
-    return experimental_attention_variant in linear_attention_variants
+    # gated_delta_net and its upstream alias gdn; see transformer_config.is_gated_delta_net_variant.
+    return is_gated_delta_net_variant(experimental_attention_variant)
 
 
 def _validate_dsa_index_share_pipeline_split(config: TransformerConfig, local_layer_ids) -> None:

megatron/core/transformer/experimental_attention_variant/absorbed_mla.py

@@ -36,6 +36,7 @@
 )
 from megatron.core.transformer.attention import Attention
 from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_config import MLATransformerConfig
 from megatron.core.utils import deprecate_inference_params, get_pg_size, is_te_min_version
@@ -167,6 +168,10 @@ def __init__(
             name=name,
         )
 
+        # Resolve which classes to use for Q and KV linear up projections and norms, based on
+        # QK-norm selection.
+        layer_classes = QKNormConfigResolver(self.config, submodules).resolve()
+
         assert not config.add_bias_linear, "add_bias_linear is not supported for AbsorbedMLA"
         assert not (
             config.tensor_model_parallel_size > 1 and not config.sequence_parallel
@@ -266,7 +271,7 @@ def __init__(
         if self.config.q_lora_rank is None:
             # Not projecting query
             self.linear_q_proj = build_module(
-                submodules.linear_q_proj,
+                layer_classes["linear_q_proj"],
                 self.config.hidden_size,
                 self.config.num_attention_heads * self.q_head_dim,
                 config=self.config,
@@ -312,7 +317,7 @@ def __init__(
             )
 
             self.linear_q_up_proj = build_module(
-                submodules.linear_q_up_proj,
+                layer_classes["linear_q_up_proj"],
                 self.config.q_lora_rank,
                 self.config.num_attention_heads * self.q_head_dim,
                 config=self.config,
@@ -373,7 +378,7 @@ def __init__(
 
         if self._uses_combined_kv_up_projection:
             self.linear_kv_up_proj = build_module(
-                submodules.linear_kv_up_proj,
+                layer_classes["linear_kv_up_proj"],
                 self.config.kv_lora_rank,
                 self.config.num_attention_heads
                 * (self.config.qk_head_dim + self.config.v_head_dim),
@@ -419,14 +424,14 @@ def __init__(
 
         if self.config.q_lora_rank is not None:
             self.q_layernorm = build_module(
-                submodules.q_layernorm,
+                layer_classes["q_layernorm"],
                 hidden_size=self.config.q_lora_rank,
                 config=self.config,
                 eps=self.config.layernorm_epsilon,
             )
 
         self.kv_layernorm = build_module(
-            submodules.kv_layernorm,
+            layer_classes["kv_layernorm"],
             hidden_size=self.config.kv_lora_rank,
             config=self.config,
             eps=self.config.layernorm_epsilon,