class_docstr: This class serves as a micro library for ChatGLM model forwards under pipeline parallelism.
class_docstr_tok: ["This","class","serves","as","a","micro","library","for","ChatGLM","model","forwards","under","pipeline","parallelism","."]
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,
      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__)
      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
      )
      # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future.
      if past_key_values:
          logger.warning_once(
              "Non-empty past_key_values is not supported for pipeline models at the moment."
          )
          past_key_values = None
      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 use_cache:
          logger.warning_once(
              "use_cache=True is not supported for pipeline models at the moment."
          )
          use_cache = False
      if stage_manager.is_first_stage():
          batch_size, seq_length = input_ids.shape
          if inputs_embeds is None:
              inputs_embeds = self.embedding(input_ids)
          hidden_states = inputs_embeds
      else:
          seq_length, batch_size = hidden_states.shape[:2]
      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 = self.get_masks(
                  input_ids,
                  past_key_values,
                  padding_mask=attention_mask,
              )
      # Rotary positional embeddings
      rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
      if position_ids is not None:
          rotary_pos_emb = rotary_pos_emb[position_ids]
      else:
          rotary_pos_emb = rotary_pos_emb[None, :seq_length]
      rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
      if not past_key_values:
          past_key_values = [None for _ in range(self.num_layers)]
      presents = () if use_cache else None
      if self.encoder.gradient_checkpointing and self.encoder.training:
          if use_cache:
              logger.warning_once(
                  "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
              )
              use_cache = False
      all_self_attentions = None
      all_hidden_states = () if output_hidden_states else None
      start_idx, end_idx = stage_index[0], stage_index[1]
  
      if shard_config.enable_sequence_parallelism:
          hidden_states = split_forward_gather_backward(
              hidden_states,
              dim=0,
              process_group=shard_config.tensor_parallel_process_group,
          )
      for idx in range(start_idx, end_idx):
          layer = self.encoder._get_layer(idx)
          if output_hidden_states:
              all_hidden_states = all_hidden_states + (hidden_states,)
          if (
              self.encoder.gradient_checkpointing
              and self.encoder.training
          ):
              layer_ret = torch.utils.checkpoint.checkpoint(
                  layer,
                  hidden_states,
                  attention_mask,
                  rotary_pos_emb,
                  past_key_values[idx],
                  use_cache,
              )
          else:
              layer_ret = layer(
                  hidden_states,
                  full_attention_mask,
                  rotary_pos_emb,
                  kv_cache=past_key_values[idx],
                  use_cache=use_cache,
              )
          hidden_states, kv_cache = layer_ret
          if use_cache:
              presents = presents + (kv_cache,)
  
      if shard_config.enable_sequence_parallelism:
          hidden_states = gather_forward_split_backward(
              hidden_states,
              dim=0,
              process_group=shard_config.tensor_parallel_process_group,
          )
      if output_hidden_states:
          all_hidden_states = all_hidden_states + (hidden_states,)
      if stage_manager.is_last_stage():
          # final layer_norm
          if self.encoder.post_layer_norm:
              hidden_states = self.encoder.final_layernorm(
                  hidden_states
              )
          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,
          )
      else:
          return {"hidden_states": hidden_states}
  
file_path: ColossalAI/colossalai/shardformer/modeling/chatglm2.py
imports: from typing import List, Optional, Tuple
  import torch
  import torch.utils.checkpoint
  from torch.nn import CrossEntropyLoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig
  from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
label_desc: Use image node_id 1 for calling the ChatGLMPipelineForwards 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, stage_manager, hidden_states, stage_index, shard_config) and returns: BaseModelOutputWithPast, dict, tuple