From 0d5124e09263b0f8e3ceff490dfbc750e6f7abce Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Wed, 10 Dec 2025 06:57:51 +0000 Subject: [PATCH] Qwen2_5_VL text encoder has 0 numerical diff with transformers --- fastvideo/configs/models/encoders/__init__.py | 3 +- fastvideo/configs/models/encoders/qwen2_5.py | 81 ++++ fastvideo/models/encoders/qwen2_5.py | 387 ++++++++++++++++++ fastvideo/models/mask_utils.py | 201 +++++++++ fastvideo/models/registry.py | 1 + .../tests/encoders/test_qwen2_5_encoder.py | 145 +++++++ 6 files changed, 817 insertions(+), 1 deletion(-) create mode 100644 fastvideo/configs/models/encoders/qwen2_5.py create mode 100644 fastvideo/models/encoders/qwen2_5.py create mode 100644 fastvideo/models/mask_utils.py create mode 100644 fastvideo/tests/encoders/test_qwen2_5_encoder.py diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 29b5d6113..04afb17e8 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -6,9 +6,10 @@ from fastvideo.configs.models.encoders.clip import ( CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig) from fastvideo.configs.models.encoders.llama import LlamaConfig from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig +from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig __all__ = [ "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", - "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig" + "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig" ] diff --git a/fastvideo/configs/models/encoders/qwen2_5.py b/fastvideo/configs/models/encoders/qwen2_5.py new file mode 100644 index 000000000..250ed739a --- /dev/null +++ b/fastvideo/configs/models/encoders/qwen2_5.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field +from typing import Optional + +from fastvideo.configs.models.encoders.base import (TextEncoderArchConfig, + TextEncoderConfig) + +def _is_transformer_layer(n: str, m) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embed_tokens") + + +def _is_final_norm(n: str, m) -> bool: + return n.endswith("norm") + + +@dataclass +class Qwen2_5_VLArchConfig(TextEncoderArchConfig): + vocab_size: int = 152064 + hidden_size: int = 8192 + intermediate_size: int = 29568 + num_hidden_layers: int = 80 + num_attention_heads: int = 64 + num_key_value_heads: int = 8 + hidden_act: str = "silu" + max_position_embeddings: int = 32768 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-05 + use_cache: bool = True + tie_word_embeddings: bool = False + rope_theta: float = 1000000.0 + use_sliding_window: bool = False + sliding_window: int = 4096 + max_window_layers: int = 80 + layer_types: list = field(default_factory=list) + attention_dropout: float = 0.0 + rope_scaling: Optional[dict] = None + bos_token_id: Optional[int] = None + eos_token_id: Optional[int] = None + pad_token_id: Optional[int] = None + vision_token_id: int = 151654 + model_type: str = "qwen2_5_vl_text" + dtype: str = "bfloat16" + + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ]) + _fsdp_shard_conditions: list = field( + default_factory=lambda: + [_is_transformer_layer, _is_embeddings, _is_final_norm]) + + def __post_init__(self): + super().__post_init__() + self.sliding_window = self.sliding_window if self.use_sliding_window else None + # for backward compatibility + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + if self.layer_types is None: + self.layer_types = [ + "sliding_attention" + if self.sliding_window is not None and i >= self.max_window_layers + else "full_attention" + for i in range(self.num_hidden_layers) + ] + if self.rope_scaling is not None and "type" in self.rope_scaling: + if self.rope_scaling["type"] == "mrope": + self.rope_scaling["type"] = "default" + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + +@dataclass +class Qwen2_5_VLConfig(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=Qwen2_5_VLArchConfig) + prefix: str = "qwen2_5_vl" \ No newline at end of file diff --git a/fastvideo/models/encoders/qwen2_5.py b/fastvideo/models/encoders/qwen2_5.py new file mode 100644 index 000000000..bd7c92f72 --- /dev/null +++ b/fastvideo/models/encoders/qwen2_5.py @@ -0,0 +1,387 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from transformers: https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py + +import math +from typing import Any, Optional, Tuple, Union, List, Callable, Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.configs.models.encoders import BaseEncoderOutput, Qwen2_5_VLConfig +from fastvideo.distributed import get_tp_rank, get_tp_world_size +from fastvideo.layers.activation import get_act_fn, SiluAndMul +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.linear import MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.loader.weight_utils import default_weight_loader +from fastvideo.models.mask_utils import sdpa_mask +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + +def sdpa_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + dropout: float = 0.0, + scaling: Optional[float] = None, + is_causal: Optional[bool] = None, + **kwargs, +) -> tuple[torch.Tensor, None]: + if kwargs.get("output_attentions", False) or kwargs.get("head_mask") is not None: + logger.warning_once( + "`sdpa` attention does not support `output_attentions=True` or `head_mask`." + " Please set your attention to `eager` if you want any of these features." + ) + + if hasattr(module, "num_key_value_groups"): + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + if attention_mask is not None and attention_mask.ndim == 4: + attention_mask = attention_mask[:, :, :, : key.shape[-2]] + # If attention_mask is not None, convert it to boolean type + if attention_mask is not None and attention_mask.dtype != torch.bool: + attention_mask = attention_mask.bool() + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=dropout, + scale=scaling, + is_causal=is_causal, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, None + +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 apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): + mrope_section = [s * 2 for s in mrope_section] + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( + unsqueeze_dim + ) + sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze( + unsqueeze_dim + ) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Qwen2_5_VLRotaryEmbedding(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, device=None): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + self.config = config + + self.rope_type = config.rope_scaling.get("rope_type", "default") + self.base = config.rope_theta + + # Simplified initialization + head_dim = config.hidden_size // config.num_attention_heads + dim = head_dim + self.attention_scaling = 1.0 + + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = inv_freq + + def forward(self, x, position_ids): + # In contrast to other models, Qwen2_5_VL has different position ids for the grids + # So we expand the inv_freq to shape (3, ...) + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Qwen2_5_VLMLP(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, quant_config: QuantizationConfig | None = None, prefix: str = ""): + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + x, _ = self.gate_up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj(x) + return x + + +class Qwen2_5_VLAttention(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, quant_config: QuantizationConfig | None = None, prefix: str = ""): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + + tp_size = get_tp_world_size() + self.total_num_heads = self.num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + + self.total_num_kv_heads = self.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = 1 + + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.layer_type = config.layer_types[layer_idx] if config.layer_types else "full_attention" + self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + output_attentions: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len, _ = hidden_states.size() + + qkv, _ = self.qkv_proj(hidden_states) + query_states, key_states, value_states = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.config.rope_scaling["mrope_section"] + ) + + attn_output = sdpa_attention_forward(self, query_states, key_states, value_states, attention_mask, dropout=self.config.attention_dropout, scaling=self.scaling, is_causal=False)[0].reshape(bsz, q_len, -1) + attn_output, _ = self.o_proj(attn_output) + + return attn_output + +class Qwen2_5_VLDecoderLayer(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, quant_config: QuantizationConfig | None = None, prefix: str = ""): + super().__init__() + self.self_attn = Qwen2_5_VLAttention(config, layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn") + self.mlp = Qwen2_5_VLMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + position_embeddings=position_embeddings, + output_attentions=output_attentions, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + return outputs + +class Qwen2_5_VLTextModel(TextEncoder): + def __init__(self, config: Qwen2_5_VLConfig): + super().__init__(config) + + quant_config = None + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size + ) + + self.layers = nn.ModuleList([ + Qwen2_5_VLDecoderLayer(config, layer_idx, quant_config=quant_config, prefix=f"{config.prefix}.layers.{layer_idx}") + for layer_idx in range(config.num_hidden_layers) + ]) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + output_hidden_states: Optional[bool] = None, + **kwargs, + ) -> BaseEncoderOutput: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings(input_ids) + + hidden_states = inputs_embeds + + if position_ids is None: + seq_length = hidden_states.shape[1] + cache_position = torch.arange(seq_length, device=hidden_states.device) + position_ids = cache_position.view(1, 1, -1).expand(3, hidden_states.shape[0], -1) + + mask_kwargs = { + "batch_size": hidden_states.shape[0], + "cache_position": cache_position, + "kv_length": attention_mask.shape[-1], + "kv_offset": 0, + "attention_mask": attention_mask + } + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + all_hidden_states = () if output_hidden_states else None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=sdpa_mask(**mask_kwargs), + position_ids=position_ids, + position_embeddings=position_embeddings, + output_attentions=False, + ) + hidden_states = layer_outputs[0] + + hidden_states = self.norm(hidden_states) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + return BaseEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in self.config.stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + # if name.endswith(".bias") and name not in params_dict: + # continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + # if name.endswith(".bias") and name not in params_dict: + # continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/fastvideo/models/mask_utils.py b/fastvideo/models/mask_utils.py new file mode 100644 index 000000000..ec944c8d9 --- /dev/null +++ b/fastvideo/models/mask_utils.py @@ -0,0 +1,201 @@ +import torch +from typing import Callable, Optional + +def and_masks(*mask_functions: Callable) -> Callable: + """Returns a mask function that is the intersection of provided mask functions""" + if not all(callable(arg) for arg in mask_functions): + raise RuntimeError(f"All inputs should be callable mask_functions: {mask_functions}") + + def and_mask(batch_idx, head_idx, q_idx, kv_idx): + result = q_idx.new_ones((), dtype=torch.bool) + for mask in mask_functions: + result = result & mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device) + return result + + return and_mask + +def causal_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + """ + This creates a basic lower-diagonal causal mask. + """ + return kv_idx <= q_idx + +def padding_mask_function(padding_mask: torch.Tensor) -> Callable: + """ + This return the mask_function function corresponding to a 2D padding mask. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + # Note that here the mask should ALWAYS be at least of the max `kv_index` size in the dimension 1. This is because + # we cannot pad it here in the mask_function as we don't know the final size, and we cannot try/except, as it is not + # vectorizable on accelerator devices + return padding_mask[batch_idx, kv_idx] + + return inner_mask + +def prepare_padding_mask( + attention_mask: Optional[torch.Tensor], kv_length: int, kv_offset: int +) -> Optional[torch.Tensor]: + """ + From the 2D attention mask, prepare the correct padding mask to use by potentially padding it. + """ + local_padding_mask = attention_mask + if attention_mask is not None: + # Pad it if necessary + if (padding_length := kv_length + kv_offset - attention_mask.shape[-1]) > 0: + local_padding_mask = torch.nn.functional.pad(attention_mask, (0, padding_length)) + return local_padding_mask + +def _non_vmap_expansion_sdpa( + batch_indices: torch.Tensor, head_indices: torch.Tensor, q_indices: torch.Tensor, kv_indices: torch.Tensor +): + """ + Used to broadcast our mask_functions over the all 4 dimensions (b_idx, h_idx, q_idx, kv_idx) of the inputs. + Allows the usage of any index-based mask function without relying on vmap. + + NOTE: This is limited to index based functions only and is not guaranteed to work otherwise. + + Reference: + - https://github.com/huggingface/optimum-onnx/blob/c123e8f4fab61b54a8e0e31ce74462bcacca576e/optimum/exporters/onnx/model_patcher.py#L362-L365 + """ + batch_indices = batch_indices[:, None, None, None] + head_indices = head_indices[None, :, None, None] + q_indices = q_indices[None, None, :, None] + kv_indices = kv_indices[None, None, None, :] + return batch_indices, head_indices, q_indices, kv_indices + +def sdpa_mask( + batch_size: int, + cache_position: torch.Tensor, + kv_length: int, + kv_offset: int = 0, + mask_function: Callable = causal_mask_function, + attention_mask: Optional[torch.Tensor] = None, + local_size: Optional[int] = None, + allow_is_causal_skip: bool = True, + allow_is_bidirectional_skip: bool = False, + allow_torch_fix: bool = True, + use_vmap: bool = False, + **kwargs, +) -> Optional[torch.Tensor]: + """ + Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that + the element should take part in the attention computation, and False that it should not. + This function can only be used with torch>=2.5, as the context manager is otherwise not available. + + Args: + batch_size (`int`): + The batch size of the input sequence. + cache_position (`torch.Tensor`): + A tensor of shape (query_length,) indicating the current indices of the input sequence elements. + kv_length (`int`): + The size that the key and value states will have during the attention computation. + kv_offset (`int`, optional): + An optional offset to indicate at which first position the key and values states will refer to. + mask_function (`Callable`): + The mask factory function describing the mask pattern. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length) + local_size (`int`, optional): + The size of the local attention, if we do not use full attention. This is used only if `allow_is_causal_skip=True` + to try to skip mask creation if possible. + allow_is_causal_skip (`bool`, optional): + Whether to allow to return `None` for the mask under conditions where we can use the `is_causal` argument in + `torch.sdpa` instead. Default to `True`. + allow_is_bidirectional_skip (`bool`, optional): + Whether to allow to return `None` for the mask under conditions where we do not have to add any bias, + i.e. full attention without any padding. Default to `False`. + allow_torch_fix (`bool`, optional): + Whether to update the mask in case a query is not attending to any tokens, to solve a bug in torch's older + versions. We need an arg to skip it when using eager. By default `True`. + use_vmap (`bool`, optional): + Whether to use `vmap` during the mask construction or not. Allows powerful custom patterns that may not be + index-based (for the cost of speed performance). By default `False`. + + + ## Creating a simple causal mask: + + To create the following causal mask: + + 0 ■ ⬚ ⬚ ⬚ ⬚ + 1 ■ ■ ⬚ ⬚ ⬚ + 2 ■ ■ ■ ⬚ ⬚ + 3 ■ ■ ■ ■ ⬚ + 4 ■ ■ ■ ■ ■ + + You can do + + ```python + >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5) + >>> tensor([[[[ True, False, False, False, False], + [ True, True, False, False, False], + [ True, True, True, False, False], + [ True, True, True, True, False], + [ True, True, True, True, True]]]]) + ``` + + ## Creating a sliding window mask: + + To create the following sliding window mask (`sliding_window=3`): + + 0 ■ ⬚ ⬚ ⬚ ⬚ + 1 ■ ■ ⬚ ⬚ ⬚ + 2 ■ ■ ■ ⬚ ⬚ + 3 ⬚ ■ ■ ■ ⬚ + 4 ⬚ ⬚ ■ ■ ■ + + You can do + + ```python + >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5, mask_function=sliding_window_causal_mask_function(3)) + >>> tensor([[[[ True, False, False, False, False], + [ True, True, False, False, False], + [ True, True, True, False, False], + [False, True, True, True, False], + [False, False, True, True, True]]]]) + ``` + + ## Creating a chunked attention mask + + To create the following chunked attention mask (`chunk_size=3`): + + 0 ■ ⬚ ⬚ ⬚ ⬚ + 1 ■ ■ ⬚ ⬚ ⬚ + 2 ■ ■ ■ ⬚ ⬚ + 3 ⬚ ⬚ ⬚ ■ ⬚ + 4 ⬚ ⬚ ⬚ ■ ■ + + You can do + + ```python + >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5, mask_function=chunked_causal_mask_function(3, torch.zeros(1, dtype=int))) + >>> tensor([[[[ True, False, False, False, False], + [ True, True, False, False, False], + [ True, True, True, False, False], + [False, False, False, True, False], + [False, False, False, True, True]]]]) + ``` + + """ + q_length = cache_position.shape[0] + + # Potentially pad the 2D mask + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) + + # Potentially add the padding 2D mask + if padding_mask is not None: + mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) + + batch_arange = torch.arange(batch_size, device=cache_position.device) + head_arange = torch.arange(1, device=cache_position.device) + # Similar to `kv_arange = torch.arange(start=kv_offset, end=kv_offset + kv_length, device=cache_position.device)` + # but without data-dependent slicing (i.e. torch.compile friendly) + kv_arange = torch.arange(kv_length, device=cache_position.device) + kv_offset + + # Actual mask creation + # Apply mask function element-wise through broadcasting + attention_mask = mask_function(*_non_vmap_expansion_sdpa(batch_arange, head_arange, cache_position, kv_arange)) + # Expand the mask to match batch size and query length if they weren't used in the mask function + attention_mask = attention_mask.expand(batch_size, -1, q_length, kv_length) + + return attention_mask \ No newline at end of file diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 04eb6936d..290744400 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -45,6 +45,7 @@ _TEXT_ENCODER_MODELS = { "T5EncoderModel": ("encoders", "t5", "T5EncoderModel"), "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"), "BertModel": ("encoders", "clip", "CLIPTextModel"), + "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"), } _IMAGE_ENCODER_MODELS: dict[str, tuple] = { diff --git a/fastvideo/tests/encoders/test_qwen2_5_encoder.py b/fastvideo/tests/encoders/test_qwen2_5_encoder.py new file mode 100644 index 000000000..2b10c74a7 --- /dev/null +++ b/fastvideo/tests/encoders/test_qwen2_5_encoder.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +import pytest +import torch +from torch.distributed.tensor import DTensor +from torch.testing import assert_close +from transformers import AutoConfig, AutoTokenizer, Qwen2_5_VLTextModel + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TextEncoderLoader +from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.configs.models.encoders import Qwen2_5_VLConfig + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29505" + +@pytest.fixture +def qwen_model_path(): + base_model_path = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + model_path = maybe_download_model(base_model_path) + text_encoder_path = os.path.join(model_path, "text_encoder") + tokenizer_path = os.path.join(model_path, "tokenizer") + return text_encoder_path, tokenizer_path + +@pytest.mark.usefixtures("distributed_setup") +def test_qwen2_5_encoder(qwen_model_path): + text_encoder_path, tokenizer_path = qwen_model_path + hf_config = AutoConfig.from_pretrained(text_encoder_path) + print(hf_config) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + # Qwen2.5-VL default dtype is usually bf16 + precision_str = "bf16" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "fp16" + precision = PRECISION_TO_TYPE[precision_str] + + logger.info(f"Using precision: {precision_str}") + + # Load HF model (Base model) + model1 = Qwen2_5_VLTextModel.from_pretrained(text_encoder_path).to( + precision).to(device).eval() + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + + # Load FastVideo model + args = FastVideoArgs(model_path=text_encoder_path, + pipeline_config=PipelineConfig(text_encoder_configs=(Qwen2_5_VLConfig(),), + text_encoder_precisions=(precision_str,)), + pin_cpu_memory=False) + + loader = TextEncoderLoader() + model2 = loader.load(text_encoder_path, args) + model2 = model2.to(precision) + model2.eval() + + # Sanity check weights + logger.info("Comparing model weights for sanity check...") + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + logger.info("Model1 has %s parameters", len(params1)) + logger.info("Model2 has %s parameters", len(params2)) + + # Check common layers like Norms which are likely not merged/sharded in a way that changes name significantly + # or simple linear layers if names match. + # Note: FastVideo uses QKVParallelLinear, so q_proj, k_proj, v_proj are merged. + # HF Qwen2_5_VL uses separate projections? No, usually they are separate nn.Linear in HF. + + weights_to_check = [ + "norm.weight", + "layers.{}.self_attn.o_proj.weight", + "layers.{}.input_layernorm.weight", + "layers.{}.post_attention_layernorm.weight", + "layers.{}.mlp.down_proj.weight" + ] + + for idx in range(hf_config.num_hidden_layers): + for w in weights_to_check: + name1 = w.format(idx) + name2 = w.format(idx) + p1 = params1[name1] + p2 = params2[name2] + p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1) + + # Check shape + assert p1.shape == p2.shape, f"Shape mismatch for {w}: {p1.shape} vs {p2.shape}" + + # Check values + assert_close(p1, p2, atol=1e-7, rtol=1e-7, msg=f"Weight mismatch for {w}") + + # Test with sample prompts + prompts = [ + "Hello world", + "The quick brown fox jumps over the lazy dog." + ] + + logger.info("Testing with sample prompts") + + with torch.no_grad(): + for prompt in prompts: + logger.info(f"Prompt: {prompt}") + tokens = tokenizer(prompt, return_tensors="pt", padding="max_length", max_length=1000, truncation=True).to(device) + + # HF Forward + # AutoModel for Qwen2.5-VL usually returns BaseModelOutputWithPast + outputs1 = model1(**tokens).last_hidden_state + + # FastVideo Forward + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs2 = model2( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask + ).last_hidden_state + + # Compare + # Filter padding for comparison if needed, but here we just check raw output matching + + # Check shapes + assert outputs1.shape == outputs2.shape, f"Output shape mismatch: {outputs1.shape} vs {outputs2.shape}" + + diff = torch.abs(outputs1 - outputs2) + max_diff = diff.max().item() + mean_diff = diff.mean().item() + + logger.info(f"Max diff: {max_diff}") + logger.info(f"Mean diff: {mean_diff}") + + # Thresholds + # Qwen2.5-VL RoPE is complex, if our implementation is slightly off (e.g. float32 conversion logic in RoPE), + # differences might appear. But should be small. + if precision_str == "bf16": + atol = 5e-2 # relaxed for bf16 + else: + atol = 1e-3 + + if max_diff > atol: + logger.warning(f"Max diff {max_diff} > {atol}. Checking if it's acceptable...") + # If mean diff is small, maybe just outliers + assert mean_diff < atol, f"Mean diff {mean_diff} too high" + else: + logger.info("Outputs match within tolerance.") +