decilm.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. # coding=utf-8
  2. # Adapted from
  3. # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
  4. # Copyright 2023 DeciAI Research Team. All rights reserved.
  5. # Copyright 2023 The vLLM team.
  6. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
  7. #
  8. # This code is based on MistralAI GPT-NeoX library and the GPT-NeoX
  9. # and OPT implementations in this library. It has been modified from its
  10. # original forms to accommodate minor architectural differences compared
  11. # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
  12. #
  13. # Licensed under the Apache License, Version 2.0 (the "License");
  14. # you may not use this file except in compliance with the License.
  15. # You may obtain a copy of the License at
  16. #
  17. # http://www.apache.org/licenses/LICENSE-2.0
  18. #
  19. # Unless required by applicable law or agreed to in writing, software
  20. # distributed under the License is distributed on an "AS IS" BASIS,
  21. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  22. # See the License for the specific language governing permissions and
  23. # limitations under the License.
  24. """Inference-only DeciLM model compatible with HuggingFace weights."""
  25. from typing import Iterable, Optional, Tuple
  26. import torch
  27. from transformers import LlamaConfig
  28. from aphrodite.common.config import CacheConfig, LoRAConfig
  29. from aphrodite.modeling.model_loader.weight_utils import default_weight_loader
  30. from aphrodite.modeling.models.llama import LlamaForCausalLM
  31. from aphrodite.quantization.base_config import QuantizationConfig
  32. class DeciLMForCausalLM(LlamaForCausalLM):
  33. """
  34. Implementation for https://huggingface.co/Deci/DeciLM-7b-instruct.
  35. Based on the llama executor.
  36. The main difference is that DeciLM uses Variable Grouped Query Attention.
  37. The constant number of GQA heads in the decoder is overridden with a value
  38. per layer.
  39. Usually, in the HuggingFace implementation, instead of
  40. "config.num_key_value_heads", we use
  41. "config.num_key_value_heads_per_layer[i]" which varies.
  42. Currently, PagedAttention does not work well with variable GQA, so we
  43. normalize the weights upon loading, and use uniform GQA with the max value
  44. instead.
  45. """
  46. def __init__(
  47. self,
  48. config: LlamaConfig,
  49. cache_config: Optional[CacheConfig] = None,
  50. quant_config: Optional[QuantizationConfig] = None,
  51. lora_config: Optional[LoRAConfig] = None,
  52. ) -> None:
  53. config.num_key_value_heads = max(config.num_key_value_heads_per_layer)
  54. delattr(config, "num_key_value_heads_per_layer")
  55. super().__init__(config=config,
  56. cache_config=cache_config,
  57. quant_config=quant_config,
  58. lora_config=lora_config)
  59. def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
  60. stacked_params_mapping = [
  61. # (param_name, shard_name, shard_id)
  62. ("qkv_proj", "q_proj", "q"),
  63. ("qkv_proj", "k_proj", "k"),
  64. ("qkv_proj", "v_proj", "v"),
  65. ("gate_up_proj", "gate_proj", 0),
  66. ("gate_up_proj", "up_proj", 1),
  67. ]
  68. params_dict = dict(self.named_parameters())
  69. for name, loaded_weight in weights:
  70. if "rotary_emb.inv_freq" in name:
  71. continue
  72. if "k_proj" in name or "v_proj" in name:
  73. loaded_weight = self._degroup_weight(loaded_weight)
  74. for (param_name, weight_name, shard_id) in stacked_params_mapping:
  75. if weight_name not in name:
  76. continue
  77. name = name.replace(weight_name, param_name)
  78. # Skip loading extra bias for GPTQ models.
  79. if name.endswith(".bias") and name not in params_dict:
  80. continue
  81. param = params_dict[name]
  82. weight_loader = param.weight_loader
  83. weight_loader(param, loaded_weight, shard_id)
  84. break
  85. else:
  86. # Skip loading extra bias for GPTQ models.
  87. if name.endswith(".bias") and name not in params_dict:
  88. continue
  89. param = params_dict[name]
  90. weight_loader = getattr(param, "weight_loader",
  91. default_weight_loader)
  92. weight_loader(param, loaded_weight)
  93. def _degroup_weight(self, loaded_weight: torch.Tensor) -> torch.Tensor:
  94. hidden_size = self.config.hidden_size
  95. head_size = self.config.hidden_size // self.config.num_attention_heads
  96. target_num_kv_heads = self.config.num_key_value_heads
  97. num_kv_heads = loaded_weight.shape[0] // head_size
  98. n_repeats = target_num_kv_heads / num_kv_heads
  99. assert n_repeats == int(n_repeats)
  100. n_repeats = int(n_repeats)
  101. loaded_weight = loaded_weight.view(num_kv_heads, head_size,
  102. hidden_size)
  103. loaded_weight = torch.repeat_interleave(loaded_weight,
  104. repeats=n_repeats,
  105. dim=0)
  106. loaded_weight = loaded_weight.reshape(target_num_kv_heads * head_size,
  107. hidden_size)
  108. return loaded_weight