stablelm.py 13.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
# Copyright 2023 Stability AI, EleutherAI, and The HuggingFace Inc. team.
# All rights reserved.
Hyunsung Lee's avatar
Hyunsung Lee committed
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#
# 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
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.
#
# This code is based off the following work:
# https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/modeling_stablelm_epoch.py
# https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/config.json
22
23
"""Inference-only StabeLM (https://github.com/Stability-AI/StableLM)
model compatible with HuggingFace weights."""
24

25
from collections.abc import Iterable
26
from itertools import islice
27
from typing import Optional, Union
Hyunsung Lee's avatar
Hyunsung Lee committed
28

29
30
import torch
from torch import nn
31
from transformers import StableLmConfig
Hyunsung Lee's avatar
Hyunsung Lee committed
32

33
from vllm.attention import Attention
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.model_executor.layers.activation import SiluAndMul
37
38
39
40
41
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
42
from vllm.model_executor.layers.logits_processor import LogitsProcessor
43
from vllm.model_executor.layers.quantization import QuantizationConfig
44
from vllm.model_executor.layers.rotary_embedding import get_rope
45
from vllm.model_executor.layers.vocab_parallel_embedding import (
46
47
48
    ParallelLMHead,
    VocabParallelEmbedding,
)
49
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
50
from vllm.sequence import IntermediateTensors
Hyunsung Lee's avatar
Hyunsung Lee committed
51

52
from .interfaces import SupportsPP
53
54
55
56
57
58
59
from .utils import (
    AutoWeightsLoader,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
60

61
62

class StablelmMLP(nn.Module):
63
64
65
66
67
68
    def __init__(
        self,
        config: StableLmConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
69
70
71
72
73
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_up_proj = MergedColumnParallelLinear(
74
75
            config.hidden_size,
            [config.intermediate_size] * 2,
76
            bias=False,
77
            quant_config=quant_config,
78
79
80
81
82
83
84
85
86
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            config.intermediate_size,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
87
88
89
90
91
92
93
94
95
96
        self.act_fn = SiluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class StablelmAttention(nn.Module):
97
98
99
100
101
102
103
    def __init__(
        self,
        config: StableLmConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = config.num_attention_heads
        self.num_heads = self.total_num_heads // tp_size

        self.total_num_key_value_heads = config.num_key_value_heads
        if self.total_num_key_value_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_key_value_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_key_value_heads == 0
120
        self.num_key_value_heads = max(1, self.total_num_key_value_heads // tp_size)
121
122
        self.head_dim = self.hidden_size // self.total_num_heads
        self.max_position_embeddings = config.max_position_embeddings
123
        self.partial_rotary_factor = getattr(
124
125
            config, "rope_pct", getattr(config, "partial_rotary_factor", 1)
        )
126
127
128
129
130
        self.scaling = self.head_dim**-0.5
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_key_value_heads * self.head_dim
        self.qkv_bias = getattr(config, "use_qkv_bias", False)
        if (self.head_dim * self.num_heads * tp_size) != self.hidden_size:
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
            raise ValueError(
                f"hidden_size must be divisible by num_heads "
                f"(got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )

        self.qkv_proj = QKVParallelLinear(
            self.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_key_value_heads,
            self.qkv_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )
153
154
        self.rotary_emb = get_rope(
            self.head_dim,
155
            rotary_dim=self.head_dim,
156
157
            max_position=self.config.max_position_embeddings,
            base=self.config.rope_theta,
158
            partial_rotary_factor=self.partial_rotary_factor,
159
        )
160
161
162
163
164
165
166
167
168
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_key_value_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
        )
169
170
171
172
173
174
175
176
177

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
178
        attn_output = self.attn(q, k, v)
179
180
181
182
183
        output, _ = self.o_proj(attn_output)
        return output


class StablelmDecoderLayer(nn.Module):
Hyunsung Lee's avatar
Hyunsung Lee committed
184
185
    def __init__(
        self,
186
        config: StableLmConfig,
187
        cache_config: Optional[CacheConfig] = None,
188
        quant_config: Optional[QuantizationConfig] = None,
189
        prefix: str = "",
Hyunsung Lee's avatar
Hyunsung Lee committed
190
    ) -> None:
191
        super().__init__()
192
193
194
        self.self_attn = StablelmAttention(
            config, cache_config, quant_config, prefix=f"{prefix}.self_attn"
        )
195
        self.mlp = StablelmMLP(config, quant_config, prefix=f"{prefix}.mlp")
196
        norm_eps = getattr(config, "norm_eps", getattr(config, "layer_norm_eps", 1e-05))
Roy's avatar
Roy committed
197
        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=norm_eps)
198
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=norm_eps)
199
200
201
202
203

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
204
    ) -> tuple[torch.Tensor, torch.Tensor]:
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
        # Self Attention
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states, residual


class StableLMEpochModel(nn.Module):
224
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
225
        super().__init__()
226
227
228
229
230

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

231
232
233
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
234
235
            quant_config=quant_config,
            prefix=f"{prefix}.embed_tokens",
236
        )
237
238
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
239
            lambda prefix: StablelmDecoderLayer(
240
241
                config, cache_config, quant_config, prefix=prefix
            ),
242
243
            prefix=f"{prefix}.layers",
        )
244
        norm_eps = getattr(config, "norm_eps", getattr(config, "layer_norm_eps", 1e-05))
Roy's avatar
Roy committed
245
        self.norm = nn.LayerNorm(config.hidden_size, eps=norm_eps)
246
247
248
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states"], config.hidden_size
        )
249

250
251
252
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

253
254
255
256
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
257
        intermediate_tensors: Optional[IntermediateTensors],
258
        inputs_embeds: Optional[torch.Tensor] = None,
259
260
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
261
262
263
264
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
265
266
267
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
268
        for layer in islice(self.layers, self.start_layer, self.end_layer):
269
            hidden_states, residual = layer(positions, hidden_states)
270
271
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
272
273
274
        hidden_states = self.norm(hidden_states)
        return hidden_states

275
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
276
277
278
279
280
281
282
283
284
        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),
        ]
        params_dict = dict(self.named_parameters())
285
        loaded_params: set[str] = set()
286
        for name, loaded_weight in weights:
287
            for param_name, weight_name, shard_id in stacked_params_mapping:
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
                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
                if is_pp_missing_parameter(name, self):
                    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
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
307
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
308
309
310
311
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

312

313
class StablelmForCausalLM(nn.Module, SupportsPP):
314
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
315
        super().__init__()
316
317
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
318
        self.config = config
319
        self.quant_config = quant_config
320
321
322
323
324
325
326
327
328
        self.model = StableLMEpochModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
        self.lm_head = ParallelLMHead(
            config.vocab_size,
            config.hidden_size,
            quant_config=quant_config,
            prefix=f"{prefix}.lm_head",
        )
329
330
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
331
        self.logits_processor = LogitsProcessor(config.vocab_size)
332
        self.make_empty_intermediate_tensors = (
333
334
            self.model.make_empty_intermediate_tensors
        )
335

336
337
338
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

339
340
341
342
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
343
        intermediate_tensors: Optional[IntermediateTensors] = None,
344
        inputs_embeds: Optional[torch.Tensor] = None,
345
    ) -> Union[torch.Tensor, IntermediateTensors]:
346
347
348
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
349
350
        return hidden_states

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

358
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
359
        loader = AutoWeightsLoader(self)
360
        return loader.load_weights(weights)