| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2023 The vLLM team. |
| 5 | # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. |
| 6 | # |
| 7 | # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX |
| 8 | # and OPT implementations in this library. It has been modified from its |
| 9 | # original forms to accommodate minor architectural differences compared |
| 10 | # to GPT-NeoX and OPT used by the Meta AI team that trained the model. |
| 11 | # |
| 12 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 13 | # you may not use this file except in compliance with the License. |
| 14 | # You may obtain a copy of the License at |
| 15 | # |
| 16 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 17 | # |
| 18 | # Unless required by applicable law or agreed to in writing, software |
| 19 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 20 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 21 | # See the License for the specific language governing permissions and |
| 22 | # limitations under the License. |
| 23 | from collections.abc import Iterable |
| 24 | |
| 25 | import torch |
| 26 | import torch.nn as nn |
| 27 | |
| 28 | from vllm.config import VllmConfig |
| 29 | from vllm.model_executor.models.llama import LlamaForCausalLM, LlamaModel |
| 30 | |
| 31 | from .llama import LlamaAttention, LlamaDecoderLayer, LlamaMLP |
| 32 | from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper |
| 33 | |
| 34 | |
| 35 | class TeleChat2Model(LlamaModel): |
| 36 | hf_to_vllm_mapper = WeightsMapper( |
| 37 | orig_to_new_stacked={ |
| 38 | ".query": (".qkv_proj", "q"), |
| 39 | ".key": (".qkv_proj", "k"), |
| 40 | ".value": (".qkv_proj", "v"), |
| 41 | ".gate_proj": (".gate_up_proj", 0), |
| 42 | ".up_proj": (".gate_up_proj", 1), |
| 43 | } |
| 44 | ) |
| 45 | |
| 46 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 47 | hf_config = vllm_config.model_config.hf_config |
| 48 | |
| 49 | vllm_config.model_config.hf_config.attribute_map = { # type: ignore[misc] |
| 50 | "num_hidden_layers": "n_layer", |
| 51 | "num_attention_heads": "n_head", |
| 52 | "intermediate_size": "ffn_hidden_size", |
| 53 | "rms_norm_eps": "layer_norm_epsilon", |
| 54 | } |
| 55 | vllm_config.model_config.hf_config.hidden_act = "silu" |
| 56 | |
| 57 | # 1. Initialize the LlamaModel with bias |
| 58 | hf_config.bias = True |
| 59 | hf_config.mlp_bias = True |
| 60 | |
| 61 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 62 | # 2. Remove the bias from the qkv_proj and gate_up_proj based on config |
| 63 | # Telechat2's gate_up_proj and qkv_proj don't have bias |
| 64 | # see: https://github.com/vllm-project/vllm/pull/10311#issuecomment-2490297566 |
| 65 | for layer in self.layers: |
| 66 | if isinstance(layer, PPMissingLayer): |
| 67 | continue |
| 68 | self_attn = layer.self_attn |
| 69 | mlp = layer.mlp |
| 70 | assert isinstance(self_attn, LlamaAttention) |
| 71 | assert isinstance(mlp, LlamaMLP) |
| 72 | self_attn.qkv_proj.register_parameter("bias", None) |
| 73 | self_attn.qkv_proj.skip_bias_add = True |
| 74 | mlp.gate_up_proj.register_parameter("bias", None) |
| 75 | mlp.gate_up_proj.skip_bias_add = True |
| 76 | |
| 77 | def _split_key_value( |
| 78 | self, weights: Iterable[tuple[str, torch.Tensor]] |
| 79 | ) -> Iterable[tuple[str, torch.Tensor]]: |
| 80 | # TeleChat2 stores k/v as a single per-head-interleaved `key_value` |
| 81 | # tensor. De-interleave it into separate k/v so the qkv_proj mapper |
| 82 | # can stack them. |
| 83 | total_num_heads = self.config.n_head |
| 84 | head_dim = self.config.hidden_size // total_num_heads |
| 85 | for name, loaded_weight in weights: |
| 86 | if "self_attn.key_value" in name: |
| 87 | starts = [i * head_dim * 2 for i in range(total_num_heads)] |
| 88 | k_weight = torch.cat( |
| 89 | [loaded_weight[s : s + head_dim, :] for s in starts], dim=0 |
| 90 | ) |
| 91 | v_weight = torch.cat( |
| 92 | [loaded_weight[s + head_dim : s + 2 * head_dim, :] for s in starts], |
| 93 | dim=0, |
| 94 | ) |
| 95 | yield name.replace("key_value", "key"), k_weight |
| 96 | yield name.replace("key_value", "value"), v_weight |
| 97 | else: |
| 98 | yield name, loaded_weight |
| 99 | |
| 100 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 101 | loader = AutoWeightsLoader(self) |
| 102 | return loader.load_weights( |
| 103 | self._split_key_value(weights), mapper=self.hf_to_vllm_mapper |
| 104 | ) |
| 105 | |
| 106 | |
| 107 | class TeleChat2ForCausalLM(LlamaForCausalLM): |
| 108 | hf_to_vllm_mapper = WeightsMapper( |
| 109 | orig_to_new_prefix={ |
| 110 | "transformer.": "model.", |
| 111 | }, |
| 112 | orig_to_new_substr={ |
| 113 | ".h.": ".layers.", |
| 114 | ".self_attention.": ".self_attn.", |
| 115 | ".word_embeddings.": ".embed_tokens.", |
| 116 | ".dense.": ".o_proj.", |
| 117 | ".ln_f.": ".norm.", |
| 118 | }, |
| 119 | ) |
| 120 | |
| 121 | def _init_model( |
| 122 | self, |
| 123 | vllm_config: VllmConfig, |
| 124 | prefix: str = "", |
| 125 | layer_type: type[nn.Module] = LlamaDecoderLayer, |
| 126 | ): |
| 127 | return TeleChat2Model(vllm_config=vllm_config, prefix=prefix) |
| 128 | |
| 129 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 130 | loader = AutoWeightsLoader(self) |
| 131 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |