Skip to content

vllm.model_executor.models.aimv2

AIMv2Attention

Bases: Module

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

    def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
                 prefix: str):
        super().__init__()
        dim = config.hidden_size

        # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
        self.num_heads = config.num_attention_heads
        self.qkv = ReplicatedLinear(dim, dim * 3, bias=config.qkv_bias)
        # self.qkv = QKVParallelLinear(
        #               hidden_size=dim,
        #               head_size=dim // config.num_attention_heads,
        #               total_num_heads=config.num_attention_heads,
        #               bias=config.qkv_bias,
        #               quant_config=quant_config,
        #               prefix=f"{prefix}.qkv")
        self.proj = ReplicatedLinear(dim, dim, bias=config.use_bias)
        # self.proj = RowParallelLinear(input_size=dim,
        #                  output_size=dim,
        #                  bias = config.use_bias,
        #                  quant_config=quant_config,
        #                  prefix=f"{prefix}.proj")

    def forward(  # todo might implement multiple attn implementations
            self,
            x: torch.Tensor,
            mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        B, N, C = x.shape
        qkv, _ = self.qkv(x)

        qkv = qkv.reshape(B, N, 3, self.num_heads,
                          C // self.num_heads).permute(2, 0, 3, 1, 4)

        q, k, v = qkv.unbind(0)

        x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        x = x.transpose(1, 2).contiguous().reshape(B, N, C)
        x, _ = self.proj(x)
        return x

num_heads instance-attribute 

num_heads = num_attention_heads

proj instance-attribute 

proj = ReplicatedLinear(dim, dim, bias=use_bias)

qkv instance-attribute 

qkv = ReplicatedLinear(dim, dim * 3, bias=qkv_bias)

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    prefix: str,
)
Source code in vllm/model_executor/models/aimv2.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
             prefix: str):
    super().__init__()
    dim = config.hidden_size

    # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
    self.num_heads = config.num_attention_heads
    self.qkv = ReplicatedLinear(dim, dim * 3, bias=config.qkv_bias)
    # self.qkv = QKVParallelLinear(
    #               hidden_size=dim,
    #               head_size=dim // config.num_attention_heads,
    #               total_num_heads=config.num_attention_heads,
    #               bias=config.qkv_bias,
    #               quant_config=quant_config,
    #               prefix=f"{prefix}.qkv")
    self.proj = ReplicatedLinear(dim, dim, bias=config.use_bias)

forward 

