Spaces:
Running on Zero
Running on Zero
File size: 2,614 Bytes
6784f71 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """Reference ragged-cache batching; positions remain local to each seat."""
import torch
from transformers import DynamicCache
from transformers.cache_utils import DynamicLayer
def cache_from_owned_tensors(layers):
"""Install freshly owned tensors without concatenating them to empty caches.
DynamicCache(ddp_cache_data=...) calls update(), which copies every tensor
once more. These inputs are already owned by this cache, so initialization
can install them directly using the pinned DynamicLayer representation.
"""
cache=DynamicCache()
for keys,values in layers:
layer=DynamicLayer()
layer.keys,layer.values=keys,values
layer.dtype,layer.device=keys.dtype,keys.device
layer.is_initialized=True
cache.layers.append(layer)
return cache
def pack_caches(caches, new_tokens, device):
lengths=[c.layers[0].keys.shape[-2] if c is not None else 0 for c in caches]
maximum=max(lengths)
if maximum==0:
return None,None,lengths
exemplar=next(c for c in caches if c is not None)
layers=[]
for li in range(len(exemplar.layers)):
kv=[]
for name in ('keys','values'):
sample=getattr(exemplar.layers[li],name)
stacked=sample.new_zeros(len(caches),sample.shape[1],maximum,sample.shape[-1])
for i,(c,length) in enumerate(zip(caches,lengths)):
if length:
stacked[i:i+1,:,-length:]=getattr(c.layers[li],name)
kv.append(stacked)
layers.append(tuple(kv))
mask=torch.arange(maximum+new_tokens,device=device)[None]>=torch.tensor(
[maximum-l for l in lengths],device=device)[:,None]
return cache_from_owned_tensors(layers),mask,lengths
def split_cache(cache,lengths,keep=None,*,copy=True):
"""Split cache rows, optionally borrowing immutable views for experiments.
The default gives independent mutable storage. copy=False permits only
DynamicLayer append/crop operations, which replace tensors rather than
writing existing storage. A borrowed row retains the whole packed batch
until all views die; measure memory before using it in a serving path.
"""
def take(tensor,i,length):
view=tensor[i:i+1,:,-length:]
return view.clone(memory_format=torch.contiguous_format) if copy else view
result=[]
for i,length in enumerate(lengths):
length=min(length,keep) if keep is not None else length
result.append(cache_from_owned_tensors([
(take(l.keys,i,length),take(l.values,i,length)) for l in cache.layers]))
return result
|