Skip to content
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 49 additions & 1 deletion src/llamafactory/data/collator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Performance Bottleneck: GPU-CPU Synchronization & Loop Overhead

In the current implementation of _fallback_rope_position_ids, there are two major performance issues that can slow down training throughput:

  1. GPU-CPU Synchronization: Calling .item() on a GPU tensor (inside attention_mask[batch_idx].bool().sum().item()) forces the CPU to block and wait for the GPU to finish its computation. Since this is executed inside a loop for every batch in the data collator, it can severely bottleneck the training pipeline and lower GPU utilization.
  2. Redundant Tensor Creation & Loop Overhead: Re-creating the positions tensor bsz times inside a Python loop is inefficient.

Solution

We can completely vectorize this method to run loop-free and avoid any GPU-CPU synchronization by leveraging PyTorch's native tensor operations (expand, torch.where, and sum along dimensions).

    @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
        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)

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, applied the vectorized version in 2e67cbb. The Python loop and .item() call have been removed.


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,
Expand Down