From 3a1d51e94902ec34909e7419011a474169c6bb4f Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 14 Apr 2026 01:04:12 +0000 Subject: [PATCH] Copy PR #4230 to main --- megatron/core/ssm/gated_delta_net.py | 127 +++++++++++++--- tests/unit_tests/ssm/test_gated_delta_net.py | 146 +++++++++++++++++++ 2 files changed, 253 insertions(+), 20 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 57315ec48d9..77fd9ecbeba 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -297,9 +297,37 @@ def forward( # TODO: support inference raise NotImplementedError("GDN does not support inference for now.") - if packed_seq_params is not None: - # TODO: support packed sequence - raise NotImplementedError("GDN does not support packed sequence for now.") + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + assert batch == 1, "Packed sequence expects batch dimension to be 1" + assert ( + not self.config.deterministic_mode + ), "Packed sequence does not support deterministic mode." + + # Resolve cu_seqlens with alignment padding handling. + cu_seqlens_q = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_q_padded, + packed_seq_params.cu_seqlens_q, + seq_len, + "cu_seqlens_q", + ) + cu_seqlens_kv = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, + seq_len, + "cu_seqlens_kv", + ) + assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( + "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " + f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" + ) + num_packed_seqs = cu_seqlens_q.shape[0] - 1 + assert num_packed_seqs > 0, ( + "Number of packed sequences must be greater than 0, " + f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" + ) + else: + cu_seqlens_q = None + cu_seqlens_kv = None # Input projection nvtx_range_push(suffix="in_proj") @@ -307,20 +335,41 @@ def forward( nvtx_range_pop(suffix="in_proj") # CP All to All: CP to HP - qkvzba = tensor_a2a_cp2hp( - qkvzba, - seq_dim=0, - head_dim=-1, - cp_group=self.pg_collection.cp, - split_sections=[ - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - self.v_dim_local_tp, - self.num_value_heads // self.tp_size, - self.num_value_heads // self.tp_size, - ], - ) + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + unpacked_qkvzba = _unpack_sequence(qkvzba, cu_seqlens_q // self.cp_size, dim=0) + outputs = [] + for qkvzba_i in unpacked_qkvzba: + qkvzba_i = tensor_a2a_cp2hp( + qkvzba_i, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + ) + outputs.append(qkvzba_i) + qkvzba = torch.cat(outputs, dim=0) + else: + qkvzba = tensor_a2a_cp2hp( + qkvzba, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + ) # Transpose: s b x --> b s x # From sbhd to bshd format @@ -430,9 +479,19 @@ def forward( norm_out = norm_out.transpose(0, 1).contiguous() # CP all to all: HP to CP - norm_out = tensor_a2a_hp2cp( - norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + unpacked_norm_out = _unpack_sequence(norm_out, cu_seqlens_q, dim=0) + outputs = [] + for norm_out_i in unpacked_norm_out: + norm_out_i = tensor_a2a_hp2cp( + norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + outputs.append(norm_out_i) + norm_out = torch.cat(outputs, dim=0) + else: + norm_out = tensor_a2a_hp2cp( + norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) # Output projection nvtx_range_push(suffix="out_proj") @@ -504,6 +563,23 @@ def _compute_g_and_beta(self, A_log_local_cp, dt_bias_local_cp, alpha, beta): beta = beta.sigmoid() return g, beta + def _resolve_cu_seqlens(self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name): + """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" + if cu_seqlens_padded is not None: + cu_seqlens = cu_seqlens_padded + else: + cu_seqlens = cu_seqlens_actual + + total_cu = cu_seqlens[-1].item() + if total_cu != total_seq_len: + raise ValueError( + f"GDN: {name}[-1]={total_cu} does not match " + f"total_sequence_length={total_seq_len}. " + f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." + ) + + return cu_seqlens + 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 @@ -602,6 +678,17 @@ def _backward_out_proj(self): self.out_proj.backward_dw() +def _unpack_sequence(x, cu_seqlens, dim=1): + unpacked_x = [] + num_seqs = cu_seqlens.shape[0] - 1 + for i in range(num_seqs): + idx_start = cu_seqlens[i].item() + idx_end = cu_seqlens[i + 1].item() + chunked_index = [slice(None)] * dim + [slice(idx_start, idx_end)] + unpacked_x.append(x[tuple(chunked_index)]) + return unpacked_x + + #################### # Sharded state dict utilities #################### diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index ec4d7a86ecf..cbf8feafff1 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -31,6 +31,10 @@ ) from tests.unit_tests.test_utilities import Utils from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness +from tests.unit_tests.transformer.test_multi_latent_attention import ( + make_test_packed_seq_params, + make_test_packed_seq_params_with_padding, +) try: import fla @@ -202,6 +206,148 @@ def test_jit_compiled_helpers(self): assert beta_sig.shape == beta.shape + def test_gpu_forward_thd_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + + # Input shape + sequence_length = 32 + micro_batch_size = 4 + cu_seqlens = [0, 32, 64, 96, 128] + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size) + ) + attention_mask_sbhd = None + hidden_states_sbhd = hidden_states_sbhd.cuda().bfloat16() + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + attention_mask_thd = None + packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + + # THD format + output_thd, _ = self.gdn( + hidden_states_thd, attention_mask_thd, packed_seq_params=packed_seq_params + ) + # SBHD format + output_sbhd, _ = self.gdn(hidden_states_sbhd, attention_mask_sbhd) + output_sbhd_T = output_sbhd.transpose(0, 1).contiguous().view(*output_thd.shape) + + rank = torch.distributed.get_rank() + assert output_thd.shape[0] == sub_sequence_length * micro_batch_size + assert output_thd.shape[1] == 1 + assert output_thd.shape[2] == self.gdn.config.hidden_size + torch.testing.assert_close( + output_sbhd_T, + output_thd, + atol=atol, + rtol=rtol, + msg=lambda msg: f"Output mismatch ({rank=}): {msg}", + ) + + def test_gpu_forward_thd_padding_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + sequence_length = 32 + micro_batch_size = 4 + + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + output_sbhd, _ = self.gdn(hidden_states_sbhd, None) + + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + output_bshd = output_sbhd.transpose(0, 1).contiguous() + + rank = torch.distributed.get_rank() + + # A) padded branch: prefer *_padded when available. + padded_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] + ) + output_thd_padded, _ = self.gdn(hidden_states_thd, None, packed_seq_params=padded_params) + output_thd2bshd = output_thd_padded.view(*output_bshd.shape) + torch.testing.assert_close( + output_bshd[..., :30], + output_thd2bshd[..., :30], + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD padded output mismatch ({rank=}): {msg}", + ) + + # B) no-padded branch: use actual cu_seqlens when it matches total_sequence_length. + no_padding_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 128]) + output_thd_no_padding, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + assert output_thd_no_padding.shape == output_thd_padded.shape + + # C) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. + padded_mismatch_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 126] + ) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) + + # D) actual mismatch branch without *_padded: should raise. + actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) + + +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGDNCuSeqlensResolve: + + @pytest.fixture + def mock_gdn(self): + class MockGDN: + cp_size = 2 + _resolve_cu_seqlens = GatedDeltaNet._resolve_cu_seqlens + + return MockGDN() + + def test_padded_preferred_when_available(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q") + assert torch.equal(result, padded) + + def test_actual_used_when_no_padding(self, mock_gdn): + actual = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q") + assert torch.equal(result, actual) + + def test_raises_when_padding_mismatch(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q") + + def test_raises_when_padded_mismatches_total(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1004], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q") + + def test_cp1_still_validates_total(self, mock_gdn): + mock_gdn.cp_size = 1 + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q") + + @pytest.mark.parametrize( ("tp", "sp", "cp"), [