Skip to content

vllm.model_executor.models.bloom

Inference-only BLOOM model compatible with HuggingFace weights.

BloomAttention

Bases: Module

Source code in vllm/model_executor/models/bloom.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 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
128
129
130
131
132
133
134
135
136
class BloomAttention(nn.Module):

    def __init__(
        self,
        config: BloomConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.total_num_heads = config.n_head
        self.head_dim = self.hidden_size // self.total_num_heads
        assert self.head_dim * self.total_num_heads == self.hidden_size

        tp_world_size = get_tensor_model_parallel_world_size()
        assert self.total_num_heads % tp_world_size == 0
        self.num_heads = self.total_num_heads // tp_world_size

        self.query_key_value = QKVParallelLinear(
            self.hidden_size,
            self.head_dim,
            self.total_num_heads,
            bias=True,
            quant_config=quant_config,
        )
        self.dense = RowParallelLinear(
            self.hidden_size,
            self.hidden_size,
            bias=True,
            quant_config=quant_config,
        )

        # Create the alibi slopes and slice them.
        tp_rank = get_tensor_model_parallel_rank()
        head_start = tp_rank * self.num_heads
        head_end = (tp_rank + 1) * self.num_heads
        alibi_slopes = _get_alibi_slopes(self.total_num_heads)
        alibi_slopes = alibi_slopes[head_start:head_end].tolist()

        scaling = self.head_dim**-0.5
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              scaling,
                              alibi_slopes=alibi_slopes,
                              cache_config=cache_config,
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        del position_ids  # Unused.
        qkv, _ = self.query_key_value(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
        attn_output = self.attn(q, k, v)
        output, _ = self.dense(attn_output)
        return output

attn instance-attribute 

attn = Attention(
    num_heads,
    head_dim,
    scaling,
    alibi_slopes=alibi_slopes,
    cache_config=cache_config,
    quant_config=quant_config,
    prefix=f"{prefix}.attn",
)

dense instance-attribute 

dense = RowParallelLinear(
    hidden_size,
    hidden_size,
    bias=True,
    quant_config=quant_config,
)

head_dim instance-attribute 

head_dim = hidden_size // total_num_heads

hidden_size instance-attribute 

hidden_size = hidden_size

num_heads instance-attribute 

num_heads = total_num_heads // tp_world_size

query_key_value instance-attribute 

query_key_value = QKVParallelLinear(
    hidden_size,
    head_dim,
    total_num_heads,
    bias=True,
    quant_config=quant_config,
)

total_num_heads instance-attribute 

total_num_heads = n_head

__init__ 

__init__(
    config: BloomConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
)
Source code in vllm/model_executor/models/bloom.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 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
def __init__(
    self,
    config: BloomConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
):
    super().__init__()
    self.hidden_size = config.hidden_size
    self.total_num_heads = config.n_head
    self.head_dim = self.hidden_size // self.total_num_heads
    assert self.head_dim * self.total_num_heads == self.hidden_size

    tp_world_size = get_tensor_model_parallel_world_size()
    assert self.total_num_heads % tp_world_size == 0
    self.num_heads = self.total_num_heads // tp_world_size

    self.query_key_value = QKVParallelLinear(
        self.hidden_size,
        self.head_dim,
        self.total_num_heads,
        bias=True,
        quant_config=quant_config,
    )
    self.dense = RowParallelLinear(
        self.hidden_size,
        self.hidden_size,
        bias=True,
        quant_config=quant_config,
    )

    # Create the alibi slopes and slice them.
    tp_rank = get_tensor_model_parallel_rank()
    head_start = tp_rank * self.num_heads
    head_end = (tp_rank + 1) * self.num_heads
    alibi_slopes = _get_alibi_slopes(self.total_num_heads)
    alibi_slopes = alibi_slopes[head_start:head_end].tolist()

    scaling = self.head_dim**-0.5
    self.attn = Attention(self.num_heads,
                          self.head_dim,
                          scaling,
                          alibi_slopes=alibi_slopes,
                          cache_config=cache_config,
                          quant_config=quant_config,
                          prefix=f"{prefix}.attn")

forward 

forward(
    position_ids: Tensor, hidden_states: Tensor
) -> Tensor
Source code in vllm/model_executor/models/bloom.py 
126
127
128
129
130
131
132
133
134
135
136
def forward(
    self,
    position_ids: torch.Tensor,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    del position_ids  # Unused.
    qkv, _ = self.query_key_value(hidden_states)
    q, k, v = qkv.chunk(chunks=3, dim=-1)
    attn_output = self.attn(q, k, v)
    output, _ = self.dense(attn_output)
    return output

BloomBlock

Bases: Module

Source code in vllm/model_executor/models/bloom.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
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
class BloomBlock(nn.Module):

    def __init__(
        self,
        config: BloomConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        hidden_size = config.hidden_size

        self.input_layernorm = nn.LayerNorm(hidden_size,
                                            eps=config.layer_norm_epsilon)
        self.self_attention = BloomAttention(config,
                                             cache_config,
                                             quant_config,
                                             prefix=f"{prefix}.self_attention")
        self.post_attention_layernorm = nn.LayerNorm(
            hidden_size, eps=config.layer_norm_epsilon)
        self.mlp = BloomMLP(config, quant_config)
        self.apply_residual_connection_post_layernorm = (
            config.apply_residual_connection_post_layernorm)

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        # Layer norm at the beginning of the transformer layer.
        layernorm_output = self.input_layernorm(hidden_states)

        # Layer norm post the self attention.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = hidden_states

        # Self attention.
        attention_output = self.self_attention(
            position_ids=position_ids,
            hidden_states=layernorm_output,
        )
        attention_output = attention_output + residual
        layernorm_output = self.post_attention_layernorm(attention_output)

        # Get residual
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = attention_output

        # MLP.
        output = self.mlp(layernorm_output) + residual
        return output

apply_residual_connection_post_layernorm instance-attribute 

apply_residual_connection_post_layernorm = (
    apply_residual_connection_post_layernorm
)

input_layernorm instance-attribute 

input_layernorm = LayerNorm(
    hidden_size, eps=layer_norm_epsilon
)

mlp instance-attribute 

mlp = BloomMLP(config, quant_config)

post_attention_layernorm instance-attribute 

post_attention_layernorm = LayerNorm(
    hidden_size, eps=layer_norm_epsilon
)

self_attention instance-attribute 

self_attention = BloomAttention(
    config,
    cache_config,
    quant_config,
    prefix=f"{prefix}.self_attention",
)

__init__ 

__init__(
    config: BloomConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
)
Source code in vllm/model_executor/models/bloom.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def __init__(
    self,
    config: BloomConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
):
    super().__init__()
    hidden_size = config.hidden_size

    self.input_layernorm = nn.LayerNorm(hidden_size,
                                        eps=config.layer_norm_epsilon)
    self.self_attention = BloomAttention(config,
                                         cache_config,
                                         quant_config,
                                         prefix=f"{prefix}.self_attention")
    self.post_attention_layernorm = nn.LayerNorm(
        hidden_size, eps=config.layer_norm_epsilon)
    self.mlp = BloomMLP(config, quant_config)
    self.apply_residual_connection_post_layernorm = (
        config.apply_residual_connection_post_layernorm)

forward 

forward(
    position_ids: Tensor, hidden_states: Tensor
) -> Tensor
Source code in vllm/model_executor/models/bloom.py 
191
192
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
def forward(
    self,
    position_ids: torch.Tensor,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    # Layer norm at the beginning of the transformer layer.
    layernorm_output = self.input_layernorm(hidden_states)

    # Layer norm post the self attention.
    if self.apply_residual_connection_post_layernorm:
        residual = layernorm_output
    else:
        residual = hidden_states

    # Self attention.
    attention_output = self.self_attention(
        position_ids=position_ids,
        hidden_states=layernorm_output,
    )
    attention_output = attention_output + residual
    layernorm_output = self.post_attention_layernorm(attention_output)

    # Get residual
    if self.apply_residual_connection_post_layernorm:
        residual = layernorm_output
    else:
        residual = attention_output

    # MLP.
    output = self.mlp(layernorm_output) + residual
    return output

BloomForCausalLM

Bases: Module, SupportsPP, SupportsV0Only, SupportsQuant

Source code in vllm/model_executor/models/bloom.py 
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
class BloomForCausalLM(nn.Module, SupportsPP, SupportsV0Only, SupportsQuant):

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.config = config
        self.quant_config = quant_config
        self.transformer = BloomModel(vllm_config=vllm_config,
                                      prefix=maybe_prefix(
                                          prefix, "transformer"))
        if self.config.tie_word_embeddings:
            self.lm_head = self.transformer.word_embeddings
        else:
            self.lm_head = ParallelLMHead(self.config.vocab_size,
                                          self.config.hidden_size)

        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)

    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.transformer.get_input_embeddings(input_ids)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        hidden_states = self.transformer(input_ids, positions,
                                         intermediate_tensors, inputs_embeds)
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        logits = self.logits_processor(self.lm_head, hidden_states,
                                       sampling_metadata)
        return logits

    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self, skip_prefixes=["lm_head.weight"])
        weights = _add_transformer_prefix(weights)
        return loader.load_weights(weights)

config instance-attribute 

config = config

lm_head instance-attribute 

lm_head = word_embeddings

logits_processor instance-attribute 

logits_processor = LogitsProcessor(vocab_size)

make_empty_intermediate_tensors instance-attribute 

make_empty_intermediate_tensors = (
    make_empty_intermediate_tensors
)

quant_config instance-attribute 

quant_config = quant_config

transformer instance-attribute 

transformer = BloomModel(
    vllm_config=vllm_config,
    prefix=maybe_prefix(prefix, "transformer"),
)

__init__ 

__init__(*, vllm_config: VllmConfig, prefix: str = '')
Source code in vllm/model_executor/models/bloom.py 
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
    super().__init__()
    config = vllm_config.model_config.hf_config
    quant_config = vllm_config.quant_config
    self.config = config
    self.quant_config = quant_config
    self.transformer = BloomModel(vllm_config=vllm_config,
                                  prefix=maybe_prefix(
                                      prefix, "transformer"))
    if self.config.tie_word_embeddings:
        self.lm_head = self.transformer.word_embeddings
    else:
        self.lm_head = ParallelLMHead(self.config.vocab_size,
                                      self.config.hidden_size)

    self.logits_processor = LogitsProcessor(config.vocab_size)
    self.make_empty_intermediate_tensors = (
        self.transformer.make_empty_intermediate_tensors)

compute_logits 

compute_logits(
    hidden_states: Tensor,
    sampling_metadata: SamplingMetadata,
) -> Optional[Tensor]
Source code in vllm/model_executor/models/bloom.py 
351
352
353
354
355
356
357
358
def compute_logits(
    self,
    hidden_states: torch.Tensor,
    sampling_metadata: SamplingMetadata,
) -> Optional[torch.Tensor]:
    logits = self.logits_processor(self.lm_head, hidden_states,
                                   sampling_metadata)
    return logits

forward 

forward(
    input_ids: Tensor,
    positions: Tensor,
    intermediate_tensors: Optional[
        IntermediateTensors
    ] = None,
    inputs_embeds: Optional[Tensor] = None,
) -> Union[Tensor, IntermediateTensors]
Source code in vllm/model_executor/models/bloom.py 
340
341
342
343
344
345
346
347
348
349
def forward(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    intermediate_tensors: Optional[IntermediateTensors] = None,
    inputs_embeds: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, IntermediateTensors]:
    hidden_states = self.transformer(input_ids, positions,
                                     intermediate_tensors, inputs_embeds)
    return hidden_states

get_input_embeddings 

get_input_embeddings(input_ids: Tensor) -> Tensor
Source code in vllm/model_executor/models/bloom.py 
337
338
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
    return self.transformer.get_input_embeddings(input_ids)

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]
Source code in vllm/model_executor/models/bloom.py 
360
361
362
363
364
def load_weights(self, weights: Iterable[tuple[str,
                                               torch.Tensor]]) -> set[str]:
    loader = AutoWeightsLoader(self, skip_prefixes=["lm_head.weight"])
    weights = _add_transformer_prefix(weights)
    return loader.load_weights(weights)

BloomMLP

Bases: Module

Source code in vllm/model_executor/models/bloom.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
class BloomMLP(nn.Module):

    def __init__(
        self,
        config: BloomConfig,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        hidden_size = config.hidden_size
        self.dense_h_to_4h = ColumnParallelLinear(
            hidden_size,
            4 * hidden_size,
            quant_config=quant_config,
        )
        self.gelu_impl = get_act_fn("gelu")
        self.dense_4h_to_h = RowParallelLinear(
            4 * hidden_size,
            hidden_size,
            quant_config=quant_config,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.dense_h_to_4h(x)
        x = self.gelu_impl(x)
        x, _ = self.dense_4h_to_h(x)
        return x

dense_4h_to_h instance-attribute 

dense_4h_to_h = RowParallelLinear(
    4 * hidden_size, hidden_size, quant_config=quant_config
)

dense_h_to_4h instance-attribute 

dense_h_to_4h = ColumnParallelLinear(
    hidden_size, 4 * hidden_size, quant_config=quant_config
)

gelu_impl instance-attribute 

gelu_impl = get_act_fn('gelu')

__init__ 

__init__(
    config: BloomConfig,
    quant_config: Optional[QuantizationConfig] = None,
)
Source code in vllm/model_executor/models/bloom.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def __init__(
    self,
    config: BloomConfig,
    quant_config: Optional[QuantizationConfig] = None,
):
    super().__init__()
    hidden_size = config.hidden_size
    self.dense_h_to_4h = ColumnParallelLinear(
        hidden_size,
        4 * hidden_size,
        quant_config=quant_config,
    )
    self.gelu_impl = get_act_fn("gelu")
    self.dense_4h_to_h = RowParallelLinear(
        4 * hidden_size,
        hidden_size,
        quant_config=quant_config,
    )

forward 

forward(x: Tensor) -> Tensor
Source code in vllm/model_executor/models/bloom.py 
160
161
162
163
164
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x, _ = self.dense_h_to_4h(x)
    x = self.gelu_impl(x)
    x, _ = self.dense_4h_to_h(x)
    return x

BloomModel

Bases: Module

Source code in vllm/model_executor/models/bloom.py 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
@support_torch_compile
class BloomModel(nn.Module):

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()

        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        self.config = config

        self.embed_dim = config.hidden_size

        # Embedding + LN Embedding
        self.word_embeddings = VocabParallelEmbedding(
            config.vocab_size,
            self.embed_dim,
        )
        self.word_embeddings_layernorm = nn.LayerNorm(
            self.embed_dim, eps=config.layer_norm_epsilon)

        # Transformer blocks
        self.start_layer, self.end_layer, self.h = make_layers(
            config.num_hidden_layers,
            lambda prefix: BloomBlock(
                config, cache_config, quant_config, prefix=prefix),
            prefix=f"{prefix}.h")

        # Final Layer Norm
        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))

    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.word_embeddings_layernorm(self.word_embeddings(input_ids))

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors],
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
        for layer in self.h[self.start_layer:self.end_layer]:
            hidden_states = layer(position_ids, hidden_states)
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
        hidden_states = self.ln_f(hidden_states)
        return hidden_states

    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            if is_pp_missing_parameter(name, self):
                continue
            param = params_dict[name]

            if "query_key_value" in name:
                # NOTE: BLOOM's fused QKV's output_dim has the shape of
                # (num_heads * 3 * head_size), while the
                # required shape is (3 * num_heads * head_size).
                # Thus, we need weight conversion.
                output_dim = getattr(param, "output_dim", None)
                num_heads = self.config.num_attention_heads
                if output_dim is not None:
                    loaded_weight_shape = loaded_weight.shape
                    loaded_weight = loaded_weight.view(
                        loaded_weight_shape[:output_dim] + (num_heads, 3, -1) +
                        loaded_weight_shape[output_dim + 1:])
                    loaded_weight = loaded_weight.transpose(
                        output_dim, output_dim + 1)
                    loaded_weight = loaded_weight.reshape(loaded_weight_shape)

            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)
            loaded_params.add(name)

        return loaded_params

