Introduce GDN to Mamba

megatron/core/inference/contexts/dynamic_context.py

@@ -329,9 +329,13 @@ 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
+            mamba_layer_map, gdn_layer_map, attention_layer_map, _, _ = (
+                get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)
             )
+
+            if len(gdn_layer_map) > 0:
+                raise NotImplementedError("GDN layers are not supported for inference.")
+
             self.num_attention_layers = len(attention_layer_map)
             self.num_mamba_layers = len(mamba_layer_map)
             self.layer_map = attention_layer_map | mamba_layer_map

megatron/core/models/mamba/mamba_layer_specs.py

@@ -12,6 +12,7 @@
     get_inference_optimized_moe_spec,
     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
@@ -83,6 +84,20 @@
                 mamba_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,
+            ),
+        ),
         # Started with spec from gpt_layer_specs.py (with MLP removed)
         # Using the TE spec because we had problems getting the non-TE spec
         # working

megatron/core/ssm/mamba_block.py

@@ -39,6 +39,7 @@ class MambaStackSubmodules:
     """
 
     mamba_layer: Union[ModuleSpec, type] = IdentityOp
+    gdn_layer: Union[ModuleSpec, type] = IdentityOp
     attention_layer: Union[ModuleSpec, type] = IdentityOp
     mlp_layer: Union[ModuleSpec, type] = IdentityOp
     moe_layer: Union[ModuleSpec, type] = IdentityOp
@@ -150,6 +151,15 @@ def __init__(
                         pg_collection=pg_collection,
                         add_layer_offset=False,
                     )
+                elif layer_type == LayerSymbols.GDN:
+                    layer = build_module(
+                        submodules.gdn_layer,
+                        config=self.config,
+                        layer_number=layer_number,
+                        pg_collection=pg_collection,
+                        # Set to False as we do not want to change offset.
+                        add_layer_offset=False,
+                    )
                 else:
                     assert False, "unexpected layer_type"
             self.layers.append(layer)

megatron/core/ssm/mamba_hybrid_layer_allocation.py

@@ -15,12 +15,13 @@ class Symbols:
     """Symbols for different layer types and pattern separators."""
 
     MAMBA = "M"
+    GDN = 'G'
     ATTENTION = "*"
     MLP = "-"
     MOE = 'E'
     PIPE = '|'
     MTP_SEPARATOR = "/"
-    VALID_LAYERS = {MAMBA, ATTENTION, MLP, MOE}
+    VALID_LAYERS = {MAMBA, GDN, ATTENTION, MLP, MOE}
 
 
 @dataclass
@@ -154,18 +155,24 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]:
         pattern: Full hybrid layer pattern string.
 
     Returns:
-        Dictionary mapping layer symbol to count. Keys are Symbols.ATTENTION,
-        Symbols.MAMBA, Symbols.MLP, and Symbols.MOE.
+        Dictionary mapping layer symbol to count. Keys are Symbols.MAMBA,
+        Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, and Symbols.MOE.
 
     Examples:
         >>> get_hybrid_layer_counts("M*M*")
-        {'*': 2, 'M': 2, '-': 0, 'E': 0}
+        {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0}
 
         >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM")
-        {'*': 1, 'M': 8, '-': 4, 'E': 0}
+        {'M': 8, 'G': 0, '*': 1, '-': 4, 'E': 0}
     """
     parsed = parse_hybrid_pattern(pattern)
-    counts = {Symbols.ATTENTION: 0, Symbols.MAMBA: 0, Symbols.MLP: 0, Symbols.MOE: 0}
+    counts = {
+        Symbols.MAMBA: 0,
+        Symbols.GDN: 0,
+        Symbols.ATTENTION: 0,
+        Symbols.MLP: 0,
+        Symbols.MOE: 0,
+    }
 
     # Count main decoder layers (skip '|' pipe separators)
     if parsed.main_pattern:
