[port] Add GDN2 (Gated DeltaNet-2) attention variant (upstream #5765)

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -4,7 +4,7 @@
 
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.backends import BackendSpecProvider
-from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
+from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNet2, GatedDeltaNetSubmodules
 from megatron.core.transformer.enums import AttnMaskType, LayerType
 from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
     AbsorbedMLASelfAttention,
@@ -80,8 +80,12 @@ def get_gated_delta_net_module_spec(
         backend = _get_backend_spec_provider(config=config)
 
     rms_norm = config.normalization == "RMSNorm"
+    # gdn2 reuses the GDN submodules and spec structure with the GatedDeltaNet2 module.
+    gdn_module = (
+        GatedDeltaNet2 if config.experimental_attention_variant == "gdn2" else GatedDeltaNet
+    )
     attention = ModuleSpec(
-        module=GatedDeltaNet,
+        module=gdn_module,
         submodules=GatedDeltaNetSubmodules(
             in_proj=backend.column_parallel_layer_norm_linear(),
             out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False),
@@ -410,8 +414,11 @@ 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."""
-    # gated_delta_net and its upstream alias gdn; see transformer_config.is_gated_delta_net_variant.
+    """Check if the experimental attention variant is a linear attention variant.
+
+    gated_delta_net (and its upstream alias gdn, and gdn2) are all served by the gated
+    delta net family; see transformer_config.is_gated_delta_net_variant.
+    """
     return is_gated_delta_net_variant(experimental_attention_variant)
 
 

megatron/core/models/hybrid/hybrid_block.py

@@ -742,8 +742,16 @@ def __init__(
                         name=(name + f".layers.{i}") if name is not None else None,
                     )
                 elif layer_type == LayerSymbols.GDN:
+                    gdn_layer_spec = submodules.gdn_layer
+                    if self.config.experimental_attention_variant == "gdn2":
+                        # 'G' layers build the GDN2 variant when the gdn2 experimental
+                        # attention variant is selected.
+                        from megatron.core.ssm.gated_delta_net import GatedDeltaNet2
+
+                        gdn_layer_spec = copy.deepcopy(gdn_layer_spec)
+                        gdn_layer_spec.submodules.self_attention.module = GatedDeltaNet2
                     layer = build_module(
-                        submodules.gdn_layer,
+                        gdn_layer_spec,
                         config=self.config,
                         layer_number=layer_number,
                         pg_collection=pg_collection,

megatron/core/ssm/gated_delta_net/__init__.py

@@ -20,17 +20,27 @@
     tensor_a2a_hp2cp,
 )
 from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet, torch_chunk_gated_delta_rule
+from megatron.core.ssm.gated_delta_net.gdn2 import (
+    HAVE_FLA_GDN2,
+    GatedDeltaNet2,
+    chunk_gdn2,
+    torch_chunk_gdn2,
+)
 
 __all__ = [
     "HAVE_FLA",
+    "HAVE_FLA_GDN2",
     "GatedDeltaNet",
+    "GatedDeltaNet2",
     "GatedDeltaNetSubmodules",
     "_split_tensor_factory",
     "causal_conv1d",
     "chunk_gated_delta_rule",
+    "chunk_gdn2",
     "get_parameter_local_cp",
     "l2norm",
     "tensor_a2a_cp2hp",
     "tensor_a2a_hp2cp",
     "torch_chunk_gated_delta_rule",
+    "torch_chunk_gdn2",
 ]