[Fix] Align GLM53 multimodal inputs with the native HF protocol

ci/scripts/CI_ENV.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 export GLM5_2_TINY_MOE_PATH=${CI_SHARE_MODEL}/GLM-5.2-30B
 export GLM5_2_MODEL_PATH=${CI_SHARE_MODEL}/GLM-5.2
+export GLM_5_3_FLASH_PATH=${CI_SHARE_MODEL}/GLM-5.3-Flash-25B
 export QWEN3_VL_MOE_PATH=${CI_SHARE_MODEL}/Qwen3-VL-30B-A3B-Instruct_MOE
 export QWEN3_VL_DENSE_PATH=${CI_SHARE_MODEL}/Qwen3-VL-4B-Instruct
 export QWEN3_VL_2B_PATH=${CI_SHARE_MODEL}/models--Qwen--Qwen3-VL-2B-Instruct/snapshots/89644892e4d85e24eaac8bacfd4f463576704203

examples/v1/config/sft_glm5p2.py

@@ -70,10 +70,10 @@ def _get_float8_config() -> Float8Config | None:
     if "INDEXER_BACKEND" in os.environ:
         # ``deep_gemm_fp8`` uses DeepGEMM's FP8 MQA path to match LMDeploy's
         # Indexer scoring contract.
-        indexer_backend = os.environ["INDEXER_BACKEND"].strip().lower()
-        model_cfg.attention.indexer_backend = indexer_backend
-    else:
-        indexer_backend = sparse_mla_backend
+        model_cfg.attention.indexer_backend = os.environ["INDEXER_BACKEND"].strip().lower()
+    # Otherwise the indexer keeps the config default; it no longer follows
+    # SPARSE_MLA_BACKEND, whose vocabulary is not the indexer's.
+    indexer_backend = model_cfg.attention.indexer_backend
     # Keep the historical one-shot path for the PyTorch selector.  TileLang
     # (including the cuDNN DSA and FlashMLA adapters) uses Slime's 8K query
     # block by default; set INDEXER_TOPK_QUERY_CHUNK_SIZE=0/none to disable it

tests/datasets/test_glm53_chat.py

