[model] fix: clamp the first CP chunk slice in preprocess_packed_seqs (#5614)

Signed-off-by: gaohongkui <gaohongkui1021@163.com>
Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com>

Author: gaohongkui

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

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

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",