[Refactor] Reuse single-media pack tensors without torch.cat copies

tests/datasets/test_glm53_vl_collator.py

@@ -15,6 +15,12 @@
     test_multi_item_pack_cast_preserves_order_and_values  拼包逐段数值一致
 TestGlm53PlaceholderPackWalksLossPath
     test_glm53_placeholder_pack_walks_ce_loss_path  占位 pack 走完 CE loss 路径：0 loss、可 backward
+TestGlm53VlCollatorSingleMediaReuse
+    test_single_image_pack_reuses_item_tensor        单图 pack 复用 item 张量（零拷贝）
+    test_single_video_pack_reuses_item_tensor        单视频 pack 复用 item 张量（零拷贝）
+    test_multi_item_pack_still_concatenates          多 item pack 仍走 cat 且顺序保持
+    test_truncated_pack_reuses_survivor_media        截断丢尾后幸存者媒体被复用
+    test_mixed_image_and_video_items_reuse_independently  图+视频 pack 各模态独立复用
 TestGlm53SequenceContextSplitKeepsMmAlignment
     test_glm53_sequence_context_split_keeps_mm_alignment  SP 切分保持模态标记对齐
 """
@@ -361,6 +367,68 @@ def test_glm53_placeholder_pack_walks_ce_loss_path(self):
         assert torch.isfinite(hidden.grad).all()
 
 
+class TestGlm53VlCollatorSingleMediaReuse:
+    IMAGE_TOKEN_ID = 999
+    MERGE_UNIT = 4
+
+    def _collate(self, pack, pack_max_length=64):
+        # media_dtype=fp32 隔离 #2207 的运输 cast（默认 bf16 会复制张量），
+        # 让本类的 `is` 断言只检验单媒体复用本身。
+        return glm53_vl_sft_collator(
+            [pack],
+            pack_max_length=pack_max_length,
+            padding_token_idx=0,
+            image_token_id=self.IMAGE_TOKEN_ID,
+            merge_unit=self.MERGE_UNIT,
+            media_dtype=torch.float32,
+        )[0]["seq_ctx"]
+
+    def test_single_image_pack_reuses_item_tensor(self):
+        # 单媒体 pack（最常见情形）必须零拷贝：直接复用 item 里已有的张量（引用相等）。
+        item = _glm53_item(text_len=10, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=4)
+        seq_ctx = self._collate([item])
+        assert seq_ctx.pixel_values is item["pixel_values"]
+        assert seq_ctx.image_grid_thw is item["image_grid_thw"]
+        assert seq_ctx.pixel_values.is_contiguous()
+
+    def test_single_video_pack_reuses_item_tensor(self):
+        item = _glm53_video_item(text_len=10, image_token_id=self.IMAGE_TOKEN_ID, num_frames=4)
+        seq_ctx = self._collate([item])
+        assert seq_ctx.pixel_values_videos is item["pixel_values_videos"]
+        assert seq_ctx.video_grid_thw is item["video_grid_thw"]
+        assert seq_ctx.pixel_values_videos.is_contiguous()
+
+    def test_multi_item_pack_still_concatenates(self):
+        # 多 item pack 语义不变：真 cat、结果与参考一致、新缓冲、顺序保持。
+        item_a = _glm53_item(text_len=8, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=2)
+        item_b = _glm53_item(text_len=6, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=1)
+        seq_ctx = self._collate([item_a, item_b])
+        expected = torch.cat([item_a["pixel_values"], item_b["pixel_values"]], dim=0)
+        assert torch.equal(seq_ctx.pixel_values, expected)
+        assert seq_ctx.pixel_values.data_ptr() != item_a["pixel_values"].data_ptr()
+        assert seq_ctx.pixel_values.data_ptr() != item_b["pixel_values"].data_ptr()
+        assert torch.equal(seq_ctx.pixel_values[: item_a["pixel_values"].shape[0]], item_a["pixel_values"])
+        assert seq_ctx.image_grid_thw.shape[0] == 2
+
+    def test_truncated_pack_reuses_survivor_media(self):
+        # 超长 pack 丢尾后只剩单个幸存 item：被丢弃 item 的媒体不进列表，幸存者被复用。
+        survivor = _glm53_item(text_len=10, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=4)
+        dropped = _glm53_item(text_len=30, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=2)
+        seq_ctx = self._collate([survivor, dropped], pack_max_length=16)
+        assert seq_ctx.pixel_values is survivor["pixel_values"]
+        assert seq_ctx.image_grid_thw is survivor["image_grid_thw"]
+
+    def test_mixed_image_and_video_items_reuse_independently(self):
+        # 图+视频 pack：两模态各自单元素，各自独立复用。
+        image_item = _glm53_item(text_len=8, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=2)
+        video_item = _glm53_video_item(text_len=8, image_token_id=self.IMAGE_TOKEN_ID, num_frames=1)
+        seq_ctx = self._collate([image_item, video_item])
+        assert seq_ctx.pixel_values is image_item["pixel_values"]
+        assert seq_ctx.image_grid_thw is image_item["image_grid_thw"]
+        assert seq_ctx.pixel_values_videos is video_item["pixel_values_videos"]
+        assert seq_ctx.video_grid_thw is video_item["video_grid_thw"]
+
+
 @pytest.fixture
 def single_rank_sp_mesh():
     """A 1-rank CPU SP mesh, torn down afterwards so the default process group does not leak

