From 0697bd2aa1b7ed12ab704e29cc27d3e81be0fb79 Mon Sep 17 00:00:00 2001 From: Haoming Date: Wed, 17 Sep 2025 11:02:40 +0800 Subject: [PATCH] qwen 2.5 --- backend/nn/llm/llama.py | 279 ++++++++++++++++++ backend/nn/llm/qwen_vl.py | 379 +++++++++++++++++++++++++ backend/text_processing/qwen_engine.py | 144 ++++++++++ 3 files changed, 802 insertions(+) create mode 100644 backend/nn/llm/llama.py create mode 100644 backend/nn/llm/qwen_vl.py create mode 100644 backend/text_processing/qwen_engine.py diff --git a/backend/nn/llm/llama.py b/backend/nn/llm/llama.py new file mode 100644 index 00000000..9c41f070 --- /dev/null +++ b/backend/nn/llm/llama.py @@ -0,0 +1,279 @@ +import math +from dataclasses import asdict, dataclass +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from backend.attention import attention_function + +from . import qwen_vl + + +@dataclass +class Qwen25_7BVLI_Config: + vocab_size: int = 152064 + hidden_size: int = 3584 + intermediate_size: int = 18944 + num_hidden_layers: int = 28 + num_attention_heads: int = 28 + num_key_value_heads: int = 4 + max_position_embeddings: int = 128000 + rms_norm_eps: float = 1e-6 + rope_theta: float = 1000000.0 + transformer_type: str = "llama" + head_dim = 128 + rms_norm_add = False + mlp_activation = "silu" + qkv_bias = True + rope_dims = [16, 24, 24] + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def precompute_freqs_cis(head_dim, position_ids, theta, rope_dims=None, device=None): + theta_numerator = torch.arange(0, head_dim, 2, device=device).float() + inv_freq = 1.0 / (theta ** (theta_numerator / head_dim)) + + inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + if rope_dims is not None and position_ids.shape[0] > 1: + mrope_section = rope_dims * 2 + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(0) + sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(0) + else: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + return (cos, sin) + + +def apply_rope(xq, xk, freqs_cis): + org_dtype = xq.dtype + cos = freqs_cis[0] + sin = freqs_cis[1] + q_embed = (xq * cos) + (rotate_half(xq) * sin) + k_embed = (xk * cos) + (rotate_half(xk) * sin) + return q_embed.to(org_dtype), k_embed.to(org_dtype) + + +class Attention(nn.Module): + def __init__(self, config: Qwen25_7BVLI_Config): + super().__init__() + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.hidden_size = config.hidden_size + + self.head_dim = config.head_dim + self.inner_size = self.num_heads * self.head_dim + + self.q_proj = nn.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias) + self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias) + self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias) + self.o_proj = nn.Linear(self.inner_size, config.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + freqs_cis: Optional[torch.Tensor] = None, + optimized_attention=None, + ): + batch_size, seq_length, _ = hidden_states.shape + xq = self.q_proj(hidden_states) + xk = self.k_proj(hidden_states) + xv = self.v_proj(hidden_states) + + xq = xq.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + xk = xk.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2) + xv = xv.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2) + + xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis) + + xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) + xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) + + output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True) + return self.o_proj(output) + + +class MLP(nn.Module): + def __init__(self, config: Qwen25_7BVLI_Config): + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + if config.mlp_activation == "silu": + self.activation = F.silu + elif config.mlp_activation == "gelu_pytorch_tanh": + self.activation = lambda a: F.gelu(a, approximate="tanh") + + def forward(self, x): + return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x)) + + +class TransformerBlock(nn.Module): + def __init__(self, config: Qwen25_7BVLI_Config): + super().__init__() + self.self_attn = Attention(config) + self.mlp = MLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + x: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + freqs_cis: Optional[torch.Tensor] = None, + optimized_attention=None, + ): + # Self Attention + residual = x + x = self.input_layernorm(x) + x = self.self_attn( + hidden_states=x, + attention_mask=attention_mask, + freqs_cis=freqs_cis, + optimized_attention=optimized_attention, + ) + x = residual + x + + # MLP + residual = x + x = self.post_attention_layernorm(x) + x = self.mlp(x) + x = residual + x + + return x + + +class Llama2_(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.normalize_in = False + + self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)]) + self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, position_ids=None, embeds_info=[]): + if embeds is not None: + x = embeds + else: + x = self.embed_tokens(x) + + if self.normalize_in: + x *= self.config.hidden_size**0.5 + + if position_ids is None: + position_ids = torch.arange(0, x.shape[1], device=x.device).unsqueeze(0) + + freqs_cis = precompute_freqs_cis(self.config.head_dim, position_ids, self.config.rope_theta, self.config.rope_dims, device=x.device) + + mask = None + if attention_mask is not None: + mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1]) + mask = mask.masked_fill(mask.to(torch.bool), float("-inf")) + + causal_mask = torch.empty(x.shape[1], x.shape[1], dtype=x.dtype, device=x.device).fill_(float("-inf")).triu_(1) + if mask is not None: + mask += causal_mask + else: + mask = causal_mask + + intermediate = None + all_intermediate = None + if intermediate_output is not None: + if intermediate_output == "all": + all_intermediate = [] + intermediate_output = None + elif intermediate_output < 0: + intermediate_output = len(self.layers) + intermediate_output + + for i, layer in enumerate(self.layers): + if all_intermediate is not None: + all_intermediate.append(x.unsqueeze(1).clone()) + x = layer( + x=x, + attention_mask=mask, + freqs_cis=freqs_cis, + optimized_attention=attention_function, + ) + if i == intermediate_output: + intermediate = x.clone() + + x = self.norm(x) + if all_intermediate is not None: + all_intermediate.append(x.unsqueeze(1).clone()) + + if all_intermediate is not None: + intermediate = torch.cat(all_intermediate, dim=1) + + if intermediate is not None and final_layer_norm_intermediate: + intermediate = self.norm(intermediate) + + return x, intermediate + + +class Qwen25_7BVLI(nn.Module): + def __init__(self, config_dict): + super().__init__() + config = Qwen25_7BVLI_Config() + + _config_dict = asdict(config) + for key, value in _config_dict.items(): + if key in config_dict: + assert value == config_dict[key] + + self.num_layers = config.num_hidden_layers + + self.model = Llama2_(config) + self.visual = qwen_vl.Qwen2VLVisionTransformer(hidden_size=1280, output_hidden_size=config.hidden_size) + + def preprocess_embed(self, embed, device): + if embed["type"] == "image": + image, grid = qwen_vl.process_qwen2vl_images(embed["data"]) + return self.visual(image.to(device, dtype=torch.float32), grid), grid + return None, None + + def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[]): + grid = None + for e in embeds_info: + if e.get("type") == "image": + grid = e.get("extra", None) + position_ids = torch.zeros((3, embeds.shape[1]), device=embeds.device) + start = e.get("index") + position_ids[:, :start] = torch.arange(0, start, device=embeds.device) + end = e.get("size") + start + len_max = int(grid.max()) // 2 + start_next = len_max + start + position_ids[:, end:] = torch.arange(start_next, start_next + (embeds.shape[1] - end), device=embeds.device) + position_ids[0, start:end] = start + max_d = int(grid[0][1]) // 2 + position_ids[1, start:end] = torch.arange(start, start + max_d, device=embeds.device).unsqueeze(1).repeat(1, math.ceil((end - start) / max_d)).flatten(0)[: end - start] + max_d = int(grid[0][2]) // 2 + position_ids[2, start:end] = torch.arange(start, start + max_d, device=embeds.device).unsqueeze(0).repeat(math.ceil((end - start) / max_d), 1).flatten(0)[: end - start] + + if grid is None: + position_ids = None + + return self.model(x, attention_mask=attention_mask, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=final_layer_norm_intermediate, dtype=dtype, position_ids=position_ids) + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, embeddings): + self.model.embed_tokens = embeddings diff --git a/backend/nn/llm/qwen_vl.py b/backend/nn/llm/qwen_vl.py new file mode 100644 index 00000000..c98286aa --- /dev/null +++ b/backend/nn/llm/qwen_vl.py @@ -0,0 +1,379 @@ +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from backend.attention import attention_function + + +def process_qwen2vl_images( + images: torch.Tensor, + min_pixels: int = 3136, + max_pixels: int = 12845056, + patch_size: int = 14, + temporal_patch_size: int = 2, + merge_size: int = 2, + image_mean: list = None, + image_std: list = None, +): + if image_mean is None: + image_mean = [0.48145466, 0.4578275, 0.40821073] + if image_std is None: + image_std = [0.26862954, 0.26130258, 0.27577711] + + batch_size, height, width, channels = images.shape + device = images.device + # dtype = images.dtype + + images = images.permute(0, 3, 1, 2) + + grid_thw_list = [] + img = images[0] + + factor = patch_size * merge_size + + h_bar = round(height / factor) * factor + w_bar = round(width / factor) * factor + + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = max(factor, math.floor(height / beta / factor) * factor) + w_bar = max(factor, math.floor(width / beta / factor) * factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = math.ceil(height * beta / factor) * factor + w_bar = math.ceil(width * beta / factor) * factor + + img_resized = F.interpolate(img.unsqueeze(0), size=(h_bar, w_bar), mode="bilinear", align_corners=False).squeeze(0) + + normalized = img_resized.clone() + for c in range(3): + normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c] + + grid_h = h_bar // patch_size + grid_w = w_bar // patch_size + grid_thw = torch.tensor([1, grid_h, grid_w], device=device, dtype=torch.long) + + pixel_values = normalized + grid_thw_list.append(grid_thw) + image_grid_thw = torch.stack(grid_thw_list) + + grid_t = 1 + channel = pixel_values.shape[0] + pixel_values = pixel_values.unsqueeze(0).repeat(2, 1, 1, 1) + + patches = pixel_values.reshape( + grid_t, + temporal_patch_size, + channel, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + + patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8) + flatten_patches = patches.reshape(grid_t * grid_h * grid_w, channel * temporal_patch_size * patch_size * patch_size) + + return flatten_patches, image_grid_thw + + +class VisionPatchEmbed(nn.Module): + def __init__( + self, + patch_size: int = 14, + temporal_patch_size: int = 2, + in_channels: int = 3, + embed_dim: int = 3584, + ): + super().__init__() + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.in_channels = in_channels + self.embed_dim = embed_dim + + kernel_size = [temporal_patch_size, patch_size, patch_size] + self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.view(-1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size) + hidden_states = self.proj(hidden_states) + return hidden_states.view(-1, self.embed_dim) + + +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_vision(q, k, cos, sin): + cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class VisionRotaryEmbedding(nn.Module): + def __init__(self, dim: int, theta: float = 10000.0): + super().__init__() + self.dim = dim + self.theta = theta + + def forward(self, seqlen: int, device) -> torch.Tensor: + inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float, device=device) / self.dim)) + seq = torch.arange(seqlen, device=inv_freq.device, dtype=inv_freq.dtype) + freqs = torch.outer(seq, inv_freq) + return freqs + + +class PatchMerger(nn.Module): + def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2): + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + self.ln_q = nn.RMSNorm(context_dim, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.ln_q(x).reshape(-1, self.hidden_size) + x = self.mlp(x) + return x + + +class VisionAttention(nn.Module): + def __init__(self, hidden_size: int, num_heads: int): + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scaling = self.head_dim**-0.5 + + self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=True) + self.proj = nn.Linear(hidden_size, hidden_size, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + cu_seqlens=None, + optimized_attention=None, + ) -> torch.Tensor: + if hidden_states.dim() == 2: + seq_length, _ = hidden_states.shape + batch_size = 1 + hidden_states = hidden_states.unsqueeze(0) + else: + batch_size, seq_length, _ = hidden_states.shape + + qkv = self.qkv(hidden_states) + qkv = qkv.reshape(batch_size, seq_length, 3, self.num_heads, self.head_dim) + query_states, key_states, value_states = qkv.reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)] + + attn_outputs = [optimized_attention(q, k, v, self.num_heads, skip_reshape=True) for q, k, v in zip(*splits)] + attn_output = torch.cat(attn_outputs, dim=1) + attn_output = attn_output.reshape(seq_length, -1) + attn_output = self.proj(attn_output) + + return attn_output + + +class VisionMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=True) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=True) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=True) + self.act_fn = nn.SiLU() + + def forward(self, hidden_state): + return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + + +class VisionBlock(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, num_heads: int): + super().__init__() + self.norm1 = nn.RMSNorm(hidden_size, eps=1e-6) + self.norm2 = nn.RMSNorm(hidden_size, eps=1e-6) + self.attn = VisionAttention(hidden_size, num_heads) + self.mlp = VisionMLP(hidden_size, intermediate_size) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + cu_seqlens=None, + optimized_attention=None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.attn(hidden_states, position_embeddings, cu_seqlens, optimized_attention) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class Qwen2VLVisionTransformer(nn.Module): + def __init__(self, hidden_size: int = 3584, output_hidden_size: int = 3584, intermediate_size: int = 3420, num_heads: int = 16, num_layers: int = 32, patch_size: int = 14, temporal_patch_size: int = 2, spatial_merge_size: int = 2, window_size: int = 112): + super().__init__() + self.hidden_size = hidden_size + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.window_size = window_size + self.fullatt_block_indexes = [7, 15, 23, 31] + + self.patch_embed = VisionPatchEmbed( + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + in_channels=3, + embed_dim=hidden_size, + ) + + head_dim = hidden_size // num_heads + self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList([VisionBlock(hidden_size, intermediate_size, num_heads) for _ in range(num_layers)]) + + self.merger = PatchMerger( + dim=output_hidden_size, + context_dim=hidden_size, + spatial_merge_size=spatial_merge_size, + ) + + def get_window_index(self, grid_thw): + window_index = [] + cu_window_seqlens = [0] + window_index_id = 0 + vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size + + for grid_t, grid_h, grid_w in grid_thw: + llm_grid_h = grid_h // self.spatial_merge_size + llm_grid_w = grid_w // self.spatial_merge_size + + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) + + pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size + pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size + num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size + + index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100) + index_padded = index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape( + grid_t, + num_windows_h * num_windows_w, + vit_merger_window_size, + vit_merger_window_size, + ) + + seqlens = (index_padded != -100).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != -100] + window_index.append(index_new + window_index_id) + + cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_size * self.spatial_merge_size + cu_window_seqlens[-1] + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += (grid_t * llm_grid_h * llm_grid_w).item() + + window_index = torch.cat(window_index, dim=0) + return window_index, cu_window_seqlens + + def get_position_embeddings(self, grid_thw, device): + pos_ids = [] + + for t, h, w in grid_thw: + hpos_ids = torch.arange(h, device=device).unsqueeze(1).expand(-1, w) + hpos_ids = hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten() + + wpos_ids = torch.arange(w, device=device).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3).flatten() + + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size, device) + return rotary_pos_emb_full[pos_ids].flatten(1) + + def forward( + self, + pixel_values: torch.Tensor, + image_grid_thw: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + hidden_states = self.patch_embed(pixel_values) + + window_index, cu_window_seqlens = self.get_window_index(image_grid_thw) + cu_window_seqlens = torch.tensor(cu_window_seqlens, device=hidden_states.device) + cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) + + position_embeddings = self.get_position_embeddings(image_grid_thw, hidden_states.device) + + seq_len, _ = hidden_states.size() + spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + hidden_states = hidden_states.reshape(seq_len // spatial_merge_unit, spatial_merge_unit, -1) + hidden_states = hidden_states[window_index, :, :] + hidden_states = hidden_states.reshape(seq_len, -1) + + position_embeddings = position_embeddings.reshape(seq_len // spatial_merge_unit, spatial_merge_unit, -1) + position_embeddings = position_embeddings[window_index, :, :] + position_embeddings = position_embeddings.reshape(seq_len, -1) + position_embeddings = torch.cat((position_embeddings, position_embeddings), dim=-1) + position_embeddings = (position_embeddings.cos(), position_embeddings.sin()) + + cu_seqlens = torch.repeat_interleave(image_grid_thw[:, 1] * image_grid_thw[:, 2], image_grid_thw[:, 0]).cumsum( + dim=0, + dtype=torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + for i, block in enumerate(self.blocks): + if i in self.fullatt_block_indexes: + cu_seqlens_now = cu_seqlens + else: + cu_seqlens_now = cu_window_seqlens + hidden_states = block(hidden_states, position_embeddings, cu_seqlens_now, optimized_attention=attention_function) + + hidden_states = self.merger(hidden_states) + return hidden_states diff --git a/backend/text_processing/qwen_engine.py b/backend/text_processing/qwen_engine.py new file mode 100644 index 00000000..532d3bcf --- /dev/null +++ b/backend/text_processing/qwen_engine.py @@ -0,0 +1,144 @@ +import torch + +from backend import memory_management +from backend.text_processing import emphasis, parsing +from modules.shared import opts + + +class PromptChunk: + def __init__(self): + self.tokens = [] + self.multipliers = [] + + +class QwenTextProcessingEngine: + def __init__(self, text_encoder, tokenizer): + super().__init__() + + self.text_encoder = text_encoder + self.tokenizer = tokenizer + + self.emphasis = emphasis.get_current_option(opts.emphasis)() + self.max_length = 99999999 + self.min_length = 1 + self.id_pad = 151643 + + self.llama_template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + + def tokenize(self, texts): + llama_texts = [self.llama_template.format(text) for text in texts] + return self.tokenizer(llama_texts)["input_ids"] + + def encode_with_transformers(self, tokens): + device = memory_management.text_encoder_device() + tokens = tokens.to(device) + self.text_encoder.to(device=device) + return self.text_encoder(x=tokens) + + def tokenize_line(self, line): + parsed = parsing.parse_prompt_attention(line, self.emphasis.name) + + tokenized = self.tokenize([text for text, _ in parsed]) + + chunks = [] + chunk = PromptChunk() + token_count = 0 + + def next_chunk(): + nonlocal token_count + nonlocal chunk + + current_chunk_length = len(chunk.tokens) + token_count += current_chunk_length + remaining_count = self.min_length - current_chunk_length + + if self.min_length > 0 and remaining_count > 0: + chunk.tokens += [self.id_pad] * remaining_count + chunk.multipliers += [1.0] * remaining_count + + chunks.append(chunk) + chunk = PromptChunk() + + for tokens, (text, weight) in zip(tokenized, parsed): + if text == "BREAK" and weight == -1: + next_chunk() + continue + + position = 0 + while position < len(tokens): + token = tokens[position] + chunk.tokens.append(token) + chunk.multipliers.append(weight) + position += 1 + + if chunk.tokens or not chunks: + next_chunk() + + return chunks, token_count + + def __call__(self, texts): + zs = [] + cache = {} + + self.emphasis = emphasis.get_current_option(opts.emphasis)() + + for line in texts: + if line in cache: + line_z_values = cache[line] + else: + chunks, token_count = self.tokenize_line(line) + line_z_values = [] + + # pad all chunks to length of longest chunk + max_tokens = 0 + for chunk in chunks: + max_tokens = max(len(chunk.tokens), max_tokens) + + for chunk in chunks: + tokens = chunk.tokens + multipliers = chunk.multipliers + + remaining_count = max_tokens - len(tokens) + if remaining_count > 0: + tokens += [self.id_pad] * remaining_count + multipliers += [1.0] * remaining_count + + z = self.process_tokens([tokens], [multipliers])[0] + z = self.postprocess_tokens(z, tokens) + line_z_values.append(z) + cache[line] = line_z_values + + zs.extend(line_z_values) + + return torch.stack(zs) + + def postprocess_tokens(self, out, tokens): + """strip the llama_template""" + template_end = 0 + count_im_start = 0 + + for i, v in enumerate(tokens): + elem = int(v) + if elem == 151644 and count_im_start < 2: + template_end = i + count_im_start += 1 + + if out.shape[1] > (template_end + 3): + if int(tokens[template_end + 1]) == 872: + if int(tokens[template_end + 2]) == 198: + template_end += 3 + + return out[template_end:] + + def process_tokens(self, batch_tokens, batch_multipliers): + tokens = torch.asarray(batch_tokens) + + z, _ = self.encode_with_transformers(tokens) + + self.emphasis.tokens = batch_tokens + self.emphasis.multipliers = torch.asarray(batch_multipliers).to(z) + self.emphasis.z = z + self.emphasis.after_transformers() + z = self.emphasis.z + + return z