Skip to content

vllm.model_executor.models.aimv2

AIMv2Attention

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
 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
class AIMv2Attention(nn.Module):

    def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
                 prefix: str):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                "embed_dim must be divisible by num_heads "
                f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
                f" {self.num_heads}).")
        self.scale = self.head_dim**-0.5

        self.qkv = QKVParallelLinear(
            hidden_size=self.embed_dim,
            head_size=self.head_dim,
            total_num_heads=self.num_heads,
            bias=config.qkv_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv",
        )

        self.proj = RowParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            bias=config.use_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.proj",
        )

        self.tp_size = get_tensor_model_parallel_world_size()
        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

        self.attn = MultiHeadAttention(self.num_heads_per_partition,
                                       self.head_dim, self.scale)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        qkv, _ = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)

        x = self.attn(q, k, v)
        x, _ = self.proj(x)
        return x

attn instance-attribute 

attn = MultiHeadAttention(
    num_heads_per_partition, head_dim, scale
)

config instance-attribute 

config = config

embed_dim instance-attribute 

embed_dim = hidden_size

head_dim instance-attribute 

head_dim = embed_dim // num_heads

num_heads instance-attribute 

num_heads = num_attention_heads

num_heads_per_partition instance-attribute 

num_heads_per_partition = divide(num_heads, tp_size)

proj instance-attribute 

proj = RowParallelLinear(
    input_size=embed_dim,
    output_size=embed_dim,
    bias=use_bias,
    quant_config=quant_config,
    prefix=f"{prefix}.proj",
)

qkv instance-attribute 

qkv = QKVParallelLinear(
    hidden_size=embed_dim,
    head_size=head_dim,
    total_num_heads=num_heads,
    bias=qkv_bias,
    quant_config=quant_config,
    prefix=f"{prefix}.qkv",
)

scale instance-attribute 

scale = head_dim ** -0.5

tp_size instance-attribute 

tp_size = get_tensor_model_parallel_world_size()

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    prefix: str,
)
Source code in vllm/model_executor/models/aimv2.py 
 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
def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
             prefix: str):
    super().__init__()
    self.config = config
    self.embed_dim = config.hidden_size
    self.num_heads = config.num_attention_heads
    self.head_dim = self.embed_dim // self.num_heads
    if self.head_dim * self.num_heads != self.embed_dim:
        raise ValueError(
            "embed_dim must be divisible by num_heads "
            f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
            f" {self.num_heads}).")
    self.scale = self.head_dim**-0.5

    self.qkv = QKVParallelLinear(
        hidden_size=self.embed_dim,
        head_size=self.head_dim,
        total_num_heads=self.num_heads,
        bias=config.qkv_bias,
        quant_config=quant_config,
        prefix=f"{prefix}.qkv",
    )

    self.proj = RowParallelLinear(
        input_size=self.embed_dim,
        output_size=self.embed_dim,
        bias=config.use_bias,
        quant_config=quant_config,
        prefix=f"{prefix}.proj",
    )

    self.tp_size = get_tensor_model_parallel_world_size()
    self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

    self.attn = MultiHeadAttention(self.num_heads_per_partition,
                                   self.head_dim, self.scale)

forward 

forward(x: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
133
134
135
136
137
138
139
def forward(self, x: torch.Tensor) -> torch.Tensor:
    qkv, _ = self.qkv(x)
    q, k, v = qkv.chunk(3, dim=-1)

    x = self.attn(q, k, v)
    x, _ = self.proj(x)
    return x

AIMv2Block

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class AIMv2Block(nn.Module):

    def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
                 prefix: str):
        super().__init__()
        self.attn = AIMv2Attention(config,
                                   quant_config=quant_config,
                                   prefix=f"{prefix}.attn")
        self.norm_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = AIMv2SwiGLUFFN(config,
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.mlp")
        self.norm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm_1.forward_native(x))
        x = x + self.mlp(self.norm_2.forward_native(x))
        return x

attn instance-attribute 

attn = AIMv2Attention(
    config,
    quant_config=quant_config,
    prefix=f"{prefix}.attn",
)

mlp instance-attribute 

mlp = AIMv2SwiGLUFFN(
    config,
    quant_config=quant_config,
    prefix=f"{prefix}.mlp",
)

norm_1 instance-attribute 

norm_1 = RMSNorm(hidden_size, eps=rms_norm_eps)

norm_2 instance-attribute 

norm_2 = RMSNorm(hidden_size, eps=rms_norm_eps)

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    prefix: str,
)
Source code in vllm/model_executor/models/aimv2.py 
144
145
146
147
148
149
150
151
152
153
154
def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
             prefix: str):
    super().__init__()
    self.attn = AIMv2Attention(config,
                               quant_config=quant_config,
                               prefix=f"{prefix}.attn")
    self.norm_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
    self.mlp = AIMv2SwiGLUFFN(config,
                              quant_config=quant_config,
                              prefix=f"{prefix}.mlp")
    self.norm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

