code: def get_chatglm_sequence_parallel_forward_fn(
      shard_config: ShardConfig,
  ):
      def forward(
          self,
          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,
      ):
          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 = 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()
  
          # Run encoder.
          # [seq_len, batch_size, hidden_size] -> [seq_len/TP_size, batch_size, hidden_size]
          inputs_embeds = split_forward_gather_backward(
              inputs_embeds,
              dim=0,
              process_group=shard_config.tensor_parallel_process_group,
          )
          (
              hidden_states,
              presents,
              all_hidden_states,
              all_self_attentions,
          ) = self.encoder(
              inputs_embeds,
              full_attention_mask,
              rotary_pos_emb=rotary_pos_emb,
              kv_caches=past_key_values,
              use_cache=use_cache,
              output_hidden_states=output_hidden_states,
          )
  
          hidden_states = gather_forward_split_backward(
              hidden_states,
              dim=0,
              process_group=shard_config.tensor_parallel_process_group,
          )
  
          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,
          )
  
      return forward
  
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 3 for calling a global function with example usage: get_chatglm_sequence_parallel_forward_fn(shard_config) and returns: forward, BaseModelOutputWithPast, tuple