From c2ca4d958ead7610e10c8c9ef8e7f2485508899e Mon Sep 17 00:00:00 2001 From: Ottohere-Mourn <3311436628@qq.com> Date: Wed, 17 Jun 2026 12:14:39 +0800 Subject: [PATCH 1/2] fix: add defensive fallback for Qwen VL mrope get_rope_index shape mismatch When video samples have inconsistent grid metadata between the tokenizer and the forward pass (e.g. due to odd-frame padding logic in _get_qwen_video_grid_metadata vs. the native video_processor), the placeholder token count in input_ids disagrees with the feature count, causing a RuntimeError in get_rope_index: "value tensor of shape [3, A] cannot be broadcast to indexing result of shape [3, B]" This affects a tiny fraction of samples (~0.001% in a 9.16M sample mix) and is normally invisible in smoke tests with small max_samples. The fix catches this specific shape mismatch at the collator level, degrades to a flat 3-axis positional encoding for the offending batch (allowing training to continue without interruption), and prints a one-line diagnostic so the user can later locate and clean those samples at the data level. Verified on a 9.16M-sample multimodal training run. No false positives detected. --- src/llamafactory/data/collator.py | 50 ++++++++++++++++++++++++++++++- 1 file changed, 49 insertions(+), 1 deletion(-) diff --git a/src/llamafactory/data/collator.py b/src/llamafactory/data/collator.py index af234d99bf..0fc29cc148 100644 --- a/src/llamafactory/data/collator.py +++ b/src/llamafactory/data/collator.py @@ -199,7 +199,55 @@ def _compute_rope_position_ids(self, features: dict[str, "torch.Tensor"], mm_inp features["position_ids"], rope_deltas = self.get_rope_func(**rope_index_kwargs) features["rope_deltas"] = rope_deltas - (1 - rope_index_kwargs["attention_mask"]).sum(dim=-1).unsqueeze(-1) else: # for qwen vl - features["position_ids"], features["rope_deltas"] = self.get_rope_func(**rope_index_kwargs) + try: + features["position_ids"], features["rope_deltas"] = self.get_rope_func(**rope_index_kwargs) + except RuntimeError as err: + # Defensive fallback for the mrope get_rope_index shape mismatch + # ("value tensor of shape [3, A] cannot be broadcast to indexing result of shape [3, B]") + # which fires when the vision placeholder token count in input_ids disagrees with the + # image_grid_thw / video_grid_thw declared by the processor. This happens on a tiny + # fraction of samples (e.g. a truncated or irregular-grid video) and would otherwise + # kill the whole run. Degrade to a flat positional encoding (3 identical axes) for the + # offending batch so training continues, and log it for later data-level cleanup. + if "shape mismatch" not in str(err) and "cannot be broadcast" not in str(err): + raise + self._log_rope_mismatch(err, features, mm_inputs) + features["position_ids"], features["rope_deltas"] = self._fallback_rope_position_ids( + features["input_ids"], rope_index_kwargs["attention_mask"] + ) + + + @staticmethod + def _fallback_rope_position_ids(input_ids: "torch.Tensor", attention_mask: "torch.Tensor"): + r"""Flat positional fallback: every non-pad token gets a monotonic position; all 3 mrope axes identical.""" + bsz, seq_len = input_ids.shape + position_ids = torch.zeros(3, bsz, seq_len, dtype=input_ids.dtype, device=input_ids.device) + rope_deltas = torch.zeros(bsz, dtype=input_ids.dtype, device=input_ids.device) + for batch_idx in range(bsz): + valid_len = int(attention_mask[batch_idx].bool().sum().item()) + positions = torch.arange(seq_len, device=input_ids.device) + position_ids[:, batch_idx] = positions.view(1, -1).expand(3, -1) + rope_deltas[batch_idx] = positions.max() + 1 - valid_len if valid_len > 0 else 0 + return position_ids, rope_deltas + + def _log_rope_mismatch(self, err: Exception, features: dict[str, "torch.Tensor"], mm_inputs: dict[str, Any]) -> None: + r"""Print a one-line diagnostic when the mrope fallback path is taken.""" + input_ids = features.get("input_ids") + image_grid = mm_inputs.get("image_grid_thw") + video_grid = mm_inputs.get("video_grid_thw") + image_token_id = getattr(self.model.config, "image_token_id", None) if self.model is not None else None + video_token_id = getattr(self.model.config, "video_token_id", None) if self.model is not None else None + n_img = int((input_ids == image_token_id).sum().item()) if image_token_id is not None and input_ids is not None else -1 + n_vid = int((input_ids == video_token_id).sum().item()) if video_token_id is not None and input_ids is not None else -1 + print( + "rope_fallback: get_rope_index shape mismatch caught -> " + f"input_ids={tuple(input_ids.shape) if input_ids is not None else None} " + f"image_tokens={n_img} video_tokens={n_vid} " + f"image_grid_thw={image_grid.tolist() if image_grid is not None else None} " + f"video_grid_thw={video_grid.tolist() if video_grid is not None else None} " + f"err={err}", + flush=True, + ) def _compute_rope_position_ids_with_packing( self, From 2e67cbb3a43e5ec8ce6848515dba9884ed86a150 Mon Sep 17 00:00:00 2001 From: Ottohere-Mourn <3311436628@qq.com> Date: Wed, 17 Jun 2026 12:24:46 +0800 Subject: [PATCH 2/2] perf: vectorize _fallback_rope_position_ids (remove GPU-CPU sync) --- src/llamafactory/data/collator.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src/llamafactory/data/collator.py b/src/llamafactory/data/collator.py index 0fc29cc148..6473e9af1e 100644 --- a/src/llamafactory/data/collator.py +++ b/src/llamafactory/data/collator.py @@ -221,14 +221,11 @@ def _compute_rope_position_ids(self, features: dict[str, "torch.Tensor"], mm_inp def _fallback_rope_position_ids(input_ids: "torch.Tensor", attention_mask: "torch.Tensor"): r"""Flat positional fallback: every non-pad token gets a monotonic position; all 3 mrope axes identical.""" bsz, seq_len = input_ids.shape - position_ids = torch.zeros(3, bsz, seq_len, dtype=input_ids.dtype, device=input_ids.device) - rope_deltas = torch.zeros(bsz, dtype=input_ids.dtype, device=input_ids.device) - for batch_idx in range(bsz): - valid_len = int(attention_mask[batch_idx].bool().sum().item()) - positions = torch.arange(seq_len, device=input_ids.device) - position_ids[:, batch_idx] = positions.view(1, -1).expand(3, -1) - rope_deltas[batch_idx] = positions.max() + 1 - valid_len if valid_len > 0 else 0 - return position_ids, rope_deltas + positions = torch.arange(seq_len, device=input_ids.device, dtype=input_ids.dtype) + position_ids = positions.view(1, 1, seq_len).expand(3, bsz, seq_len).contiguous() + valid_len = attention_mask.bool().sum(dim=-1) + rope_deltas = torch.where(valid_len > 0, seq_len - valid_len, torch.zeros_like(valid_len)) + return position_ids, rope_deltas.to(dtype=input_ids.dtype) def _log_rope_mismatch(self, err: Exception, features: dict[str, "torch.Tensor"], mm_inputs: dict[str, Any]) -> None: r"""Print a one-line diagnostic when the mrope fallback path is taken."""