class_docstr: This class serves as a micro library for forward function substitution of Llama models
  under pipeline setting.
code: def openmoe_model_forward(
      self: OpenMoeModel,
      input_ids: torch.LongTensor = None,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[List[torch.FloatTensor]] = None,
      inputs_embeds: 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,
      past_router_aux_loss: Optional[torch.FloatTensor] = None,
      past_router_z_loss: Optional[torch.FloatTensor] = None,
  ):
      # reset moe loss for different data
      MOE_MANAGER.reset_loss()
  
      logger = logging.get_logger(__name__)
  
      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
      )
  
      # retrieve input_ids and inputs_embeds
      if stage_manager.is_first_stage():
          if input_ids is not None and inputs_embeds is not None:
              raise ValueError(
                  "You cannot specify both decoder_input_ids and decoder_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 decoder_input_ids or decoder_inputs_embeds"
              )
          device = (
              input_ids.device
              if input_ids is not None
              else inputs_embeds.device
          )
          if inputs_embeds is None:
              inputs_embeds = self.embed_tokens(input_ids)
          hidden_states = inputs_embeds
      else:
          input_shape = hidden_states.shape[:-1]
          batch_size, seq_length = input_shape
          device = hidden_states.device
  
      seq_length_with_past = seq_length
      past_key_values_length = 0
  
      # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future.
      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 past_key_values is not None:
          past_key_values_length = past_key_values[0][0].shape[2]
          seq_length_with_past = (
              seq_length_with_past + past_key_values_length
          )
  
      if position_ids is None:
          position_ids = torch.arange(
              past_key_values_length,
              seq_length + past_key_values_length,
              dtype=torch.long,
              device=device,
          )
          position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
      else:
          position_ids = position_ids.view(-1, seq_length).long()
  
      # embed positions, for the first stage, hidden_states is the input embeddings,
      # for the other stages, hidden_states is the output of the previous stage
      if attention_mask is None:
          attention_mask = torch.ones(
              (batch_size, seq_length_with_past),
              dtype=torch.bool,
              device=hidden_states.device,
          )
      attention_mask = self._prepare_decoder_attention_mask(
          attention_mask,
          (batch_size, seq_length),
          hidden_states,
          past_key_values_length,
      )
  
      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
  
      # decoder layers
      all_hidden_states = () if output_hidden_states else None
      all_self_attns = () if output_attentions else None
      next_decoder_cache = () if use_cache else None
  
      start_idx, end_idx = stage_index[0], stage_index[1]
      for idx, decoder_layer in enumerate(
          self.layers[start_idx:end_idx], start=start_idx
      ):
          if output_hidden_states:
              all_hidden_states += (hidden_states,)
  
          past_key_value = (
              past_key_values[idx]
              if past_key_values is not None
              else None
          )
  
          if self.gradient_checkpointing and self.training:
  
              def create_custom_forward(module):
                  def custom_forward(*inputs):
                      # None for past_key_value
                      return module(*inputs, output_attentions, None)
  
                  return custom_forward
  
              layer_outputs = torch.utils.checkpoint.checkpoint(
                  create_custom_forward(decoder_layer),
                  hidden_states,
                  attention_mask,
                  position_ids,
                  None,
              )
          else:
              layer_outputs = decoder_layer(
                  hidden_states,
                  attention_mask=attention_mask,
                  position_ids=position_ids,
                  past_key_value=past_key_value,
                  output_attentions=output_attentions,
                  use_cache=use_cache,
              )
  
          hidden_states = layer_outputs[0]
  
          if use_cache:
              next_decoder_cache += (
                  layer_outputs[2 if output_attentions else 1],
              )
          if output_attentions:
              all_self_attns += (layer_outputs[1],)
  
      if stage_manager.is_last_stage():
          hidden_states = self.norm(hidden_states)
  
      # add hidden states from the last decoder layer
      if output_hidden_states:
          all_hidden_states += (hidden_states,)
      next_cache = next_decoder_cache if use_cache else None
  
      # concat past losses with current ones
      router_aux_loss, router_z_loss = MOE_MANAGER.get_loss()
      if (
          past_router_aux_loss is not None
          and past_router_z_loss is not None
      ):
          router_aux_loss = past_router_aux_loss + router_aux_loss
          router_z_loss = past_router_z_loss + router_z_loss
  
      if stage_manager.is_last_stage():
          return tuple(
              [
                  hidden_states,
                  next_cache,
                  all_hidden_states,
                  all_self_attns,
                  router_aux_loss,
                  router_z_loss,
              ]
          )
      # always return dict for imediate stage
      return {
          "hidden_states": hidden_states,
          "router_aux_loss": router_aux_loss,
          "router_z_loss": router_z_loss,
      }
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Optional, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.nn import Module
  from transformers.modeling_outputs import CausalLMOutputWithPast
  from transformers.utils import logging
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
  from .modeling_openmoe import OpenMoeDecoderLayer, OpenMoeForCausalLM, OpenMoeModel
label_desc: Use image node_id 1 for calling the OpenMoePipelineForwards obj's underlying member method code with example usage: obj.openmoe_model_forward(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, stage_manager, hidden_states, stage_index, past_router_aux_loss, past_router_z_loss) and returns: dict, tuple, custom_forward, module