olmo.py 14.9 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, 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
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
        return hidden_states


224
@support_torch_compile
Isotr0py's avatar
Isotr0py committed
225
226
227
class OlmoModel(nn.Module):

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

235
236
        self.embed_tokens = VocabParallelEmbedding(config.vocab_size,
                                                   config.hidden_size)
237
238
239
240
241
        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")
242
243
244
        self.norm = nn.LayerNorm(config.hidden_size,
                                 elementwise_affine=False,
                                 bias=False)
245
246
247
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))
Isotr0py's avatar
Isotr0py committed
248
249
250
251
252

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
253
254
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
255
256
        intermediate_tensors: Optional[IntermediateTensors],
    ) -> Union[torch.Tensor, IntermediateTensors]:
Isotr0py's avatar
Isotr0py committed
257
258
259
        """
        :param input_ids: A tensor of shape `(batch_size, seq_len)`.
        """
260
261
262
263
        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)
264

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

        # Apply blocks one-by-one.
272
        for i in range(self.start_layer, self.end_layer):
Isotr0py's avatar
Isotr0py committed
273
            # shape: (batch_size, seq_len, d_model)
274
            hidden_states = self.layers[i](
Isotr0py's avatar
Isotr0py committed
275
                positions,
276
                hidden_states,
277
                kv_caches[i - self.start_layer],
278
                attn_metadata,
Isotr0py's avatar
Isotr0py committed
279
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
class OlmoForCausalLM(nn.Module, SupportsPP):
Isotr0py's avatar
Isotr0py committed
290
291
292
293
    """
    Extremely barebones HF model wrapper.
    """

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

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

337
338
339
340
341
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
342
        logits = self.logits_processor(self.lm_head, hidden_states,
343
344
345
                                       sampling_metadata)
        return logits

Isotr0py's avatar
Isotr0py committed
346
347
    def sample(
        self,
348
        logits: torch.Tensor,
Isotr0py's avatar
Isotr0py committed
349
350
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
351
        next_tokens = self.sampler(logits, sampling_metadata)
Isotr0py's avatar
Isotr0py committed
352
353
        return next_tokens

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