forward 

forward(x: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
156
157
158
159
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x = x + self.attn(self.norm_1.forward_native(x))
    x = x + self.mlp(self.norm_2.forward_native(x))
    return x

AIMv2Model

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
class AIMv2Model(torch.nn.Module):

    def __init__(self,
                 config: AIMv2Config,
                 quant_config: QuantizationConfig,
                 *,
                 require_post_norm: Optional[bool] = None,
                 prefix: str = ""):
        super().__init__()
        self.preprocessor = AIMv2ViTPreprocessor(config)
        self.trunk = AIMv2Transformer(config,
                                      quant_config=quant_config,
                                      require_post_norm=require_post_norm,
                                      prefix=f"{prefix}.trunk")

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:

        x = self.preprocessor(pixel_values)
        x = self.trunk(x)

        return x

    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".fc13", ".fc1", 0),
            (".fc13", ".fc3", 1),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()

        for name, loaded_weight in weights:
            # post_layernorm is optional in SiglipVisionModel
            if (name.startswith("trunk.post_trunk_norm")
                    and self.trunk.post_trunk_norm is None):
                continue

            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                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

preprocessor instance-attribute 

preprocessor = AIMv2ViTPreprocessor(config)

trunk instance-attribute 

trunk = AIMv2Transformer(
    config,
    quant_config=quant_config,
    require_post_norm=require_post_norm,
    prefix=f"{prefix}.trunk",
)

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    *,
    require_post_norm: Optional[bool] = None,
    prefix: str = "",
)
Source code in vllm/model_executor/models/aimv2.py 
195
196
197
198
199
200
201
202
203
204
205
206
def __init__(self,
             config: AIMv2Config,
             quant_config: QuantizationConfig,
             *,
             require_post_norm: Optional[bool] = None,
             prefix: str = ""):
    super().__init__()
    self.preprocessor = AIMv2ViTPreprocessor(config)
    self.trunk = AIMv2Transformer(config,
                                  quant_config=quant_config,
                                  require_post_norm=require_post_norm,
                                  prefix=f"{prefix}.trunk")

forward 

forward(pixel_values: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
208
209
210
211
212
213
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:

    x = self.preprocessor(pixel_values)
    x = self.trunk(x)

    return x

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]
Source code in vllm/model_executor/models/aimv2.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
def load_weights(self, weights: Iterable[tuple[str,
                                               torch.Tensor]]) -> set[str]:
    stacked_params_mapping = [
        # (param_name, shard_name, shard_id)
        (".fc13", ".fc1", 0),
        (".fc13", ".fc3", 1),
    ]
    params_dict = dict(self.named_parameters())
    loaded_params: set[str] = set()

    for name, loaded_weight in weights:
        # post_layernorm is optional in SiglipVisionModel
        if (name.startswith("trunk.post_trunk_norm")
                and self.trunk.post_trunk_norm is None):
            continue

        for (param_name, weight_name, shard_id) in stacked_params_mapping:
            if weight_name not in name:
                continue
            name = name.replace(weight_name, param_name)

            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
            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

AIMv2PatchEmbed

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class AIMv2PatchEmbed(nn.Module):

    def __init__(self, config: AIMv2Config):
        super().__init__()
        self.proj = nn.Conv2d(
            config.num_channels,
            config.hidden_size,
            kernel_size=(config.patch_size, config.patch_size),
            stride=(config.patch_size, config.patch_size),
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj(x).flatten(2).transpose(1, 2)
        x = self.norm.forward_native(x)
        return x

norm instance-attribute 

norm = RMSNorm(hidden_size, eps=rms_norm_eps)

proj instance-attribute 

proj = Conv2d(
    num_channels,
    hidden_size,
    kernel_size=(patch_size, patch_size),
    stride=(patch_size, patch_size),
)

__init__ 

__init__(config: AIMv2Config)
Source code in vllm/model_executor/models/aimv2.py 
60
61
62
63
64
65
66
67
68
def __init__(self, config: AIMv2Config):
    super().__init__()
    self.proj = nn.Conv2d(
        config.num_channels,
        config.hidden_size,
        kernel_size=(config.patch_size, config.patch_size),
        stride=(config.patch_size, config.patch_size),
    )
    self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

forward 

forward(x: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
70
71
72
73
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x = self.proj(x).flatten(2).transpose(1, 2)
    x = self.norm.forward_native(x)
    return x

AIMv2SwiGLUFFN

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
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
class AIMv2SwiGLUFFN(nn.Module):

    def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
                 prefix: str):
        super().__init__()
        hidden_features = config.intermediate_size
        in_features = config.hidden_size
        bias = config.use_bias

        self.fc13 = MergedColumnParallelLinear(
            in_features,
            [hidden_features] * 2,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.fc13",
        )
        self.fc2 = RowParallelLinear(
            input_size=hidden_features,
            output_size=in_features,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.fc2",
        )
        self.act_fn = SiluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.fc13(x)
        x = self.act_fn(x)
        x, _ = self.fc2(x)
        return x

act_fn instance-attribute 

act_fn = SiluAndMul()

fc13 instance-attribute 

fc13 = MergedColumnParallelLinear(
    in_features,
    [hidden_features] * 2,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.fc13",
)

fc2 instance-attribute 

fc2 = RowParallelLinear(
    input_size=hidden_features,
    output_size=in_features,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.fc2",
)

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    prefix: str,
)
Source code in vllm/model_executor/models/aimv2.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
             prefix: str):
    super().__init__()
    hidden_features = config.intermediate_size
    in_features = config.hidden_size
    bias = config.use_bias

    self.fc13 = MergedColumnParallelLinear(
        in_features,
        [hidden_features] * 2,
        bias=bias,
        quant_config=quant_config,
        prefix=f"{prefix}.fc13",
    )
    self.fc2 = RowParallelLinear(
        input_size=hidden_features,
        output_size=in_features,
        bias=bias,
        quant_config=quant_config,
        prefix=f"{prefix}.fc2",
    )
    self.act_fn = SiluAndMul()

