class_docstr: Engine class for tensor parallel inference.
  
  Args:
      model (Module): original model, e.g. huggingface CausalLM
      shard_config (ShardConfig): The config for sharding original model
      max_batch_size (int): maximum batch size
      max_input_len (int): maximum input length of sequence
      max_output_len (int): maximum output length of output tokens
      dtype (torch.dtype): datatype used to init KV cache space
      device (str): device the KV cache of engine to be initialized on
  
  Examples:
      >>> # define model and shard config for your inference
      >>> model = ...
      >>> generate_kwargs = ...
      >>> shard_config = ShardConfig(enable_tensor_parallelism=True, extra_kwargs={"inference_only": True})
      >>> infer_engine = TPInferEngine(model, shard_config, MAX_BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
      >>> outputs = infer_engine.generate(input_ids, **generate_kwargs)
class_docstr_tok: ["Engine","class","for","tensor","parallel","inference",".","Args",":","model","(","Module",")",":","original","model",",","e.g",".","huggingface","CausalLM","shard_config","(","ShardConfig",")",":","The","config","for","sharding","original","model","max_batch_size","(","int",")",":","maximum","batch","size","max_input_len","(","int",")",":","maximum","input","length","of","sequence","max_output_len","(","int",")",":","maximum","output","length","of","output","tokens","dtype","(","torch.dtype",")",":","datatype","used","to","init","KV","cache","space","device","(","str",")",":","device","the","KV","cache","of","engine","to","be","initialized","on","Examples",":",">",">",">","#","define","model","and","shard","config","for","your","inference",">",">",">","model","=","...",">",">",">","generate_kwargs","=","...",">",">",">","shard_config","=","ShardConfig","(","enable_tensor_parallelism=True",",","extra_kwargs=","{","``","inference_only","''",":","True","}",")",">",">",">","infer_engine","=","TPInferEngine","(","model",",","shard_config",",","MAX_BATCH_SIZE",",","MAX_INPUT_LEN",",","MAX_OUTPUT_LEN",")",">",">",">","outputs","=","infer_engine.generate","(","input_ids",",","*","*","generate_kwargs",")"]
docstr: Shard original model by the given ShardFormer and store the sharded model.
docstr_tok: ["Shard","original","model","by","the","given","ShardFormer","and","store","the","sharded","model","."]
code: def _shard_model_by(
      self, shardformer: ShardFormer, model: nn.Module
  ) -> None:
      """Shard original model by the given ShardFormer and store the sharded model."""
      assert (
          self.tp_size == shardformer.shard_config.tensor_parallel_size
      ), "Discrepancy between the tp size of TPInferEngine and the tp size of shard config"
      model_name = model.__class__.__name__
      assert (
          model_name in self.supported_models
      ), f"Unsupported model cls {model_name} for TP inference."
      if self.shard_config.extra_kwargs.get("inference_gptq", False):
          model = model.model
      policy = get_autopolicy(model, shard_config=self.shard_config)
      self.model, _ = shardformer.optimize(model, policy)
      if self.shard_config.extra_kwargs.get("inference_gptq", False):
          self._post_init_gptq_buffer(self.model)
  
      self.model = self.model.cuda()
  
imports: from typing import Any, Callable, List, Optional, Union
  import torch
  import torch.nn
  from transformers import BloomForCausalLM, LlamaForCausalLM
  from transformers.generation import GenerationConfig
  from transformers.generation.stopping_criteria import StoppingCriteriaList
  from transformers.tokenization_utils_base import BatchEncoding
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.shardformer.policies.auto_policy import get_autopolicy
  from .batch_infer_state import BatchInferState
  from .kvcache_manager import MemoryManager
label_desc: Use image node_id 6 for calling the TPInferEngine obj's underlying member method code with example usage: obj._shard_model_by(shardformer, model) without return types