olmo.py 13.1 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
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_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
39
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
40
from vllm.model_executor.layers.rotary_embedding import get_rope
Isotr0py's avatar
Isotr0py committed
41
from vllm.model_executor.layers.sampler import Sampler
42
from vllm.model_executor.layers.vocab_parallel_embedding import (
43
    ParallelLMHead, VocabParallelEmbedding)
44
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
Isotr0py's avatar
Isotr0py committed
45
46
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import SamplerOutput
47

Isotr0py's avatar
Isotr0py committed
48
49
50

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

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

        assert self.hidden_size % self.total_num_heads == 0
Isotr0py's avatar
Isotr0py committed
70
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
71

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

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

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

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

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


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

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

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

        # Activation function.
153
        self.act_fn = SiluAndMul()
Isotr0py's avatar
Isotr0py committed
154
155

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

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


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

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

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

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

        # MLP block.
214
215
216
217
        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
218
219
220
221
222
223
        return hidden_states


class OlmoModel(nn.Module):

    def __init__(self,
224
                 config: OlmoConfig,
225
                 cache_config: Optional[CacheConfig] = None,
226
                 quant_config: Optional[QuantizationConfig] = None):
Isotr0py's avatar
Isotr0py committed
227
228
229
        super().__init__()
        self.config = config

230
231
232
        self.embed_tokens = VocabParallelEmbedding(config.vocab_size,
                                                   config.hidden_size)
        self.layers = nn.ModuleList([
233
            OlmoDecoderLayer(config, cache_config, quant_config)
234
235
236
237
238
            for layer_idx in range(config.num_hidden_layers)
        ])
        self.norm = nn.LayerNorm(config.hidden_size,
                                 elementwise_affine=False,
                                 bias=False)
Isotr0py's avatar
Isotr0py committed
239
240
241
242
243

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
244
245
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
Isotr0py's avatar
Isotr0py committed
246
247
248
249
250
251
    ) -> torch.Tensor:
        """
        :param input_ids: A tensor of shape `(batch_size, seq_len)`.
        """
        # Get embeddings of input.
        # shape: (batch_size, seq_len, d_model)
252
253
254
255
        inputs_embeds = self.embed_tokens(input_ids)

        # embed positions
        hidden_states = inputs_embeds
Isotr0py's avatar
Isotr0py committed
256
257

        # Apply blocks one-by-one.
258
        for layer_idx, decoder_layer in enumerate(self.layers):
Isotr0py's avatar
Isotr0py committed
259
            # shape: (batch_size, seq_len, d_model)
260
            hidden_states = decoder_layer(
Isotr0py's avatar
Isotr0py committed
261
                positions,
262
263
                hidden_states,
                kv_caches[layer_idx],
264
                attn_metadata,
Isotr0py's avatar
Isotr0py committed
265
266
267
268
            )

        # Apply final layer norm.
        # shape: (batch_size, seq_len or 1, d_model)
269
270
        hidden_states = self.norm(hidden_states)
        return hidden_states
Isotr0py's avatar
Isotr0py committed
271
272


273
class OlmoForCausalLM(nn.Module):
Isotr0py's avatar
Isotr0py committed
274
275
276
277
278
    """
    Extremely barebones HF model wrapper.
    """

    def __init__(self,
279
                 config: OlmoConfig,
280
                 cache_config: Optional[CacheConfig] = None,
281
                 quant_config: Optional[QuantizationConfig] = None):
Isotr0py's avatar
Isotr0py committed
282
283
        super().__init__()
        self.config = config
284
        self.model = OlmoModel(config, cache_config, quant_config)
285
286
287
288
289
290
291
292
293
294
        if config.tie_word_embeddings:
            self.lm_head_weight = self.model.embed_tokens.weight
        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,
            )
            self.lm_head_weight = self.lm_head.weight
295
296
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
Isotr0py's avatar
Isotr0py committed
297
298
299
300
301

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
302
303
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
Isotr0py's avatar
Isotr0py committed
304
305
306
307
308
    ) -> torch.Tensor:
        hidden_states = self.model(
            input_ids=input_ids,
            positions=positions,
            kv_caches=kv_caches,
309
            attn_metadata=attn_metadata,
Isotr0py's avatar
Isotr0py committed
310
311
312
        )
        return hidden_states

313
314
315
316
317
318
    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        logits = self.logits_processor(self.lm_head_weight, hidden_states,
                                       sampling_metadata)
        return logits

Isotr0py's avatar
Isotr0py committed
319
320
    def sample(
        self,
321
        logits: torch.Tensor,
Isotr0py's avatar
Isotr0py committed
322
323
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
324
        next_tokens = self.sampler(logits, sampling_metadata)
Isotr0py's avatar
Isotr0py committed
325
326
        return next_tokens

327
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
328
329
330
331
332
333
334
335
        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
336
        params_dict = dict(self.named_parameters(remove_duplicate=False))
337
        for name, loaded_weight in weights:
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
            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
            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
                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
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)