olmo.py 14.5 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

Isotr0py's avatar
Isotr0py committed
3
# Adapted from
4
5
6
# https://github.com/huggingface/transformers/blob/v4.40.1/src/transformers/models/olmo/modeling_olmo.py
# Copyright 2024 The vLLM team.
# Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved.
Isotr0py's avatar
Isotr0py committed
7
#
8
9
10
11
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
Isotr0py's avatar
Isotr0py committed
12
#
13
14
15
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
Isotr0py's avatar
Isotr0py committed
16
#
17
#     http://www.apache.org/licenses/LICENSE-2.0
Isotr0py's avatar
Isotr0py committed
18
#
19
20
21
22
23
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Isotr0py's avatar
Isotr0py committed
24
"""Inference-only OLMo model compatible with HuggingFace weights."""
25
26
from collections.abc import Iterable
from typing import Optional, Union
Isotr0py's avatar
Isotr0py committed
27
28
29

import torch
from torch import nn
30
from transformers import OlmoConfig
Isotr0py's avatar
Isotr0py committed
31

32
from vllm.attention import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.model_executor.layers.activation import SiluAndMul
37
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
38
39
                                               QKVParallelLinear,
                                               RowParallelLinear)
40
from vllm.model_executor.layers.logits_processor import LogitsProcessor
41
from vllm.model_executor.layers.quantization import QuantizationConfig
42
from vllm.model_executor.layers.rotary_embedding import get_rope
43
from vllm.model_executor.layers.vocab_parallel_embedding import (
44
    ParallelLMHead, VocabParallelEmbedding)
45
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
Isotr0py's avatar
Isotr0py committed
46
from vllm.model_executor.sampling_metadata import SamplingMetadata
47
from vllm.sequence import IntermediateTensors
48

49
from .interfaces import SupportsPP
50
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
51
52
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
53

Isotr0py's avatar
Isotr0py committed
54
55
56

class OlmoAttention(nn.Module):
    """
57
58
    This is the attention block where the output is computed as
    ``Attention(LN(x))`` in ``MLP(LN(x + Attention(LN(x))))``
Isotr0py's avatar
Isotr0py committed
59
60
61
62
63
    (plus another skip connection).
    """

    def __init__(
        self,
64
        config: OlmoConfig,
65
        cache_config: Optional[CacheConfig] = None,
66
        quant_config: Optional[QuantizationConfig] = None,
67
        prefix: str = "",
Isotr0py's avatar
Isotr0py committed
68
69
70
    ):
        super().__init__()
        self.config = config
71
        self.hidden_size = config.hidden_size
72
73
        tensor_model_parallel_world_size = (
            get_tensor_model_parallel_world_size())
74
75
76
        self.total_num_heads = config.num_attention_heads

        assert self.hidden_size % self.total_num_heads == 0
Isotr0py's avatar
Isotr0py committed
77
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
78

79
80
        self.num_heads = (self.total_num_heads //
                          tensor_model_parallel_world_size)
Isotr0py's avatar
Isotr0py committed
81
        self.head_dim = self.hidden_size // self.total_num_heads
82
83
84
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = config.rope_theta
        self.clip_qkv = config.clip_qkv
Isotr0py's avatar
Isotr0py committed
85
86

        # Attention input projection. Projects x -> (q, k, v)
87
88
        self.qkv_proj = QKVParallelLinear(
            self.hidden_size,
Isotr0py's avatar
Isotr0py committed
89
90
            self.head_dim,
            self.total_num_heads,
91
            bias=config.attention_bias,
92
            quant_config=quant_config,
Isotr0py's avatar
Isotr0py committed
93
94
95
        )

        # Rotary embeddings.
96
97
98
99
100
101
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=self.max_position_embeddings,
            base=self.rope_theta,
        )
Isotr0py's avatar
Isotr0py committed
102
        self.scaling = self.head_dim**-0.5
