TheDrummer/Valkyrie-49B-v2

1# coding=utf-82# Copyright 2024 Nvidia Corporation. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16from copy import deepcopy17from typing import Optional, Dict, Any, Tuple18 19import torch20from transformers.cache_utils import Cache  # used to let GenerationMixin know that we use a Cache object21 22from .configuration_decilm import DeciLMConfig23from .transformers_4_44_2__cache_utils import Cache as Cache_4_44_2, SinkCache, StaticCache, SlidingWindowCache24 25 26class VariableCache(Cache_4_44_2, Cache):27    """28    A Cache object that supports a different Cache implementation for every layer,29    including layers without any kv-cache.30    Implemented using a list of Cache objects, each represents a "model" with 1 layer.31    The default implementation for the layer caches is StaticCache.32    The cache of each layer is allocated to the same gpu as the layer itself.33    """34 35    def __init__(36            self,37            *,  # key-word only, no positional args allowed to avoid mix-ups with newer transformers versions38            config: DeciLMConfig,39            batch_size: int = None,40            max_cache_len: int = None,41            dtype: torch.dtype = torch.float32,42            max_batch_size: Optional[int] = None,43            **kwargs,44    ) -> None:45        Cache_4_44_2.__init__(self)46 47        self.config = deepcopy(config)48        self.max_batch_size = batch_size or max_batch_size49        self.batch_size = self.max_batch_size50        self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len51        self.dtype = dtype52 53        self.layer_caches: list[Cache_4_44_2 | None] = [None] * config.num_hidden_layers54        self.layer_devices: list[torch.device | None] = [None] * config.num_hidden_layers55 56    def update(57            self,58            key_states: torch.Tensor,59            value_states: torch.Tensor,60            layer_idx: int,61            cache_kwargs: Optional[Dict[str, Any]] = None,62    ) -> Tuple[torch.Tensor, torch.Tensor]:63        if self.layer_caches[layer_idx] is None:64            self.layer_devices[layer_idx] = key_states.device65            self._init_layer_cache(layer_idx)66 67        layer_cache = self.layer_caches[layer_idx]68        assert layer_cache is not None, f"Trying to update the cache of a cache-less layer: {layer_idx=}"69 70        k_out, v_out = layer_cache.update(key_states=key_states,71                                          value_states=value_states,72                                          layer_idx=0,73                                          cache_kwargs=cache_kwargs)74        seq_len = self.get_seq_length(layer_idx)75        k_out = k_out[:, :, :seq_len, :]76        v_out = v_out[:, :, :seq_len, :]77        return k_out, v_out78 79    def _init_layer_cache(self, layer_idx: int) -> None:80        block_config = self.config.block_configs[layer_idx]81        attention_config = block_config.attention82 83        if attention_config.no_op or attention_config.replace_with_linear:84            return None85 86        device = self.layer_devices[layer_idx]87        assert device is not None, f"Trying to init layer cache for {layer_idx=} without device"88 89        config = deepcopy(self.config)90        config.num_hidden_layers = 191        config.num_key_value_heads = self.config.num_attention_heads // attention_config.n_heads_in_group92 93        if attention_config.window_length is not None:94            if not attention_config.is_sink:95                config.sliding_window = attention_config.window_length96                self.layer_caches[layer_idx] = SlidingWindowCache(config=config,97                                                                  max_batch_size=self.max_batch_size,98                                                                  max_cache_len=self.max_cache_len,99                                                                  device=device,100                                                                  dtype=self.dtype)101                return102            elif not attention_config.unshifted_sink:103                self.layer_caches[layer_idx] = SinkCache(window_length=attention_config.window_length,104                                                         num_sink_tokens=attention_config.num_sink_tokens)105                return106 107        self.layer_caches[layer_idx] = StaticCache(config=config,108                                                   max_batch_size=self.max_batch_size,109                                                   max_cache_len=self.max_cache_len,110                                                   device=device,111                                                   dtype=self.dtype)112 113    def _get_first_real_cache(self) -> Cache:114        for layer_cache in self.layer_caches:115            if layer_cache is not None:116                return layer_cache117        raise ValueError(f"No real cache found, all layer caches are None.")118 119    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:120        if layer_idx == 0 and self.layer_caches[0] is None:121            try:122                layer_cache = self._get_first_real_cache()123            except ValueError:124                return 0125        else:126            layer_cache = self.layer_caches[layer_idx]127        return layer_cache.get_seq_length()128 129    def get_max_length(self) -> Optional[int]:130        """Returns the maximum sequence length of the cached states."""131        return self.max_cache_len132 133    def reset(self):134        for layer_idx in range(len(self.layer_caches)):135            layer_cache = self.layer_caches[layer_idx]136            if hasattr(layer_cache, "reset"):137                layer_cache.reset()138            else:139                self._init_layer_cache(layer_idx)140