class_docstr: This class holds forwards for Chatglm2 inference.
  We intend to replace the forward methods for ChatGLMModel, ChatGLMEecoderLayer, and ChatGLMAttention.
class_docstr_tok: ["This","class","holds","forwards","for","Chatglm2","inference",".","We","intend","to","replace","the","forward","methods","for","ChatGLMModel",",","ChatGLMEecoderLayer",",","and","ChatGLMAttention","."]
code: def chatglm_model_forward(
      self: ChatGLMModel,
      input_ids,
      position_ids: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.BoolTensor] = None,
      full_attention_mask: Optional[torch.BoolTensor] = None,
      past_key_values: Optional[
          Tuple[Tuple[torch.Tensor, torch.Tensor], ...]
      ] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      use_cache: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      infer_state: BatchInferState = None,
  ):
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.config.output_hidden_states
      )
      use_cache = (
          use_cache if use_cache is not None else self.config.use_cache
      )
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
      batch_size, seq_length = input_ids.shape
  
      if inputs_embeds is None:
          inputs_embeds = self.embedding(input_ids)
  
      if self.pre_seq_len is not None:
          if past_key_values is None:
              past_key_values = self.get_prompt(
                  batch_size=batch_size,
                  device=input_ids.device,
                  dtype=inputs_embeds.dtype,
              )
          if attention_mask is not None:
              attention_mask = torch.cat(
                  [
                      attention_mask.new_ones(
                          (batch_size, self.pre_seq_len)
                      ),
                      attention_mask,
                  ],
                  dim=-1,
              )
      if full_attention_mask is None:
          if (
              attention_mask is not None and not attention_mask.all()
          ) or (past_key_values and seq_length != 1):
              full_attention_mask = get_masks(
                  self,
                  input_ids,
                  infer_state.cache_manager.past_key_values_length,
                  padding_mask=attention_mask,
              )
  
      # Run encoder.
      (
          hidden_states,
          presents,
          all_hidden_states,
          all_self_attentions,
      ) = self.encoder(
          inputs_embeds,
          full_attention_mask,
          kv_caches=past_key_values,
          use_cache=use_cache,
          output_hidden_states=output_hidden_states,
          infer_state=infer_state,
      )
  
      # update indices
      # infer_state.block_loc[:, infer_state.max_len_in_batch-1] = infer_state.total_token_num + torch.arange(0, batch_size, dtype=torch.int32, device="cuda")
      infer_state.start_loc = infer_state.start_loc + torch.arange(
          0, batch_size, dtype=torch.int32, device="cuda"
      )
      infer_state.seq_len += 1
      infer_state.max_len_in_batch += 1
  
      if not return_dict:
          return tuple(
              v
              for v in [
                  hidden_states,
                  presents,
                  all_hidden_states,
                  all_self_attentions,
              ]
              if v is not None
          )
  
      return BaseModelOutputWithPast(
          last_hidden_state=hidden_states,
          past_key_values=presents,
          hidden_states=all_hidden_states,
          attentions=all_self_attentions,
      )
  
file_path: ColossalAI/colossalai/legacy/inference/tensor_parallel/modeling/chatglm2.py
imports: import os
  from typing import Optional, Tuple
  import torch
  from torch.nn import CrossEntropyLoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel, GLMBlock, GLMTransformer, SelfAttention, split_tensor_along_last_dim
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 2 for calling the ChatGLM2InferenceForwards obj's underlying member method code with example usage: obj.chatglm_model_forward(input_ids, position_ids, attention_mask, full_attention_mask, past_key_values, inputs_embeds, use_cache, output_hidden_states, return_dict, infer_state) and returns: BaseModelOutputWithPast, tuple