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] = True,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      return_last_logit: Optional[bool] = False,
      infer_state: Optional[BatchInferState] = None,
      stage_manager: Optional[PipelineStageManager] = None,
      hidden_states: Optional[torch.FloatTensor] = None,
      stage_index: Optional[List[int]] = None,
      shard_config: ShardConfig = None,
  ):
      logger = logging.get_logger(__name__)
  
      if output_attentions:
          logger.warning_once(
              "output_attentions=True is not supported for pipeline models at the moment."
          )
          output_attentions = False
      if output_hidden_states:
          logger.warning_once(
              "output_hidden_states=True is not supported for pipeline models at the moment."
          )
          output_hidden_states = False
  
      # If is first stage and hidden_states is not None, go throught lm_head first
      if stage_manager.is_first_stage() and hidden_states is not None:
          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()
          return {"logits": lm_logits}
  
      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,
          stage_manager=stage_manager,
          hidden_states=hidden_states,
          stage_index=stage_index,
          shard_config=shard_config,
      )
  
      return outputs
  
file_path: colossalai/colossalai/inference/engine/modeling/chatglm2.py
imports: from typing import List, Optional, Tuple
  import torch
  from transformers.utils import logging
  from colossalai.inference.kv_cache import BatchInferState
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig
  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, infer_state, stage_manager, hidden_states, stage_index, shard_config) and returns: outputs, dict