olmo.py 15.6 KB
Newer Older
Isotr0py's avatar
Isotr0py committed
1
# Adapted from
2
3
4
# 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
5
#
6
7
8
9
# 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
10
#
11
12
13
# 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
14
#
15
#     http://www.apache.org/licenses/LICENSE-2.0
Isotr0py's avatar
Isotr0py committed
16
#
17
18
19
20
21
# 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
22
"""Inference-only OLMo model compatible with HuggingFace weights."""
23
from typing import Iterable, List, Optional, Set, Tuple, Union
Isotr0py's avatar
Isotr0py committed
24
25
26

import torch
from torch import nn
27
from transformers import OlmoConfig
Isotr0py's avatar
Isotr0py committed
28

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

47
48
from .interfaces import SupportsPP
from .utils import (is_pp_missing_parameter,
49
50
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
51

Isotr0py's avatar
Isotr0py committed
52
53
54

class OlmoAttention(nn.Module):
    """
55
56
    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
57
58
59
60
61
    (plus another skip connection).
    """

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

        assert self.hidden_size % self.total_num_heads == 0
Isotr0py's avatar
Isotr0py committed
75
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
76

77
78
        self.num_heads = (self.total_num_heads //
                          tensor_model_parallel_world_size)
Isotr0py's avatar
Isotr0py committed
79
        self.head_dim = self.hidden_size // self.total_num_heads
80
81
82
        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
83
84

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

        # Rotary embeddings.
94
95
96
97
98
99
        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
100
        self.scaling = self.head_dim**-0.5
101
102
        self.attn = Attention(self.num_heads,
                              self.head_dim,
103
                              scale=self.scaling,
104
                              cache_config=cache_config,
105
106
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
Isotr0py's avatar
Isotr0py committed
107
108

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

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
120
121
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Isotr0py's avatar
Isotr0py committed
122
    ) -> 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, kv_cache, attn_metadata)
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
214
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Isotr0py's avatar
Isotr0py committed
215
216
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        # Attention block.
217
218
219
220
221
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(positions, hidden_states, kv_cache,
                                       attn_metadata)
        hidden_states = hidden_states + residual
Isotr0py's avatar
Isotr0py committed
222
223

        # MLP block.
224
225
226
227
        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
228
229
230
        return hidden_states


231
@support_torch_compile
Isotr0py's avatar
Isotr0py committed
232
233
class OlmoModel(nn.Module):

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

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

Isotr0py's avatar
Isotr0py committed
241
242
        self.config = config

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

257
258
259
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

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

        # Apply blocks one-by-one.
282
        for i in range(self.start_layer, self.end_layer):
Isotr0py's avatar
Isotr0py committed
283
            # shape: (batch_size, seq_len, d_model)
284
            hidden_states = self.layers[i](
Isotr0py's avatar
Isotr0py committed
285
                positions,
286
                hidden_states,
287
                kv_caches[i - self.start_layer],
288
                attn_metadata,
Isotr0py's avatar
Isotr0py committed
289
290
            )

291
292
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
Isotr0py's avatar
Isotr0py committed
293
294
        # Apply final layer norm.
        # shape: (batch_size, seq_len or 1, d_model)
295
296
        hidden_states = self.norm(hidden_states)
        return hidden_states
Isotr0py's avatar
Isotr0py committed
297
298


299
class OlmoForCausalLM(nn.Module, SupportsPP):
Isotr0py's avatar
Isotr0py committed
300
301
302
303
    """
    Extremely barebones HF model wrapper.
    """

304
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Isotr0py's avatar
Isotr0py committed
305
        super().__init__()
306
307
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
Isotr0py's avatar
Isotr0py committed
308
        self.config = config
309
310
        self.model = OlmoModel(vllm_config=vllm_config,
                               prefix=maybe_prefix(prefix, "model"))
311
        if config.tie_word_embeddings:
312
            self.lm_head = self.model.embed_tokens
313
314
315
316
317
318
        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,
319
                quant_config=quant_config,
320
            )
321
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
322
        self.sampler = get_sampler()
323
324
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Isotr0py's avatar
Isotr0py committed
325

326
327
328
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

Isotr0py's avatar
Isotr0py committed
329
330
331
332
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
333
334
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
335
        intermediate_tensors: Optional[IntermediateTensors] = None,
336
        inputs_embeds: Optional[torch.Tensor] = None,
337
    ) -> Union[torch.Tensor, IntermediateTensors]:
Isotr0py's avatar
Isotr0py committed
338
339
340
341
        hidden_states = self.model(
            input_ids=input_ids,
            positions=positions,
            kv_caches=kv_caches,
342
            attn_metadata=attn_metadata,
343
            intermediate_tensors=intermediate_tensors,
344
            inputs_embeds=inputs_embeds,
Isotr0py's avatar
Isotr0py committed
345
346
347
        )
        return hidden_states

348
349
350
351
352
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
353
        logits = self.logits_processor(self.lm_head, hidden_states,
354
355
356
                                       sampling_metadata)
        return logits

Isotr0py's avatar
Isotr0py committed
357
358
    def sample(
        self,
359
        logits: torch.Tensor,
Isotr0py's avatar
Isotr0py committed
360
361
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
362
        next_tokens = self.sampler(logits, sampling_metadata)
Isotr0py's avatar
Isotr0py committed
363
364
        return next_tokens

365
366
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
367
368
369
370
371
372
373
374
        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),
        ]
Isotr0py's avatar
Isotr0py committed
375
        params_dict = dict(self.named_parameters(remove_duplicate=False))
376
        loaded_params: Set[str] = set()
377
        for name, loaded_weight in weights:
378
379
380
381
382
383
384
            if "rotary_emb.inv_freq" in name:
                continue
            if ("rotary_emb.cos_cached" in name
                    or "rotary_emb.sin_cached" in name):
                # Models trained using ColossalAI may include these tensors in
                # the checkpoint. Skip them.
                continue
385
386
387
388
389
            # With tie_word_embeddings, we can skip lm_head.weight
            # The weight might appear unnecessarily in the files if the model is
            # processed with quantization, LoRA, fine-tuning, etc.
            if self.config.tie_word_embeddings and "lm_head.weight" in name:
                continue
390
391
392
393
394
395
396
            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
397
398
                if is_pp_missing_parameter(name, self):
                    continue
399
400
401
402
403
404
405
406
                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
407
408
                if is_pp_missing_parameter(name, self):
                    continue
409
410
411
412
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
413
414
            loaded_params.add(name)
        return loaded_params