mirror of https://github.com/hpcaitech/ColossalAI
249 lines
8.7 KiB
Python
249 lines
8.7 KiB
Python
from enum import Enum
|
|
from typing import Dict
|
|
|
|
import torch
|
|
|
|
from ..kv_cache import BatchInferState, MemoryManager
|
|
|
|
__all__ = "MicroBatchManager"
|
|
|
|
|
|
class Status(Enum):
|
|
PREFILL = 1
|
|
GENERATE = 2
|
|
DONE = 3
|
|
COOLDOWN = 4
|
|
|
|
|
|
class MicroBatchDescription:
|
|
"""
|
|
This is the class to record the information of each microbatch, and also do some update operation.
|
|
This class is the base class of `HeadMicroBatchDescription` and `BodyMicroBatchDescription`, for more
|
|
details, please refer to the doc of these two classes blow.
|
|
|
|
Args:
|
|
inputs_dict (Dict[str, torch.Tensor]): the inputs of current stage. The key should have `input_ids` and `attention_mask`.
|
|
output_dict (Dict[str, torch.Tensor]): the outputs of previous stage. The key should have `hidden_states` and `past_key_values`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
inputs_dict: Dict[str, torch.Tensor],
|
|
max_input_len: int,
|
|
max_output_len: int,
|
|
cache_manager: MemoryManager,
|
|
) -> None:
|
|
self.mb_length = inputs_dict["input_ids"].shape[-1]
|
|
self.target_length = self.mb_length + max_output_len
|
|
self.infer_state = BatchInferState.init_from_batch(
|
|
batch=inputs_dict, max_input_len=max_input_len, max_output_len=max_output_len, cache_manager=cache_manager
|
|
)
|
|
# print(f"[init] {inputs_dict}, {max_input_len}, {max_output_len}, {cache_manager}, {self.infer_state}")
|
|
|
|
def update(self, *args, **kwargs):
|
|
pass
|
|
|
|
@property
|
|
def state(self):
|
|
"""
|
|
Return the state of current micro batch, when current length is equal to target length,
|
|
the state is DONE, otherwise GENERATE
|
|
|
|
"""
|
|
# TODO: add the condition for early stopping
|
|
if self.cur_length == self.target_length:
|
|
return Status.DONE
|
|
elif self.cur_length == self.target_length - 1:
|
|
return Status.COOLDOWN
|
|
else:
|
|
return Status.GENERATE
|
|
|
|
@property
|
|
def cur_length(self):
|
|
"""
|
|
Return the current sequence length of micro batch
|
|
|
|
"""
|
|
|
|
|
|
class HeadMicroBatchDescription(MicroBatchDescription):
|
|
"""
|
|
This class is used to record the information of the first stage of pipeline, the first stage should have attributes `input_ids` and `attention_mask`
|
|
and `new_tokens`, and the `new_tokens` is the tokens generated by the first stage. Also due to the schedule of pipeline, the operation to update the
|
|
information and the condition to determine the state is different from other stages.
|
|
|
|
Args:
|
|
inputs_dict (Dict[str, torch.Tensor]): the inputs of current stage. The key should have `input_ids` and `attention_mask`.
|
|
output_dict (Dict[str, torch.Tensor]): the outputs of previous stage. The key should have `hidden_states` and `past_key_values`.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
inputs_dict: Dict[str, torch.Tensor],
|
|
max_input_len: int,
|
|
max_output_len: int,
|
|
cache_manager: MemoryManager,
|
|
) -> None:
|
|
super().__init__(inputs_dict, max_input_len, max_output_len, cache_manager)
|
|
assert inputs_dict is not None
|
|
assert inputs_dict.get("input_ids") is not None and inputs_dict.get("attention_mask") is not None
|
|
self.input_ids = inputs_dict["input_ids"]
|
|
self.attn_mask = inputs_dict["attention_mask"]
|
|
self.new_tokens = None
|
|
|
|
def update(self, new_token: torch.Tensor = None):
|
|
if new_token is not None:
|
|
self._update_newtokens(new_token)
|
|
if self.state is not Status.DONE and new_token is not None:
|
|
self._update_attnmask()
|
|
|
|
def _update_newtokens(self, new_token: torch.Tensor):
|
|
if self.new_tokens is None:
|
|
self.new_tokens = new_token
|
|
else:
|
|
self.new_tokens = torch.cat([self.new_tokens, new_token], dim=-1)
|
|
|
|
def _update_attnmask(self):
|
|
self.attn_mask = torch.cat(
|
|
(self.attn_mask, torch.ones((self.attn_mask.shape[0], 1), dtype=torch.int64, device="cuda")), dim=-1
|
|
)
|
|
|
|
@property
|
|
def cur_length(self):
|
|
"""
|
|
When there is no new_token, the length is mb_length, otherwise the sequence length is `mb_length` plus the length of new_token
|
|
|
|
"""
|
|
if self.new_tokens is None:
|
|
return self.mb_length
|
|
else:
|
|
return self.mb_length + len(self.new_tokens[0])
|
|
|
|
|
|
class BodyMicroBatchDescription(MicroBatchDescription):
|
|
"""
|
|
This class is used to record the information of the stages except the first stage of pipeline, the stages should have attributes `hidden_states` and `past_key_values`,
|
|
|
|
Args:
|
|
inputs_dict (Dict[str, torch.Tensor]): will always be `None`. Other stages only receive hiddenstates from previous stage.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
inputs_dict: Dict[str, torch.Tensor],
|
|
max_input_len: int,
|
|
max_output_len: int,
|
|
cache_manager: MemoryManager,
|
|
) -> None:
|
|
super().__init__(inputs_dict, max_input_len, max_output_len, cache_manager)
|
|
|
|
@property
|
|
def cur_length(self):
|
|
"""
|
|
When there is no kv_cache, the length is mb_length, otherwise the sequence length is `kv_cache[0][0].shape[-2]` plus 1
|
|
|
|
"""
|
|
return self.infer_state.seq_len.max().item()
|
|
|
|
|
|
class MicroBatchManager:
|
|
"""
|
|
MicroBatchManager is a class that manages the micro batch.
|
|
|
|
Args:
|
|
stage (int): stage id of current stage.
|
|
micro_batch_size (int): the micro batch size.
|
|
micro_batch_buffer_size (int): the buffer size for micro batch. Normally, it should be the same as the number of pipeline stages.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
stage: int,
|
|
micro_batch_size: int,
|
|
micro_batch_buffer_size: int,
|
|
max_input_len: int,
|
|
max_output_len: int,
|
|
cache_manager_list: MemoryManager,
|
|
):
|
|
self.stage = stage
|
|
self.micro_batch_size = micro_batch_size
|
|
self.buffer_size = micro_batch_buffer_size
|
|
self.max_input_len = max_input_len
|
|
self.max_output_len = max_output_len
|
|
self.cache_manager_list = cache_manager_list
|
|
self.mb_description_buffer = {}
|
|
self.new_tokens_buffer = {}
|
|
self.idx = 0
|
|
|
|
def add_description(self, inputs_dict: Dict[str, torch.Tensor]):
|
|
if self.stage == 0:
|
|
self.mb_description_buffer[self.idx] = HeadMicroBatchDescription(
|
|
inputs_dict, self.max_input_len, self.max_output_len, self.cache_manager_list[self.idx]
|
|
)
|
|
else:
|
|
self.mb_description_buffer[self.idx] = BodyMicroBatchDescription(
|
|
inputs_dict, self.max_input_len, self.max_output_len, self.cache_manager_list[self.idx]
|
|
)
|
|
|
|
def step(self, new_token: torch.Tensor = None):
|
|
"""
|
|
Update the state if microbatch manager, 2 conditions.
|
|
1. For first stage in PREFILL, receive inputs and outputs, `_add_description` will save its inputs.
|
|
2. For other condition, only receive the output of previous stage, and update the description.
|
|
|
|
Args:
|
|
inputs_dict (Dict[str, torch.Tensor]): the inputs of current stage. The key should have `input_ids` and `attention_mask`.
|
|
output_dict (Dict[str, torch.Tensor]): the outputs of previous stage. The key should have `hidden_states` and `past_key_values`.
|
|
new_token (torch.Tensor): the new token generated by current stage.
|
|
"""
|
|
# Add description first if the description is None
|
|
self.cur_description.update(new_token)
|
|
return self.cur_state
|
|
|
|
def export_new_tokens(self):
|
|
new_tokens_list = []
|
|
for i in self.mb_description_buffer.values():
|
|
new_tokens_list.extend(i.new_tokens.tolist())
|
|
return new_tokens_list
|
|
|
|
def is_micro_batch_done(self):
|
|
if len(self.mb_description_buffer) == 0:
|
|
return False
|
|
for mb in self.mb_description_buffer.values():
|
|
if mb.state != Status.DONE:
|
|
return False
|
|
return True
|
|
|
|
def clear(self):
|
|
self.mb_description_buffer.clear()
|
|
for cache in self.cache_manager_list:
|
|
cache.free_all()
|
|
|
|
def next(self):
|
|
self.idx = (self.idx + 1) % self.buffer_size
|
|
|
|
def _remove_description(self):
|
|
self.mb_description_buffer.pop(self.idx)
|
|
|
|
@property
|
|
def cur_description(self) -> MicroBatchDescription:
|
|
return self.mb_description_buffer.get(self.idx)
|
|
|
|
@property
|
|
def cur_infer_state(self):
|
|
if self.cur_description is None:
|
|
return None
|
|
return self.cur_description.infer_state
|
|
|
|
@property
|
|
def cur_state(self):
|
|
"""
|
|
Return the state of current micro batch, when current description is None, the state is PREFILL
|
|
|
|
"""
|
|
if self.cur_description is None:
|
|
return Status.PREFILL
|
|
return self.cur_description.state
|