@@ -463,12 +470,12 @@ def select_pipeline_segment(
 
 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 [Mamba, GDN, Attention, MLP, MoE] given a layer type list.
     """
-    layer_types = [Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE]
+    layer_types = [Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE]
     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]

megatron/training/arguments.py

@@ -3201,10 +3201,10 @@ def _add_experimental_args(parser):
                        'For more details, see the model class, '
                        '`transformer_block.py`, or `transformer_layer.py`')
     group.add_argument('--hybrid-layer-pattern', type=str, default=None,
-                       help='Specify a hybrid layer pattern using M (mamba), * (attention), '
-                       '- (mlp), E (moe). Use | to define pipeline stage boundaries for '
-                       'flexible virtual pipeline parallel (fVPP). Use / to separate MTP '
-                       'patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". '
+                       help='Specify a hybrid layer pattern using M (mamba), G (gdn), '
+                       '* (attention), - (mlp), E (moe). Use | to define pipeline stage '
+                       'boundaries for flexible virtual pipeline parallel (fVPP). Use / to '
+                       'separate MTP patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". '
                        'When this flag is used, it is the sole indicator that a hybrid model '
                        'is being run.')
     group.add_argument('--hybrid-override-pattern', type=str, default=None,

megatron/training/training.py

@@ -335,15 +335,39 @@ def mamba_layer_flops(batch_size, seq_len, hidden_size, state_dim=16,
             + (2 * batch_size * seq_len * d_in * hidden_size)  # out_proj
         )
 
+    def gdn_layer_flops(batch_size, seq_len, hidden_size,
+                        qk_head_dim=128, v_head_dim=128,
+                        num_qk_heads=16, num_v_heads=32,
+                        conv_kernel_dim=4):
+        """Calculate FLOPs for a Gated Delta Net (GDN) layer."""
+        qk_dim = qk_head_dim * num_qk_heads
+        v_dim = v_head_dim * num_v_heads
+        return (
+            2 * batch_size * seq_len * (
+                # in_proj: hidden_size -> (2*qk_dim + 2*v_dim + 2*num_v_heads)
+                hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads)
+                # conv1d
+                + conv_kernel_dim * (2 * qk_dim + v_dim)
+                # gated delta rule: KK^T, VK^T, S(a(I-bKK^T)), and SQ
+                + num_v_heads * (v_head_dim ** 2) * 4
+                # out_proj: v_dim -> hidden_size
+                + hidden_size * v_dim
+            )
+        )
+
     def hybrid_flops(batch_size, seq_len, hidden_size,
                      num_attn_layers, num_mamba_layers, num_mlp_layers, num_moe_layers,
+                     num_gdn_layers=0,
                      mamba_state_dim=128, mamba_head_dim=64,
                      mamba_num_groups=8, mamba_num_heads=128,
                      num_attn_heads=32, gqa=True,
                      gqa_groups=8, kv_channels=None,
                      mlp_expansion=4.0, swiglu=False,
                      moe_latent_size=None,
                      moe_ffn_hidden_size=2048, shared_expert_ffn_hidden_size=2048, num_experts_routed_to=1,
+                     gdn_qk_head_dim=128, gdn_v_head_dim=128,
+                     gdn_num_qk_heads=16, gdn_num_v_heads=32,
+                     gdn_conv_kernel_dim=4,
                      vocab_size=256000, mtp_num_layers=0):
         """Calculate total FLOPs for the hybrid model."""
         flops_fwd = (
@@ -357,6 +381,10 @@ def hybrid_flops(batch_size, seq_len, hidden_size,
                 num_moe_layers * moe_layer_flops(batch_size, seq_len, hidden_size, moe_ffn_hidden_size,
                                                  shared_expert_ffn_hidden_size, num_experts_routed_to,
                                                  moe_latent_size, swiglu) +
+                num_gdn_layers * gdn_layer_flops(batch_size, seq_len, hidden_size,
+                                                  gdn_qk_head_dim, gdn_v_head_dim,
+                                                  gdn_num_qk_heads, gdn_num_v_heads,
+                                                  gdn_conv_kernel_dim) +
                 (2 * batch_size * seq_len * hidden_size * vocab_size * (1 + mtp_num_layers))  # logits computation
         )
         return flops_fwd * 3
@@ -637,9 +665,11 @@ def transformer_flops():
         from operator import itemgetter
 
         from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, get_hybrid_layer_counts
-        num_attn_layers, num_mamba_layers, num_mlp_layers, num_moe_layers = itemgetter(
-            Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE
-        )(get_hybrid_layer_counts(args.hybrid_layer_pattern))
+        num_mamba_layers, num_gdn_layers, num_attn_layers, num_mlp_layers, num_moe_layers = (
+            itemgetter(Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE)(
+                get_hybrid_layer_counts(args.hybrid_layer_pattern)
+            )
+        )
 
         mtp_num_layers = args.mtp_num_layers
         if mtp_num_layers is None:
@@ -653,6 +683,7 @@ def transformer_flops():
             num_mamba_layers=num_mamba_layers,
             num_mlp_layers=num_mlp_layers,
             num_moe_layers=num_moe_layers,
+            num_gdn_layers=num_gdn_layers,
             mamba_state_dim=args.mamba_state_dim,
             mamba_head_dim=args.mamba_head_dim,
             mamba_num_groups=args.mamba_num_groups,
@@ -669,6 +700,11 @@ def transformer_flops():
             shared_expert_ffn_hidden_size=(0 if args.moe_shared_expert_intermediate_size is None
                                            else args.moe_shared_expert_intermediate_size),
             num_experts_routed_to=args.moe_router_topk,
+            gdn_qk_head_dim=args.linear_key_head_dim or 128,
+            gdn_v_head_dim=args.linear_value_head_dim or 128,
+            gdn_num_qk_heads=args.linear_num_key_heads or 16,
+            gdn_num_v_heads=args.linear_num_value_heads or 32,
+            gdn_conv_kernel_dim=args.linear_conv_kernel_dim or 4,
             vocab_size=args.padded_vocab_size,
             mtp_num_layers=mtp_num_layers,
         )

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, validate_segment_layers
 from megatron.core.ssm.mamba_layer import MambaLayer
@@ -89,3 +90,51 @@ def test_invalid_layer_types_cause_failure(self):
         # validate_segment_layers() in mamba_hybrid_layer_allocation.py throws a ValueError.
         with pytest.raises(ValueError):
             block = self.get_mamba_block(layer_pattern)
+
+    def test_gdn_layer_types(self):
+        """
+        Make sure that G creates a TransformerLayer wrapping GatedDeltaNet,
+        while * creates a TransformerLayer wrapping SelfAttention.
+        """
+        layer_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA
+        block = self.get_mamba_block(layer_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."""
+        layer_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA
+        layer_type_list = validate_segment_layers(layer_pattern)
+        transformer_config = TransformerConfig(
+            hidden_size=256,
+            num_layers=len(layer_type_list),
+            num_attention_heads=4,
+            use_cpu_initialization=True,
+            activation_func=torch.nn.functional.silu,
+        )
+        modules = mamba_stack_spec.submodules
+        block = MambaStack(
+            transformer_config,
+            modules,
+            layer_type_list=layer_type_list,
+            pp_layer_offset=0,
+            pg_collection=self.get_pg_collection(),
+        )
+        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

@@ -73,6 +73,8 @@ def test_valid_patterns(self):
             ("MM*-MM*-", ['M', 'M', '*', '-', 'M', 'M', '*', '-']),
             ("E", ['E']),
             ("", []),
+            ("GGG*GGG*", ['G', 'G', 'G', '*', 'G', 'G', 'G', '*']),
+            ("GEGEGE*E", ['G', 'E', 'G', 'E', 'G', 'E', '*', 'E']),
         ]
         for pattern, expected in test_cases:
             result = validate_segment_layers(pattern)
