Spaces:
Running on Zero
Running on Zero
File size: 7,749 Bytes
414b4fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """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,
)
|