103
104
        self.attn = Attention(self.num_heads,
                              self.head_dim,
105
                              scale=self.scaling,
106
                              cache_config=cache_config,
107
108
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
Isotr0py's avatar
Isotr0py committed
109
110

        # Attention output projection.
111
112
113
114
        self.o_proj = RowParallelLinear(
            self.hidden_size,
            self.hidden_size,
            bias=config.attention_bias,
115
            quant_config=quant_config,
Isotr0py's avatar
Isotr0py committed
116
117
118
119
120
121
122
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
123
124
125
        qkv, _ = self.qkv_proj(hidden_states)
        if self.clip_qkv is not None:
            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
Isotr0py's avatar
Isotr0py committed
126
        q, k, v = qkv.chunk(chunks=3, dim=-1)
127
        q, k = self.rotary_emb(positions, q, k)
128
        attn_output = self.attn(q, k, v)
129
        output, _ = self.o_proj(attn_output)
Isotr0py's avatar
Isotr0py committed
130
131
132
133
134
        return output


class OlmoMLP(nn.Module):
    """
135
136
    This is the MLP block where the output is computed as
    ``MLP(LN(x))`` in ``MLP(LN(x + Attention(LN(x))))``
Isotr0py's avatar
Isotr0py committed
137
138
139
140
141
    (plus another skip connection).
    """

    def __init__(
        self,
142
        config: OlmoConfig,
143
        quant_config: Optional[QuantizationConfig] = None,
Isotr0py's avatar
Isotr0py committed
144
145
146
    ):
        super().__init__()
        self.config = config
147
148
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
Isotr0py's avatar
Isotr0py committed
149
150

        # Feed-forward input projection.
151
152
153
154
        self.gate_up_proj = MergedColumnParallelLinear(
            self.hidden_size,
            [self.intermediate_size] * 2,
            bias=False,
155
            quant_config=quant_config,
Isotr0py's avatar
Isotr0py committed
156
157
158
        )

        # Activation function.
159
        self.act_fn = SiluAndMul()
Isotr0py's avatar
Isotr0py committed
160
161

        # Feed-forward output projection.
162
163
164
165
        self.down_proj = RowParallelLinear(
            self.intermediate_size,
            self.hidden_size,
            bias=False,
166
            quant_config=quant_config,
Isotr0py's avatar
Isotr0py committed
167
168
169
170
171
172
        )

    def forward(
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
173
174
175
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
Isotr0py's avatar
Isotr0py committed
176
177
178
        return x


179
class OlmoDecoderLayer(nn.Module):
Isotr0py's avatar
Isotr0py committed
180
    """
181
182
    This is a typical transformer block where the output is
    computed as ``MLP(LN(x + Attention(LN(x))))``
Isotr0py's avatar
Isotr0py committed
183
184
185
186
    (plus another skip connection).
    """

    def __init__(self,
187
                 config: OlmoConfig,
188
                 cache_config: Optional[CacheConfig] = None,
189
190
                 quant_config: Optional[QuantizationConfig] = None,
                 prefix: str = ""):
Isotr0py's avatar
Isotr0py committed
191
192
        super().__init__()
        # Attention block.
193
194
195
196
        self.self_attn = OlmoAttention(config,
                                       cache_config,
                                       quant_config,
                                       prefix=f"{prefix}.self_attn")
Isotr0py's avatar
Isotr0py committed
197
198

        # MLP block.
199
        self.mlp = OlmoMLP(config, quant_config)
Isotr0py's avatar
Isotr0py committed
200

201
202
203
204
205
206
207
208
        # LayerNorm
        self.input_layernorm = nn.LayerNorm(config.hidden_size,
                                            elementwise_affine=False,
                                            bias=False)
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size,
                                                     elementwise_affine=False,
                                                     bias=False)

Isotr0py's avatar
Isotr0py committed
209
210
211
212
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
213
    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
Isotr0py's avatar
Isotr0py committed
214
        # Attention block.
215
216
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
217
        hidden_states = self.self_attn(positions, hidden_states)
218
        hidden_states = hidden_states + residual
Isotr0py's avatar
Isotr0py committed
219
220

        # MLP block.
221
222
223
224
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
Isotr0py's avatar
Isotr0py committed
225
226
227
        return hidden_states


228
@support_torch_compile
Isotr0py's avatar
Isotr0py committed
229
230
class OlmoModel(nn.Module):

231
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Isotr0py's avatar
Isotr0py committed
232
        super().__init__()
233
234
235
236
237

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

Isotr0py's avatar
Isotr0py committed
238
239
        self.config = config

240
241
        self.embed_tokens = VocabParallelEmbedding(config.vocab_size,
                                                   config.hidden_size)
242
243
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
244
245
            lambda prefix: OlmoDecoderLayer(
                config, cache_config, quant_config, prefix=prefix),
246
            prefix=f"{prefix}.layers")
247
248
249
        self.norm = nn.LayerNorm(config.hidden_size,
                                 elementwise_affine=False,
                                 bias=False)
250
251
252
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))
Isotr0py's avatar
Isotr0py committed
253

