class_docstr: InferenceEngine 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.
      dtype (str): the data type of the model, should be one of 'fp16', 'fp32', 'bf16'.
      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. It will be determined by the model type if not provided.
      micro_batch_size (int): the micro batch size. Only useful when `pp_size` > 1.
      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.
      quant (str): the quantization method, should be one of 'smoothquant', 'gptq', None.
      verbose (bool): whether to return the time cost of each step.
class_docstr_tok: ["InferenceEngine","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",".","dtype","(","str",")",":","the","data","type","of","the","model",",","should","be","one","of","'fp16","'",",","'fp32","'",",","'bf16","'",".","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",".","It","will","be","determined","by","the","model","type","if","not","provided",".","micro_batch_size","(","int",")",":","the","micro","batch","size",".","Only","useful","when","`","pp_size","`",">","1.","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",".","quant","(","str",")",":","the","quantization","method",",","should","be","one","of","'smoothquant","'",",","'gptq","'",",","None",".","verbose","(","bool",")",":","whether","to","return","the","time","cost","of","each","step","."]
code: 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
  
file_path: ColossalAI/colossalai/inference/engine/engine.py
imports: from typing import Union
  import torch
  import torch.distributed
  import torch.nn
  from transformers.utils import logging
  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 ..kv_cache import MemoryManager
  from .microbatch_manager import MicroBatchManager
  from .policies import model_policy_map
label_desc: Use image node_id 4 for calling the InferenceEngine obj's underlying member method code with example usage: obj._init_manager(model, max_batch_size, max_input_len, max_output_len) and returns: cache_manager