wehub-resource-sync
269 lines
10 KiB
Python
269 lines
10 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import os
|
|
import json
|
|
import pickle
|
|
from typing import Iterable, Tuple
|
|
|
|
import torch
|
|
|
|
import deepspeed.comm as dist
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.comm.comm import init_distributed
|
|
|
|
from .model_implementations import InferenceV2Policy
|
|
from .logging import inference_logger
|
|
from .ragged import DSStateManager, RaggedBatchWrapper, PlaceholderSequenceDescriptor
|
|
from .scheduling_utils import SchedulingError, SchedulingResult
|
|
from .model_implementations.flat_model_helpers import make_param_filename, make_metadata_filename
|
|
from .model_implementations.inference_model_base import DSInferenceModelBase
|
|
|
|
from .config_v2 import RaggedInferenceEngineConfig
|
|
|
|
INFERENCE_MODEL_TIMER = "model-forward-inference"
|
|
|
|
|
|
class InferenceEngineV2:
|
|
|
|
_config: RaggedInferenceEngineConfig
|
|
"""
|
|
Configuration of the inference engine.
|
|
"""
|
|
|
|
_model: DSInferenceModelBase
|
|
"""
|
|
Inference model supporting ragged inference.
|
|
"""
|
|
|
|
_state_manager: DSStateManager
|
|
"""
|
|
Persistent state manager for sequences and KV-cache.
|
|
"""
|
|
|
|
@property
|
|
def free_blocks(self) -> torch.Tensor:
|
|
"""
|
|
Number of free KV blocks. This is a tensor of shape [n_kv_cache_groups] where each
|
|
element is the number of free blocks in the corresponding KV cache group.
|
|
"""
|
|
return self._state_manager.free_blocks
|
|
|
|
@property
|
|
def n_kv_cache_groups(self) -> int:
|
|
"""
|
|
Number of KV cache groups.
|
|
"""
|
|
return self._state_manager.n_kv_cache_groups
|
|
|
|
def model(self) -> DSInferenceModelBase:
|
|
"""
|
|
The model implementation.
|
|
"""
|
|
return self._model
|
|
|
|
def __init__(self, policy: InferenceV2Policy, engine_config: RaggedInferenceEngineConfig) -> None:
|
|
"""
|
|
Create the Inference V2 engine.
|
|
|
|
Arguments:
|
|
policy (InferenceV2Policy): Policy for the model implementation. This policy object
|
|
will be used to build the model and load the checkpoint associated with it.
|
|
engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
|
|
"""
|
|
self._config = engine_config
|
|
self._policy = policy
|
|
self._base_mp_group = self._initialize_tp_group()
|
|
|
|
# Build model from policy
|
|
inference_logger().info("Building model...")
|
|
self._model = self._policy.build_model(self._config, self._base_mp_group)
|
|
inference_logger().info("Model built.")
|
|
|
|
# Create state manager
|
|
self._batch = RaggedBatchWrapper(self._config.state_manager)
|
|
self._state_manager = DSStateManager(self._config.state_manager,
|
|
self._model.kv_cache_config(),
|
|
base_mp_group=self._base_mp_group)
|
|
self._model.set_state_manager(self._state_manager)
|
|
|
|
def _initialize_tp_group(self):
|
|
"""
|
|
Implementation of our TP group initialization.
|
|
"""
|
|
init_distributed()
|
|
local_rank = int(os.getenv("LOCAL_RANK", 0))
|
|
get_accelerator().set_device(local_rank)
|
|
|
|
if local_rank >= self._config.tensor_parallel.tp_size:
|
|
raise RuntimeError("Local rank is greater than TP size, ensure that the TP config is correct.")
|
|
|
|
ranks = list(range(self._config.tensor_parallel.tp_size))
|
|
return dist.new_group(ranks=ranks)
|
|
|
|
def put(self,
|
|
batch_uids: Iterable[int],
|
|
batch_tokens: Iterable[torch.Tensor],
|
|
do_checks: bool = True) -> torch.Tensor:
|
|
"""
|
|
Put a ragged batch onto the inference engine. This will perform one forward and return
|
|
a Tensor of the shape [len(batch_uids), *output_shape]. Logits for the non-final tokens
|
|
are not calculated.
|
|
|
|
Arguments:
|
|
batch_uids: Iterable of uids for the batch on the host
|
|
batch_tokens: Iterable of token tensors for the batch on the host
|
|
do_checks: Check schedulability when it is set to True. You can skip this check for better performance when it has already been completed.
|
|
"""
|
|
|
|
if do_checks:
|
|
token_lens = [len(tokens) for tokens in batch_tokens]
|
|
schedule_check = self.can_schedule(batch_uids, token_lens)
|
|
if schedule_check != SchedulingResult.Success:
|
|
raise SchedulingError(schedule_check)
|
|
|
|
self._batch.clear()
|
|
for uid, tokens in zip(batch_uids, batch_tokens):
|
|
|
|
host_seq_desc = self._state_manager.get_or_create_sequence(uid)
|
|
self._model.maybe_allocate_kv(host_seq_desc, tokens.numel())
|
|
host_seq_desc.pre_forward(tokens.numel())
|
|
|
|
# We can disable checks since we already validated schedulability.
|
|
self._batch.insert_sequence(host_seq_desc, tokens, do_checks=do_checks)
|
|
|
|
# Send all metadata to the device
|
|
self._batch.finalize()
|
|
|
|
# Prep all data structures for the actual forward (in anticipation of CG in the future)
|
|
# and also to amortize some of the costs in a more straightforward way.
|
|
self._model.prepare_batch(self._batch)
|
|
|
|
# Model implementation will pick up in the forward.
|
|
logits = self._model.forward(self._batch)
|
|
|
|
# We return one set of logits per sequence in the batch (saves cost on unembedding)
|
|
assert logits.shape[0] == self._batch.current_sequences
|
|
|
|
for uid in batch_uids:
|
|
host_seq_desc = self._state_manager.get_sequence(uid)
|
|
host_seq_desc.post_forward() # Updates sequence metadata.
|
|
self._model.maybe_free_kv(host_seq_desc)
|
|
|
|
return logits
|
|
|
|
def query(self, uid: int, max_request_tokens: int, max_request_blocks) -> Tuple[int, torch.Tensor]:
|
|
"""
|
|
Determine the number of tokens and KV blocks to reserve for a given request. Given a UID
|
|
(this UID may not be recognized by the model yet), this will return the number of tokens
|
|
and blocks to reserve for the request.
|
|
|
|
Arguments:
|
|
uid (int): The UID of the sequence (as tracked by the scheduling entity). If
|
|
this is a new sequence (with a UID unknown to the inference engine), then
|
|
an empty placeholder is created to pass to the occupancy logic.
|
|
n_tokens (int): The number of tokens to hypothetically send.
|
|
|
|
Returns:
|
|
Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks
|
|
required to schedule the sequence.
|
|
"""
|
|
seq_desc = self._state_manager.get_sequence(uid)
|
|
if seq_desc is None:
|
|
if (self._state_manager.n_tracked_sequences == self._config.state_manager.max_tracked_sequences):
|
|
return (0, 0)
|
|
seq_desc = PlaceholderSequenceDescriptor()
|
|
|
|
req_tokens, req_blocks = self._model.get_kv_requirements(seq_desc, max_request_tokens, max_request_blocks)
|
|
|
|
return (req_tokens, req_blocks)
|
|
|
|
def can_schedule(self, uids: Iterable[int], lengths: Iterable[int]) -> SchedulingResult:
|
|
"""
|
|
Dry run a batch to determine if it can be scheduled. Placeholder sequences will be
|
|
created for any UIDs that are unknown to the inference engine.
|
|
|
|
Arguments:
|
|
uids (Iterable[int]): Iterable of UIDs for the batch
|
|
lengths (Iterable[int]): Iterable of lengths for each sequence of the batch. This lengths
|
|
corresponds to the number of tokens to send in the hypothetical forward; history
|
|
tokens will be determined via UID lookup and future tokens are disregarded.
|
|
|
|
Returns:
|
|
bool: True if the batch can be scheduled, False otherwise.
|
|
"""
|
|
|
|
cur_seqs = self._state_manager.n_tracked_sequences
|
|
free_blocks = self._state_manager.free_blocks
|
|
req_blocks = 0
|
|
batch_len = 0
|
|
|
|
if len(uids) > self._config.state_manager.max_ragged_sequence_count:
|
|
# Can only compose a batch from a limited number of sequences
|
|
return SchedulingResult.BatchSequenceLimitExceeded
|
|
|
|
for uid, length in zip(uids, lengths):
|
|
seq_desc = self._state_manager.get_sequence(uid)
|
|
if seq_desc is None:
|
|
cur_seqs += 1
|
|
seq_desc = PlaceholderSequenceDescriptor()
|
|
|
|
sched_len, sched_blocks = self._model.get_kv_requirements(seq_desc, length, free_blocks)
|
|
|
|
if sched_len != length:
|
|
# We ran out of KV cache
|
|
return SchedulingResult.KVCacheLimitExceeded
|
|
|
|
batch_len += length
|
|
free_blocks -= sched_blocks
|
|
|
|
if cur_seqs > self._config.state_manager.max_tracked_sequences:
|
|
# Would run out of tracking metadata
|
|
return SchedulingResult.EngineSequenceLimitExceeded
|
|
|
|
if batch_len > self._config.state_manager.max_ragged_batch_size:
|
|
# Would exceed the maximum batch size
|
|
return SchedulingResult.BatchTokenLimitExceeded
|
|
|
|
return SchedulingResult.Success
|
|
|
|
def get_remaining_block_capacity(self, uid: int) -> int:
|
|
"""
|
|
Get the remaining capacity of the last block already allocated.
|
|
"""
|
|
seq_desc = self._state_manager.get_sequence(uid)
|
|
if seq_desc is None:
|
|
return 0
|
|
return self._model.get_remaining_block_capacity(seq_desc)
|
|
|
|
def flush(self, uid: int) -> None:
|
|
"""
|
|
Remove all state associated with a sequence from the inference engine.
|
|
|
|
Arguments:
|
|
uid (int): The UID of the sequence to flush.
|
|
"""
|
|
self._state_manager.flush_sequence(uid)
|
|
|
|
def serialize(self, save_path: str) -> None:
|
|
"""
|
|
Serialize the model to a file.
|
|
|
|
Arguments:
|
|
path (str): Path to the file to serialize to.
|
|
"""
|
|
param_file_name = make_param_filename(save_path, self._model.tp_rank, self._model.tp_size)
|
|
metadata_file_name = make_metadata_filename(save_path, self._model.tp_rank, self._model.tp_size)
|
|
|
|
# Save the flattened parameters
|
|
|
|
torch.save(self._model.flattened_params, param_file_name)
|
|
|
|
json.dump(self._model.flattened_param_metadata.json(), open(metadata_file_name, "w"))
|
|
|
|
if self._model.tp_rank == 0:
|
|
pickle.dump(self._model._config, open(os.path.join(save_path, "ds_model_config.pkl"), "wb"))
|