tests/datasets/test_glm53_vl_tokenize_fn.py

@@ -20,6 +20,9 @@
 TestGlm53MixedMediaAndTruncation
     test_glm53_mixed_media_is_discarded             图文混合视频未支持：cache 警告丢弃、runtime 占位，不中断训练
     test_glm53_visual_truncation_is_dropped         截断会切断视觉跨度时 cache 阶段就丢弃
+TestGlm53SingleMediaTokenizeReuse
+    test_single_image_reuses_processor_output       单图样本复用 processor 输出（零拷贝）
+    test_single_video_reuses_processor_output       单视频样本同上
 TestGlm53MediaRoleContract
     test_media_is_collected_from_all_roles           system/assistant/tool 角色的图片也会被收集
     test_media_collection_preserves_input_order     媒体按原始消息顺序收集，与 HF 一致
@@ -477,6 +480,59 @@ def test_glm53_visual_truncation_is_dropped(self, ckpt_path, tokenizer, image_fi
             short_fn(_image_item(image_fixture))
 
 
+class TestGlm53SingleMediaTokenizeReuse:
+    def _recording_processor(self, original, recorded):
+        # Spy that records each call's output and forwards everything else -- attributes
+        # included, since the tokenize fn also reads geometry attrs off the processor.
+        class Spy:
+            def __call__(self, **kwargs):
+                out = original(**kwargs)
+                recorded.append(out)
+                return out
+
+            def __getattr__(self, name):
+                return getattr(original, name)
+
+        return Spy()
+
+    def test_single_image_reuses_processor_output(self, tokenize_fn, image_fixture, monkeypatch):
+        # 单图样本（最常见情形）的 pixel_values/image_grid_thw 应是 processor 原输出对象本身，
+        # 而不是 cat 出来的全量拷贝（引用相等 = 零拷贝），且值与真实 processor 输出一致。
+        original = tokenize_fn.image_processor
+        recorded = []
+        monkeypatch.setattr(tokenize_fn, "image_processor", self._recording_processor(original, recorded))
+        tokenize_fn.state = "runtime"
+        actual = tokenize_fn(_image_item(image_fixture))
+
+        assert len(recorded) == 1
+        assert actual["pixel_values"] is recorded[0]["pixel_values"]
+        assert actual["image_grid_thw"] is recorded[0]["image_grid_thw"]
+
+        monkeypatch.undo()
+        from xtuner.v1.datasets.mllm_tokenize_fn import glm53_vl_tokenize_fn as module
+
+        image = module.apply_exif_orientation(module.load_image(image_fixture))
+        expected = original(images=[image], return_tensors="pt")
+        image.close()
+        assert torch.equal(actual["pixel_values"], expected["pixel_values"])
+        assert torch.equal(actual["image_grid_thw"], expected["image_grid_thw"])
+
+    def test_single_video_reuses_processor_output(self, tokenize_fn, video_fixture, monkeypatch):
+        original = tokenize_fn.video_processor
+        recorded = []
+        monkeypatch.setattr(tokenize_fn, "video_processor", self._recording_processor(original, recorded))
+        tokenize_fn.state = "runtime"
+        actual = tokenize_fn(_video_item(video_fixture))
+
+        assert len(recorded) == 1
+        assert actual["pixel_values_videos"] is recorded[0]["pixel_values_videos"]
+
+        # 去掉 spy 再跑一遍真实路径，值必须逐位一致。
+        monkeypatch.undo()
+        second = tokenize_fn(_video_item(video_fixture))
+        assert torch.equal(actual["pixel_values_videos"], second["pixel_values_videos"])
+
+
 class TestGlm53MediaRoleContract:
     @pytest.mark.parametrize("role", ["system", "assistant", "tool"])
     def test_media_is_collected_from_all_roles(self, tokenize_fn: Glm53VLTokenizeFunction, role: str) -> None:

tests/model/test_glm53_compose.py

@@ -325,3 +325,62 @@ def test_pack_with_an_image_sample_and_a_video_sample(self):
         reference = model.language_model(ref_ctx, None).logits
 
         torch.testing.assert_close(logits, reference, rtol=0, atol=0)
+
+
+@pytest.mark.gpu
+class TestGlm53ComposeSingleModalityNoCopy:
+    """单模态 pack 进视觉塔前不再做全量 cat 拷贝（复用唯一张量）；图+视频 pack 仍统一
+    cat 后一次塔调用（保持各 rank 的 FSDP all-gather 序列一致，见上面的测试）。"""
+
+    def test_image_only_pack_feeds_tower_without_copy(self):
+        model = _build_model()
+        captured = []
+        handle = model.vision_tower.register_forward_pre_hook(lambda module, args: captured.append(args))
+        try:
+            num_placeholders = 4  # merge_unit = MERGE**2 = 4
+            raw_patches = num_placeholders * MERGE * MERGE
+            grid_side = int(raw_patches**0.5)
+            pixel_values = torch.randn(raw_patches, _patch_dim(model), device="cuda", dtype=torch.bfloat16)
+            image_grid_thw = torch.tensor([[1, grid_side, grid_side]], device="cuda")
+            input_ids = torch.randint(2, 200, (1, SEQ_LEN), device="cuda")
+            mm_type = torch.zeros(1, SEQ_LEN, dtype=torch.long, device="cuda")
+            mm_type[0, 10 : 10 + num_placeholders] = 1
+            seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cuda")
+            seq_ctx.mm_token_type_ids = mm_type
+            seq_ctx.pixel_values = pixel_values
+            seq_ctx.image_grid_thw = image_grid_thw
+            model(seq_ctx=seq_ctx, loss_ctx=None)
+        finally:
+            handle.remove()
+
+        # 塔收到的就是 pack 里的张量本体（引用相等 = 零拷贝）。
+        assert len(captured) == 1
+        assert captured[0][0] is pixel_values
+        assert captured[0][1] is image_grid_thw
+
+    def test_mixed_pack_still_concatenates_for_the_tower(self):
+        model = _build_model()
+        captured = []
+        handle = model.vision_tower.register_forward_pre_hook(lambda module, args: captured.append(args))
+        try:
+            input_ids = torch.randint(2, 200, (1, SEQ_LEN), device="cuda")
+            mm_type = torch.zeros(1, SEQ_LEN, dtype=torch.long, device="cuda")
+            mm_type[0, 10] = 1  # 一个图像 token：[1, 2, 2] 的 4 个 patch 合并成 1 个
+            mm_type[0, 60] = 2  # 一个视频 token
+            pixel_values = torch.randn(4, _patch_dim(model), device="cuda", dtype=torch.bfloat16)
+            pixel_values_videos = torch.randn(4, _patch_dim(model), device="cuda", dtype=torch.bfloat16)
+            grid = torch.tensor([[1, 2, 2]], device="cuda")
+            seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cuda")
+            seq_ctx.mm_token_type_ids = mm_type
+            seq_ctx.pixel_values, seq_ctx.image_grid_thw = pixel_values, grid
+            seq_ctx.pixel_values_videos, seq_ctx.video_grid_thw = pixel_values_videos, grid
+            model(seq_ctx=seq_ctx, loss_ctx=None)
+        finally:
+            handle.remove()
+
+        # 图+视频 pack 仍收到 cat 之后的新缓冲（值一致、地址独立于两个源）。
+        assert len(captured) == 1
+        tower_patches = captured[0][0]
+        assert torch.equal(tower_patches, torch.cat([pixel_values, pixel_values_videos], dim=0))
+        assert tower_patches.data_ptr() != pixel_values.data_ptr()
+        assert tower_patches.data_ptr() != pixel_values_videos.data_ptr()

xtuner/v1/datasets/collator.py

@@ -310,6 +310,14 @@ def qwen3_vl_sft_collator(
     return ret
 
 
+def _reuse_single_or_cat(tensors: list[torch.Tensor]) -> torch.Tensor | None:
+    # A single-media pack (the common case) already holds the final tensor; reuse it instead of
+    # paying a full torch.cat copy of every patch. Multi-item packs still cat; empty stays None.
+    if len(tensors) == 1:
+        return tensors[0]
+    return torch.cat(tensors, dim=0) if tensors else None
+
+
 def glm53_vl_sft_collator(
     instances: list[list[Glm53VLDataItem]],
     pack_max_length: int,
@@ -387,23 +395,14 @@ def glm53_vl_sft_collator(
         # before cat and the cross-process IPC: the vision tower's patch_embed casts patches to
         # its parameter dtype anyway, so fp32 media would only double the bytes crossing shm,
         # main-process residency, and H2D. No-op when the dtypes already match.
-        pixel_values: list | torch.Tensor | None
-        pixel_values = [i["pixel_values"].to(media_dtype) for i in instance if "pixel_values" in i]
-        pixel_values = torch.cat(pixel_values, dim=0) if pixel_values else None
-
-        image_grid_thw: list | torch.Tensor | None
-        image_grid_thw = [i["image_grid_thw"] for i in instance if "image_grid_thw" in i]
-        image_grid_thw = torch.cat(image_grid_thw, dim=0) if image_grid_thw else None
-
-        pixel_values_videos: list | torch.Tensor | None
-        pixel_values_videos = [
-            i["pixel_values_videos"].to(media_dtype) for i in instance if "pixel_values_videos" in i
-        ]
-        pixel_values_videos = torch.cat(pixel_values_videos, dim=0) if pixel_values_videos else None
-
-        video_grid_thw: list | torch.Tensor | None
-        video_grid_thw = [i["video_grid_thw"] for i in instance if "video_grid_thw" in i]
-        video_grid_thw = torch.cat(video_grid_thw, dim=0) if video_grid_thw else None
+        pixel_values = _reuse_single_or_cat(
+            [i["pixel_values"].to(media_dtype) for i in instance if "pixel_values" in i]
+        )
+        image_grid_thw = _reuse_single_or_cat([i["image_grid_thw"] for i in instance if "image_grid_thw" in i])
+        pixel_values_videos = _reuse_single_or_cat(
+            [i["pixel_values_videos"].to(media_dtype) for i in instance if "pixel_values_videos" in i]
+        )
+        video_grid_thw = _reuse_single_or_cat([i["video_grid_thw"] for i in instance if "video_grid_thw" in i])
 
         # mm_token_type_ids must track input_ids through the exact same [:, :-1] drop + pad
         # (design doc §8.3). Slice each item to its own (possibly pack-truncated) input_ids

xtuner/v1/datasets/mllm_tokenize_fn/glm53_vl_tokenize_fn.py

@@ -358,8 +358,14 @@ def multi_modal_get_item(self, data_item: dict, media_root: str = "") -> Glm53VL
             image_grid_list.append(processed_i["image_grid_thw"])
             del image, processed_i
 
-        pixel_values = torch.cat(pixel_values_list, dim=0)
-        image_grid_thw = torch.cat(image_grid_list, dim=0)
+        # A single-image sample already holds the processor's final tensor; reuse it
+        # instead of paying a full torch.cat copy of every patch.
+        if len(pixel_values_list) == 1:
+            pixel_values = pixel_values_list[0]
+            image_grid_thw = image_grid_list[0]
+        else:
+            pixel_values = torch.cat(pixel_values_list, dim=0)
+            image_grid_thw = torch.cat(image_grid_list, dim=0)
         num_images = len(pixel_values_list)
         del pixel_values_list, image_grid_list
 
@@ -488,7 +494,12 @@ def video_get_item(self, data_item: dict, media_root: str = "") -> Glm53VLDataIt
             grid_list.append(out["video_grid_thw"][0])
             metadata_list.append(metadata)
 
-        pixel_values_videos = torch.cat(pixel_values_list, dim=0)
+        # A single-video sample already holds the processor's final tensor; reuse it
+        # instead of paying a full torch.cat copy of every patch.
+        if len(pixel_values_list) == 1:
+            pixel_values_videos = pixel_values_list[0]
+        else:
+            pixel_values_videos = torch.cat(pixel_values_list, dim=0)
         video_grid_thw = torch.stack(grid_list)
 
         text, loss_mask = self._render_base(data_item)

xtuner/v1/model/compose/glm53/modeling_glm53.py

@@ -200,7 +200,11 @@ def _prepare_llm_inputs(self, seq_ctx: SequenceContext) -> torch.Tensor:
             pixel_values.append(seq_ctx.pixel_values_videos)  # type: ignore[arg-type]
             grid_thw.append(flat_grid_thw)
             num_features[2] = int((flat_grid_thw.prod(-1) // merge_unit).sum().item())
-        features = self.get_visual_features(torch.cat(pixel_values), torch.cat(grid_thw), sp_mesh)
+        # A single-modality pack (the common case) already holds the final tensor; reuse it
+        # instead of paying a full torch.cat copy of every patch.
+        flat_pixel_values = pixel_values[0] if len(pixel_values) == 1 else torch.cat(pixel_values)
+        flat_grid_thw = grid_thw[0] if len(grid_thw) == 1 else torch.cat(grid_thw)
+        features = self.get_visual_features(flat_pixel_values, flat_grid_thw, sp_mesh)
         return self._splice(inputs_embeds, mm_token_type_ids, features, num_features, sp_mesh)
 
     def forward(