[data] fix: pack VLM batches directly (#4507)

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/data/datasets/sft.py

@@ -1240,9 +1240,8 @@ def collate_fn(self, batch):
                 if x_tensor.sum().item() == 0:
                     logger.warning(
                         "Due to truncation to max_seq_length, no assistant tokens are found in sample. "
-                        "Setting loss_mask to all ones."
+                        "Keeping loss_mask empty to avoid supervising non-assistant tokens."
                     )
-                    loss_mask[i] = [1] * self.max_seq_length
 
             contexts = [x[: self.max_seq_length] for x in contexts]
             answers = [x[: self.max_seq_length] for x in answers]

src/megatron/bridge/data/datasets/utils.py

@@ -879,6 +879,17 @@ def _convert_to_openai_messages(source: dict) -> list[dict]:
     return chat
 
 
+def _chat_template_input_ids(tokenized_chat: Any) -> list[int]:
+    input_ids = tokenized_chat.get("input_ids") if hasattr(tokenized_chat, "get") else tokenized_chat
+    if isinstance(input_ids, torch.Tensor):
+        input_ids = input_ids.detach().cpu().tolist()
+    if isinstance(input_ids, (list, tuple)) and input_ids and isinstance(input_ids[0], (list, tuple)):
+        if len(input_ids) != 1:
+            raise ValueError("Expected a single tokenized chat sequence from apply_chat_template.")
+        input_ids = input_ids[0]
+    return [int(token_id) for token_id in input_ids]
+
+
 def _chat_preprocess(source: dict, tokenizer: MegatronTokenizer, tool_schemas: Optional[list[Any]] = None) -> dict:
     """
     Preprocess messages to apply chat template and tokenize. Returns a dictionary of tokens.
@@ -931,33 +942,69 @@ def _chat_preprocess(source: dict, tokenizer: MegatronTokenizer, tool_schemas: O
     if getattr(tokenizer, "legacy", False):
         tokenizer = tokenizer._tokenizer
 
-    # assistant mask only works if chat template has generation keyword
     template_has_generation_kwd = GENERATION_REGEX.search(tokenizer.chat_template) is not None
+    from megatron.bridge.data.vlm_processing import build_assistant_loss_mask, infer_assistant_mask_boundary_config
 
-    if not template_has_generation_kwd:
-        raise ValueError(
-            "The tokenizer's chat_template does not contain a {% generation %} block, which is required "
-            "for HF's apply_chat_template to produce assistant-only loss masks via "
-            "return_assistant_tokens_mask=True. Without it, the loss mask would silently fall back to "
-            "all-ones (loss computed on the entire conversation including system/user tokens). "
-            "To fix this, either: (1) patch the chat_template to wrap assistant content with "
-            "{% generation %}...{% endgeneration %}, or (2) use the legacy special-tokens preprocessing "
-            "path instead of use_hf_tokenizer_chat_template=True."
+    known_boundary_markers = ("<|im_start|>assistant", "<|turn>model", "<start_of_turn>model")
+    boundary_config = (
+        infer_assistant_mask_boundary_config(tokenizer)
+        if any(marker in tokenizer.chat_template for marker in known_boundary_markers)
+        else None
+    )
+
+    if template_has_generation_kwd:
+        tokenized_chat = tokenizer.apply_chat_template(
+            chat,
+            tools=tools,
+            tokenize=True,
+            return_dict=True,
+            return_assistant_tokens_mask=True,
         )
+        input_ids = _chat_template_input_ids(tokenized_chat)
+        if boundary_config is None:
+            mask = tokenized_chat["assistant_masks"]
+        else:
+            chat_example = {"conversation": chat}
+            if tools is not None:
+                chat_example["tools"] = tools
+            mask = (
+                build_assistant_loss_mask(
+                    chat_example,
+                    torch.LongTensor(input_ids),
+                    tokenizer,
+                    boundary_config=boundary_config,
+                )
+                .to(dtype=torch.bool)
+                .tolist()
+            )
+    else:
+        if boundary_config is None:
+            raise ValueError(
+                "The tokenizer's chat_template does not contain a {% generation %} block and Bridge could not "
+                "infer assistant boundary markers for an assistant-only loss mask. Add a generation block to the "
+                "chat_template or use a model collate path that passes AssistantMaskBoundaryConfig explicitly."
+            )
 
-    tokenized_chat = tokenizer.apply_chat_template(
-        chat,
-        tools=tools,
-        tokenize=True,
-        return_dict=True,
-        return_assistant_tokens_mask=True,
-    )
+        tokenized_chat = tokenizer.apply_chat_template(
+            chat,
+            tools=tools,
+            tokenize=True,
+            return_dict=True,
+        )
+        input_ids = _chat_template_input_ids(tokenized_chat)
+        mask = (
+            build_assistant_loss_mask(
+                chat,
+                torch.LongTensor(input_ids),
+                tokenizer,
+                boundary_config=boundary_config,
+            )
+            .to(dtype=torch.bool)
+            .tolist()
+        )
 
     # Choose the last conversation as answer other history are context by finding the last masked token
     # which indicates end of context and beginning of answer
-    input_ids = tokenized_chat.get("input_ids")
-    mask = tokenized_chat["assistant_masks"]
-
     if 0 in mask:
         # traverse the list backward for first occurrence of masked token
         context_end_idx = len(mask) - mask[::-1].index(0)

src/megatron/bridge/data/energon/hf_task_encoder.py

@@ -59,11 +59,12 @@ class HFEnergonBatch(Batch):
     position_ids: torch.Tensor = field(default_factory=lambda: torch.empty(0))  # [B, seq_len]
     visual_inputs: GenericVisualInputs | None = None
     attention_mask: torch.Tensor | None = None
-    cu_seqlens: torch.Tensor | None = None
-    cu_seqlens_unpadded: torch.Tensor | None = None
-    cu_seqlens_argmin: torch.Tensor | None = None
-    cu_seqlens_unpadded_argmin: torch.Tensor | None = None
-    max_seqlen: torch.Tensor | None = None
+    cu_seqlens_q: torch.Tensor | None = None
+    cu_seqlens_kv: torch.Tensor | None = None
+    cu_seqlens_q_padded: torch.Tensor | None = None
+    cu_seqlens_kv_padded: torch.Tensor | None = None
+    max_seqlen_q: torch.Tensor | None = None
+    max_seqlen_kv: torch.Tensor | None = None
 
 
 class HFTaskEncoder(DefaultTaskEncoder[ChatMLSample, HFEnergonSample, HFEnergonBatch, dict]):
@@ -181,8 +182,8 @@ def batch(self, samples: List[HFEnergonSample]) -> HFEnergonBatch:
         examples = [sample.example for sample in samples]
         collated = self.collate_fn(examples)
         collated_seq_len = (
-            int(collated["max_seqlen"].max().item())
-            if collated.get("max_seqlen") is not None
+            int(collated["max_seqlen_q"].max().item())
+            if collated.get("max_seqlen_q") is not None
             else collated["input_ids"].shape[1]
         )
         if collated_seq_len > self.seq_length:
@@ -202,11 +203,12 @@ def batch(self, samples: List[HFEnergonSample]) -> HFEnergonBatch:
             attention_mask=collated.get("attention_mask"),
             position_ids=collated["position_ids"],
             visual_inputs=collated.get("visual_inputs"),
-            cu_seqlens=collated.get("cu_seqlens"),
-            cu_seqlens_unpadded=collated.get("cu_seqlens_unpadded"),
-            cu_seqlens_argmin=collated.get("cu_seqlens_argmin"),
-            cu_seqlens_unpadded_argmin=collated.get("cu_seqlens_unpadded_argmin"),
-            max_seqlen=collated.get("max_seqlen"),
+            cu_seqlens_q=collated.get("cu_seqlens_q"),
+            cu_seqlens_kv=collated.get("cu_seqlens_kv"),
+            cu_seqlens_q_padded=collated.get("cu_seqlens_q_padded"),
+            cu_seqlens_kv_padded=collated.get("cu_seqlens_kv_padded"),
+            max_seqlen_q=collated.get("max_seqlen_q"),
+            max_seqlen_kv=collated.get("max_seqlen_kv"),
         )
 
         return HFEnergonBatch(**batch_kwargs)

src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py

@@ -37,7 +37,7 @@
     find_pattern_indices,
     get_ltor_masks_and_position_ids,
 )
-from megatron.bridge.data.sequence_batching import pad_or_pack_sequence
+from megatron.bridge.data.sequence_batching import prepare_padded_or_packed_sequence_batch
 from megatron.bridge.training.utils.visual_inputs import GenericVisualInputs
 
 
@@ -86,11 +86,12 @@ class NemotronOmniTaskBatch(Batch):
     num_frames: Optional[torch.Tensor] = None  # [num_media_items]
     num_image_tiles: Optional[torch.Tensor] = None  # [total_images] LM-side token count per image
     # Packed-sequence metadata (only populated when enable_in_batch_packing=True).
-    cu_seqlens: Optional[torch.Tensor] = None
-    cu_seqlens_unpadded: Optional[torch.Tensor] = None
-    cu_seqlens_argmin: Optional[torch.Tensor] = None
-    cu_seqlens_unpadded_argmin: Optional[torch.Tensor] = None
-    max_seqlen: Optional[torch.Tensor] = None
+    cu_seqlens_q: Optional[torch.Tensor] = None
+    cu_seqlens_kv: Optional[torch.Tensor] = None
+    cu_seqlens_q_padded: Optional[torch.Tensor] = None
+    cu_seqlens_kv_padded: Optional[torch.Tensor] = None
+    max_seqlen_q: Optional[torch.Tensor] = None
+    max_seqlen_kv: Optional[torch.Tensor] = None
 
 
 # ---------------------------------------------------------------------------
@@ -500,50 +501,62 @@ def encode_sample(self, sample: ChatMLSample) -> NemotronOmniTaskSample:
 
     def batch(self, samples: List[NemotronOmniTaskSample]) -> NemotronOmniTaskBatch:
         """Pad-and-collate (default) OR pack samples along the seq dim when
-        ``enable_in_batch_packing=True``. Packing emits ``cu_seqlens`` / ``cu_seqlens_unpadded``
-        / ``max_seqlen`` so TE's THD kernels handle cross-sample masking (and CP
-        partitioning via ``thd_get_partitioned_indices``) without an attention mask.
+        ``enable_in_batch_packing=True``. Packing emits current MCore
+        packed-sequence metadata so TE's THD kernels handle cross-sample masking
+        without an attention mask.
         """
         pad_id = self._pad_token_id
         batch_size = len(samples)
 
-        cu_seqlens_t: Optional[torch.Tensor] = None
-        cu_seqlens_unpadded_t: Optional[torch.Tensor] = None
-        cu_seqlens_argmin_t: Optional[torch.Tensor] = None
-        cu_seqlens_unpadded_argmin_t: Optional[torch.Tensor] = None
-        max_seqlen_t: Optional[torch.Tensor] = None
+        cu_seqlens_q_t: Optional[torch.Tensor] = None
+        cu_seqlens_kv_t: Optional[torch.Tensor] = None
+        cu_seqlens_q_padded_t: Optional[torch.Tensor] = None
+        cu_seqlens_kv_padded_t: Optional[torch.Tensor] = None
+        max_seqlen_q_t: Optional[torch.Tensor] = None
+        max_seqlen_kv_t: Optional[torch.Tensor] = None
 
         if self.enable_in_batch_packing:
             # Concatenate samples along the seq dim into a single [1, total_len]
             # microbatch. TE attention kernels use cu_seqlens for per-sample
             # masking; no attention_mask needed.
+            if self.in_batch_packing_pad_to_multiple_of < 1:
+                raise ValueError("in_batch_packing_pad_to_multiple_of must be >= 1.")
             lengths = [int(s.input_ids.size(0)) for s in samples]
+            padded_lengths = [
+                ((length + self.in_batch_packing_pad_to_multiple_of - 1) // self.in_batch_packing_pad_to_multiple_of)
+                * self.in_batch_packing_pad_to_multiple_of
+                for length in lengths
+            ]
             cu_seqlens = [0]
-            for L in lengths:
-                cu_seqlens.append(cu_seqlens[-1] + L)
-
-            tokens_flat = torch.cat([s.input_ids for s in samples], dim=0)
-            labels_flat = torch.cat([s.labels for s in samples], dim=0)
-            loss_mask_flat = torch.cat([s.loss_mask for s in samples], dim=0)
-            # Per-sample resetting position ids: [0..L1-1, 0..L2-1, ...]
-            position_ids_flat = torch.cat([torch.arange(L, dtype=torch.long) for L in lengths], dim=0)
-
-            tokens = tokens_flat.unsqueeze(0)
-            tokens[tokens == pad_id] = 0
-            labels = labels_flat.unsqueeze(0)
-            loss_mask_t = loss_mask_flat.unsqueeze(0)
-            position_ids = position_ids_flat.unsqueeze(0)
+            cu_seqlens_padded = [0]
+            for length, padded_length in zip(lengths, padded_lengths):
+                cu_seqlens.append(cu_seqlens[-1] + length)
+                cu_seqlens_padded.append(cu_seqlens_padded[-1] + padded_length)
+
+            total_len = cu_seqlens_padded[-1]
+            tokens = torch.zeros((1, total_len), dtype=samples[0].input_ids.dtype)
+            labels = torch.full((1, total_len), IGNORE_INDEX, dtype=samples[0].labels.dtype)
+            loss_mask_t = torch.zeros((1, total_len), dtype=samples[0].loss_mask.dtype)
+            position_ids = torch.zeros((1, total_len), dtype=torch.long)
+
+            offset = 0
+            for sample, length, padded_length in zip(samples, lengths, padded_lengths):
+                sample_tokens = sample.input_ids.clone()
+                sample_tokens[sample_tokens == pad_id] = 0
+                tokens[0, offset : offset + length] = sample_tokens
+                labels[0, offset : offset + length] = sample.labels
+                loss_mask_t[0, offset : offset + length] = sample.loss_mask
+                position_ids[0, offset : offset + padded_length] = torch.arange(padded_length, dtype=torch.long)
+                offset += padded_length
             attention_mask = None  # TE derives the causal+padding mask from cu_seqlens.
 
-            cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32)
-            cu_seqlens_unpadded_t = cu_seqlens_t.clone()
-            # get_packed_seq_params truncates cu_seqlens_padded[: argmin.item()]; the
-            # trick in the fixed-size-batched case is sentinel=-1 padding with argmin
-            # pointing at the first sentinel. Here we emit an unpadded cu_seqlens and
-            # set argmin = len(cu_seqlens) so the slice is a no-op (keeps every entry).
-            cu_seqlens_argmin_t = torch.tensor(len(cu_seqlens), dtype=torch.int32)
-            cu_seqlens_unpadded_argmin_t = torch.tensor(len(cu_seqlens), dtype=torch.int32)
-            max_seqlen_t = torch.tensor(max(lengths), dtype=torch.int32)
+            cu_seqlens_q_t = torch.tensor(cu_seqlens, dtype=torch.int32)
+            cu_seqlens_kv_t = cu_seqlens_q_t
+            if self.in_batch_packing_pad_to_multiple_of > 1:
+                cu_seqlens_q_padded_t = torch.tensor(cu_seqlens_padded, dtype=torch.int32)
+                cu_seqlens_kv_padded_t = cu_seqlens_q_padded_t
+            max_seqlen_q_t = torch.tensor(max(padded_lengths), dtype=torch.int32)
+            max_seqlen_kv_t = max_seqlen_q_t
         else:
             max_seq_len = max(s.input_ids.size(0) for s in samples)
             input_ids_mat = np.full((batch_size, max_seq_len), pad_id, dtype=np.int64)
@@ -575,7 +588,7 @@ def batch(self, samples: List[NemotronOmniTaskSample]) -> NemotronOmniTaskBatch:
                 "position_ids": position_ids,
                 "attention_mask": attention_mask,
             }
-            pad_or_pack_sequence(
+            prepare_padded_or_packed_sequence_batch(
                 text_batch,
                 sequence_length=self.seq_length,
                 pad_to_max_length=self.pad_to_max_length,
@@ -664,11 +677,12 @@ def batch(self, samples: List[NemotronOmniTaskSample]) -> NemotronOmniTaskBatch:
             imgs_sizes=imgs_sizes_batch,
             num_frames=num_frames_batch,
             num_image_tiles=num_image_tiles_batch,
-            cu_seqlens=cu_seqlens_t,
-            cu_seqlens_unpadded=cu_seqlens_unpadded_t,
-            cu_seqlens_argmin=cu_seqlens_argmin_t,
-            cu_seqlens_unpadded_argmin=cu_seqlens_unpadded_argmin_t,
-            max_seqlen=max_seqlen_t,
+            cu_seqlens_q=cu_seqlens_q_t,
+            cu_seqlens_kv=cu_seqlens_kv_t,
+            cu_seqlens_q_padded=cu_seqlens_q_padded_t,
+            cu_seqlens_kv_padded=cu_seqlens_kv_padded_t,
+            max_seqlen_q=max_seqlen_q_t,
+            max_seqlen_kv=max_seqlen_kv_t,
         )
 
         return NemotronOmniTaskBatch(**batch_kwargs)
@@ -691,11 +705,12 @@ def encode_batch(self, batch: NemotronOmniTaskBatch) -> dict:
             "imgs_sizes": batch.imgs_sizes,
             "num_frames": batch.num_frames,
             "num_image_tiles": batch.num_image_tiles,
-            "cu_seqlens": batch.cu_seqlens,
-            "cu_seqlens_unpadded": batch.cu_seqlens_unpadded,
-            "cu_seqlens_argmin": batch.cu_seqlens_argmin,
-            "cu_seqlens_unpadded_argmin": batch.cu_seqlens_unpadded_argmin,
-            "max_seqlen": batch.max_seqlen,
+            "cu_seqlens_q": batch.cu_seqlens_q,
+            "cu_seqlens_kv": batch.cu_seqlens_kv,
+            "cu_seqlens_q_padded": batch.cu_seqlens_q_padded,
+            "cu_seqlens_kv_padded": batch.cu_seqlens_kv_padded,
+            "max_seqlen_q": batch.max_seqlen_q,
+            "max_seqlen_kv": batch.max_seqlen_kv,
         }
 
         vt = batch.visual_tensors if batch.visual_tensors else {}