diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 0c172d0b47f..a602758b937 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/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 diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py index 39f7fab0266..d2a85d004ef 100755 --- a/megatron/core/models/mamba/mamba_layer_specs.py +++ b/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 diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py index 65e95d038c5..f42f3542c3d 100644 --- a/megatron/core/ssm/mamba_block.py +++ b/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) diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py index 92dde8132ce..1cb1e4a31d1 100644 --- a/megatron/core/ssm/mamba_hybrid_layer_allocation.py +++ b/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] diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index a4d938f2e7b..0d1391d8b97 100644 --- a/megatron/training/arguments.py +++ b/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, diff --git a/megatron/training/training.py b/megatron/training/training.py index fb3d85d3a60..471cf422841 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -335,8 +335,29 @@ 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, @@ -344,6 +365,9 @@ def hybrid_flops(batch_size, seq_len, hidden_size, 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, ) diff --git a/tests/unit_tests/ssm/test_mamba_block.py b/tests/unit_tests/ssm/test_mamba_block.py index c65623e08e0..7b743afbfad 100644 --- a/tests/unit_tests/ssm/test_mamba_block.py +++ b/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 diff --git a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py index a0a0401d113..440c843bc27 100644 --- a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py +++ b/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