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

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

30
from vllm.attention import Attention, AttentionMetadata
31
from vllm.config import CacheConfig
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
40
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
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
49
50
from .interfaces import SupportsPP
from .utils import (is_pp_missing_parameter,
                    make_empty_intermediate_tensors_factory, make_layers)

Isotr0py's avatar
Isotr0py committed
51
52
53

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

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

        assert self.hidden_size % self.total_num_heads == 0
Isotr0py's avatar
Isotr0py committed
73
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
74

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

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

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

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

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
117
118
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Isotr0py's avatar
Isotr0py committed
119
    ) -> torch.Tensor:
120
121
122
        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
123
        q, k, v = qkv.chunk(chunks=3, dim=-1)
124
        q, k = self.rotary_emb(positions, q, k)
125
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
126
        output, _ = self.o_proj(attn_output)
Isotr0py's avatar
Isotr0py committed
127
128
129
130
131
        return output


class OlmoMLP(nn.Module):
    """
132
133
    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
134
135
136
137
138
    (plus another skip connection).
    """

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

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

        # Activation function.
156
        self.act_fn = SiluAndMul()
Isotr0py's avatar
Isotr0py committed
157
158

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

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


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

    def __init__(self,
184
                 config: OlmoConfig,
185
                 cache_config: Optional[CacheConfig] = None,
186
                 quant_config: Optional[QuantizationConfig] = None):
Isotr0py's avatar
Isotr0py committed
187
188
        super().__init__()
        # Attention block.
189
        self.self_attn = OlmoAttention(config, cache_config, quant_config)
Isotr0py's avatar
Isotr0py committed
190
191

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

194
195
196
197
198
199
200
201
        # 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
202
203
204
205
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
206
207
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Isotr0py's avatar
Isotr0py committed
208
209
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        # Attention block.
210
211
212
213
214
        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
215
216

        # MLP block.
217
218
219
220
        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
221
222
223
224
225
226
        return hidden_states


class OlmoModel(nn.Module):

    def __init__(self,
227
                 config: OlmoConfig,
228
                 cache_config: Optional[CacheConfig] = None,
229
230
                 quant_config: Optional[QuantizationConfig] = None,
                 prefix: str = ""):
Isotr0py's avatar
Isotr0py committed
231
232
233
        super().__init__()
        self.config = config

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

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
252
253
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
254
255
        intermediate_tensors: Optional[IntermediateTensors],
    ) -> Union[torch.Tensor, IntermediateTensors]:
Isotr0py's avatar
Isotr0py committed
256
257
258
        """
        :param input_ids: A tensor of shape `(batch_size, seq_len)`.
        """
259
260
261
262
        if get_pp_group().is_first_rank:
            # Get embeddings of input.
            # shape: (batch_size, seq_len, d_model)
            inputs_embeds = self.embed_tokens(input_ids)
263

264
265
266
267
268
            # embed positions
            hidden_states = inputs_embeds
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
Isotr0py's avatar
Isotr0py committed
269
270

        # Apply blocks one-by-one.
271
        for i in range(self.start_layer, self.end_layer):
Isotr0py's avatar
Isotr0py committed
272
            # shape: (batch_size, seq_len, d_model)
273
            hidden_states = self.layers[i](
Isotr0py's avatar
Isotr0py committed
274
                positions,
275
                hidden_states,
276
                kv_caches[i - self.start_layer],
277
                attn_metadata,
Isotr0py's avatar
Isotr0py committed
278
279
            )

280
281
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
Isotr0py's avatar
Isotr0py committed
282
283
        # Apply final layer norm.
        # shape: (batch_size, seq_len or 1, d_model)
284
285
        hidden_states = self.norm(hidden_states)
        return hidden_states
Isotr0py's avatar
Isotr0py committed
286
287


288
class OlmoForCausalLM(nn.Module, SupportsPP):
Isotr0py's avatar
Isotr0py committed
289
290
291
292
293
    """
    Extremely barebones HF model wrapper.
    """

    def __init__(self,
294
                 config: OlmoConfig,
295
                 cache_config: Optional[CacheConfig] = None,
296
                 quant_config: Optional[QuantizationConfig] = None):
Isotr0py's avatar
Isotr0py committed
297
298
        super().__init__()
        self.config = config
299
        self.model = OlmoModel(config, cache_config, quant_config)
300
        if config.tie_word_embeddings:
301
            self.lm_head = self.model.embed_tokens
302
303
304
305
306
307
        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,
308
                quant_config=quant_config,
309
            )
310
311
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
312
313
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Isotr0py's avatar
Isotr0py committed
314
315
316
317
318

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
319
320
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
321
        intermediate_tensors: Optional[IntermediateTensors] = None,
322
    ) -> Union[torch.Tensor, IntermediateTensors]:
Isotr0py's avatar
Isotr0py committed
323
324
325
326
        hidden_states = self.model(
            input_ids=input_ids,
            positions=positions,
            kv_caches=kv_caches,
327
            attn_metadata=attn_metadata,
328
            intermediate_tensors=intermediate_tensors,
Isotr0py's avatar
Isotr0py committed
329
330
331
        )
        return hidden_states

332
333
334
335
336
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
337
        logits = self.logits_processor(self.lm_head, hidden_states,
338
339
340
                                       sampling_metadata)
        return logits

Isotr0py's avatar
Isotr0py committed
341
342
    def sample(
        self,
343
        logits: torch.Tensor,
Isotr0py's avatar
Isotr0py committed
344
345
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
346
        next_tokens = self.sampler(logits, sampling_metadata)
Isotr0py's avatar
Isotr0py committed
347
348
        return next_tokens

349
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
350
351
352
353
354
355
356
357
        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
358
        params_dict = dict(self.named_parameters(remove_duplicate=False))
359
        for name, loaded_weight in weights:
360
361
362
363
364
365
366
            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
367
368
369
370
371
            # 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
372
373
374
375
376
377
378
            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
379
380
                if is_pp_missing_parameter(name, self):
                    continue
381
382
383
384
385
386
387
388
                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
389
390
                if is_pp_missing_parameter(name, self):
                    continue
391
392
393
394
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)