Introduce GDN to Mamba

megatron/core/inference/contexts/dynamic_context.py

@@ -298,14 +298,14 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
             # For hybrid models, the layer map converts the global layer index to the
             # corresponding attention layer index or Mamba layer index depending on the
             # layer type.
-            attention_layer_map, mamba_layer_map, _, _ = get_layer_maps_from_layer_type_list(
-                mamba_inference_state_config.layer_type_list
+            attention_layer_map, mamba_layer_map, _, _, gdn_layer_map = (
+                get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)
             )
             self.num_attention_layers = len(attention_layer_map)
             self.num_mamba_layers = len(mamba_layer_map)
             self.mamba_conv_states_shape = mamba_conv_states_shape
             self.mamba_ssm_states_shape = mamba_ssm_states_shape
-            self.layer_map = attention_layer_map | mamba_layer_map
+            self.layer_map = attention_layer_map | mamba_layer_map | gdn_layer_map
         else:
             # The layer map is the identity function for pure Transformer models.
             self.num_attention_layers = model_config.num_layers // pp_size

megatron/core/models/mamba/mamba_layer_specs.py

@@ -9,6 +9,7 @@
 )
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
+from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
 from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
 from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules
 from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules
@@ -117,6 +118,20 @@
                 pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
             ),
         ),
+        gdn_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                self_attention=ModuleSpec(
+                    module=GatedDeltaNet,
+                    submodules=GatedDeltaNetSubmodules(
+                        in_proj=TELayerNormColumnParallelLinear,
+                        out_norm=TENorm,
+                        out_proj=TERowParallelLinear,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
         mtp_block_spec=_mamba_mtp_block_spec,
     ),
 )
@@ -179,6 +194,20 @@
                 pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
             ),
         ),
+        gdn_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                self_attention=ModuleSpec(
+                    module=GatedDeltaNet,
+                    submodules=GatedDeltaNetSubmodules(
+                        in_proj=TELayerNormColumnParallelLinear,
+                        out_norm=TENorm,
+                        out_proj=TERowParallelLinear,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
         mtp_block_spec=_mamba_mtp_block_spec,
     ),
 )

megatron/core/ssm/mamba_block.py

@@ -43,6 +43,7 @@ class MambaStackSubmodules:
     attention_layer: Union[ModuleSpec, type] = IdentityOp
     mlp_layer: Union[ModuleSpec, type] = IdentityOp
     moe_layer: Union[ModuleSpec, type] = IdentityOp
+    gdn_layer: Union[ModuleSpec, type] = IdentityOp
     mtp_block_spec: Optional[ModuleSpec] = None
 
 
@@ -176,6 +177,14 @@ def __init__(
                         layer_number=i + 1,
                         pg_collection=pg_collection,
                     )
+                elif layer_type == LayerSymbols.GDN:
+                    layer = build_module(
+                        submodules.gdn_layer,
+                        config=self.config,
+                        layer_number=i + 1,
+                        pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
+                    )
                 else:
                     assert False, "unexpected layer_type"
             self.layers.append(layer)

megatron/core/ssm/mamba_hybrid_layer_allocation.py

@@ -30,8 +30,9 @@ class Symbols:
     ATTENTION = "*"
     MLP = "-"
     MOE = 'E'
+    GDN = 'G'
     MTP_SEPARATOR = "/"
-    VALID = {MAMBA, ATTENTION, MLP, MOE}
+    VALID = {MAMBA, ATTENTION, MLP, MOE, GDN}
 
 
 @dataclass
@@ -264,7 +265,8 @@ def allocate_layers(
             logging.INFO,
             f"Hybrid allocation ({Symbols.MAMBA} is mamba, "
             f"{Symbols.ATTENTION} is attention, "
-            f"{Symbols.MLP} is mlp):",
+            f"{Symbols.MLP} is mlp, "
+            f"{Symbols.GDN} is gdn):",
         )
         maybe_log_single_rank(logger, logging.INFO, allocation_string)
         maybe_log_single_rank(
@@ -290,17 +292,24 @@ def allocate_layers(
             f"Target mlp ratio: {target_mlp_ratio:.2f}. "
             f"Actual mlp ratio: {actual_mlp_ratio:.2f}.",
         )
+        actual_gdn_layers_count = layer_type_list.count(Symbols.GDN)
+        actual_gdn_ratio = actual_gdn_layers_count / total_layers_count
+        maybe_log_single_rank(
+            logger,
+            logging.INFO,
+            f"{actual_gdn_layers_count} gdn layers in " f"{total_layers_count} total layers.",
+        )
     return layer_type_list
 
 
 def get_layer_maps_from_layer_type_list(
     layer_type_list: List[str],
-) -> Tuple[Dict[int, int], Dict[int, int], Dict[int, int]]:
+) -> Tuple[Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int]]:
     """
     Returns maps from global layer index to the corresponding layer index
-    for each layer type in [Attention, Mamba, MLP, MoE] given a layer type list.
+    for each layer type in [Attention, Mamba, MLP, MoE, GDN] given a layer type list.
     """