@@ -0,0 +1,95 @@
+"""GLM-5.3-Flash 的 chat template 与 loss mask，见设计文档 F1。
+
+TestGlm53Rendering
+    test_reasoning_effort_domain_matches_jinja  reasoning_effort 取值域与 jinja 一致
+    test_all_generation_eos_are_supervised_at_assistant_boundaries  停止 token 计入 loss
+    test_media_placeholder_matches_jinja        媒体占位符与 jinja 一致
+"""
+
+import os
+
+import pytest
+
+from transformers import AutoTokenizer
+from xtuner.v1.data_proto.messages import Glm53ChatMessages
+from xtuner.v1.data_proto.messages.glm53_chat import glm53_tokenize_fn_slowspeed
+
+
+GLM_5_3_FLASH_PATH = os.environ["GLM_5_3_FLASH_PATH"]
+GLM53_TEMPLATE_DEFAULTS = {"enable_thinking": True, "reasoning_effort": "max", "clear_thinking": False}
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    if not os.path.isdir(GLM_5_3_FLASH_PATH):
+        pytest.fail(f"GLM_5_3_FLASH_PATH not found: {GLM_5_3_FLASH_PATH}")
+    return AutoTokenizer.from_pretrained(GLM_5_3_FLASH_PATH, trust_remote_code=True)
+
+
+def _render_from_hf(tokenizer, messages, **kwargs):
+    return tokenizer.apply_chat_template(messages, tokenize=False, **{**GLM53_TEMPLATE_DEFAULTS, **kwargs})
+
+
+class TestGlm53Rendering:
+    @pytest.mark.parametrize("reasoning_effort", ["low", "high", "max", "bogus"])
+    @pytest.mark.parametrize("enable_thinking", [False, True])
+    def test_reasoning_effort_domain_matches_jinja(self, tokenizer, reasoning_effort, enable_thinking):
+        # jinja: reasoning_effort in ['low', 'high'] else 'max' -- "bogus" exercises the fallback.
+        messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]
+        tokenized = Glm53ChatMessages(messages=messages).tokenize(
+            tokenizer, reasoning_effort=reasoning_effort, enable_thinking=enable_thinking
+        )
+        hf_rendered = _render_from_hf(
+            tokenizer, messages, reasoning_effort=reasoning_effort, enable_thinking=enable_thinking
+        )
+        rendered = tokenizer.decode(tokenized["input_ids"], skip_special_tokens=False)
+        assert rendered == hf_rendered + tokenizer.eos_token
+        slow_ids, slow_labels = glm53_tokenize_fn_slowspeed(
+            tokenizer, messages, reasoning_effort=reasoning_effort, enable_thinking=enable_thinking
+        )
+        assert tokenized["input_ids"] == slow_ids
+        assert tokenized["labels"] == slow_labels
+
+    def test_all_generation_eos_are_supervised_at_assistant_boundaries(self, tokenizer):
+        # 助手轮边界上的停止 token 必须计入 loss，否则模型学不会停。
+        messages = [
+            {"role": "user", "content": "First question"},
+            {"role": "assistant", "content": "First answer"},
+            {"role": "user", "content": "Final question"},
+            {"role": "assistant", "content": "Final answer"},
+        ]
+        tokenized = Glm53ChatMessages(messages=messages).tokenize(tokenizer)
+        slow_ids, slow_labels = glm53_tokenize_fn_slowspeed(tokenizer, messages)
+        rendered = tokenizer.decode(tokenized["input_ids"], skip_special_tokens=False)
+        hf_rendered = _render_from_hf(tokenizer, messages)
+
+        assert rendered == hf_rendered + tokenizer.eos_token
+        assert tokenized["input_ids"] == slow_ids
+        assert tokenized["labels"] == slow_labels
+        user_id = tokenizer.convert_tokens_to_ids("<|user|>")
+        first_user_label = tokenized["labels"][tokenized["input_ids"].index(user_id)]
+        assert first_user_label == -100  # first <|user|> has no preceding assistant turn
+        assert tokenized["labels"][-1] == tokenizer.convert_tokens_to_ids("<|endoftext|>")
+
+    def test_media_placeholder_matches_jinja(self, tokenizer):
+        # 媒体占位符必须与官方 jinja 的两段式标记一致。
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                    {"type": "video"},
+                    {"type": "audio"},
+                    {"type": "text", "text": "describe these"},
+                ],
+            },
+            {"role": "assistant", "content": "ok"},
+        ]
+        tokenized = Glm53ChatMessages(messages=messages).tokenize(tokenizer)
+        rendered = tokenizer.decode(tokenized["input_ids"], skip_special_tokens=False)
+        hf_rendered = _render_from_hf(tokenizer, messages)
+
+        assert rendered == hf_rendered + tokenizer.eos_token
+        assert "<|begin_of_image|><|image|><|end_of_image|>" in rendered
+        assert "<|begin_of_video|><|video|><|end_of_video|>" in rendered
+        assert "<|begin_of_audio|><|end_of_audio|>" in rendered

tests/datasets/test_glm53_vl_collator.py

@@ -0,0 +1,216 @@
+"""GLM-5.3-Flash 的 VL collator，见 doc/xtuner_glm5p3flash_design.md F1 §8.4。
+
+TestSharedCollatorsStillWork
+    test_qwen3_vl_collator_still_carries_media   共用 build_text_ctx_labels 的 qwen3_vl 回归
+    test_intern_s1_collator_still_carries_media  同上，intern_s1
+TestGlm53VlSftCollator
+    test_glm53_collator_keeps_media_with_instance    截断丢样本时媒体同步丢弃
+    test_glm53_collator_image_video_batch            图像/视频混批各自聚合
+    test_glm53_collator_pure_text_sample_has_no_media 纯文本样本不带媒体字段
+TestGlm53SequenceContextSplitKeepsMmAlignment
+    test_glm53_sequence_context_split_keeps_mm_alignment  SP 切分保持模态标记对齐
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+
+from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.datasets.collator import (
+    glm53_vl_sft_collator,
+    intern_s1_vl_sft_collator,
+    qwen3_vl_sft_collator,
+)
+from xtuner.v1.datasets.config import DataloaderConfig
+from xtuner.v1.utils import IGNORE_INDEX
+
+
+def _text_item(text_len: int):
+    input_ids = list(range(1, text_len + 1))
+    labels = list(input_ids)
+    return {"input_ids": input_ids, "labels": labels, "num_tokens": text_len}
+
+
+class TestSharedCollatorsStillWork:
+    def test_qwen3_vl_collator_still_carries_media(self):
+        # qwen3_vl 与 glm53 共用 build_text_ctx_labels，改动后它仍要正常携带媒体。
+        item = _text_item(4)
+        item["pixel_values"] = torch.randn(4, 8)
+        item["image_grid_thw"] = torch.tensor([[1, 2, 2]])
+        item["num_img_tokens"] = [4]
+        item["num_imgs"] = [1]
+        result = qwen3_vl_sft_collator([[item]], pack_max_length=16, padding_token_idx=0)
+        seq_ctx = result[0]["seq_ctx"]
+        assert seq_ctx.pixel_values.shape == (4, 8)
+        assert seq_ctx.num_img_tokens == [[4]]
+
+    def test_intern_s1_collator_still_carries_media(self):
+        # intern_s1 同上。
+        item = _text_item(4)
+        item["pixel_values"] = torch.randn(4, 8)
+        item["num_img_tokens"] = [4]
+        item["num_imgs"] = [1]
+        result = intern_s1_vl_sft_collator([[item]], pack_max_length=16, padding_token_idx=0)
+        seq_ctx = result[0]["seq_ctx"]
+        assert seq_ctx.pixel_values.shape == (4, 8)
+
+
+def _glm53_item(text_len: int, image_token_id: int, num_placeholders: int, merge_unit: int = 4):
+    input_ids = [100] * text_len
+    for i in range(num_placeholders):
+        input_ids[i] = image_token_id
+    labels = list(input_ids)
+    labels[0] = IGNORE_INDEX
+    mm_token_type_ids = torch.zeros(1, text_len, dtype=torch.long)
+    mm_token_type_ids[0, :num_placeholders] = 1
+    item = {
+        "input_ids": input_ids,
+        "labels": labels,
+        "num_tokens": text_len,
+        "num_imgs": [1] if num_placeholders else [0],
+        "num_img_tokens": [num_placeholders * merge_unit] if num_placeholders else [0],
+        "mm_token_type_ids": mm_token_type_ids,
+    }
+    if num_placeholders:
+        item["pixel_values"] = torch.randn(num_placeholders * merge_unit, 8)
+        item["image_grid_thw"] = torch.tensor([[1, 2, 2 * num_placeholders]])
+    return item
+
+
+class TestGlm53VlSftCollator:
+    IMAGE_TOKEN_ID = 999
+    MERGE_UNIT = 4
+
+    def test_oversized_media_sample_cannot_bypass_tokenize_validation(self) -> None:
+        item = _glm53_item(text_len=10, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=4)
+        with pytest.raises(ValueError, match="truncate at tokenize time"):
+            glm53_vl_sft_collator(
+                [[item]],
+                pack_max_length=4,
+                padding_token_idx=0,
+                image_token_id=self.IMAGE_TOKEN_ID,
+                merge_unit=self.MERGE_UNIT,
+            )
+
+    def test_oversized_text_sample_keeps_text_truncation(self) -> None:
+        item = _glm53_item(text_len=10, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=0)
+        expected_ids = item["input_ids"][:3]
+        expected_labels = item["labels"][1:4]
+        result = glm53_vl_sft_collator(
+            [[item]],
+            pack_max_length=4,
+            padding_token_idx=0,
+            image_token_id=self.IMAGE_TOKEN_ID,
+            merge_unit=self.MERGE_UNIT,
+        )[0]
+        assert result["seq_ctx"].input_ids[0, :3].tolist() == expected_ids
+        assert result["shifted_labels"][0, :3].tolist() == expected_labels
+        assert result["shifted_labels"][0, 3].item() == -100
+        assert result["seq_ctx"].pixel_values is None
+
+    def test_glm53_collator_keeps_media_with_instance(self):
+        # 4 placeholders * merge_unit(4) = 16 raw patches.
+        item = _glm53_item(text_len=10, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=4)
+        result = glm53_vl_sft_collator(
+            [[item]],
+            pack_max_length=32,
+            padding_token_idx=0,
+            image_token_id=self.IMAGE_TOKEN_ID,
+            merge_unit=self.MERGE_UNIT,
+        )
+        seq_ctx = result[0]["seq_ctx"]
+        assert seq_ctx.pixel_values.shape == (16, 8)
+        assert seq_ctx.num_img_tokens == [[16]]
+        assert seq_ctx.mm_token_type_ids.shape[-1] == seq_ctx.input_ids.shape[-1]
+        # instance text length 10 -> [:-1] drop -> 9 real tokens, padded to 32.
+        assert (seq_ctx.mm_token_type_ids[0, :4] == 1).all()
+        assert (seq_ctx.mm_token_type_ids[0, 9:] == 0).all()
+
+    def test_glm53_collator_image_video_batch(self):
+        # Two image-only samples packed together; each keeps its own media and placeholder count.
+        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)
+        result = glm53_vl_sft_collator(
+            [[item_a, item_b]],
+            pack_max_length=32,
+            padding_token_idx=0,
+            image_token_id=self.IMAGE_TOKEN_ID,
+            merge_unit=self.MERGE_UNIT,
+        )
+        seq_ctx = result[0]["seq_ctx"]
+        # 2*4 + 1*4 = 12 raw patches total, preserving sample order.
+        assert seq_ctx.pixel_values.shape == (12, 8)
+        assert seq_ctx.num_img_tokens == [[8], [4]]
+        assert seq_ctx.image_grid_thw.shape[0] == 2
+
+    def test_glm53_collator_pure_text_sample_has_no_media(self):
+        # 纯文本样本不应带出任何媒体字段。
+        item = _glm53_item(text_len=6, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=0)
+        result = glm53_vl_sft_collator(
+            [[item]],
+            pack_max_length=16,
+            padding_token_idx=0,
+            image_token_id=self.IMAGE_TOKEN_ID,
+            merge_unit=self.MERGE_UNIT,
+        )
+        seq_ctx = result[0]["seq_ctx"]
+        assert seq_ctx.pixel_values is None
+        assert seq_ctx.image_grid_thw is None
+        assert (seq_ctx.mm_token_type_ids == 0).all()
+
+    def test_glm53_collator_is_reachable_by_name_from_dataloader_config(self):
+        # 训练侧只能按名字选 collator；`glm53_vl_sft_collator` 比其它 VL collator 多两个必填
+        # 参数，要靠 `collator_kwargs` 绑定，否则 dataloader 一构建就 TypeError。
+        cfg = DataloaderConfig(
+            collator="glm53_vl_sft_collator",
+            collator_kwargs={"image_token_id": self.IMAGE_TOKEN_ID, "merge_unit": self.MERGE_UNIT},
+        )
+        collator = cfg.build_collator()
+        item = _glm53_item(text_len=10, image_token_id=self.IMAGE_TOKEN_ID, num_placeholders=4)
+        # 只补 dataloader 自己会传的那几个参数，额外参数必须已经绑好。
+        result = collator([[item]], pack_max_length=32, padding_token_idx=0)
+        assert result[0]["seq_ctx"].pixel_values.shape == (16, 8)
+
+
+@pytest.fixture
+def single_rank_sp_mesh():
+    """A 1-rank CPU SP mesh, torn down afterwards so the default process group does not leak
+    into the rest of the session (a later test creating its own would hit "trying to initialize
+    the default process group twice")."""
+    owned = not dist.is_initialized()
+    if owned:
+        os.environ.setdefault("RANK", "0")
+        os.environ.setdefault("WORLD_SIZE", "1")
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29556")
+        dist.init_process_group(backend="gloo", rank=0, world_size=1)
+    try:
+        yield init_device_mesh("cpu", (1,), mesh_dim_names=("sp",))
+    finally:
+        if owned:
+            dist.destroy_process_group()
+
+
+class TestGlm53SequenceContextSplitKeepsMmAlignment:
+    def test_glm53_sequence_context_split_keeps_mm_alignment(self, single_rank_sp_mesh):
+        # SP 切分后 mm_token_type_ids 必须与 input_ids 走同一套 pad/split。
+        sp_mesh = single_rank_sp_mesh
+
+        seq_len = 10
+        input_ids = torch.arange(seq_len).view(1, -1)
+        ctx = SequenceContext.from_input_ids((input_ids,), device="cpu")
+        mm_token_type_ids = torch.zeros(1, seq_len, dtype=torch.long)
+        mm_token_type_ids[0, 2:5] = 1
+        ctx.mm_token_type_ids = mm_token_type_ids
+
+        split_ctx = ctx.split(sequence_parallel_mesh=sp_mesh)
+
+        # world_size=1: the split is a structural no-op, but it must still run mm_token_type_ids
+        # through the exact same pad_to_multiple_of + split_for_sequence_parallel path as
+        # input_ids (design doc §8.3), not silently pass it through untouched.
+        assert split_ctx.mm_token_type_ids is not None
+        assert split_ctx.mm_token_type_ids.shape[-1] == split_ctx.input_ids.shape[-1]
+        assert torch.equal(split_ctx.mm_token_type_ids[0, :seq_len], mm_token_type_ids[0])