class_docstr: This class serves as a micro library for forward function substitution of GPT2 models
  under pipeline setting.
class_docstr_tok: ["This","class","serves","as","a","micro","library","for","forward","function","substitution","of","GPT2","models","under","pipeline","setting","."]
code: def gpt2_model_forward(
      self: GPT2Model,
      input_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
      attention_mask: Optional[torch.FloatTensor] = None,
      token_type_ids: Optional[torch.LongTensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      head_mask: Optional[torch.FloatTensor] = None,
      inputs_embeds: Optional[torch.FloatTensor] = None,
      encoder_hidden_states: Optional[torch.Tensor] = None,
      encoder_attention_mask: Optional[torch.FloatTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: 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,
  ) -> Union[Dict, Tuple, BaseModelOutputWithPastAndCrossAttentions]:
      # This function is modified on the basis of transformers.models.gpt2.modeling_gpt2.GPT2Model.forward.
      # Please refer to original code of transformers for more details.
  
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      logger = logging.get_logger(__name__)
  
      # Preprocess passed in arguments
      # TODO(baizhou): 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_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 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():
          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:
              input_shape = input_ids.size()
              input_ids = input_ids.view(-1, input_shape[-1])
              batch_size = input_ids.shape[0]
          elif inputs_embeds is not None:
              input_shape = inputs_embeds.size()[:-1]
              batch_size = inputs_embeds.shape[0]
          else:
              raise ValueError(
                  "You have to specify either input_ids or inputs_embeds"
              )
  
          device = (
              input_ids.device
              if input_ids is not None
              else inputs_embeds.device
          )
          if token_type_ids is not None:
              token_type_ids = token_type_ids.view(-1, input_shape[-1])
      else:
          if hidden_states is None:
              raise ValueError(
                  "hidden_states shouldn't be None for stages other than the first stage."
              )
          input_shape = hidden_states.size()[:-1]
          device = hidden_states.device
          hidden_states = hidden_states.view(
              (-1,) + hidden_states.shape[-2:]
          )
          batch_size = hidden_states.shape[0]
  
      # GPT2Attention mask.
      if attention_mask is not None:
          if batch_size <= 0:
              raise ValueError("batch_size has to be defined and > 0")
          attention_mask = attention_mask.view(batch_size, -1)
          # We create a 3D attention mask from a 2D tensor mask.
          # Sizes are [batch_size, 1, 1, to_seq_length]
          # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
          # this attention mask is more simple than the triangular masking of causal attention
          # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
          attention_mask = attention_mask[:, None, None, :]
  
          # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
          # masked positions, this operation will create a tensor which is 0.0 for
          # positions we want to attend and the dtype's smallest value for masked positions.
          # Since we are adding it to the raw scores before the softmax, this is
          # effectively the same as removing these entirely.
          attention_mask = attention_mask.to(
              dtype=self.dtype
          )  # fp16 compatibility
          attention_mask = (1.0 - attention_mask) * torch.finfo(
              self.dtype
          ).min
  
      # If a 2D or 3D attention mask is provided for the cross-attention
      # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
      if (
          self.config.add_cross_attention
          and encoder_hidden_states is not None
      ):
          (
              encoder_batch_size,
              encoder_sequence_length,
              _,
          ) = encoder_hidden_states.size()
          encoder_hidden_shape = (
              encoder_batch_size,
              encoder_sequence_length,
          )
          if encoder_attention_mask is None:
              encoder_attention_mask = torch.ones(
                  encoder_hidden_shape, device=device
              )
          encoder_attention_mask = self.invert_attention_mask(
              encoder_attention_mask
          )
      else:
          encoder_attention_mask = None
  
      # Prepare head mask if needed
      # 1.0 in head_mask indicate we keep the head
      # attention_probs has shape bsz x n_heads x N x N
      # head_mask has shape n_layer x batch x n_heads x N x N
      head_mask = self.get_head_mask(head_mask, self.config.n_layer)
  
      if stage_manager.is_first_stage():
          if position_ids is not None:
              position_ids = position_ids.view(-1, input_shape[-1])
          else:
              position_ids = torch.arange(
                  0, input_shape[-1], dtype=torch.long, device=device
              )
              position_ids = position_ids.unsqueeze(0).view(
                  -1, input_shape[-1]
              )
  
          if inputs_embeds is None:
              inputs_embeds = self.wte(input_ids)
          position_embeds = self.wpe(position_ids)
          hidden_states = inputs_embeds + position_embeds
          if token_type_ids is not None:
              token_type_embeds = self.wte(token_type_ids)
              hidden_states = hidden_states + token_type_embeds
          hidden_states = self.drop(hidden_states)
  
      output_shape = input_shape + (hidden_states.size(-1),)
  
      if self.gradient_checkpointing and self.training:
          if use_cache:
              logger.warning_once(
                  "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
              )
              use_cache = False
      presents = () if use_cache else None
      all_self_attentions = () if output_attentions else None
      all_cross_attentions = (
          ()
          if output_attentions and self.config.add_cross_attention
          else None
      )
      all_hidden_states = () if output_hidden_states else None
  
      # split the input tensor along sequence dimension
      # [batch_size, seq_len, hidden_size] -> [batch_size, seq_len/TP_size, hidden_size]
      if shard_config.enable_sequence_parallelism:
          hidden_states = split_forward_gather_backward(
              hidden_states,
              dim=1,
              process_group=shard_config.tensor_parallel_process_group,
          )
  
      # Going through held blocks.
      start_idx, end_idx = stage_index[0], stage_index[1]
      for i in range(start_idx, end_idx):
          block = self.h[i]
          torch.cuda.set_device(hidden_states.device)
          # Ensure that attention_mask is always on the same device as hidden_states
          if attention_mask is not None:
              attention_mask = attention_mask.to(hidden_states.device)
          if isinstance(head_mask, torch.Tensor):
              head_mask = head_mask.to(hidden_states.device)
          if output_hidden_states:
              all_hidden_states = all_hidden_states + (hidden_states,)
  
          if self.gradient_checkpointing and self.training:
  
              def create_custom_forward(module):
                  def custom_forward(*inputs):
                      # None for past_key_value
                      return module(
                          *inputs, use_cache, output_attentions
                      )
  
                  return custom_forward
  
              outputs = torch.utils.checkpoint.checkpoint(
                  create_custom_forward(block),
                  hidden_states,
                  None,
                  attention_mask,
                  head_mask[i],
                  encoder_hidden_states,
                  encoder_attention_mask,
              )
          else:
              outputs = block(
                  hidden_states,
                  layer_past=None,
                  attention_mask=attention_mask,
                  head_mask=head_mask[i],
                  encoder_hidden_states=encoder_hidden_states,
                  encoder_attention_mask=encoder_attention_mask,
                  use_cache=use_cache,
                  output_attentions=output_attentions,
              )
  
          hidden_states = outputs[0]
          if use_cache is True:
              presents = presents + (outputs[1],)
  
          if output_attentions:
              all_self_attentions = all_self_attentions + (
                  outputs[2 if use_cache else 1],
              )
              if self.config.add_cross_attention:
                  all_cross_attentions = all_cross_attentions + (
                      outputs[3 if use_cache else 2],
                  )
  
      # When sequence parallelism done, gather the output tensor in forward and split it in backward
      if shard_config.enable_sequence_parallelism:
          hidden_states = gather_forward_split_backward(
              hidden_states,
              dim=1,
              process_group=shard_config.tensor_parallel_process_group,
          )
  
      if stage_manager.is_last_stage():
          hidden_states = self.ln_f(hidden_states)
  
      hidden_states = hidden_states.view(output_shape)
  
      # Add last hidden state
      if output_hidden_states:
          all_hidden_states = all_hidden_states + (hidden_states,)
  
      if stage_manager.is_last_stage():
          if not return_dict:
              return tuple(
                  v
                  for v in [
                      hidden_states,
                      presents,
                      all_hidden_states,
                      all_self_attentions,
                      all_cross_attentions,
                  ]
                  if v is not None
              )
  
          return BaseModelOutputWithPastAndCrossAttentions(
              last_hidden_state=hidden_states,
              past_key_values=presents,
              hidden_states=all_hidden_states,
              attentions=all_self_attentions,
              cross_attentions=all_cross_attentions,
          )
      else:
          # always return dict for intermediate stage
          return {"hidden_states": hidden_states}
  
file_path: ColossalAI/colossalai/shardformer/modeling/gpt2.py
imports: from typing import Dict, List, Optional, Tuple, Union
  import torch
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from transformers.models.gpt2.modeling_gpt2 import GPT2DoubleHeadsModel, GPT2DoubleHeadsModelOutput, GPT2ForQuestionAnswering, GPT2ForSequenceClassification, GPT2ForTokenClassification, GPT2LMHeadModel, GPT2Model
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
  from colossalai.shardformer.shard import ShardConfig
label_desc: Use image node_id 1 for calling the GPT2PipelineForwards obj's underlying member method code with example usage: obj.gpt2_model_forward(input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions, output_hidden_states, return_dict, stage_manager, hidden_states, stage_index, shard_config) and returns: BaseModelOutputWithPastAndCrossAttentions, dict, tuple, custom_forward, module