wehub-resource-sync
209 lines
8.4 KiB
Python
209 lines
8.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import operator
|
|
from functools import reduce
|
|
from typing import Any, Iterable, Optional, Tuple
|
|
|
|
import torch
|
|
|
|
import deepspeed.comm as dist
|
|
from deepspeed.comm.reduce_op import ReduceOp
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from ..inference_utils import elem_size
|
|
from ..logging import inference_logger
|
|
from .blocked_allocator import BlockedAllocator
|
|
from .manager_configs import AllocationMode, KVCacheConfig, MemoryConfig
|
|
|
|
|
|
def split_kv(kv_cache: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Split a KV cache instance into its key and value components.
|
|
|
|
Parameters:
|
|
kv_cache (torch.Tensor): The KV-cache to split. This should be a 5D tensor with the
|
|
following shape: [num_blocks, block_size, 2, num_heads, head_size]
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]: The key and value components of the KV-cache. Both
|
|
tensors will have the shape [num_blocks, block_size, num_heads, head_size].
|
|
"""
|
|
if kv_cache.ndim != 5:
|
|
raise ValueError(f"KV-cache must have 5 dimensions, got {kv_cache.ndim}.")
|
|
|
|
return kv_cache[:, :, 0, :, :], kv_cache[:, :, 1, :, :]
|
|
|
|
|
|
class BlockedKVCache:
|
|
|
|
_caches: Tuple[torch.Tensor, ...]
|
|
"""
|
|
Backing storage for all KV caches. This is a 6D tensor with the following shape:
|
|
(num_caches, num_blocks, block_size, 2, num_heads, head_size)
|
|
"""
|
|
|
|
_allocators: Tuple[BlockedAllocator, ...]
|
|
"""
|
|
Block allocator for tracking cache usage. This manages the GPU cache.
|
|
"""
|
|
|
|
_configs: Tuple[KVCacheConfig, ...]
|
|
"""
|
|
Configuration of the KV cache(s). See ``KVCacheConfig`` for more details. This enables the support
|
|
for different types/shapes of KV-caches (i.e. the alternating local and global attention in
|
|
GPT-Neo).
|
|
"""
|
|
|
|
def __init__(self,
|
|
configs: Tuple[KVCacheConfig, ...],
|
|
memory_config: MemoryConfig,
|
|
mp_group: Optional[Any] = None,
|
|
offload: bool = False) -> None:
|
|
"""
|
|
Create a container that will maintain the storage and allocations for a set of
|
|
blocked KV-caches.
|
|
|
|
Parameters:
|
|
config (KVCacheConfig): The configuration of the KV-cache.
|
|
slack (int): The amount of slack space to reserve in GPU memory for the cache.
|
|
enable_offload (bool): Whether to enable offloading of the cache to the host.
|
|
blocks (int): The number of blocks to pre-allocate for the cache. If this is set,
|
|
slack will be ignored.
|
|
"""
|
|
self._configs = configs
|
|
self._memory_config = memory_config
|
|
self._enable_offload = offload
|
|
|
|
if self._enable_offload:
|
|
raise NotImplementedError("Offloading of KV-caches is not yet supported.")
|
|
|
|
if AllocationMode(self._memory_config.mode) is AllocationMode.RESERVE:
|
|
# TODO(cmikeh2): Change the weighting based on the type of the KV-cache
|
|
|
|
total_per_block_footprint = 0
|
|
for config in self._configs:
|
|
per_block_footprint = reduce(operator.mul, config.cache_shape, config.block_size)
|
|
per_block_footprint *= 2 # for key and value
|
|
total_per_block_footprint += per_block_footprint * elem_size(config.cache_dtype)
|
|
|
|
# Perform a dummy nccl call before calculating available memory, on some systems (H100) we've observed higher memory allocations from NCCL
|
|
if dist.get_world_size(group=mp_group) > 1:
|
|
dummy_tensor = torch.tensor(0, dtype=torch.int32, device=get_accelerator().current_device())
|
|
dist.all_reduce(dummy_tensor, op=ReduceOp.MIN, group=mp_group)
|
|
|
|
get_accelerator().empty_cache()
|
|
available_kv_memory = get_accelerator().available_memory() - self._memory_config.size
|
|
total_memory = get_accelerator().total_memory()
|
|
|
|
inference_logger().debug(
|
|
f"Memory usage before KV-cache allocation: total_memory={total_memory}, available_kv_memory={available_kv_memory}, total_per_block_footprint={total_per_block_footprint}"
|
|
)
|
|
|
|
if available_kv_memory < total_per_block_footprint:
|
|
raise ValueError(
|
|
f"Insufficient memory to allocate KV-caches. Required: {total_per_block_footprint}, Available: {available_kv_memory}"
|
|
)
|
|
|
|
num_blocks = available_kv_memory // total_per_block_footprint
|
|
|
|
# In a multi-process setting, we need to ensure that all processes have the same
|
|
# KV cache capacity to ensure scheduling guarantees are equivalent on all ranks.
|
|
if dist.get_world_size(group=mp_group) > 1:
|
|
reduce_tensor = torch.tensor(num_blocks, dtype=torch.int32, device=get_accelerator().current_device())
|
|
dist.all_reduce(reduce_tensor, op=ReduceOp.MIN, group=mp_group)
|
|
num_blocks = reduce_tensor.item()
|
|
|
|
# This is ugly but don't want the fragmentation of the 8 byte Tensor maybe
|
|
# hanging around.
|
|
del reduce_tensor
|
|
get_accelerator().empty_cache()
|
|
else: # AllocationMode.ALLOCATE
|
|
num_blocks = self._memory_config.size
|
|
|
|
caches = []
|
|
allocators = []
|
|
|
|
for cache_group_id, config in enumerate(self._configs):
|
|
num_caches = config.cache_shape[0]
|
|
num_heads = config.cache_shape[1]
|
|
head_size = config.cache_shape[2]
|
|
|
|
alloc_shape = (num_caches, num_blocks, config.block_size, 2, num_heads, head_size)
|
|
inference_logger().info(
|
|
f"Allocating KV-cache {cache_group_id} with shape: {alloc_shape} consisting of {num_blocks} blocks.")
|
|
caches.append(torch.empty(alloc_shape, dtype=config.cache_dtype,
|
|
device=get_accelerator().current_device()))
|
|
allocators.append(BlockedAllocator(num_blocks))
|
|
|
|
self._caches = tuple(caches)
|
|
self._allocators = tuple(allocators)
|
|
|
|
def reserve(self, num_blocks: int, cache_group: int = 0) -> torch.Tensor:
|
|
"""
|
|
Reserve a number of blocks from the cache. This will return a 1D tensor of
|
|
block_ids that have been marked as reserved.
|
|
|
|
Parameters:
|
|
num_blocks (int): The number of blocks to reserve.
|
|
cache_group (int): The cache group to reserve from. Default is 0.
|
|
"""
|
|
return self._allocators[cache_group].allocate(num_blocks)
|
|
|
|
def free(self, blocks: Iterable[int], cache_group: int = 0) -> None:
|
|
"""
|
|
Free a set of blocks from the cache. This will mark the blocks as free in the
|
|
allocator.
|
|
|
|
Parameters:
|
|
blocks (Iterable[int]): The blocks to free.
|
|
cache_group (int): The cache group to free from. Default is 0.
|
|
"""
|
|
self._allocators[cache_group].free(blocks)
|
|
|
|
def offload(self, blocks: Iterable[int], cache_group: int = 0) -> torch.Tensor:
|
|
"""
|
|
Offload KV-cache blocks from accelerator memory to the host.
|
|
|
|
Parameters:
|
|
blocks (Iterable[int]): The blocks to offload.
|
|
cache_group (int): The cache group to offload from. Default is 0.
|
|
"""
|
|
raise NotImplementedError("Offloading is not yet supported.")
|
|
|
|
def restore(self, blocks: Iterable[int], cache_group: int = 0) -> torch.Tensor:
|
|
"""
|
|
Restore KV-cache blocks from the host to accelerator memory.
|
|
|
|
Parameters:
|
|
blocks (Iterable[int]): The blocks to restore.
|
|
cache_group (int): The cache group to restore to. Default is 0.
|
|
"""
|
|
raise NotImplementedError("Offloading is not yet supported.")
|
|
|
|
def get_cache(self, cache_id: int, cache_group: int = 0) -> torch.Tensor:
|
|
"""
|
|
Get the tensor associated with the given cache ID.
|
|
|
|
Parameters:
|
|
cache_id (int): The ID of the cache tensor to get.
|
|
cache_group (int): The cache group to get from. Default is 0.
|
|
"""
|
|
return self._caches[cache_group][cache_id]
|
|
|
|
@property
|
|
def free_blocks(self) -> torch.Tensor:
|
|
"""
|
|
Return the number of free blocks in each cache
|
|
"""
|
|
return [allocator.free_blocks for allocator in self._allocators]
|
|
|
|
@property
|
|
def num_caches(self) -> int:
|
|
"""
|
|
Return the number of caches
|
|
"""
|
|
return len(self._caches)
|