|
from collections import namedtuple |
|
from dataclasses import dataclass |
|
import torch |
|
from typing import Tuple, Optional |
|
|
|
|
|
@dataclass |
|
class LongLlamaMemConfig: |
|
""" |
|
Class for configuring memory caches for LongLlama model. |
|
|
|
Args: |
|
positionals (`boolean`) |
|
Whether to use positional embeddings in memory layer |
|
cache_dtype (`torch.dtype`) |
|
Specifies storing type for keys and values |
|
attention_grouping (`Tuple[int, int]`, *optional*) |
|
One can trade speed for memory by performing attention |
|
in memory layers sequentially. |
|
When equal to `(4, 128)` the memory layers will process at most 4 heads and 128 queries |
|
from each head at once. That is at most 512 queries at once. |
|
""" |
|
|
|
positionals: bool = True |
|
cache_dtype: torch.dtype = torch.bfloat16 |
|
attention_grouping: Optional[Tuple[int, int]] = None |
|
|
|
|
|
@dataclass |
|
class LongLlamaMemCache: |
|
""" |
|
Class with LongLlama's memory cache |
|
|
|
Args: |
|
keys (`torch.FloatTensor` of shape `(batch_size, num_heads, mem_length, embed_size_per_head)`) |
|
values (`torch.FloatTensor` of shape `(batch_size, num_heads, mem_length, embed_size_per_head)`) |
|
masks (`torch.FloatTensor` of shape `(batch_size, 1, mem_length, 1)`) |
|
For masking out parts of memory |
|
""" |
|
|
|
keys: torch.FloatTensor |
|
values: torch.FloatTensor |
|
masks: torch.FloatTensor |
|
|
|
|
|
def mem_apply_update( |
|
prev_mem_cache: LongLlamaMemCache, new_mem_content: LongLlamaMemCache, mem_config: LongLlamaMemConfig |
|
): |
|
def update_one(prev, new): |
|
if len(prev.shape) != 4 or len(new.shape) != 4: |
|
raise ValueError(f"Memory cache content should be consistent in shape got {prev.shape} {new.shape}") |
|
|
|
return torch.concat([prev, new], dim=-2) |
|
|
|
insert_size = new_mem_content.keys.shape[-2] |
|
|
|
if new_mem_content.values.shape[-2] != insert_size or new_mem_content.masks.shape[-2] != insert_size: |
|
raise ValueError(f"Inconsistent mem_length in new_mem_content") |
|
|
|
return LongLlamaMemCache( |
|
keys=update_one(prev_mem_cache.keys, new_mem_content.keys), |
|
values=update_one(prev_mem_cache.values, new_mem_content.values), |
|
masks=update_one(prev_mem_cache.masks, new_mem_content.masks), |
|
) |
|
|