config instance-attribute 

config = config

embed_dim instance-attribute 

embed_dim = hidden_size

ln_f instance-attribute 

ln_f = LayerNorm(embed_dim, eps=layer_norm_epsilon)

make_empty_intermediate_tensors instance-attribute 

make_empty_intermediate_tensors = (
    make_empty_intermediate_tensors_factory(
        ["hidden_states"], hidden_size
    )
)

word_embeddings instance-attribute 

word_embeddings = VocabParallelEmbedding(
    vocab_size, embed_dim
)

word_embeddings_layernorm instance-attribute 

word_embeddings_layernorm = LayerNorm(
    embed_dim, eps=layer_norm_epsilon
)

__init__ 

__init__(*, vllm_config: VllmConfig, prefix: str = '')
Source code in vllm/model_executor/models/bloom.py 
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
    super().__init__()

    config = vllm_config.model_config.hf_config
    cache_config = vllm_config.cache_config
    quant_config = vllm_config.quant_config
    self.config = config

    self.embed_dim = config.hidden_size

    # Embedding + LN Embedding
    self.word_embeddings = VocabParallelEmbedding(
        config.vocab_size,
        self.embed_dim,
    )
    self.word_embeddings_layernorm = nn.LayerNorm(
        self.embed_dim, eps=config.layer_norm_epsilon)

    # Transformer blocks
    self.start_layer, self.end_layer, self.h = make_layers(
        config.num_hidden_layers,
        lambda prefix: BloomBlock(
            config, cache_config, quant_config, prefix=prefix),
        prefix=f"{prefix}.h")

    # Final Layer Norm
    self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
    self.make_empty_intermediate_tensors = (
        make_empty_intermediate_tensors_factory(["hidden_states"],
                                                config.hidden_size))