-    layer_types = [Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE]
+    layer_types = [Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE, Symbols.GDN]
     layer_maps = {layer_type: {} for layer_type in layer_types}
     for global_layer_idx, layer_type in enumerate(layer_type_list):
         layer_map = layer_maps[layer_type]

tests/unit_tests/ssm/test_mamba_block.py

@@ -5,6 +5,7 @@
 
 from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.ssm.gated_delta_net import GatedDeltaNet
 from megatron.core.ssm.mamba_block import MambaStack
 from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
 from megatron.core.ssm.mamba_layer import MambaLayer
@@ -87,3 +88,36 @@ def test_invalid_layer_types_cause_failure(self):
         # _allocate_override() in mamba_hybrid_layer_allocation.py throws a ValueError.
         with pytest.raises(ValueError):
             block = self.get_mamba_block(hybrid_override_pattern)
+
+    def test_gdn_layer_types(self):
+        """
+        Make sure that G creates a TransformerLayer wrapping GatedDeltaNet,
+        while * creates a TransformerLayer wrapping SelfAttention.
+        """
+        hybrid_override_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA
+        block = self.get_mamba_block(hybrid_override_pattern)
+        layers = block.layers
+        assert isinstance(layers[0], TransformerLayer)
+        assert isinstance(layers[0].self_attention, GatedDeltaNet)
+        assert isinstance(layers[1], TransformerLayer)
+        assert isinstance(layers[1].self_attention, SelfAttention)
+        assert isinstance(layers[2], MambaLayer)
+
+    def test_gdn_gpu_forward(self):
+        """Test GPU forward pass with GDN, attention, and Mamba layers."""
+        hybrid_override_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA
+        block = self.get_mamba_block(hybrid_override_pattern)
+        block.cuda()
+        micro_batch_size = 2
+        sequence_length = 32
+        hidden_states = torch.ones((sequence_length, micro_batch_size, block.config.hidden_size))
+        hidden_states = hidden_states.cuda()
+        attention_mask = torch.ones(
+            (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool
+        )
+        attention_mask = attention_mask.cuda()
+        output = block(hidden_states, attention_mask=attention_mask)
+        assert output.shape[0] == sequence_length
+        assert output.shape[1] == micro_batch_size
+        assert output.shape[2] == block.config.hidden_size
+        assert output.dtype == torch.float32

tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py

@@ -27,6 +27,8 @@ def test_hybrid_layer_allocation(self):
             (8, 0.25, 0.25, "MM*-MM*-"),
             (8, 0.5, 0.25, "M**-M**-"),
             (48, 0.5, 0.2, None),
+            (8, 0.0, 0.0, "GGG*GGG*"),
+            (8, 0.0, 0.0, "GEGEGE*E"),
         ]
         for test in test_cases:
             (layers_count, attention_ratio, mlp_ratio, override_pattern) = test
@@ -101,6 +103,8 @@ def test_main_pattern_only(self):
             ("*M*M", "*M*M"),
             ("MM-*", "MM-*"),
             ("E", "E"),
+            ("GGG*GGG*", "GGG*GGG*"),
+            ("GEGEGE*E", "GEGEGE*E"),
         ]
         for pattern, expected_main in test_cases:
             result = parse_hybrid_pattern(pattern)
@@ -200,6 +204,8 @@ def test_complex_patterns(self):
             ("*****/M/M/M/M", "*****", "M", 4),
             # MoE in main pattern
             ("MEME/MM/MM", "MEME", "MM", 2),
+            # GDN+MoE main pattern with GDN MTP
+            ("GEGEGE*E/GG/GG", "GEGEGE*E", "GG", 2),
         ]
         for pattern, expected_main, expected_mtp, expected_depths in test_cases:
             result = parse_hybrid_pattern(pattern)