# coding=utf-8 # adapted from https://github.com/huggingface/transformers/blob/v4.39.3/src/transformers/models/persimmon/modeling_persimmon.py # Copyright 2023 The PygmalionAI team. # Copyright 2023 The vLLM team. # Copyright 2023 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # 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. # # 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. """Inference-only persimmon model compatible with HuggingFace weights.""" from typing import Iterable, List, Optional, Tuple import torch from torch import nn from transformers import PersimmonConfig from transformers.activations import ReLUSquaredActivation from aphrodite.attention import Attention, AttentionMetadata from aphrodite.common.config import CacheConfig from aphrodite.common.sequence import IntermediateTensors, SamplerOutput from aphrodite.common.utils import progress_bar from aphrodite.distributed import get_tensor_model_parallel_world_size from aphrodite.modeling.layers.linear import (ColumnParallelLinear, QKVParallelLinear, RowParallelLinear) from aphrodite.modeling.layers.logits_processor import LogitsProcessor from aphrodite.modeling.layers.rotary_embedding import get_rope from aphrodite.modeling.layers.sampler import Sampler from aphrodite.modeling.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding) from aphrodite.modeling.model_loader.weight_utils import default_weight_loader from aphrodite.modeling.sampling_metadata import SamplingMetadata from aphrodite.quantization.base_config import QuantizationConfig class PersimmonMLP(nn.Module): def __init__(self, config: PersimmonConfig, quant_config: Optional[QuantizationConfig] = None): super().__init__() self.dense_h_to_4h = ColumnParallelLinear(config.hidden_size, config.intermediate_size, quant_config=quant_config) self.dense_4h_to_h = RowParallelLinear(config.intermediate_size, config.hidden_size, quant_config=quant_config) self.act = ReLUSquaredActivation() def forward(self, hidden_states) -> torch.Tensor: hidden_states, _ = self.dense_h_to_4h(hidden_states) hidden_states = self.act(hidden_states) hidden_states, _ = self.dense_4h_to_h(hidden_states) return hidden_states class PersimmonAttention(nn.Module): def __init__(self, config: PersimmonConfig, cache_config: Optional[CacheConfig] = None, quant_config: Optional[QuantizationConfig] = None): super().__init__() self.config = config tensor_parallel_world_size = get_tensor_model_parallel_world_size() self.hidden_size = config.hidden_size self.total_num_heads = config.num_attention_heads self.num_heads = self.total_num_heads // tensor_parallel_world_size self.head_dim = self.hidden_size // self.total_num_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.partial_rotary_factor = config.partial_rotary_factor self.is_causal = True assert (self.head_dim * self.total_num_heads) == self.hidden_size assert self.total_num_heads % tensor_parallel_world_size == 0 self.query_key_value = QKVParallelLinear( self.hidden_size, self.head_dim, self.total_num_heads, bias=True, quant_config=quant_config, ) self.dense = RowParallelLinear( self.num_heads * self.head_dim, self.hidden_size, bias=True, quant_config=quant_config, ) self.is_qk_layernorm = config.qk_layernorm if self.is_qk_layernorm: self.q_layernorm = nn.LayerNorm(self.head_dim) self.k_layernorm = nn.LayerNorm(self.head_dim) self.rotary_emb = get_rope( self.head_dim, rotary_dim=int(self.partial_rotary_factor * self.head_dim), max_position=self.max_position_embeddings, base=self.rope_theta, ) self.scaling = self.head_dim**-0.5 self.attn = Attention(self.num_heads, self.head_dim, scale=self.scaling, cache_config=cache_config, quant_config=quant_config) def _split_heads(self, x: torch.Tensor) -> torch.Tensor: # [seq_length, hidden_size] -> [seq_length, num_heads, head_dim] seq_length = x.shape[0] return x.view(seq_length, self.num_heads, self.head_dim) def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: # [seq_length, num_heads, head_dim] -> [seq_length, hidden_size] seq_length = x.shape[0] return x.view(seq_length, self.num_heads * self.head_dim) def forward( self, position_ids: torch.Tensor, hidden_states: torch.Tensor, kv_cache: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: # [seq_length, 3 x hidden_size] qkv, _ = self.query_key_value(hidden_states) q, k, v = qkv.chunk(chunks=3, dim=-1) if self.is_qk_layernorm: # [seq_length, num_heads, head_dim] q = self._split_heads(q) k = self._split_heads(k) q = self.q_layernorm(q) k = self.k_layernorm(k) q = self._merge_heads(q) k = self._merge_heads(k) q, k = self.rotary_emb(position_ids, q, k) attn_output = self.attn(q, k, v, kv_cache, attn_metadata) output, _ = self.dense(attn_output) return output class PersimmonDecoderLayer(nn.Module): def __init__(self, config: PersimmonConfig, cache_config: Optional[CacheConfig] = None, quant_config: Optional[QuantizationConfig] = None): super().__init__() self.hidden_size = config.hidden_size self.self_attn = PersimmonAttention(config=config, cache_config=cache_config, quant_config=quant_config) self.mlp = PersimmonMLP(config, quant_config=quant_config) self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward( self, position_ids: torch.Tensor, hidden_states: torch.Tensor, kv_cache: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states = self.self_attn( position_ids=position_ids, hidden_states=hidden_states, kv_cache=kv_cache, attn_metadata=attn_metadata, ) 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 = hidden_states + residual outputs = hidden_states return outputs class PersimmonModel(nn.Module): def __init__(self, config: PersimmonConfig, cache_config: Optional[CacheConfig] = None, quant_config: Optional[QuantizationConfig] = None): super().__init__() self.vocab_size = config.vocab_size self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList([ PersimmonDecoderLayer(config, cache_config=cache_config, quant_config=quant_config) for _ in range(config.num_hidden_layers) ]) self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, kv_caches: List[torch.Tensor], attn_metadata: AttentionMetadata, inputs_embeds: Optional[torch.Tensor] = None, ) -> torch.Tensor: if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.embed_tokens(input_ids) for i in range(len(self.layers)): hidden_states = self.layers[i]( positions, hidden_states, kv_caches[i], attn_metadata, ) hidden_states = self.final_layernorm(hidden_states) return hidden_states class PersimmonForCausalLM(nn.Module): def __init__(self, config, cache_config: Optional[CacheConfig] = None, quant_config: Optional[QuantizationConfig] = None): super().__init__() self.config = config self.vocab_size = config.vocab_size self.model = PersimmonModel(config, cache_config=cache_config, quant_config=quant_config) self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, bias=False) self.logits_processor = LogitsProcessor(config.vocab_size) self.sampler = Sampler() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, kv_caches: List[torch.Tensor], attn_metadata: AttentionMetadata, intermediate_tensors: Optional[IntermediateTensors] = None, inputs_embeds: Optional[torch.Tensor] = None, ): hidden_states = self.model( input_ids=input_ids, positions=positions, kv_caches=kv_caches, attn_metadata=attn_metadata, inputs_embeds=inputs_embeds, ) return hidden_states def compute_logits( self, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata, ) -> Optional[torch.Tensor]: logits = self.logits_processor(self.lm_head, hidden_states, sampling_metadata) return logits def sample( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, ) -> Optional[SamplerOutput]: next_tokens = self.sampler(logits, sampling_metadata) return next_tokens def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): params_dict = dict(self.named_parameters(remove_duplicate=False)) weights_list = list(weights) for name, loaded_weight in progress_bar(weights_list, desc="Loading modules..."): 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 param = params_dict[name] if "query_key_value" in name: # copy from vllm/model_executor/models/bloom.py # NOTE: Persimmon's fused QKV's output_dim has the shape of # (num_heads * 3 * head_size), while the # required shape is (3 * num_heads * head_size). # Thus, we need weight conversion. output_dim = getattr(param, "output_dim", None) num_heads = self.config.num_attention_heads if output_dim is not None: loaded_weight_shape = loaded_weight.shape loaded_weight = loaded_weight.view( loaded_weight_shape[:output_dim] + (num_heads, 3, -1) + loaded_weight_shape[output_dim + 1:]) loaded_weight = loaded_weight.transpose( output_dim, output_dim + 1) loaded_weight = loaded_weight.reshape(loaded_weight_shape) weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight)