fix(data): support GLM-4.5V assistant masking

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/models/glm_vl/data/collate_fn.py

@@ -14,23 +14,76 @@
 
 """GLM VL collator implementations."""
 
+from typing import Any
+
 import torch
 
 from megatron.bridge.data.collators.sequence import prepare_sequence_batch
 from megatron.bridge.data.collators.sequence_padding import use_processor_right_padding
 from megatron.bridge.data.collators.visual import THW_GRID_VISUAL_KEYS
 from megatron.bridge.data.conversation_processing import (
+    AssistantMaskBoundaryConfig,
+    assistant_mask_boundary_config_from_markers,
     build_assistant_loss_mask,
     chat_template_kwargs_from_example,
-    infer_assistant_mask_boundary_config,
+    get_processor_tokenizer,
     shared_chat_template_kwargs_from_examples,
+    tokenize_text_without_special_tokens,
 )
 from megatron.bridge.data.datasets.utils import IGNORE_INDEX
 from megatron.bridge.data.packing.in_batch import build_mcore_thd_sequence_batch_from_rows
 from megatron.bridge.data.token_utils import extract_skipped_token_ids
 from megatron.bridge.training.utils.visual_inputs import GenericVisualInputs
 
 
+GLM4V_ASSISTANT_START = "<|assistant|>\n"
+GLM4V_ASSISTANT_END = "<|endoftext|>"
+GLM4V_EMPTY_THINK = "<think></think>\n"
+GLM4V_NEXT_ROLE_MARKERS = ("<|system|>\n", "<|user|>\n", "<|observation|>\n")
+
+
+def _glm4v_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundaryConfig:
+    """Build GLM-4.5V role boundaries and exclude its empty thinking prefix."""
+    tokenizer = get_processor_tokenizer(processor)
+    empty_think_tokens = tokenize_text_without_special_tokens(tokenizer, GLM4V_EMPTY_THINK)
+    trim_leading_token_sequences = (empty_think_tokens,) if empty_think_tokens else ()
+
+    return assistant_mask_boundary_config_from_markers(
+        processor,
+        assistant_start=GLM4V_ASSISTANT_START,
+        assistant_end=GLM4V_ASSISTANT_END,
+        assistant_end_fallbacks=GLM4V_NEXT_ROLE_MARKERS,
+        include_end_tokens_for_roles=(),
+        trim_leading_token_sequences=trim_leading_token_sequences,
+    )
+
+
+def _build_glm4v_assistant_loss_mask(
+    example: dict,
+    input_ids: torch.Tensor,
+    processor: Any,
+    skipped_tokens: torch.Tensor,
+    boundary_config: AssistantMaskBoundaryConfig,
+) -> torch.Tensor:
+    """Build GLM-4.5V loss spans using a virtual final turn terminator."""
+    tokenizer = get_processor_tokenizer(processor)
+    eos_token_id = getattr(tokenizer, "eos_token_id", None)
+    if eos_token_id is None:
+        raise ValueError("GLM-4.5V assistant masking requires tokenizer.eos_token_id.")
+
+    # GLM role markers terminate intermediate turns, but its final assistant
+    # turn ends directly at the sequence boundary. Add EOS only to delimit that
+    # final span for masking; it is neither returned nor trained as a target.
+    terminated_input_ids = torch.cat([input_ids, input_ids.new_tensor([int(eos_token_id)])])
+    return build_assistant_loss_mask(
+        example,
+        terminated_input_ids,
+        processor,
+        skipped_tokens,
+        boundary_config=boundary_config,
+    )[:-1]
+
+
 def glm4v_collate_fn(
     examples: list,
     processor,
@@ -55,7 +108,7 @@ def glm4v_collate_fn(
     del visual_keys, min_pixels, max_pixels
 
     skipped_tokens = extract_skipped_token_ids(processor)
-    boundary_config = infer_assistant_mask_boundary_config(processor)
+    boundary_config = _glm4v_assistant_mask_boundary_config(processor)
 
     if enable_in_batch_packing:
         sequence_rows = []
@@ -80,12 +133,12 @@ def glm4v_collate_fn(
                     if position_ids is not None
                     else torch.arange(input_ids.numel(), device=input_ids.device, dtype=torch.long)
                 )
-                loss_mask = build_assistant_loss_mask(
+                loss_mask = _build_glm4v_assistant_loss_mask(
                     example,
                     input_ids,
                     processor,
                     skipped_tokens,
-                    boundary_config=boundary_config,
+                    boundary_config,
                 ).to(device=input_ids.device, dtype=torch.float32)
                 labels = torch.cat([input_ids[1:], input_ids.new_full((1,), IGNORE_INDEX)])
                 if skipped_tokens.numel() > 0:
@@ -144,12 +197,12 @@ def glm4v_collate_fn(
 
     loss_mask = torch.stack(
         [
-            build_assistant_loss_mask(
+            _build_glm4v_assistant_loss_mask(
                 example,
                 input_ids,
                 processor,
                 skipped_tokens,
-                boundary_config=boundary_config,
+                boundary_config,
             )
             for example, input_ids in zip(examples, batch["input_ids"])
         ]

tests/unit_tests/data/collators/test_model_collators.py

@@ -880,10 +880,10 @@ def apply_chat_template(self, conversations, **kwargs):
     monkeypatch.setattr(
         glm_vl_collate, "extract_skipped_token_ids", lambda processor: torch.empty(0, dtype=torch.long)
     )
-    monkeypatch.setattr(glm_vl_collate, "infer_assistant_mask_boundary_config", lambda processor: None)
+    monkeypatch.setattr(glm_vl_collate, "_glm4v_assistant_mask_boundary_config", lambda processor: None)
     monkeypatch.setattr(
         glm_vl_collate,
-        "build_assistant_loss_mask",
+        "_build_glm4v_assistant_loss_mask",
         lambda example, input_ids, *args, **kwargs: (input_ids != 0).to(dtype=torch.float32),
     )
     examples = [
@@ -906,6 +906,59 @@ def apply_chat_template(self, conversations, **kwargs):
     assert processor.tokenizer.padding_side == "left"
 
 
+@pytest.mark.parametrize(
+    ("input_ids", "expected_mask"),
+    [
+        (
+            [100, 1, 102, 55, 56, 15, 3, 4, 100, 2, 102, 55, 56, 15, 5, 6],
+            [0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1],
+        ),
+        (
+            [100, 1, 102, 55, 56, 15, 3, 4, 99, 99],
+            [0, 0, 0, 0, 0, 0, 1, 1, 0, 0],
+        ),
+        (
+            [100, 1, 102, 55, 56, 15, 70, 71, 107, 8, 102, 55, 56, 15, 3, 4],
+            [0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1],
+        ),
+    ],
+)
+def test_glm4v_assistant_mask_uses_role_boundaries_and_virtual_final_terminator(input_ids, expected_mask):
+    class _GlmProcessor:
+        class _Tokenizer:
+            chat_template = "<|user|>...<|assistant|>...<|observation|>"
+            eos_token_id = 99
+
+            def encode(self, text, add_special_tokens=False):
+                return self(text, add_special_tokens=add_special_tokens)["input_ids"]
+
+            def __call__(self, text, add_special_tokens=False):
+                mapping = {
+                    "<|assistant|>\n": [102],
+                    "<|endoftext|>": [99],
+                    "<|system|>\n": [105],
+                    "<|user|>\n": [100],
+                    "<|observation|>\n": [107],
+                    "<think></think>\n": [55, 56, 15],
+                }
+                return {"input_ids": mapping[text]}
+
+        tokenizer = _Tokenizer()
+
+    processor = _GlmProcessor()
+    boundary_config = glm_vl_collate._glm4v_assistant_mask_boundary_config(processor)
+
+    mask = glm_vl_collate._build_glm4v_assistant_loss_mask(
+        {"conversation": []},
+        torch.tensor(input_ids),
+        processor,
+        torch.empty(0, dtype=torch.long),
+        boundary_config,
+    )
+
+    assert mask.tolist() == expected_mask
+
+
 def test_expand_image_tokens_handles_multiple_images_and_temporal_grids():
     image_token_id = 163605
     input_ids = torch.tensor([11, image_token_id, 22, image_token_id, 33])