"""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, )