diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index 63f2d244840..8305eaff5c8 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -389,6 +389,7 @@ def _prepare_input_for_gated_delta_rule( batch: int, seq_len: int, *gate_feats: tuple[torch.Tensor], + use_qk_l2norm_in_kernel: bool = False, ) -> dict[str, torch.Tensor]: """ Prepare all gated delta rule kernel inputs. @@ -396,6 +397,8 @@ def _prepare_input_for_gated_delta_rule( Fuses split, reshape, L2 norm, decay/gate activations, repeat_interleave, and contiguous operations. ``gate_feats`` holds the variant-specific in_proj sections, which ``_compute_gates`` turns into the decay and gating tensors. + When ``use_qk_l2norm_in_kernel`` is true, q/k normalization is deferred to + the gated delta rule kernel to avoid materializing normalized q/k here. Returns: (dict[str, Tensor]): Kernel inputs keyed by kernel argument name (``q``, @@ -413,8 +416,9 @@ def _prepare_input_for_gated_delta_rule( query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) value = value.reshape(batch, seq_len, -1, self.value_head_dim) - # Apply L2 norm to query and key - if self.use_qk_l2norm: + # Let a supporting kernel own normalization so caller autograd does not retain an + # additional pre-split, normalized query_key activation alongside the kernel q/k. + if self.use_qk_l2norm and not use_qk_l2norm_in_kernel: query_key = l2norm(query_key.contiguous()) # Split query and key diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index c7a983841cb..194725b219e 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -248,7 +248,15 @@ def forward( # Prepare all kernel inputs (split, reshape, L2 norm, gates, contiguous) nvtx_range_push(suffix="prepare_input_for_gated_delta_rule") kernel_inputs = self._prepare_input_for_gated_delta_rule( - qkv, gate, A_log_local_cp, dt_bias_local_cp, batch, seq_len, beta, alpha + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + seq_len, + beta, + alpha, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, ) gate = kernel_inputs.pop("gate") nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule") @@ -258,7 +266,7 @@ def forward( **kernel_inputs, initial_state=None, output_final_state=False, - use_qk_l2norm_in_kernel=False, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, cu_seqlens=cu_seqlens_q, ) nvtx_range_pop(suffix="gated_delta_rule") @@ -452,8 +460,8 @@ def torch_chunk_gated_delta_rule( query, key, value = q, k, v initial_dtype = query.dtype if use_qk_l2norm_in_kernel: - query = l2norm(query, dim=-1, eps=1e-6) - key = l2norm(key, dim=-1, eps=1e-6) + query = l2norm(query) + key = l2norm(key) query, key, value, beta, g = [ x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) ] diff --git a/megatron/core/ssm/gated_delta_net/gdn2.py b/megatron/core/ssm/gated_delta_net/gdn2.py index 62ad3ef0f32..73431e6c021 100644 --- a/megatron/core/ssm/gated_delta_net/gdn2.py +++ b/megatron/core/ssm/gated_delta_net/gdn2.py @@ -377,8 +377,8 @@ def torch_chunk_gdn2( initial_dtype = q.dtype if use_qk_l2norm_in_kernel: - q = l2norm(q, dim=-1, eps=1e-6) - k = l2norm(k, dim=-1, eps=1e-6) + q = l2norm(q) + k = l2norm(k) # b s h d -> b h s d, and compute the whole recurrence in fp32 query, key, value, g, b, w = [ diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index b63420b2fa9..fc238ea91ff 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -25,6 +25,7 @@ from megatron.core.ssm.gated_delta_net.common import ( _build_head_perm_for_split_sections, _build_thd_cp_a2a_perm, + l2norm, tensor_a2a_cp2hp, tensor_a2a_hp2cp, ) @@ -54,6 +55,229 @@ def _unpack_sequence(x: torch.Tensor, cu_seqlens: torch.Tensor, dim=1) -> list[t return unpacked_x +@pytest.mark.parametrize("use_gdn2", [False, True], ids=["gdn", "gdn2"]) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +def test_torch_gdn_l2norm_flag_matches_explicit_normalization(use_gdn2): + """The torch fallback flag must match explicit q/k normalization.""" + torch.manual_seed(123) + torch.cuda.set_device(Utils.local_rank % torch.cuda.device_count()) + device = torch.cuda.current_device() + batch_size, sequence_length, num_heads, head_dim = 1, 8, 2, 16 + shape = (batch_size, sequence_length, num_heads, head_dim) + q = torch.randn(shape, device=device) + k = torch.randn(shape, device=device) + v = torch.randn(shape, device=device) + + if use_gdn2: + kernel = torch_chunk_gdn2 + kernel_kwargs = { + "g": torch.full(shape, -0.01, device=device), + "b": torch.full(shape, 0.1, device=device), + "w": torch.full(shape, 0.1, device=device), + } + else: + kernel = torch_chunk_gated_delta_rule + gate_shape = (batch_size, sequence_length, num_heads) + kernel_kwargs = { + "g": torch.full(gate_shape, -0.01, device=device), + "beta": torch.full(gate_shape, 0.1, device=device), + } + + actual_output, actual_state = kernel( + q=q, + k=k, + v=v, + chunk_size=8, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + **kernel_kwargs, + ) + expected_output, expected_state = kernel( + q=l2norm(q), + k=l2norm(k), + v=v, + chunk_size=8, + output_final_state=True, + use_qk_l2norm_in_kernel=False, + **kernel_kwargs, + ) + + torch.testing.assert_close(actual_output, expected_output, rtol=0, atol=0) + torch.testing.assert_close(actual_state, expected_state, rtol=0, atol=0) + + +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +def test_gdn_in_kernel_l2norm_matches_caller_end_to_end(): + """GDN's BF16 in-kernel route must match the former caller-normalized route.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 + ) + try: + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + config = TransformerConfig( + hidden_size=128, + linear_conv_kernel_dim=4, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=2, + linear_num_value_heads=4, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=4, + num_query_groups=2, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=1, + sequence_parallel=False, + context_parallel_size=1, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + transformer_impl="transformer_engine", + ) + gdn_spec = get_experimental_attention_variant_module_spec(config=config) + + def build_gdn(): + return ( + gdn_spec.module( + config, + submodules=gdn_spec.submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + .cuda() + .bfloat16() + ) + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + kernel_norm_gdn = build_gdn() + caller_norm_gdn = build_gdn() + caller_norm_gdn.load_state_dict(kernel_norm_gdn.state_dict()) + + def install_route_probe(module, caller_normalizes): + original_prepare = module._prepare_input_for_gated_delta_rule + original_kernel = module.gated_delta_rule + observed = {} + + def wrapped_prepare( + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + seq_len, + *gate_feats, + use_qk_l2norm_in_kernel=False, + ): + observed["prepare_flag"] = use_qk_l2norm_in_kernel + effective_flag = False if caller_normalizes else use_qk_l2norm_in_kernel + kernel_inputs = original_prepare( + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + seq_len, + *gate_feats, + use_qk_l2norm_in_kernel=effective_flag, + ) + + # The new route must pass the raw post-convolution q/k to FLA. This + # explicit check catches accidental caller normalization (and hence + # double normalization) even when BF16 rounding hides the difference. + if not caller_normalizes: + query_key, _ = torch.split( + qkv.detach(), + [ + 2 * module.qk_dim_local_tp // module.cp_size, + module.v_dim_local_tp // module.cp_size, + ], + dim=-1, + ) + query_key = query_key.reshape(batch, seq_len, -1, module.key_head_dim) + split_size = module.qk_dim_local_tp // module.key_head_dim // module.cp_size + raw_query, raw_key = torch.split(query_key, [split_size, split_size], dim=2) + repeat_factor = module.num_value_heads // module.num_key_heads + if repeat_factor > 1: + raw_query = raw_query.repeat_interleave(repeat_factor, dim=2) + raw_key = raw_key.repeat_interleave(repeat_factor, dim=2) + observed["prepare_kept_qk_raw"] = torch.equal( + kernel_inputs["q"].detach(), raw_query.contiguous() + ) and torch.equal(kernel_inputs["k"].detach(), raw_key.contiguous()) + return kernel_inputs + + def wrapped_kernel(*args, **kwargs): + observed["kernel_flag"] = kwargs.get("use_qk_l2norm_in_kernel", False) + if caller_normalizes: + kwargs = {**kwargs, "use_qk_l2norm_in_kernel": False} + return original_kernel(*args, **kwargs) + + module._prepare_input_for_gated_delta_rule = wrapped_prepare + module.gated_delta_rule = wrapped_kernel + return observed + + kernel_route = install_route_probe(kernel_norm_gdn, caller_normalizes=False) + caller_route = install_route_probe(caller_norm_gdn, caller_normalizes=True) + + torch.manual_seed(123) + hidden_states = torch.randn( + (16, 2, config.hidden_size), device=torch.cuda.current_device(), dtype=torch.bfloat16 + ) + + def run(module): + module.zero_grad(set_to_none=True) + module_input = hidden_states.detach().clone().requires_grad_(True) + output, _ = module(module_input, None) + output.float().square().mean().backward() + parameter_grads = { + name: parameter.grad.detach().clone() + for name, parameter in module.named_parameters() + if parameter.grad is not None + } + return output.detach(), module_input.grad.detach().clone(), parameter_grads + + kernel_output, kernel_input_grad, kernel_parameter_grads = run(kernel_norm_gdn) + caller_output, caller_input_grad, caller_parameter_grads = run(caller_norm_gdn) + + assert kernel_route["prepare_flag"] is True + assert kernel_route["kernel_flag"] is True + assert kernel_route["prepare_kept_qk_raw"] is True + # The compatibility route still goes through GatedDeltaNet.forward, then + # overrides only where normalization is performed to model the former code. + assert caller_route["prepare_flag"] is True + assert caller_route["kernel_flag"] is True + + atol = rtol = 2e-2 + torch.testing.assert_close( + kernel_output.float(), caller_output.float(), atol=atol, rtol=rtol + ) + torch.testing.assert_close( + kernel_input_grad.float(), caller_input_grad.float(), atol=atol, rtol=rtol + ) + assert set(kernel_parameter_grads) == set(caller_parameter_grads) + for name in kernel_parameter_grads: + torch.testing.assert_close( + kernel_parameter_grads[name].float(), + caller_parameter_grads[name].float(), + atol=atol, + rtol=rtol, + msg=lambda msg, name=name: f"Parameter grad mismatch for {name}: {msg}", + ) + finally: + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("use_gdn2", [False, True], ids=["gdn", "gdn2"]) @pytest.mark.parametrize( ("tp_size", "sp", "cp_size"),