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",")"]
code: def _post_init_gptq_buffer(self, model: nn.Module) -> None:
      from colossalai.inference.quant.gptq.cai_gptq import (
          CaiQuantLinear,
      )
  
      HAS_GPTQ_CUDA = False
      try:
          from colossalai.kernel.op_builder.gptq import GPTQBuilder
  
          gptq_cuda = GPTQBuilder().load()
          HAS_GPTQ_CUDA = True
      except ImportError:
          warnings.warn("CUDA gptq is not installed")
          HAS_GPTQ_CUDA = False
  
      for name, submodule in model.named_modules():
          if isinstance(submodule, CaiQuantLinear):
              self.max_dq_buffer_size = max(
                  self.max_dq_buffer_size, submodule.qweight.numel() * 8
              )
  
              if self.use_act_order:
                  self.max_inner_outer_dim = max(
                      self.max_inner_outer_dim,
                      submodule.infeatures,
                      submodule.outfeatures,
                  )
              self.bits = submodule.bits
      if not (HAS_GPTQ_CUDA and self.bits == 4):
          return
  
      max_input_len = 1
      if self.use_act_order:
          max_input_len = self.max_input_len
      # The temp_state buffer is required to reorder X in the act-order case.
      # The temp_dq buffer is required to dequantize weights when using cuBLAS, typically for the prefill.
      self.gptq_temp_state_buffer = torch.zeros(
          (max_input_len, self.max_inner_outer_dim),
          dtype=torch.float16,
          device=torch.cuda.current_device(),
      )
      self.gptq_temp_dq_buffer = torch.zeros(
          (1, self.max_dq_buffer_size),
          dtype=torch.float16,
          device=torch.cuda.current_device(),
      )
  
      gptq_cuda.prepare_buffers(
          torch.device(torch.cuda.current_device()),
          self.gptq_temp_state_buffer,
          self.gptq_temp_dq_buffer,
      )
      # Using the default from exllama repo here.
      matmul_recons_thd = 8
      matmul_fused_remap = False
      matmul_no_half2 = False
      gptq_cuda.set_tuning_params(
          matmul_recons_thd, matmul_fused_remap, matmul_no_half2
      )
  
      torch.cuda.empty_cache()
  
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 3 for calling the TPInferEngine obj's underlying member method code with example usage: obj._post_init_gptq_buffer(model) without return types