padoc/modeling.py

"""Model loading and the flat-sequence training wrapper."""

from __future__ import annotations

from pathlib import Path

import torch
import torch.nn as nn
from transformers import AutoModelForImageTextToText, AutoProcessor, PreTrainedModel

from .preprocess import get_fork_token_map
from .tree_flash_attention import ATTN_IMPLEMENTATION as TREE_FLASH_ATTN


def load_padoc_model(
    model_path: str | Path,
    *,
    dtype: torch.dtype = torch.bfloat16,
    device_map: str | dict | None = "auto",
    attn_implementation: str | None = None,
) -> tuple[PreTrainedModel, AutoProcessor, dict[str, str]]:
    """Load a preprocessed checkpoint and validate its atomic fork tokens."""

    kwargs: dict = {"dtype": dtype, "device_map": device_map}
    if attn_implementation == TREE_FLASH_ATTN:
        kwargs["attn_implementation"] = {
            "": "sdpa",
            "text_config": TREE_FLASH_ATTN,
            "vision_config": "flash_attention_2",
        }
    elif attn_implementation:
        kwargs["attn_implementation"] = attn_implementation

    model = AutoModelForImageTextToText.from_pretrained(str(model_path), **kwargs)
    processor = AutoProcessor.from_pretrained(str(model_path))
    fork_token_map = get_fork_token_map(model)
    for token in set(fork_token_map) | set(fork_token_map.values()):
        ids = processor.tokenizer.encode(token, add_special_tokens=False)
        if len(ids) != 1:
            raise RuntimeError(f"Checkpoint fork token {token!r} is not atomic: {ids}")
    return model, processor, fork_token_map


class ForkTokenTrainer(nn.Module):
    """Adapt tree masks and logical positions to a Qwen3-VL-style model."""

    def __init__(self, model: PreTrainedModel):
        super().__init__()
        self.model = model
        self.config = model.config

    @staticmethod
    def _float_mask(attention_mask: torch.BoolTensor, dtype: torch.dtype) -> torch.Tensor:
        if attention_mask.ndim != 3:
            raise ValueError(
                "SDPA reference attention_mask must have shape (batch, sequence, sequence)."
            )
        output = torch.zeros(
            attention_mask.shape[0],
            1,
            attention_mask.shape[1],
            attention_mask.shape[2],
            dtype=dtype,
            device=attention_mask.device,
        )
        output.masked_fill_(~attention_mask.unsqueeze(1), torch.finfo(dtype).min)
        return output

    def _model_position_ids(
        self,
        *,
        input_ids: torch.LongTensor,
        logical_position_ids: torch.LongTensor,
        prompt_lens: torch.LongTensor,
        sequence_lengths: torch.LongTensor,
        mm_token_type_ids: torch.IntTensor | None,
        image_grid_thw: torch.LongTensor | None,
        video_grid_thw: torch.LongTensor | None,
        has_multimodal: bool,
    ) -> torch.Tensor:
        batch_size, sequence_length = input_ids.shape
        if has_multimodal and mm_token_type_ids is not None:
            token_range = torch.arange(sequence_length, device=input_ids.device)
            presence_mask = token_range.unsqueeze(0) < sequence_lengths.unsqueeze(1)
            positions, _ = self.model.model.get_rope_index(
                input_ids=input_ids,
                mm_token_type_ids=mm_token_type_ids,
                image_grid_thw=image_grid_thw,
                video_grid_thw=video_grid_thw,
                attention_mask=presence_mask.long(),
            )
            for batch_index in range(batch_size):
                prompt_len = int(prompt_lens[batch_index])
                sequence_len = int(sequence_lengths[batch_index])
                if prompt_len <= 0 or prompt_len >= sequence_len:
                    continue
                delta = positions[0, batch_index, prompt_len - 1] - (prompt_len - 1)
                shifted = (
                    logical_position_ids[batch_index, prompt_len:sequence_len]
                    .to(positions.dtype)
                    .add(delta)
                )
                positions[:, batch_index, prompt_len:sequence_len] = shifted.unsqueeze(0)
            return positions

        if logical_position_ids.ndim != 2:
            raise ValueError("Logical position_ids must have shape (batch, sequence).")
        return logical_position_ids.unsqueeze(0).expand(4, -1, -1).contiguous()

    def forward(
        self,
        *,
        input_ids: torch.LongTensor,
        labels: torch.LongTensor,
        position_ids: torch.LongTensor,
        prompt_lens: torch.LongTensor,
        sequence_lengths: torch.LongTensor,
        attention_mask: torch.Tensor | None = None,
        pixel_values: torch.Tensor | None = None,
        image_grid_thw: torch.LongTensor | None = None,
        pixel_values_videos: torch.Tensor | None = None,
        video_grid_thw: torch.LongTensor | None = None,
        mm_token_type_ids: torch.IntTensor | None = None,
        tree_q_indices: torch.LongTensor | None = None,
        tree_kv_indices: torch.LongTensor | None = None,
        tree_cu_seqlens_q: torch.IntTensor | None = None,
        tree_cu_seqlens_kv: torch.IntTensor | None = None,
        tree_max_seqlen_q: int | torch.Tensor | None = None,
        tree_max_seqlen_kv: int | torch.Tensor | None = None,
    ):
        uses_tree_fa2 = tree_q_indices is not None
        uses_sdpa = attention_mask is not None and attention_mask.ndim == 3
        if uses_tree_fa2 == uses_sdpa:
            raise ValueError("Pass exactly one of tree FA2 metadata or a 3-D SDPA mask.")
        if uses_tree_fa2:
            text_config = getattr(self.model.config, "text_config", self.model.config)
            implementation = getattr(text_config, "_attn_implementation", None)
            if implementation is not None and implementation != TREE_FLASH_ATTN:
                raise ValueError(
                    "Tree metadata requires a model loaded with "
                    f"attn_implementation={TREE_FLASH_ATTN!r}."
                )

        dtype = next(self.model.parameters()).dtype
        mask_to_pass = None if uses_tree_fa2 else self._float_mask(attention_mask, dtype)
        has_multimodal = (pixel_values is not None and image_grid_thw is not None) or (
            pixel_values_videos is not None and video_grid_thw is not None
        )
        model_positions = self._model_position_ids(
            input_ids=input_ids,
            logical_position_ids=position_ids,
            prompt_lens=prompt_lens,
            sequence_lengths=sequence_lengths,
            mm_token_type_ids=mm_token_type_ids,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            has_multimodal=has_multimodal,
        )

        if pixel_values is not None:
            pixel_values = pixel_values.to(dtype=dtype)
        if pixel_values_videos is not None:
            pixel_values_videos = pixel_values_videos.to(dtype=dtype)

        tree_kwargs = {}
        if uses_tree_fa2:
            tree_kwargs = {
                "tree_q_indices": tree_q_indices,
                "tree_kv_indices": tree_kv_indices,
                "tree_cu_seqlens_q": tree_cu_seqlens_q,
                "tree_cu_seqlens_kv": tree_cu_seqlens_kv,
                "tree_max_seqlen_q": tree_max_seqlen_q,
                "tree_max_seqlen_kv": tree_max_seqlen_kv,
            }

        return self.model(
            input_ids=input_ids,
            attention_mask=mask_to_pass,
            position_ids=model_positions,
            labels=labels,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            pixel_values_videos=pixel_values_videos,
            video_grid_thw=video_grid_thw,
            mm_token_type_ids=mm_token_type_ids,
            **tree_kwargs,
        )