olmo.py 12.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
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.distributed import get_tensor_model_parallel_world_size
32
from vllm.model_executor.layers.activation import SiluAndMul
33
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
34
35
                                               QKVParallelLinear,
                                               RowParallelLinear)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
38
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
39
from vllm.model_executor.layers.rotary_embedding import get_rope
Isotr0py's avatar
Isotr0py committed
40
from vllm.model_executor.layers.sampler import 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
45
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import SamplerOutput
46

Isotr0py's avatar
Isotr0py committed
47
48
49

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

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

        assert self.hidden_size % self.total_num_heads == 0
Isotr0py's avatar
Isotr0py committed
68
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
69

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

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

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

        # Attention output projection.
99
100
101
102
        self.o_proj = RowParallelLinear(
            self.hidden_size,
            self.hidden_size,
            bias=config.attention_bias,
103
            quant_config=quant_config,
Isotr0py's avatar
Isotr0py committed
104
105
106
107
108
109
        )

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


class OlmoMLP(nn.Module):
    """
125
126
    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
127
128
129
130
131
    (plus another skip connection).
    """

    def __init__(
        self,
132
        config: OlmoConfig,
133
        quant_config: Optional[QuantizationConfig] = None,
Isotr0py's avatar
Isotr0py committed
134
135
136
    ):
        super().__init__()
        self.config = config
137
138
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
Isotr0py's avatar
Isotr0py committed
139
140

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

        # Activation function.
149
        self.act_fn = SiluAndMul()
Isotr0py's avatar
Isotr0py committed
150
151

        # Feed-forward output projection.
152
153
154
155
        self.down_proj = RowParallelLinear(
            self.intermediate_size,
            self.hidden_size,
            bias=False,
156
            quant_config=quant_config,
Isotr0py's avatar
Isotr0py committed
157
158
159
160
161
162
        )

    def forward(
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
163
164
165
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
Isotr0py's avatar
Isotr0py committed
166
167
168
        return x


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

    def __init__(self,
177
                 config: OlmoConfig,
178
                 quant_config: Optional[QuantizationConfig] = None):
Isotr0py's avatar
Isotr0py committed
179
180
        super().__init__()
        # Attention block.
181
        self.self_attn = OlmoAttention(config, quant_config)
Isotr0py's avatar
Isotr0py committed
182
183

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

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

        # MLP block.
209
210
211
212
        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
213
214
215
216
217
218
        return hidden_states


class OlmoModel(nn.Module):

    def __init__(self,
219
                 config: OlmoConfig,
220
                 quant_config: Optional[QuantizationConfig] = None):
Isotr0py's avatar
Isotr0py committed
221
222
223
        super().__init__()
        self.config = config

224
225
226
        self.embed_tokens = VocabParallelEmbedding(config.vocab_size,
                                                   config.hidden_size)
        self.layers = nn.ModuleList([
227
            OlmoDecoderLayer(config, quant_config)
228
229
230
231
232
            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
233
234
235
236
237

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

        # embed positions
        hidden_states = inputs_embeds
Isotr0py's avatar
Isotr0py committed
250
251

        # Apply blocks one-by-one.
252
        for layer_idx, decoder_layer in enumerate(self.layers):
Isotr0py's avatar
Isotr0py committed
253
            # shape: (batch_size, seq_len, d_model)
254
            hidden_states = decoder_layer(
Isotr0py's avatar
Isotr0py committed
255
                positions,
256
257
                hidden_states,
                kv_caches[layer_idx],
258
                attn_metadata,
Isotr0py's avatar
Isotr0py committed
259
260
261
262
            )

        # Apply final layer norm.
        # shape: (batch_size, seq_len or 1, d_model)
263
264
        hidden_states = self.norm(hidden_states)
        return hidden_states
Isotr0py's avatar
Isotr0py committed
265
266


267
class OlmoForCausalLM(nn.Module):
Isotr0py's avatar
Isotr0py committed
268
269
270
271
272
    """
    Extremely barebones HF model wrapper.
    """

    def __init__(self,
273
                 config: OlmoConfig,
274
                 quant_config: Optional[QuantizationConfig] = None):
Isotr0py's avatar
Isotr0py committed
275
276
        super().__init__()
        self.config = config
277
        self.model = OlmoModel(config, quant_config)
278
279
280
281
282
283
284
285
286
287
        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
288
289
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
Isotr0py's avatar
Isotr0py committed
290
291
292
293
294

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
295
296
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
Isotr0py's avatar
Isotr0py committed
297
298
299
300
301
    ) -> torch.Tensor:
        hidden_states = self.model(
            input_ids=input_ids,
            positions=positions,
            kv_caches=kv_caches,
302
            attn_metadata=attn_metadata,
Isotr0py's avatar
Isotr0py committed
303
304
305
        )
        return hidden_states

306
307
308
309
310
311
    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
312
313
    def sample(
        self,
314
        logits: torch.Tensor,
Isotr0py's avatar
Isotr0py committed
315
316
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
317
        next_tokens = self.sampler(logits, sampling_metadata)
Isotr0py's avatar
Isotr0py committed
318
319
        return next_tokens

320
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
321
322
323
324
325
326
327
328
        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
329
        params_dict = dict(self.named_parameters(remove_duplicate=False))
330
        for name, loaded_weight in weights:
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
            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)