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_for_conditional_generation_forward(
      self: ChatGLMForConditionalGeneration,
      input_ids: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.Tensor] = None,
      past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      labels: Optional[torch.Tensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      return_last_logit: Optional[bool] = False,
  ):
      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
      )
      infer_state = self.infer_state
  
      if input_ids is not None and inputs_embeds is not None:
          raise ValueError(
              "You cannot specify both input_ids and inputs_embeds at the same time"
          )
      elif input_ids is not None:
          batch_size, seq_length = input_ids.shape
      elif inputs_embeds is not None:
          batch_size, seq_length, _ = inputs_embeds.shape
      else:
          raise ValueError(
              "You have to specify either input_ids or inputs_embeds"
          )
  
      if infer_state.is_context_stage:
          past_key_values_length = 0
      else:
          past_key_values_length = infer_state.max_len_in_batch - 1
  
      seq_length_with_past = seq_length + past_key_values_length
  
      # prefill stage at first
      if use_cache and seq_length != 1:
          infer_state.is_context_stage = True
          infer_state.context_mem_index = (
              infer_state.cache_manager.alloc(
                  infer_state.total_token_num
              )
          )
          infer_state.init_block_loc(
              infer_state.block_loc,
              infer_state.seq_len,
              seq_length,
              infer_state.context_mem_index,
          )
      else:
          infer_state.is_context_stage = False
          alloc_mem = infer_state.cache_manager.alloc_contiguous(
              batch_size
          )
          if alloc_mem is not None:
              infer_state.decode_is_contiguous = True
              infer_state.decode_mem_index = alloc_mem[0]
              infer_state.decode_mem_start = alloc_mem[1]
              infer_state.decode_mem_end = alloc_mem[2]
              infer_state.block_loc[
                  :, seq_length_with_past - 1
              ] = infer_state.decode_mem_index
          else:
              print(f" *** Encountered allocation non-contiguous")
              print(
                  f"    infer_state.cache_manager.past_key_values_length: {infer_state.cache_manager.past_key_values_length}"
              )
              infer_state.decode_is_contiguous = False
              alloc_mem = infer_state.cache_manager.alloc(batch_size)
              infer_state.decode_mem_index = alloc_mem
              # infer_state.decode_key_buffer = torch.empty((batch_size, self.tp_head_num_, self.head_dim_), dtype=torch.float16, device="cuda")
              # infer_state.decode_value_buffer = torch.empty((batch_size, self.tp_head_num_, self.head_dim_), dtype=torch.float16, device="cuda")
              infer_state.block_loc[
                  :, seq_length_with_past - 1
              ] = infer_state.decode_mem_index
  
      # related to rotary embedding
      if infer_state.is_context_stage:
          infer_state.position_cos = torch.index_select(
              self._cos_cached, 0, position_ids.view(-1)
          ).view(position_ids.view(-1).shape[0], -1)
          infer_state.position_sin = torch.index_select(
              self._sin_cached, 0, position_ids.view(-1)
          ).view(position_ids.view(-1).shape[0], -1)
      else:
          seq_len = infer_state.seq_len
          infer_state.position_cos = torch.index_select(
              self._cos_cached, 0, seq_len - 1
          ).view(seq_len.shape[0], -1)
          infer_state.position_sin = torch.index_select(
              self._sin_cached, 0, seq_len - 1
          ).view(seq_len.shape[0], -1)
          infer_state.other_kv_index = infer_state.block_loc[
              0, infer_state.max_len_in_batch - 1
          ].item()
  
      transformer_outputs = self.transformer(
          input_ids=input_ids,
          position_ids=position_ids,
          attention_mask=attention_mask,
          past_key_values=past_key_values,
          inputs_embeds=inputs_embeds,
          use_cache=use_cache,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          infer_state=infer_state,
      )
  
      hidden_states = transformer_outputs[0]
      if return_last_logit:
          hidden_states = hidden_states[-1:]
      lm_logits = self.transformer.output_layer(hidden_states)
      lm_logits = lm_logits.transpose(0, 1).contiguous()
  
      loss = None
      if labels is not None:
          lm_logits = lm_logits.to(torch.float32)
  
          # Shift so that tokens < n predict n
          shift_logits = lm_logits[..., :-1, :].contiguous()
          shift_labels = labels[..., 1:].contiguous()
          # Flatten the tokens
          loss_fct = CrossEntropyLoss(ignore_index=-100)
          loss = loss_fct(
              shift_logits.view(-1, shift_logits.size(-1)),
              shift_labels.view(-1),
          )
  
          lm_logits = lm_logits.to(hidden_states.dtype)
          loss = loss.to(hidden_states.dtype)
  
      if not return_dict:
          output = (lm_logits,) + transformer_outputs[1:]
          return ((loss,) + output) if loss is not None else output
  
      return CausalLMOutputWithPast(
          loss=loss,
          logits=lm_logits,
          past_key_values=transformer_outputs.past_key_values,
          hidden_states=transformer_outputs.hidden_states,
          attentions=transformer_outputs.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 1 for calling the ChatGLM2InferenceForwards obj's underlying member method code with example usage: obj.chatglm_for_conditional_generation_forward(input_ids, position_ids, attention_mask, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, return_last_logit) and returns: CausalLMOutputWithPast, unknown