forward 

forward(
    input_ids: Tensor,
    position_ids: Tensor,
    intermediate_tensors: Optional[IntermediateTensors],
    inputs_embeds: Optional[Tensor] = None,
) -> Union[Tensor, IntermediateTensors]
Source code in vllm/model_executor/models/bloom.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
def forward(
    self,
    input_ids: torch.Tensor,
    position_ids: torch.Tensor,
    intermediate_tensors: Optional[IntermediateTensors],
    inputs_embeds: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, IntermediateTensors]:
    if get_pp_group().is_first_rank:
        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            hidden_states = self.get_input_embeddings(input_ids)
    else:
        assert intermediate_tensors is not None
        hidden_states = intermediate_tensors["hidden_states"]
    for layer in self.h[self.start_layer:self.end_layer]:
        hidden_states = layer(position_ids, hidden_states)
    if not get_pp_group().is_last_rank:
        return IntermediateTensors({"hidden_states": hidden_states})
    hidden_states = self.ln_f(hidden_states)
    return hidden_states

get_input_embeddings 

get_input_embeddings(input_ids: Tensor) -> Tensor
Source code in vllm/model_executor/models/bloom.py 
258
259
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
    return self.word_embeddings_layernorm(self.word_embeddings(input_ids))

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]
Source code in vllm/model_executor/models/bloom.py 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
def load_weights(self, weights: Iterable[tuple[str,
                                               torch.Tensor]]) -> set[str]:
    params_dict = dict(self.named_parameters(remove_duplicate=False))
    loaded_params: set[str] = set()
    for name, loaded_weight in weights:
        if is_pp_missing_parameter(name, self):
            continue
        param = params_dict[name]

        if "query_key_value" in name:
            # NOTE: BLOOM's fused QKV's output_dim has the shape of
            # (num_heads * 3 * head_size), while the
            # required shape is (3 * num_heads * head_size).
            # Thus, we need weight conversion.
            output_dim = getattr(param, "output_dim", None)
            num_heads = self.config.num_attention_heads
            if output_dim is not None:
                loaded_weight_shape = loaded_weight.shape
                loaded_weight = loaded_weight.view(
                    loaded_weight_shape[:output_dim] + (num_heads, 3, -1) +
                    loaded_weight_shape[output_dim + 1:])
                loaded_weight = loaded_weight.transpose(
                    output_dim, output_dim + 1)
                loaded_weight = loaded_weight.reshape(loaded_weight_shape)

        weight_loader = getattr(param, "weight_loader",
                                default_weight_loader)
        weight_loader(param, loaded_weight)
        loaded_params.add(name)

    return loaded_params

