| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Adapted from |
| 5 | # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py |
| 6 | # https://huggingface.co/Qwen/Qwen-7B/blob/main/modeling_qwen.py |
| 7 | # https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20 |
| 8 | # |
| 9 | # Copyright 2023 The Qwen team. |
| 10 | # Copyright 2023 The vLLM team. |
| 11 | # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. |
| 12 | # |
| 13 | # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX |
| 14 | # and OPT implementations in this library. It has been modified from its |
| 15 | # original forms to accommodate minor architectural differences compared |
| 16 | # to GPT-NeoX and OPT used by the Meta AI team that trained the model. |
| 17 | # |
| 18 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 19 | # you may not use this file except in compliance with the License. |
| 20 | # You may obtain a copy of the License at |
| 21 | # |
| 22 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 23 | # |
| 24 | # Unless required by applicable law or agreed to in writing, software |
| 25 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 26 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 27 | # See the License for the specific language governing permissions and |
| 28 | # limitations under the License. |
| 29 | """Shared resampler perceiver network used in multimodal models and |
| 30 | related helpers for sincos positional embeddings. |
| 31 | |
| 32 | Example models: Qwen (Qwen-VL), MiniCPM-V 2.0 |
| 33 | """ |
| 34 | |
| 35 | import math |
| 36 | from collections.abc import Callable |
| 37 | from functools import partial |
| 38 | |
| 39 | import numpy as np |
| 40 | import torch |
| 41 | import torch.nn.functional as F |
| 42 | from torch import nn |
| 43 | |
| 44 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 45 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 46 | |
| 47 | DEFAULT_LN = partial(nn.LayerNorm, eps=1e-6) |
| 48 | |
| 49 | |
| 50 | def get_abs_pos(abs_pos: torch.Tensor, tgt_size: torch.Tensor | int) -> torch.Tensor: |
| 51 | # abs_pos: L, C |
| 52 | # tgt_size: (H, W) |
| 53 | # return: M, C |
| 54 | src_size = int(math.sqrt(abs_pos.size(0))) |
| 55 | dtype = abs_pos.dtype |
| 56 | if isinstance(tgt_size, int): |
| 57 | tgt_size = (tgt_size, tgt_size) |
| 58 | if src_size == tgt_size[0] and src_size == tgt_size[1]: |
| 59 | return abs_pos |
| 60 | return ( |
| 61 | F.interpolate( |
| 62 | abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2), |
| 63 | size=(tgt_size[0], tgt_size[1]), |
| 64 | mode="bicubic", |
| 65 | align_corners=False, |
| 66 | ) |
| 67 | .permute(0, 2, 3, 1) |
| 68 | .flatten(0, 2) |
| 69 | .to(dtype=dtype) |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | # sin/cos positional embedding helpers are adapted from: |
| 74 | # https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20 |
| 75 | def get_1d_sincos_pos_embed_from_grid( |
| 76 | embed_dim: int, pos: np.ndarray, version: tuple[int, int] = (2, 0) |
| 77 | ) -> torch.Tensor: |
| 78 | """embed_dim: output dimension for each position |
| 79 | pos: a list of positions to be encoded: size (M,) / (H, W) |
| 80 | out: (M, D) / (H, W, D) |
| 81 | """ |
| 82 | assert embed_dim % 2 == 0 |
| 83 | omega = np.arange(embed_dim // 2, dtype=np.float32) |
| 84 | omega /= embed_dim / 2.0 |
| 85 | omega = 1.0 / 10000**omega # (D/2,) |
| 86 | |
| 87 | if version == (2, 0): |
| 88 | pos = pos.reshape(-1) # (M,) |
| 89 | out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product |
| 90 | emb_sin = np.sin(out) # (M, D/2) |
| 91 | emb_cos = np.cos(out) # (M, D/2) |
| 92 | emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) |
| 93 | else: |
| 94 | out = np.einsum("hw,d->hwd", pos, omega) # (H, W, D/2), outer product |
| 95 | emb_sin = np.sin(out) # (H, W, D/2) |
| 96 | emb_cos = np.cos(out) # (H, W, D/2) |
| 97 | emb = np.concatenate([emb_sin, emb_cos], axis=-1) # (H, W, D) |
| 98 | return emb |
| 99 | |
| 100 | |
| 101 | def get_2d_sincos_pos_embed_from_grid( |
| 102 | embed_dim: int, grid: np.ndarray, version: tuple[int, int] = (2, 0) |
| 103 | ) -> torch.Tensor: |
| 104 | assert embed_dim % 2 == 0 |
| 105 | |
| 106 | # use half of dimensions to encode grid_h |
| 107 | emb_h = get_1d_sincos_pos_embed_from_grid( |
| 108 | embed_dim // 2, grid[0], version |
| 109 | ) # (H*W, D/2) or (H, W, D/2) |
| 110 | emb_w = get_1d_sincos_pos_embed_from_grid( |
| 111 | embed_dim // 2, grid[1], version |
| 112 | ) # (H*W, D/2) or (H, W, D/2) |
| 113 | |
| 114 | if version == (2, 0): |
| 115 | emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) |
| 116 | else: |
| 117 | emb = np.concatenate([emb_h, emb_w], axis=-1) # (H, W, D) |
| 118 | return emb |
| 119 | |
| 120 | |
| 121 | def get_2d_sincos_pos_embed( |
| 122 | embed_dim: int, |
| 123 | grid_size: int | tuple[int, int], |
| 124 | cls_token: bool = False, |
| 125 | version: tuple[int, int] = (2, 0), |
| 126 | ) -> torch.Tensor: |
| 127 | """grid_size: int of the grid height and width |
| 128 | return: |
| 129 | pos_embed: [grid_size*grid_size, embed_dim] or |
| 130 | [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) |
| 131 | """ |
| 132 | if isinstance(grid_size, int): |
| 133 | grid_h_size, grid_w_size = grid_size, grid_size |
| 134 | else: |
| 135 | grid_h_size, grid_w_size = grid_size[0], grid_size[1] |
| 136 | |
| 137 | grid_h = np.arange(grid_h_size, dtype=np.float32) |
| 138 | grid_w = np.arange(grid_w_size, dtype=np.float32) |
| 139 | grid = np.meshgrid(grid_w, grid_h) # here w goes first |
| 140 | grid = np.stack(grid, axis=0) |
| 141 | assert isinstance(grid, np.ndarray) and grid.shape == (2, grid_h_size, grid_w_size) |
| 142 | |
| 143 | if version == (2, 0): |
| 144 | grid = grid.reshape([2, 1, grid_h_size, grid_w_size]) |
| 145 | pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version) |
| 146 | if cls_token: |
| 147 | pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) |
| 148 | else: |
| 149 | pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version) |
| 150 | return pos_embed |
| 151 | |
| 152 | |
| 153 | class BaseResampler(nn.Module): |
| 154 | """A 2D perceiver-resampler network with one cross attention layers by |
| 155 | (grid_size**2) learnable queries and 2d sincos pos_emb. |
| 156 | Outputs: |
| 157 | A tensor with the shape of (grid_size**2, embed_dim) |
| 158 | """ |
| 159 | |
| 160 | def __init__( |
| 161 | self, |
| 162 | num_queries: int, |
| 163 | embed_dim: int, |
| 164 | num_heads: int, |
| 165 | kv_dim: int | None = None, |
| 166 | norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN, |
| 167 | do_post_projection: bool = True, |
| 168 | quant_config: QuantizationConfig | None = None, |
| 169 | prefix: str = "", |
| 170 | ) -> None: |
| 171 | super().__init__() |
| 172 | |
| 173 | self.num_queries = num_queries |
| 174 | self.embed_dim = embed_dim |
| 175 | self.num_heads = num_heads |
| 176 | |
| 177 | self.query = nn.Parameter(torch.empty(self.num_queries, embed_dim)) |
| 178 | |
| 179 | if kv_dim is not None and kv_dim != embed_dim: |
| 180 | self.kv_proj = ReplicatedLinear( |
| 181 | kv_dim, |
| 182 | embed_dim, |
| 183 | bias=False, |
| 184 | quant_config=quant_config, |
| 185 | prefix=f"{prefix}.kv_proj", |
| 186 | ) |
| 187 | else: |
| 188 | # Maintain the same return value with ReplicatedLinear.forward |
| 189 | self.kv_proj = lambda *args, **kwargs: ( # type: ignore # noqa |
| 190 | nn.Identity()(*args, **kwargs), |
| 191 | None, |
| 192 | ) |
| 193 | self.attn = nn.MultiheadAttention(embed_dim, num_heads) |
| 194 | self.ln_q = norm_layer(embed_dim) |
| 195 | self.ln_kv = norm_layer(embed_dim) |
| 196 | self.do_post_projection = do_post_projection |
| 197 | if self.do_post_projection: |
| 198 | self.ln_post = norm_layer(embed_dim) |
| 199 | data = (embed_dim**-0.5) * torch.empty(embed_dim, embed_dim) |
| 200 | self.proj = nn.Parameter(data=data) |
| 201 | |
| 202 | def _repeat(self, query, N: int): |
| 203 | return query.unsqueeze(1).repeat(1, N, 1) |
| 204 | |
| 205 | |
| 206 | class Resampler2(BaseResampler): |
| 207 | """Resampler-perceiver network to be used for a variety of model types, |
| 208 | e.g., Qwen-vl / Minicpmv 2.0. The main difference is the addition of the |
| 209 | do_post_projection arg, which indicates whether or not there should be |
| 210 | a post layer normalization and projector after the attention. This is |
| 211 | present in minicpmv2.0, but not qwen-vl. |
| 212 | """ |
| 213 | |
| 214 | def __init__( |
| 215 | self, |
| 216 | grid_size: int, |
| 217 | embed_dim: int, |
| 218 | num_heads: int, |
| 219 | kv_dim: int | None = None, |
| 220 | norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN, |
| 221 | adaptive: bool = False, |
| 222 | do_post_projection: bool = True, |
| 223 | quant_config: QuantizationConfig | None = None, |
| 224 | prefix: str = "", |
| 225 | ) -> None: |
| 226 | super().__init__( |
| 227 | grid_size**2, |
| 228 | embed_dim, |
| 229 | num_heads, |
| 230 | kv_dim, |
| 231 | norm_layer, |
| 232 | do_post_projection=do_post_projection, |
| 233 | quant_config=quant_config, |
| 234 | prefix=prefix, |
| 235 | ) |
| 236 | |
| 237 | self.adaptive = adaptive |
| 238 | pos_embed_arr = get_2d_sincos_pos_embed(embed_dim, grid_size, version=(2, 0)) |
| 239 | |
| 240 | self.pos_embed = nn.Parameter( |
| 241 | torch.from_numpy(pos_embed_arr).requires_grad_(False) |
| 242 | ) |
| 243 | |
| 244 | def forward( |
| 245 | self, |
| 246 | x: torch.Tensor, |
| 247 | tgt_sizes: torch.Tensor | None = None, |
| 248 | attn_mask: torch.Tensor | None = None, |
| 249 | ) -> torch.Tensor: |
| 250 | if tgt_sizes is None: |
| 251 | tgt_sizes = int(math.sqrt(x.size(1))) |
| 252 | if self.adaptive: |
| 253 | pos_embed_arr = get_2d_sincos_pos_embed( |
| 254 | self.embed_dim, tgt_sizes, version=(2, 0) |
| 255 | ) |
| 256 | pos_embed = torch.from_numpy(pos_embed_arr).to( |
| 257 | device=x.device, dtype=x.dtype |
| 258 | ) |
| 259 | else: |
| 260 | pos_embed = get_abs_pos(self.pos_embed, tgt_sizes).to( |
| 261 | device=x.device, dtype=x.dtype |
| 262 | ) |
| 263 | |
| 264 | x, _ = self.kv_proj(x) |
| 265 | x = self.ln_kv(x).permute(1, 0, 2) |
| 266 | |
| 267 | N = x.shape[1] |
| 268 | q = self.ln_q(self.query) |
| 269 | out = self.attn( |
| 270 | self._repeat(q, N) + self.pos_embed.unsqueeze(1), |
| 271 | x + pos_embed.unsqueeze(1), |
| 272 | x, |
| 273 | attn_mask=attn_mask, |
| 274 | )[0] |
| 275 | x = out.permute(1, 0, 2) |
| 276 | if self.do_post_projection: |
| 277 | x = self.ln_post(x) |
| 278 | x = x @ self.proj |
| 279 | return x |