class_docstr: This class serves as a micro library for bloom pipeline forwards.
code: def bloom_model_forward(
      self: BloomModel,
      input_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[
          Tuple[Tuple[torch.Tensor, torch.Tensor], ...]
      ] = None,
      attention_mask: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.LongTensor] = None,
      inputs_embeds: Optional[torch.LongTensor] = 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,
      **deprecated_arguments,
  ) -> Union[
      Tuple[torch.Tensor, ...],
      "BaseModelOutputWithPastAndCrossAttentions",
  ]:
      logger = logging.get_logger(__name__)
  
      if deprecated_arguments.pop("position_ids", False) is not False:
          # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
          warnings.warn(
              "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
              " passing `position_ids`.",
              FutureWarning,
          )
      if len(deprecated_arguments) > 0:
          raise ValueError(
              f"Got unexpected arguments: {deprecated_arguments}"
          )
  
      output_attentions = (
          output_attentions
          if output_attentions is not None
          else self.config.output_attentions
      )
      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
      )
  
      # add warnings here
      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
      # Prepare head mask if needed
      # 1.0 in head_mask indicate we keep the head
      # attention_probs has shape batch_size x num_heads x N x N
  
      # head_mask has shape n_layer x batch x num_heads x N x N
      head_mask = self.get_head_mask(head_mask, self.config.n_layer)
  
      # case: First stage of training
      if stage_manager.is_first_stage():
          # check input_ids and inputs_embeds
          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:
              batch_size, seq_length = input_ids.shape
          elif inputs_embeds is not None:
              batch_size, seq_length, _ = inputs_embeds.shape
          else:
              raise ValueError(
                  "You have to specify either input_ids or inputs_embeds"
              )
  
          if inputs_embeds is None:
              inputs_embeds = self.word_embeddings(input_ids)
  
          hidden_states = self.word_embeddings_layernorm(inputs_embeds)
          # initialize in the first stage and then pass to the next stage
      else:
          input_shape = hidden_states.shape[:-1]
          batch_size, seq_length = input_shape
  
      # extra recording tensor should be generated in the first stage
  
      presents = () if use_cache else None
      all_self_attentions = () if output_attentions else None
      all_hidden_states = () if output_hidden_states else None
  
      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
  
      if past_key_values is None:
          past_key_values = tuple([None] * len(self.h))
      # Compute alibi tensor: check build_alibi_tensor documentation,build for every stage
      seq_length_with_past = seq_length
      past_key_values_length = 0
      if past_key_values[0] is not None:
          past_key_values_length = past_key_values[0][0].shape[
              2
          ]  # source_len
  
          seq_length_with_past = (
              seq_length_with_past + past_key_values_length
          )
      if attention_mask is None:
          attention_mask = torch.ones(
              (batch_size, seq_length_with_past),
              device=hidden_states.device,
          )
      else:
          attention_mask = attention_mask.to(hidden_states.device)
  
      alibi = self.build_alibi_tensor(
          attention_mask, self.num_heads, dtype=hidden_states.dtype
      )
  
      # causal_mask is constructed every stage and its input is passed through different stages
      causal_mask = self._prepare_attn_mask(
          attention_mask,
          input_shape=(batch_size, seq_length),
          past_key_values_length=past_key_values_length,
      )
  
      # 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,
          )
  
      start_idx, end_idx = stage_index[0], stage_index[1]
      for i, (block, layer_past) in enumerate(
          zip(
              self.h[start_idx:end_idx],
              past_key_values[start_idx:end_idx],
          ),
          start=start_idx,
      ):
          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=use_cache,
                          output_attentions=output_attentions,
                      )
  
                  return custom_forward
  
              outputs = torch.utils.checkpoint.checkpoint(
                  create_custom_forward(block),
                  hidden_states,
                  alibi,
                  causal_mask,
                  layer_past,
                  head_mask[i],
              )
          else:
              outputs = block(
                  hidden_states,
                  layer_past=layer_past,
                  attention_mask=causal_mask,
                  head_mask=head_mask[i],
                  use_cache=use_cache,
                  output_attentions=output_attentions,
                  alibi=alibi,
              )
  
          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],
              )
  
      # 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():
          # Add last hidden state
          hidden_states = self.ln_f(hidden_states)
  
      # TODO(jianghai): deal with all_hidden_states, all_self_attentions, presents
      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,
                  ]
                  if v is not None
              )
  
          # attention_mask is not returned ; presents = past_key_values
          return BaseModelOutputWithPastAndCrossAttentions(
              last_hidden_state=hidden_states,
              past_key_values=presents,
              hidden_states=all_hidden_states,
              attentions=all_self_attentions,
          )
      else:
          # always return dict for imediate stage
          return {"hidden_states": hidden_states}
  
imports: import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from torch.nn import functional
  from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from transformers.models.bloom.modeling_bloom import BloomForCausalLM, BloomForQuestionAnswering, BloomForSequenceClassification, BloomForTokenClassification, BloomModel
  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 BloomPipelineForwards obj's underlying member method code with example usage: obj.bloom_model_forward(input_ids, past_key_values, attention_mask, head_mask, inputs_embeds, 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