mirror of https://github.com/hpcaitech/ColossalAI
202 lines
8.5 KiB
Python
202 lines
8.5 KiB
Python
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from transformers.tokenization_utils_base import BatchEncoding
|
|
|
|
from colossalai.cluster import ProcessGroupMesh
|
|
from colossalai.pipeline.schedule.generate import GenerateSchedule
|
|
from colossalai.pipeline.stage_manager import PipelineStageManager
|
|
from colossalai.shardformer import ShardConfig, ShardFormer
|
|
from colossalai.shardformer.policies.base_policy import Policy
|
|
|
|
from ..kvcache_manager import MemoryManager
|
|
from .microbatch_manager import MicroBatchManager
|
|
|
|
PP_AXIS, TP_AXIS = 0, 1
|
|
|
|
_supported_models = ["LlamaForCausalLM", "BloomForCausalLM", "LlamaGPTQForCausalLM", "SmoothLlamaForCausalLM", "ChatGLMForConditionalGeneration"]
|
|
|
|
class CaiInferEngine:
|
|
"""
|
|
CaiInferEngine is a class that handles the pipeline parallel inference.
|
|
|
|
Args:
|
|
tp_size (int): the size of tensor parallelism.
|
|
pp_size (int): the size of pipeline parallelism.
|
|
model (`nn.Module`): the model not in pipeline style, and will be modified with `ShardFormer`.
|
|
model_policy (`colossalai.shardformer.policies.base_policy.Policy`): the policy to shardformer model.
|
|
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.
|
|
max_batch_size (int): the maximum batch size.
|
|
max_input_len (int): the maximum input length.
|
|
max_output_len (int): the maximum output length.
|
|
|
|
Example:
|
|
|
|
```python
|
|
from colossalai.inference import InferEngine
|
|
from colossalai.inference.pipeline.policies import LlamaModelInferPolicy
|
|
import colossalai
|
|
from transformers import LlamaForCausalLM, LlamaTokenizer
|
|
|
|
colossalai.launch_from_torch(config={})
|
|
|
|
model = LlamaForCausalLM.from_pretrained("your_path_to_model")
|
|
tokenizer = LlamaTokenizer.from_pretrained("/home/lczyh/share/models/llama-7b-hf")
|
|
# assume the model is infered with 2 pipeline stages
|
|
inferengine = CaiInferEngine(pp_size=2, model=model, model_policy=LlamaModelInferPolicy())
|
|
|
|
input = ["Introduce a landmark in China ","Introduce a landmark in China "]
|
|
data = tokenizer(input, return_tensors='pt')
|
|
output = inferengine.inference([data.to('cuda').data])
|
|
|
|
```
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
tp_size: int = 1,
|
|
pp_size: int = 1,
|
|
dtype: str = "fp16",
|
|
model: nn.Module = None,
|
|
model_policy: Policy = None,
|
|
micro_batch_size: int = 1,
|
|
micro_batch_buffer_size: int = None,
|
|
max_batch_size: int = 4,
|
|
max_input_len: int = 32,
|
|
max_output_len: int = 32,
|
|
quant: str = None,
|
|
verbose: bool = False,
|
|
# TODO: implement early_stopping, and various gerneration options
|
|
early_stopping: bool = False,
|
|
do_sample: bool = False,
|
|
num_beams: int = 1,
|
|
) -> None:
|
|
if quant == "gptq":
|
|
from ..quant.gptq import GPTQManager
|
|
|
|
self.gptq_manager = GPTQManager(model.quantize_config, max_input_len=max_input_len)
|
|
model = model.model
|
|
elif quant == "smoothquant":
|
|
model = model.model
|
|
|
|
assert model.__class__.__name__ in _supported_models, f"Model {model.__class__.__name__} is not supported."
|
|
assert (
|
|
tp_size * pp_size == dist.get_world_size()
|
|
), f"TP size({tp_size}) * PP size({pp_size}) should be equal to the global world size ({dist.get_world_size()})"
|
|
assert model and model_policy, "Model with model_policy should be provided."
|
|
assert dtype in ["fp16", "fp32", "bf16"], "dtype should be one of 'fp16', 'fp32', 'bf16'"
|
|
|
|
assert max_batch_size <= 64, "Max batch size exceeds the constraint"
|
|
assert max_input_len + max_output_len <= 4096, "Max length exceeds the constraint"
|
|
assert quant in ["smoothquant", "gptq", None], "quant should be one of 'smoothquant', 'gptq'"
|
|
self.pp_size = pp_size
|
|
self.tp_size = tp_size
|
|
self.quant = quant
|
|
|
|
if quant == "smoothquant" and dtype != "fp32":
|
|
dtype = "fp32"
|
|
print("Warning: smoothquant only support fp32 and int8 mix precision. set dtype to fp32")
|
|
|
|
if dtype == "fp16":
|
|
self.dtype = torch.float16
|
|
model.half()
|
|
elif dtype == "bf16":
|
|
self.dtype = torch.bfloat16
|
|
model.to(torch.bfloat16)
|
|
else:
|
|
self.dtype = torch.float32
|
|
|
|
# Init pg mesh
|
|
pg_mesh = ProcessGroupMesh(pp_size, tp_size)
|
|
|
|
stage_manager = PipelineStageManager(pg_mesh, PP_AXIS, True)
|
|
self.cache_manager_list = [
|
|
self._init_manager(model, max_batch_size, max_input_len, max_output_len)
|
|
for _ in range(micro_batch_buffer_size or pp_size)
|
|
]
|
|
self.mb_manager = MicroBatchManager(
|
|
stage_manager.stage,
|
|
micro_batch_size,
|
|
micro_batch_buffer_size or pp_size,
|
|
max_input_len,
|
|
max_output_len,
|
|
self.cache_manager_list,
|
|
)
|
|
self.verbose = verbose
|
|
self.schedule = GenerateSchedule(stage_manager, self.mb_manager, verbose)
|
|
|
|
self.model = self._shardformer(model, model_policy, stage_manager, pg_mesh.get_group_along_axis(TP_AXIS))
|
|
if quant == "gptq":
|
|
self.gptq_manager.post_init_gptq_buffer(self.model)
|
|
|
|
def inference(self, input_list):
|
|
"""
|
|
Args:
|
|
input_list (list): a list of input data, each element is a `BatchEncoding` or `dict`.
|
|
|
|
Returns:
|
|
out (list): a list of output data, each element is a list of token.
|
|
timestamp (float): the time cost of the inference, only return when verbose is `True`.
|
|
"""
|
|
assert isinstance(
|
|
input_list, (BatchEncoding, dict)
|
|
), f"Only accept BatchEncoding or dict as input, but get {input_list.__class__.__name__}."
|
|
if isinstance(input_list, BatchEncoding):
|
|
input_list = input_list.data
|
|
out, timestamp = self.schedule.generate_step(self.model, iter([input_list]))
|
|
if self.verbose:
|
|
return out, timestamp
|
|
else:
|
|
return out
|
|
|
|
def _shardformer(self, model, model_policy, stage_manager, tp_group):
|
|
shardconfig = ShardConfig(
|
|
tensor_parallel_process_group=tp_group,
|
|
pipeline_stage_manager=stage_manager,
|
|
enable_tensor_parallelism=True if self.tp_size > 1 else False,
|
|
enable_fused_normalization=False,
|
|
enable_all_optimization=False,
|
|
enable_flash_attention=False,
|
|
enable_jit_fused=False,
|
|
enable_sequence_parallelism=False,
|
|
quant=self.quant,
|
|
)
|
|
shardformer = ShardFormer(shard_config=shardconfig)
|
|
shard_model, _ = shardformer.optimize(model, model_policy)
|
|
return shard_model.cuda()
|
|
|
|
def _init_manager(self, model, max_batch_size: int, max_input_len: int, max_output_len: int) -> None:
|
|
max_total_token_num = max_batch_size * (max_input_len + max_output_len)
|
|
if model.config.model_type == "llama":
|
|
head_dim = model.config.hidden_size // model.config.num_attention_heads
|
|
head_num = model.config.num_key_value_heads // self.tp_size
|
|
num_hidden_layers = (
|
|
model.config.num_hidden_layers
|
|
if hasattr(model.config, "num_hidden_layers")
|
|
else model.config.num_layers
|
|
)
|
|
layer_num = num_hidden_layers // self.pp_size
|
|
elif model.config.model_type == "bloom":
|
|
head_dim = model.config.hidden_size // model.config.n_head
|
|
head_num = model.config.n_head // self.tp_size
|
|
num_hidden_layers = model.config.n_layer
|
|
layer_num = num_hidden_layers // self.pp_size
|
|
elif model.config.model_type == "chatglm":
|
|
head_dim = model.config.hidden_size // model.config.num_attention_heads
|
|
if model.config.multi_query_attention:
|
|
head_num = model.config.multi_query_group_num // self.tp_size
|
|
else:
|
|
head_num = model.config.num_attention_heads // self.tp_size
|
|
num_hidden_layers = model.config.num_layers
|
|
layer_num = num_hidden_layers // self.pp_size
|
|
else:
|
|
raise NotImplementedError("Only support llama, bloom and chatglm model.")
|
|
|
|
if self.quant == "smoothquant":
|
|
cache_manager = MemoryManager(max_total_token_num, torch.int8, head_num, head_dim, layer_num)
|
|
else:
|
|
cache_manager = MemoryManager(max_total_token_num, self.dtype, head_num, head_dim, layer_num)
|
|
return cache_manager
|