Get the memory (in bytes) required to store the key and value cache for a transformer layer in inference, given the batch size, sequence length, activation data type, and tensor parallelism size. Args: batch_size (int): batch size seq_len (int): seque
(
self,
batch_size: int,
seq_len: int,
kv_cache_dtype_bytes: int = None,
)
| 843 | return activation_memory_per_layer |
| 844 | |
| 845 | def get_memory_kv_cache_per_layer( |
| 846 | self, |
| 847 | batch_size: int, |
| 848 | seq_len: int, |
| 849 | kv_cache_dtype_bytes: int = None, |
| 850 | ) -> float: |
| 851 | """Get the memory (in bytes) required to store the key and value cache for a |
| 852 | transformer layer in inference, given the batch size, sequence length, |
| 853 | activation data type, and tensor parallelism size. |
| 854 | |
| 855 | Args: |
| 856 | batch_size (int): batch size |
| 857 | seq_len (int): sequence length |
| 858 | kv_cache_dtype_bytes (int, optional): number of bytes in the data type for the kv_cache. Defaults to None. Often has to be at least FP16 in inference to maintain model accuracy. |
| 859 | |
| 860 | Returns: |
| 861 | float: the memory (in bytes) required to store the key and value cache for a transformer layer in inference |
| 862 | """ |
| 863 | if kv_cache_dtype_bytes is None: |
| 864 | kv_cache_dtype_bytes = (self.dtype_config.activation_bits / |
| 865 | BITS_PER_BYTE) |
| 866 | head_dim = self.model_config.hidden_dim / self.model_config.n_head |
| 867 | num_heads_per_gpu = max( |
| 868 | self.model_config.num_key_value_heads / |
| 869 | self.parallelism_config.tp_size, |
| 870 | 1) # At least on attention head on each tensor-parallel GPU |
| 871 | |
| 872 | memory_kv_cache_per_layer = (2 * batch_size * seq_len * head_dim * |
| 873 | num_heads_per_gpu) * kv_cache_dtype_bytes |
| 874 | logger.debug( |
| 875 | f"memory_kv_cache_per_layer = {_num_to_string(memory_kv_cache_per_layer)} B" |
| 876 | ) |
| 877 | |
| 878 | return memory_kv_cache_per_layer |
| 879 | |
| 880 | def get_num_flops_fwd_per_layer_attn(self, batch_size: int, |
| 881 | seq_len: int) -> int: |
no test coverage detected