forward 

forward(x: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
51
52
53
54
55
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x, _ = self.fc13(x)
    x = self.act_fn(x)
    x, _ = self.fc2(x)
    return x

AIMv2Transformer

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
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
189
190
class AIMv2Transformer(nn.Module):

    def __init__(
        self,
        config: AIMv2Config,
        quant_config: QuantizationConfig,
        *,
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
    ):
        super().__init__()

        self.blocks = nn.ModuleList([
            AIMv2Block(config, quant_config, prefix=f"{prefix}.blocks.{i}")
            for i in range(config.num_hidden_layers)
        ])
        if require_post_norm:
            self.post_trunk_norm = RMSNorm(config.hidden_size,
                                           eps=config.rms_norm_eps)
        else:
            self.post_trunk_norm = None

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        # they take the -1 as the ref embeddings, like a clip skip
        for block in self.blocks:
            tokens = block(tokens)
        if self.post_trunk_norm is not None:
            tokens = self.post_trunk_norm(tokens)
        return tokens

blocks instance-attribute 

blocks = ModuleList(
    [
        AIMv2Block(
            config,
            quant_config,
            prefix=f"{prefix}.blocks.{i}",
        )
        for i in range(num_hidden_layers)
    ]
)

post_trunk_norm instance-attribute 

post_trunk_norm = RMSNorm(hidden_size, eps=rms_norm_eps)

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    *,
    require_post_norm: Optional[bool] = None,
    prefix: str = "",
)
Source code in vllm/model_executor/models/aimv2.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def __init__(
    self,
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    *,
    require_post_norm: Optional[bool] = None,
    prefix: str = "",
):
    super().__init__()

    self.blocks = nn.ModuleList([
        AIMv2Block(config, quant_config, prefix=f"{prefix}.blocks.{i}")
        for i in range(config.num_hidden_layers)
    ])
    if require_post_norm:
        self.post_trunk_norm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)
    else:
        self.post_trunk_norm = None

forward 

forward(tokens: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
184
185
186
187
188
189
190
def forward(self, tokens: torch.Tensor) -> torch.Tensor:
    # they take the -1 as the ref embeddings, like a clip skip
    for block in self.blocks:
        tokens = block(tokens)
    if self.post_trunk_norm is not None:
        tokens = self.post_trunk_norm(tokens)
    return tokens

AIMv2ViTPreprocessor

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
class AIMv2ViTPreprocessor(nn.Module):

    def __init__(self, config: AIMv2Config):
        super().__init__()
        num_patches = (config.image_size // config.patch_size)**2

        self.patchifier = AIMv2PatchEmbed(config)
        self.pos_embed = nn.Parameter(
            torch.zeros((1, num_patches, config.hidden_size)))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        tokens = self.patchifier(x)
        _, N, _ = tokens.shape
        pos_embed = self.pos_embed.to(tokens.device)
        tokens = tokens + pos_embed[:, :N]
        return tokens

patchifier instance-attribute 

patchifier = AIMv2PatchEmbed(config)

pos_embed instance-attribute 

pos_embed = Parameter(zeros((1, num_patches, hidden_size)))

__init__ 

__init__(config: AIMv2Config)
Source code in vllm/model_executor/models/aimv2.py 
78
79
80
81
82
83
84
def __init__(self, config: AIMv2Config):
    super().__init__()
    num_patches = (config.image_size // config.patch_size)**2

    self.patchifier = AIMv2PatchEmbed(config)
    self.pos_embed = nn.Parameter(
        torch.zeros((1, num_patches, config.hidden_size)))

forward 

forward(x: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
86
87
88
89
90
91
def forward(self, x: torch.Tensor) -> torch.Tensor:
    tokens = self.patchifier(x)
    _, N, _ = tokens.shape
    pos_embed = self.pos_embed.to(tokens.device)
    tokens = tokens + pos_embed[:, :N]
    return tokens