diff --git a/src/megatron/bridge/models/exaone/exaone45/modelling_exaone45/utils.py b/src/megatron/bridge/models/exaone/exaone45/modelling_exaone45/utils.py index d5e9eebe4c..2c8479d82a 100755 --- a/src/megatron/bridge/models/exaone/exaone45/modelling_exaone45/utils.py +++ b/src/megatron/bridge/models/exaone/exaone45/modelling_exaone45/utils.py @@ -633,14 +633,24 @@ def preprocess_packed_seqs( start_idx = cu_seqlens_padded_cpu[i] // cp_size # split to 2 chunks d = input_ids[i, attention_mask[i]] - input_ids_rmpad[start_idx : start_idx + half_seqlen] = d[ - half_seqlen * cp_rank : half_seqlen * (cp_rank + 1) - ] + # Chunk indices are computed in the *padded* coordinate space + # (seqlen_padded_i), while d holds only the *valid* tokens. When a row's + # valid length is shorter than the alignment padding, the first-chunk + # slice can start past the end of d and yield an empty tensor, which + # raises on assignment: + # RuntimeError: The expanded size of the tensor (2) must match the + # existing size (0) at non-singleton dimension 0 + # Clamp it the same way the remaining chunk below already is. + first_start = half_seqlen * cp_rank + first_end = min(half_seqlen * (cp_rank + 1), d.shape[0]) + first_len = max(first_end - first_start, 0) + if first_len > 0: + input_ids_rmpad[start_idx : start_idx + first_len] = d[first_start:first_end] remain_start = seqlen_padded_i - half_seqlen * (cp_rank + 1) remain_end = seqlen_padded_i - half_seqlen * cp_rank remain_end = min(remain_end, d.shape[0]) - remain_len = remain_end - remain_start + remain_len = max(remain_end - remain_start, 0) if remain_len > 0: input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[ remain_start:remain_end diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py index 4078b147cb..9f7c579322 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py @@ -767,14 +767,24 @@ def preprocess_packed_seqs( start_idx = cu_seqlens_padded_cpu[i] // cp_size # split to 2 chunks d = input_ids[i, attention_mask[i]] - input_ids_rmpad[start_idx : start_idx + half_seqlen] = d[ - half_seqlen * cp_rank : half_seqlen * (cp_rank + 1) - ] + # Chunk indices are computed in the *padded* coordinate space + # (seqlen_padded_i), while d holds only the *valid* tokens. When a row's + # valid length is shorter than the alignment padding, the first-chunk + # slice can start past the end of d and yield an empty tensor, which + # raises on assignment: + # RuntimeError: The expanded size of the tensor (2) must match the + # existing size (0) at non-singleton dimension 0 + # Clamp it the same way the remaining chunk below already is. + first_start = half_seqlen * cp_rank + first_end = min(half_seqlen * (cp_rank + 1), d.shape[0]) + first_len = max(first_end - first_start, 0) + if first_len > 0: + input_ids_rmpad[start_idx : start_idx + first_len] = d[first_start:first_end] remain_start = seqlen_padded_i - half_seqlen * (cp_rank + 1) remain_end = seqlen_padded_i - half_seqlen * cp_rank remain_end = min(remain_end, d.shape[0]) - remain_len = remain_end - remain_start + remain_len = max(remain_end - remain_start, 0) if remain_len > 0: input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[ remain_start:remain_end diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py index 933e9e3e12..e0afcb7521 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py @@ -507,6 +507,41 @@ def test_preprocess_packed_seqs(self): self.destroy_parallel_state() + @pytest.mark.skipif( + not torch.cuda.is_available() or int(os.environ.get("WORLD_SIZE", "1")) < 2, + reason="Requires at least 2 GPUs", + ) + def test_preprocess_packed_seqs_row_shorter_than_alignment(self): + """A row whose valid length is below the alignment padding must not raise. + + Chunk indices come from the alignment-padded length while ``d`` holds only the + valid tokens, so on the higher CP ranks the first-chunk slice can start past the + end of ``d``. Assigning that empty slice used to raise:: + + RuntimeError: The expanded size of the tensor (1) must match the existing + size (0) at non-singleton dimension 0 + + With tp=1/cp=2 the alignment is 4, so a single-token row is the smallest case + that reaches the out-of-range slice on cp_rank 1. + """ + tp_size = 1 + cp_size = 2 + self._setup_parallel_state(tp_size=tp_size, cp_size=cp_size) + + batch_size, seq_len = 2, 8 + input_ids = torch.randint(0, 1000, (batch_size, seq_len)) + attention_mask = torch.zeros(batch_size, seq_len, dtype=torch.bool) + attention_mask[0, :] = True # a full-length row + attention_mask[1, 0] = True # a single valid token + + ids_out, packed = preprocess_packed_seqs(input_ids, attention_mask, pre_process=True) + assert packed.max_seqlen_q == packed.max_seqlen_kv + assert ids_out.shape[0] == 1 + # 8 valid tokens need no padding; the 1-token row is padded to align_size = 4. + assert ids_out.shape[1] == (8 + 4) // cp_size + + self.destroy_parallel_state() + @pytest.mark.skipif( not torch.cuda.is_available() or int(os.environ.get("WORLD_SIZE", "1")) < 2, reason="Requires at least 2 GPUs",