_add_transformer_prefix 

_add_transformer_prefix(
    weights: Iterable[tuple[str, Tensor]],
) -> Iterable[tuple[str, Tensor]]
Source code in vllm/model_executor/models/bloom.py 
367
368
369
370
371
372
373
def _add_transformer_prefix(
    weights: Iterable[tuple[str, torch.Tensor]]
) -> Iterable[tuple[str, torch.Tensor]]:
    for name, tensor in weights:
        if not name.startswith('transformer.'):
            name = 'transformer.' + name
        yield name, tensor

_get_alibi_slopes 

_get_alibi_slopes(total_num_heads: int) -> Tensor
Source code in vllm/model_executor/models/bloom.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
    closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
    base = torch.tensor(
        2**(-(2**-(math.log2(closest_power_of_2) - 3))),
        dtype=torch.float32,
    )
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != total_num_heads:
        extra_base = torch.tensor(
            2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
            dtype=torch.float32,
        )
        num_remaining_heads = min(closest_power_of_2,
                                  total_num_heads - closest_power_of_2)
        extra_powers = torch.arange(start=1,
                                    end=1 + 2 * num_remaining_heads,
                                    step=2,
                                    dtype=torch.int32)
        slopes = torch.cat(
            [slopes, torch.pow(extra_base, extra_powers)], dim=0)
    return slopes