From 5e50ac20bbc51f79c29a257a3f191f92611a6018 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Wed, 22 Apr 2026 14:43:14 -0700 Subject: [PATCH] fix: pass total_tokens to PackedSeqParams for SSM seq_idx generation Without total_tokens, PackedSeqParams.__post_init__ skips seq_idx computation, causing Mamba/SSM state to bleed across sequence boundaries in varlen (packed sequence) training. This affects all hybrid SSM models (Nemotron Nano, NemotronH, Nemotron Nano V2, Nemotron 3 Super), not just Nemotron Nano. Fixes #3474 Co-Authored-By: Claude Opus 4.6 Signed-off-by: yaoyu-33 --- src/megatron/bridge/training/gpt_step.py | 1 + src/megatron/bridge/training/llava_step.py | 1 + .../bridge/training/utils/packed_seq_utils.py | 8 ++- src/megatron/bridge/training/vlm_step.py | 1 + .../training/utils/test_packed_seq_utils.py | 54 +++++++++++++++++++ 5 files changed, 63 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index f05981d57f..589235c9ea 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -253,6 +253,7 @@ def _forward_step_common( "max_seqlen": max_seqlen, "cu_seqlens_unpadded": cu_seqlens_unpadded, "cu_seqlens_unpadded_argmin": cu_seqlens_unpadded_argmin, + "total_tokens": tokens.size(1) if tokens is not None else labels.size(1), } forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params) diff --git a/src/megatron/bridge/training/llava_step.py b/src/megatron/bridge/training/llava_step.py index 1723cfaeaf..f3f1019556 100644 --- a/src/megatron/bridge/training/llava_step.py +++ b/src/megatron/bridge/training/llava_step.py @@ -198,6 +198,7 @@ def forward_step( "cu_seqlens": cu_seqlens, "cu_seqlens_argmin": cu_seqlens_argmin, "max_seqlen": max_seqlen, + "total_tokens": input_ids.size(1) if input_ids is not None else labels.size(1), } forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params) diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py index 1194fc7ed2..9683a17278 100644 --- a/src/megatron/bridge/training/utils/packed_seq_utils.py +++ b/src/megatron/bridge/training/utils/packed_seq_utils.py @@ -27,8 +27,9 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams: Args: batch: A dictionary containing packed-sequence metadata. Expected keys: - `cu_seqlens`, optional `cu_seqlens_unpadded`, optional argmins, and - optional `max_seqlen`. + `cu_seqlens`, optional `cu_seqlens_unpadded`, optional argmins, + optional `max_seqlen`, and optional `total_tokens` (required for + hybrid SSM/Mamba models to generate ``seq_idx``). Returns: PackedSeqParams with identical q/kv parameters and `qkv_format` set to @@ -57,6 +58,7 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams: cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded)] max_seqlen = batch["max_seqlen"].squeeze() if "max_seqlen" in batch else None + total_tokens = batch.get("total_tokens") # When cu_seqlens_unpadded is present (pad_seq_to_mult > 1), pass both unpadded and padded # for proper THD CP support. Otherwise, just use cu_seqlens_padded to avoid slower TE kernel. @@ -68,6 +70,7 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams: cu_seqlens_kv_padded=cu_seqlens_padded, max_seqlen_q=max_seqlen, max_seqlen_kv=max_seqlen, + total_tokens=total_tokens, qkv_format="thd", ) else: @@ -76,5 +79,6 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams: cu_seqlens_kv=cu_seqlens_padded, max_seqlen_q=max_seqlen, max_seqlen_kv=max_seqlen, + total_tokens=total_tokens, qkv_format="thd", ) diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py index c3622c9815..8c3fb6df6b 100644 --- a/src/megatron/bridge/training/vlm_step.py +++ b/src/megatron/bridge/training/vlm_step.py @@ -448,6 +448,7 @@ def forward_step( "cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen, "cu_seqlens_argmin": cu_seqlens_argmin, + "total_tokens": tokens.size(1) if tokens is not None else labels.size(1), } forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params) diff --git a/tests/unit_tests/training/utils/test_packed_seq_utils.py b/tests/unit_tests/training/utils/test_packed_seq_utils.py index 3f9267f69a..e63aa5127e 100644 --- a/tests/unit_tests/training/utils/test_packed_seq_utils.py +++ b/tests/unit_tests/training/utils/test_packed_seq_utils.py @@ -153,6 +153,60 @@ def test_single_sequence(self): torch.testing.assert_close(result.cu_seqlens_q, expected) assert result.cu_seqlens_q_padded is None # No unpadded, so no padded variants + def test_total_tokens_generates_seq_idx(self): + """Test that passing total_tokens causes PackedSeqParams to generate seq_idx. + + This is critical for hybrid SSM/Mamba models in varlen (packed sequence) + settings. Without total_tokens, seq_idx remains None and SSM state bleeds + across sequence boundaries. + """ + batch = { + "cu_seqlens": torch.IntTensor([0, 5, 7, 11, -1]), + "cu_seqlens_argmin": torch.tensor(4), + "max_seqlen": torch.tensor(6), + "total_tokens": 16, + } + + result = get_packed_seq_params(batch) + + assert result.total_tokens == 16 + assert result.seq_idx is not None + # seq_idx maps each token position to its sequence index: + # seq 0: tokens 0-4 (len 5), seq 1: tokens 5-6 (len 2), + # seq 2: tokens 7-10 (len 4), seq 3: tokens 11-15 (len 5) + expected_seq_idx = torch.IntTensor([[0, 0, 0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3]]) + torch.testing.assert_close(result.seq_idx, expected_seq_idx) + + def test_without_total_tokens_seq_idx_is_none(self): + """Test that omitting total_tokens leaves seq_idx as None (backward compat).""" + batch = { + "cu_seqlens": torch.IntTensor([0, 128, 256, -1]), + "cu_seqlens_argmin": torch.tensor(3), + "max_seqlen": torch.tensor(128), + } + + result = get_packed_seq_params(batch) + + assert result.total_tokens is None + assert result.seq_idx is None + + def test_total_tokens_with_cu_seqlens_unpadded(self): + """Test total_tokens flows through when cu_seqlens_unpadded is present.""" + batch = { + "cu_seqlens": torch.IntTensor([0, 128, 256, 384, -1]), + "cu_seqlens_argmin": torch.tensor(4), + "cu_seqlens_unpadded": torch.IntTensor([0, 120, 245, 370, -1]), + "cu_seqlens_unpadded_argmin": torch.tensor(4), + "max_seqlen": torch.tensor(128), + "total_tokens": 384, + } + + result = get_packed_seq_params(batch) + + assert result.total_tokens == 384 + # seq_idx should be generated from cu_seqlens_q_padded (the padded variant) + assert result.seq_idx is not None + def test_performance_no_unnecessary_padded_variants(self): """Verify that when unpadded is not provided, padded variants are None.