@@ -151,6 +153,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)
@@ -271,6 +275,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)
@@ -289,34 +295,47 @@ def test_dataclass_equality(self):
 class TestGetHybridLayerCounts:
 
     def test_simple_pattern(self):
-        assert get_hybrid_layer_counts("M*M*") == {'*': 2, 'M': 2, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("M*M*") == {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0}
 
     def test_all_layer_types(self):
-        assert get_hybrid_layer_counts("M*-E") == {'*': 1, 'M': 1, '-': 1, 'E': 1}
+        assert get_hybrid_layer_counts("MG*-E") == {'M': 1, 'G': 1, '*': 1, '-': 1, 'E': 1}
 
     def test_with_pipes(self):
         # Pipes should be skipped in counting
-        assert get_hybrid_layer_counts("M*|M*") == {'*': 2, 'M': 2, '-': 0, 'E': 0}
-        assert get_hybrid_layer_counts("M-M-|M-M*-") == {'*': 1, 'M': 4, '-': 4, 'E': 0}
+        assert get_hybrid_layer_counts("M*|M*") == {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("M-M-|M-M*-") == {'M': 4, 'G': 0, '*': 1, '-': 4, 'E': 0}
 
     def test_with_mtp(self):
         # MTP pattern "MM" repeated 2 depths -> 4 extra mamba layers
-        assert get_hybrid_layer_counts("M*M*/MM/MM") == {'*': 2, 'M': 6, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("M*M*/MM/MM") == {'M': 6, 'G': 0, '*': 2, '-': 0, 'E': 0}
 
     def test_with_pipes_and_mtp(self):
         # Main: M-M-|M-M*- -> 1 attn, 4 mamba, 4 mlp
         # MTP: MM x 2 depths -> +4 mamba
-        assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == {'*': 1, 'M': 8, '-': 4, 'E': 0}
+        assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == {
+            'M': 8,
+            'G': 0,
+            '*': 1,
+            '-': 4,
+            'E': 0,
+        }
 
     def test_moe_pattern(self):
-        assert get_hybrid_layer_counts("MEME") == {'*': 0, 'M': 2, '-': 0, 'E': 2}
+        assert get_hybrid_layer_counts("MEME") == {'M': 2, 'G': 0, '*': 0, '-': 0, 'E': 2}
 
     def test_mtp_with_attention(self):
         # MTP pattern "*M" repeated 3 depths -> 3 attn + 3 mamba from MTP
-        assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == {'*': 3, 'M': 7, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == {'M': 7, 'G': 0, '*': 3, '-': 0, 'E': 0}
 
     def test_empty_pattern(self):
-        assert get_hybrid_layer_counts("") == {'*': 0, 'M': 0, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("") == {'M': 0, 'G': 0, '*': 0, '-': 0, 'E': 0}
+
+    def test_gdn_pattern(self):
+        assert get_hybrid_layer_counts("GMGM") == {'M': 2, 'G': 2, '*': 0, '-': 0, 'E': 0}
+
+    def test_gdn_hybrid_pattern(self):
+        # GDN + Mamba + Attention
+        assert get_hybrid_layer_counts("G*GM*") == {'M': 1, 'G': 2, '*': 2, '-': 0, 'E': 0}
 
 
 @pytest.mark.internal