diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 601a72a4356..9cb50ba6953 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -388,37 +388,12 @@ def forward( ) nvtx_range_pop(suffix="conv1d") - # Split qkv into query_key, and value - query_key, value = torch.split( - qkv, - [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], - dim=-1, - ) - 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: - query_key = l2norm(query_key.contiguous()) - # Split query and key. - query, key = torch.split( - query_key, - [ - self.qk_dim_local_tp // self.key_head_dim // self.cp_size, - self.qk_dim_local_tp // self.key_head_dim // self.cp_size, - ], - dim=2, + # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous) + nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule") + query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( + qkv, gate, beta, alpha, batch, seq_len ) - if self.num_value_heads // self.num_key_heads > 1: - query = query.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) - key = key.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) - - # Make contiguous - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - gate = gate.contiguous() - beta = beta.contiguous() - alpha = alpha.contiguous() + nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule") # Calculate g and beta nvtx_range_push(suffix="g_and_beta") @@ -426,8 +401,7 @@ def forward( dt_bias_local_cp = get_parameter_local_cp( self.dt_bias, dim=0, cp_group=self.pg_collection.cp ) - g = -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp) # In fp32 - beta = beta.sigmoid() + g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") nvtx_range_push(suffix="gated_delta_rule") @@ -477,6 +451,57 @@ def _apply_gated_norm(self, x, gate): y = y.to(x_dtype) return y + @jit_fuser + def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len): + """ + Prepare query, key, value, gate, beta, alpha tensors for gated delta rule. + Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations. + """ + # Split qkv into query_key and value + query_key, value = torch.split( + qkv, + [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], + dim=-1, + ) + + # Reshape query_key and value + 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: + query_key = l2norm(query_key.contiguous()) + + # Split query and key + split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size + query, key = torch.split(query_key, [split_size, split_size], dim=2) + + # Expand query and key if needed (grouped query attention) + if self.num_value_heads // self.num_key_heads > 1: + repeat_factor = self.num_value_heads // self.num_key_heads + query = query.repeat_interleave(repeat_factor, dim=2) + key = key.repeat_interleave(repeat_factor, dim=2) + + # Make all tensors contiguous + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + gate = gate.contiguous() + beta = beta.contiguous() + alpha = alpha.contiguous() + + return query, key, value, gate, beta, alpha + + @jit_fuser + def _compute_g_and_beta(self, A_log_local_cp, dt_bias_local_cp, alpha, beta): + """ + Compute g (decay) and beta (sigmoid) for gated delta rule. + Fuses exp, softplus, mul, neg, and sigmoid operations. + """ + g = -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp) # In fp32 + beta = beta.sigmoid() + return g, beta + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): """Provide a sharded state dictionary for distributed checkpointing.""" # Guard for cases metadata is not provided diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 81f8eed0574..8f3c59b3d43 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -138,6 +138,70 @@ def test_gpu_forward(self): output.dtype == hidden_states.dtype ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + def test_jit_compiled_helpers(self): + import torch._dynamo + + gdn = self.gdn + batch = 2 + seq_len = 16 + + num_v_heads_local = gdn.num_value_heads // gdn.tp_size // gdn.cp_size + + qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // gdn.cp_size + qkv = torch.randn( + batch, seq_len, qkv_last_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 + ) + gate = torch.randn( + batch, + seq_len, + num_v_heads_local, + gdn.value_head_dim, + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + beta = torch.randn( + batch, + seq_len, + num_v_heads_local, + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + alpha = torch.randn( + batch, + seq_len, + num_v_heads_local, + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + # Disable dynamo so coverage.py can trace through the method bodies, + # which are normally wrapped by @jit_fuser (torch.compile). + with torch._dynamo.config.patch(disable=True): + query, key, value, gate_out, beta_out, alpha_out = ( + gdn._prepare_qkv_for_gated_delta_rule(qkv, gate, beta, alpha, batch, seq_len) + ) + + assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert key.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert value.shape == (batch, seq_len, num_v_heads_local, gdn.value_head_dim) + assert query.is_contiguous() + assert key.is_contiguous() + assert value.is_contiguous() + + A_log_mock = torch.randn( + num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16 + ) + dt_bias_mock = torch.randn( + num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16 + ) + + with torch._dynamo.config.patch(disable=True): + g, beta_sig = gdn._compute_g_and_beta(A_log_mock, dt_bias_mock, alpha, beta) + + assert g.dtype == torch.float32 + assert g.shape == alpha.shape + assert beta_sig.shape == beta.shape + @pytest.mark.parametrize( ("tp", "sp", "cp"),