class_docstr: This class serves as a micro library for forward function substitution of Bert models
  under pipeline setting.
docstr: encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
      Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
      the model is configured as a decoder.
  encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
      Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
      the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
  
      - 1 for tokens that are **not masked**,
      - 0 for tokens that are **masked**.
  labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
      Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
      `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
      ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
  past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
      Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
  
      If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
      don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
      `decoder_input_ids` of shape `(batch_size, sequence_length)`.
  use_cache (`bool`, *optional*):
      If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
      `past_key_values`).
docstr_tok: ["encoder_hidden_states","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","sequence_length",",","hidden_size",")","`",",","*","optional","*",")",":","Sequence","of","hidden-states","at","the","output","of","the","last","layer","of","the","encoder",".","Used","in","the","cross-attention","if","the","model","is","configured","as","a","decoder",".","encoder_attention_mask","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Mask","to","avoid","performing","attention","on","the","padding","token","indices","of","the","encoder","input",".","This","mask","is","used","in","the","cross-attention","if","the","model","is","configured","as","a","decoder",".","Mask","values","selected","in","`","[","0",",","1","]","`",":","-","1","for","tokens","that","are","*","*","not","masked","*","*",",","-","0","for","tokens","that","are","*","*","masked","*","*",".","labels","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Labels","for","computing","the","left-to-right","language","modeling","loss","(","next","word","prediction",")",".","Indices","should","be","in","`","[","-100",",","0",",","...",",","config.vocab_size","]","`","(","see","`","input_ids","`","docstring",")","Tokens","with","indices","set","to","`","-100","`","are","ignored","(","masked",")",",","the","loss","is","only","computed","for","the","tokens","with","labels","n","`","[","0",",","...",",","config.vocab_size","]","`","past_key_values","(","`","tuple","(","tuple","(","torch.FloatTensor",")",")","`","of","length","`","config.n_layers","`","with","each","tuple","having","4","tensors","of","shape","`","(","batch_size",",","num_heads",",","sequence_length","-","1",",","embed_size_per_head",")","`",")",":","Contains","precomputed","key","and","value","hidden","states","of","the","attention","blocks",".","Can","be","used","to","speed","up","decoding",".","If","`","past_key_values","`","are","used",",","the","user","can","optionally","input","only","the","last","`","decoder_input_ids","`","(","those","that","do","n't","have","their","past","key","value","states","given","to","this","model",")","of","shape","`","(","batch_size",",","1",")","`","instead","of","all","`","decoder_input_ids","`","of","shape","`","(","batch_size",",","sequence_length",")","`",".","use_cache","(","`","bool","`",",","*","optional","*",")",":","If","set","to","`","True","`",",","`","past_key_values","`","key","value","states","are","returned","and","can","be","used","to","speed","up","decoding","(","see","`","past_key_values","`",")","."]
code: def bert_lm_head_model_forward(
      self: BertLMHeadModel,
      input_ids: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.Tensor] = None,
      token_type_ids: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.Tensor] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      encoder_hidden_states: Optional[torch.Tensor] = None,
      encoder_attention_mask: Optional[torch.Tensor] = None,
      labels: Optional[torch.Tensor] = None,
      past_key_values: Optional[List[torch.Tensor]] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      hidden_states: Optional[torch.FloatTensor] = None,
      stage_manager: Optional[PipelineStageManager] = None,
      stage_index: Optional[List[int]] = None,
      shard_config: ShardConfig = None,
  ):
      r"""
      encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
          Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
          the model is configured as a decoder.
      encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
          the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
  
          - 1 for tokens that are **not masked**,
          - 0 for tokens that are **masked**.
      labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
          `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
          ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
      past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
          Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
  
          If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
          don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
          `decoder_input_ids` of shape `(batch_size, sequence_length)`.
      use_cache (`bool`, *optional*):
          If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
          `past_key_values`).
      """
      logger = logging.get_logger(__name__)
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      if labels is not None:
          use_cache = False
      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
  
      outputs = BertPipelineForwards.bert_model_forward(
          self.bert,
          input_ids,
          attention_mask=attention_mask,
          token_type_ids=token_type_ids,
          position_ids=position_ids,
          head_mask=head_mask,
          inputs_embeds=inputs_embeds,
          encoder_hidden_states=encoder_hidden_states,
          encoder_attention_mask=encoder_attention_mask,
          past_key_values=past_key_values,
          use_cache=use_cache,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          stage_manager=stage_manager,
          hidden_states=hidden_states
          if hidden_states is not None
          else None,
          stage_index=stage_index,
          shard_config=shard_config,
      )
      past_key_values = None
  
      if stage_manager.is_last_stage():
          sequence_output = outputs[0]
          prediction_scores = self.cls(sequence_output)
  
          lm_loss = None
          if labels is not None:
              # we are doing next-token prediction; shift prediction scores and input ids by one
              shifted_prediction_scores = prediction_scores[
                  :, :-1, :
              ].contiguous()
              labels = labels[:, 1:].contiguous()
              loss_fct = CrossEntropyLoss()
              lm_loss = loss_fct(
                  shifted_prediction_scores.view(
                      -1, self.config.vocab_size
                  ),
                  labels.view(-1),
              )
  
          if not return_dict:
              output = (prediction_scores,) + outputs[2:]
              return (
                  ((lm_loss,) + output)
                  if lm_loss is not None
                  else output
              )
  
          return CausalLMOutputWithCrossAttentions(
              loss=lm_loss,
              logits=prediction_scores,
              past_key_values=outputs.past_key_values,
              hidden_states=outputs.hidden_states,
              attentions=outputs.attentions,
              cross_attentions=outputs.cross_attentions,
          )
      else:
          hidden_states = outputs.get("hidden_states")
          # intermediate stage always return dict
          return {"hidden_states": hidden_states}
  
imports: import math
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions, CausalLMOutputWithCrossAttentions, MaskedLMOutput, MultipleChoiceModelOutput, NextSentencePredictorOutput, QuestionAnsweringModelOutput, SequenceClassifierOutput, TokenClassifierOutput
  from transformers.models.bert.modeling_bert import BertForMaskedLM, BertForMultipleChoice, BertForNextSentencePrediction, BertForPreTraining, BertForPreTrainingOutput, BertForQuestionAnswering, BertForSequenceClassification, BertForTokenClassification, BertLMHeadModel, BertModel
  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
label_desc: Use image node_id 3 for calling the BertPipelineForwards obj's underlying member method code with example usage: obj.bert_lm_head_model_forward(input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict, hidden_states, stage_manager, stage_index, shard_config) and returns: CausalLMOutputWithCrossAttentions, dict, unknown