Skip to content

[Copy to main] fix(ssm): handle alignment padding in GDN packed seq + CP #4289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
svcnvidia-nemo-ci wants to merge 1 commit into from
Closed

[Copy to main] fix(ssm): handle alignment padding in GDN packed seq + CP #4289

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 107 additions & 20 deletions megatron/core/ssm/gated_delta_net.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -297,30 +297,79 @@ 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")
qkvzba, _ = self.in_proj(hidden_states)
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
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
####################
Expand Down
146 changes: 146 additions & 0 deletions tests/unit_tests/ssm/test_gated_delta_net.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"),
[
Expand Down