xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 08:04:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Remove xformers requirement for Pixtral (#9597)

Browse Source 
Signed-off-by: mgoin <michael@neuralmagic.com>
This commit is contained in:
Michael Goin 2024-10-24 18:38:05 -04:00 committed by GitHub
parent 59449095ab
commit c91ed47c43
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 46 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
vllm/model_executor/models/pixtral.py
View File

@ -14,8 +14,6 @@ from transformers.models.pixtral.image_processing_pixtral import (
_num_image_tokens)
from transformers.models.pixtral.modeling_pixtral import (
PixtralRotaryEmbedding, apply_rotary_pos_emb, position_ids_in_meshgrid)
from xformers.ops.fmha import memory_efficient_attention
from xformers.ops.fmha.attn_bias import BlockDiagonalMask
from vllm.attention import AttentionMetadata
from vllm.config import CacheConfig, ModelConfig, MultiModalConfig
@ -38,6 +36,12 @@ from vllm.utils import is_list_of
from .interfaces import SupportsMultiModal, SupportsPP
from .utils import init_vllm_registered_model
try:
from xformers import ops as xops
USE_XFORMERS_OPS = True
except ImportError:
USE_XFORMERS_OPS = False
def get_max_pixtral_image_tokens(ctx: InputContext):
tokenizer = cached_get_tokenizer(
@ -416,7 +420,7 @@ class Attention(nn.Module):
def forward(
self,
x: torch.Tensor,
mask: BlockDiagonalMask,
mask: torch.Tensor,
freqs_cis: torch.Tensor,
) -> torch.Tensor:
batch, patches, _ = x.shape
@ -427,7 +431,7 @@ class Attention(nn.Module):
v = v.reshape(batch, patches, self.n_heads, self.head_dim)
q, k = apply_rotary_emb_vit(q, k, freqs_cis=freqs_cis)
out = memory_efficient_attention(q, k, v, attn_bias=mask)
out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)
out = out.reshape(batch, patches, self.n_heads * self.head_dim)
return self.wo(out)
@ -444,7 +448,7 @@ class TransformerBlock(nn.Module):
def forward(
self,
x: torch.Tensor,
mask: BlockDiagonalMask,
mask: torch.Tensor,
freqs_cis: torch.Tensor,
) -> torch.Tensor:
r = self.attention.forward(self.attention_norm(x),
@ -467,7 +471,7 @@ class Transformer(nn.Module):
def forward(
self,
x: torch.Tensor,
mask: BlockDiagonalMask,
mask: torch.Tensor,
freqs_cis: Optional[torch.Tensor],
) -> torch.Tensor:
for layer in self.layers:
@ -562,8 +566,12 @@ class VisionTransformer(nn.Module):
freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]
# pass through Transformer with a block diagonal mask delimiting images
mask = BlockDiagonalMask.from_seqlens(
[p.shape[-2] * p.shape[-1] for p in patch_embeds_list], )
if USE_XFORMERS_OPS:
mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(
[p.shape[-2] * p.shape[-1] for p in patch_embeds_list], )
else:
raise ImportError("Xformers is required for Pixtral inference "
"with the Mistral format")
out = self.transformer(patch_embeds, mask=mask, freqs_cis=freqs_cis)
# remove batch dimension of the single sequence
@ -828,7 +836,7 @@ class PixtralHFAttention(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: BlockDiagonalMask,
attention_mask: torch.Tensor,
position_embeddings: torch.Tensor,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
batch, patches, _ = hidden_states.size()
@ -843,12 +851,23 @@ class PixtralHFAttention(nn.Module):
cos, sin = position_embeddings
q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=0)
# Transpose q and k back for attention
q = q.transpose(1, 2).contiguous()
k = k.transpose(1, 2).contiguous()
v = v.reshape(batch, patches, self.n_heads, self.head_dim)
if USE_XFORMERS_OPS:
# Transpose q and k back for attention
q = q.transpose(1, 2).contiguous()
k = k.transpose(1, 2).contiguous()
v = v.reshape(batch, patches, self.n_heads, self.head_dim)
out = xops.memory_efficient_attention(q,
k,
v,
attn_bias=attention_mask)
else:
v = v.reshape(batch, patches, self.n_heads,
self.head_dim).transpose(1, 2)
out = nn.functional.scaled_dot_product_attention(
q, k, v, attn_mask=attention_mask)
out = out.transpose(1, 2)
out = memory_efficient_attention(q, k, v, attn_bias=attention_mask)
out = out.reshape(batch, patches, self.n_heads * self.head_dim)
return self.o_proj(out)
@ -877,7 +896,7 @@ class PixtralHFTransformerBlock(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: BlockDiagonalMask,
attention_mask: torch.Tensor,
position_embeddings: torch.Tensor,
) -> torch.Tensor:
r = self.attention.forward(self.attention_norm(hidden_states),
@ -916,7 +935,7 @@ class PixtralHFTransformer(nn.Module):
def forward(
self,
x: torch.Tensor,
attention_mask: BlockDiagonalMask,
attention_mask: torch.Tensor,
position_embeddings: torch.Tensor,
) -> torch.Tensor:
for layer in self.layers:
@ -1000,11 +1019,19 @@ class PixtralHFVisionModel(nn.Module):
patch_embeds_list,
max_width=self.config.image_size // self.config.patch_size).to(
self.device)
position_embedding = self.patch_positional_embedding(
patch_embeds, position_ids)
attention_mask = BlockDiagonalMask.from_seqlens(
[p.shape[-2] * p.shape[-1] for p in patch_embeds_list], )
if USE_XFORMERS_OPS:
attention_mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(
[p.shape[-2] * p.shape[-1] for p in patch_embeds_list], )
else:
from transformers.models.pixtral.modeling_pixtral import (
generate_block_attention_mask)
attention_mask = generate_block_attention_mask(
[p.shape[-2] * p.shape[-1] for p in patch_embeds_list],
patch_embeds)
out = self.transformer(patch_embeds, attention_mask,
position_embedding)