diff --git a/tests/datasets/test_glm53_vl_collator.py b/tests/datasets/test_glm53_vl_collator.py index 6855408616..e309732aea 100644 --- a/tests/datasets/test_glm53_vl_collator.py +++ b/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 diff --git a/tests/datasets/test_glm53_vl_tokenize_fn.py b/tests/datasets/test_glm53_vl_tokenize_fn.py index 563b07face..f7ab4b7382 100644 --- a/tests/datasets/test_glm53_vl_tokenize_fn.py +++ b/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: diff --git a/tests/model/test_glm53_compose.py b/tests/model/test_glm53_compose.py index 572e00d7c7..f56d2ccaba 100644 --- a/tests/model/test_glm53_compose.py +++ b/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() diff --git a/xtuner/v1/datasets/collator.py b/xtuner/v1/datasets/collator.py index 06dcb332a5..f1b4a576e1 100644 --- a/xtuner/v1/datasets/collator.py +++ b/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 diff --git a/xtuner/v1/datasets/mllm_tokenize_fn/glm53_vl_tokenize_fn.py b/xtuner/v1/datasets/mllm_tokenize_fn/glm53_vl_tokenize_fn.py index faaec537ce..90775caa2b 100644 --- a/xtuner/v1/datasets/mllm_tokenize_fn/glm53_vl_tokenize_fn.py +++ b/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) diff --git a/xtuner/v1/model/compose/glm53/modeling_glm53.py b/xtuner/v1/model/compose/glm53/modeling_glm53.py index 731d5802c7..a66248b764 100644 --- a/xtuner/v1/model/compose/glm53/modeling_glm53.py +++ b/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(