254
255
256
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Isotr0py's avatar
Isotr0py committed
257
258
259
260
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
261
        intermediate_tensors: Optional[IntermediateTensors],
262
        inputs_embeds: Optional[torch.Tensor] = None,
263
    ) -> Union[torch.Tensor, IntermediateTensors]:
Isotr0py's avatar
Isotr0py committed
264
265
266
        """
        :param input_ids: A tensor of shape `(batch_size, seq_len)`.
        """
267
        if get_pp_group().is_first_rank:
268
269
270
271
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
272
273
274
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
Isotr0py's avatar
Isotr0py committed
275
276

        # Apply blocks one-by-one.
277
        for layer in self.layers[self.start_layer:self.end_layer]:
Isotr0py's avatar
Isotr0py committed
278
            # shape: (batch_size, seq_len, d_model)
279
            hidden_states = layer(positions, hidden_states)
Isotr0py's avatar
Isotr0py committed
280

281
282
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
Isotr0py's avatar
Isotr0py committed
283
284
        # Apply final layer norm.
        # shape: (batch_size, seq_len or 1, d_model)
285
286
        hidden_states = self.norm(hidden_states)
        return hidden_states
Isotr0py's avatar
Isotr0py committed
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
314
315
316
317
318
319
320
321
322
323
324
325
326
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            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)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if is_pp_missing_parameter(name, self):
                    continue
                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

Isotr0py's avatar
Isotr0py committed
327

328
class OlmoForCausalLM(nn.Module, SupportsPP):
Isotr0py's avatar
Isotr0py committed
329
330
331
332
    """
    Extremely barebones HF model wrapper.
    """

333
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Isotr0py's avatar
Isotr0py committed
334
        super().__init__()
335
336
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
Isotr0py's avatar
Isotr0py committed
337
        self.config = config
338
339
        self.model = OlmoModel(vllm_config=vllm_config,
                               prefix=maybe_prefix(prefix, "model"))
340
        if config.tie_word_embeddings:
341
            self.lm_head = self.model.embed_tokens
342
343
344
345
346
347
        else:
            self.unpadded_vocab_size = config.vocab_size
            self.lm_head = ParallelLMHead(
                self.unpadded_vocab_size,
                config.hidden_size,
                org_num_embeddings=config.vocab_size,
348
                quant_config=quant_config,
349
            )
350
        self.logits_processor = LogitsProcessor(config.vocab_size)
351
352
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Isotr0py's avatar
Isotr0py committed
353

354
355
356
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

Isotr0py's avatar
Isotr0py committed
357
358
359
360
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
361
        intermediate_tensors: Optional[IntermediateTensors] = None,
362
        inputs_embeds: Optional[torch.Tensor] = None,
363
    ) -> Union[torch.Tensor, IntermediateTensors]:
Isotr0py's avatar
Isotr0py committed
364
365
366
        hidden_states = self.model(
            input_ids=input_ids,
            positions=positions,
367
            intermediate_tensors=intermediate_tensors,
368
            inputs_embeds=inputs_embeds,
Isotr0py's avatar
Isotr0py committed
369
370
371
        )
        return hidden_states

372
373
374
375
376
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
377
        logits = self.logits_processor(self.lm_head, hidden_states,
378
379
380
                                       sampling_metadata)
        return logits

381
382
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
383
384
        loader = AutoWeightsLoader(
            self,
385
386
            skip_prefixes=(["lm_head.weight"]
                           if self.config.tie_word_embeddings else None),
387
388
        )
        return loader.load_weights(weights)