forward(x: Tensor, mask: Optional[Tensor] = None) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def forward(  # todo might implement multiple attn implementations
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    B, N, C = x.shape
    qkv, _ = self.qkv(x)

    qkv = qkv.reshape(B, N, 3, self.num_heads,
                      C // self.num_heads).permute(2, 0, 3, 1, 4)

    q, k, v = qkv.unbind(0)

    x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
    x = x.transpose(1, 2).contiguous().reshape(B, N, C)
    x, _ = self.proj(x)
    return x

AIMv2Block

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
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,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        x = x + self.attn(self.norm_1.forward_native(x), mask)
        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 
132
133
134
135
136
137
138
139
140
141
142
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, mask: Optional[Tensor] = None) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
144
145
146
147
148
149
def forward(self,
            x: torch.Tensor,
            mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    x = x + self.attn(self.norm_1.forward_native(x), mask)
    x = x + self.mlp(self.norm_2.forward_native(x))
    return x

AIMv2Model

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
class AIMv2Model(torch.nn.Module):

    def __init__(self,
                 config: AIMv2Config,
                 quant_config: QuantizationConfig,
                 prefix: str = ""):
        super().__init__()
        self.preprocessor = AIMv2ViTPreprocessor(config)
        self.trunk = AIMv2Transformer(config,
                                      quant_config=quant_config,
                                      prefix=f"{prefix}.trunk")

    def forward(
        self,
        pixel_values: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:

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

        return x

preprocessor instance-attribute 

preprocessor = AIMv2ViTPreprocessor(config)

trunk instance-attribute 

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

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    prefix: str = "",
)
Source code in vllm/model_executor/models/aimv2.py 
180
181
182
183
184
185
186
187
188
def __init__(self,
             config: AIMv2Config,
             quant_config: QuantizationConfig,
             prefix: str = ""):
    super().__init__()
    self.preprocessor = AIMv2ViTPreprocessor(config)
    self.trunk = AIMv2Transformer(config,
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.trunk")

forward 

forward(
    pixel_values: Tensor, mask: Optional[Tensor] = None
) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
190
191
192
193
194
195
196
197
198
199
def forward(
    self,
    pixel_values: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:

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

    return x

AIMv2PatchEmbed

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
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 
54
55
56
57
58
59
60
61
62
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 
64
65
66
67
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 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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

        # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
        self.fc1 = ReplicatedLinear(in_features,
                                    hidden_features,
                                    bias=bias,
                                    quant_config=quant_config,
                                    prefix=f"{prefix}.fc1")
        self.fc2 = ReplicatedLinear(hidden_features,
                                    in_features,
                                    bias=bias,
                                    quant_config=quant_config,
                                    prefix=f"{prefix}.fc2")
        self.fc3 = ReplicatedLinear(in_features,
                                    hidden_features,
                                    bias=bias,
                                    quant_config=quant_config,
                                    prefix=f"{prefix}.fc3")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_parallel, _ = self.fc1(x)
        gate, _ = self.fc3(x)
        x_parallel = F.silu(x_parallel) * gate
        out, _ = self.fc2(x_parallel)
        return out

fc1 instance-attribute 

fc1 = ReplicatedLinear(
    in_features,
    hidden_features,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.fc1",
)

fc2 instance-attribute 

fc2 = ReplicatedLinear(
    hidden_features,
    in_features,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.fc2",
)

fc3 instance-attribute 

fc3 = ReplicatedLinear(
    in_features,
    hidden_features,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.fc3",
)

__init__ 

__init__(
    config: AIMv2Config,
    quant_config: QuantizationConfig,
    prefix: str,
)
Source code in vllm/model_executor/models/aimv2.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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

    # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
    self.fc1 = ReplicatedLinear(in_features,
                                hidden_features,
                                bias=bias,
                                quant_config=quant_config,
                                prefix=f"{prefix}.fc1")
    self.fc2 = ReplicatedLinear(hidden_features,
                                in_features,
                                bias=bias,
                                quant_config=quant_config,
                                prefix=f"{prefix}.fc2")
    self.fc3 = ReplicatedLinear(in_features,
                                hidden_features,
                                bias=bias,
                                quant_config=quant_config,
                                prefix=f"{prefix}.fc3")

forward 

forward(x: Tensor) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
44
45
46
47
48
49
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x_parallel, _ = self.fc1(x)
    gate, _ = self.fc3(x)
    x_parallel = F.silu(x_parallel) * gate
    out, _ = self.fc2(x_parallel)
    return out

AIMv2Transformer

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
class AIMv2Transformer(nn.Module):

    def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
                 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)
        ])
        self.post_trunk_norm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)

    def forward(
        self,
        tokens: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # they take the -1 as the ref embeddings, like a clip skip
        for block in self.blocks:
            tokens = block(tokens, mask)
        # NO NORM IN THE OG IMPLEMENTATION
        # 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,
    prefix: str,
)
Source code in vllm/model_executor/models/aimv2.py 
154
155
156
157
158
159
160
161
162
163
def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
             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)
    ])
    self.post_trunk_norm = RMSNorm(config.hidden_size,
                                   eps=config.rms_norm_eps)

forward 

forward(
    tokens: Tensor, mask: Optional[Tensor] = None
) -> Tensor
Source code in vllm/model_executor/models/aimv2.py 
165
166
167
168
169
170
171
172
173
174
175
def forward(
    self,
    tokens: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    # they take the -1 as the ref embeddings, like a clip skip
    for block in self.blocks:
        tokens = block(tokens, mask)
    # NO NORM IN THE OG IMPLEMENTATION
    # tokens = self.post_trunk_norm(tokens)
    return tokens

AIMv2ViTPreprocessor

Bases: Module

Source code in vllm/model_executor/models/aimv2.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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 
72
73
74
75
76
77
78
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 
80
81
82
83
84
85
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