class_docstr: This class serves as a micro library for forward function substitution of OPT models
  under pipeline setting.
class_docstr_tok: ["This","class","serves","as","a","micro","library","for","forward","function","substitution","of","OPT","models","under","pipeline","setting","."]
docstr: This function is modified on the basis of transformers.models.opt.modeling_opt.OPTForCausalLM.forward.
  Please refer to original code of transformers for more details.
docstr_tok: ["This","function","is","modified","on","the","basis","of","transformers.models.opt.modeling_opt.OPTForCausalLM.forward",".","Please","refer","to","original","code","of","transformers","for","more","details","."]
code: def opt_for_causal_lm_forward(
      self: OPTForCausalLM,
      input_ids: torch.LongTensor = None,
      attention_mask: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.Tensor] = None,
      past_key_values: Optional[List[torch.FloatTensor]] = None,
      inputs_embeds: Optional[torch.FloatTensor] = None,
      labels: 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,
  ) -> Union[Tuple, CausalLMOutputWithPast]:
      r"""
      This function is modified on the basis of transformers.models.opt.modeling_opt.OPTForCausalLM.forward.
      Please refer to original code of transformers for more details.
      """
  
      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
      )
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
      outputs = OPTPipelineForwards.opt_model_forward(
          self.model,
          input_ids=input_ids,
          attention_mask=attention_mask,
          head_mask=head_mask,
          past_key_values=past_key_values,
          inputs_embeds=inputs_embeds,
          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,
          stage_index=stage_index,
      )
      if stage_manager.is_last_stage():
          logits = self.lm_head(outputs[0]).contiguous()
          loss = None
          if labels is not None:
              # move labels to correct device to enable model parallelism
              labels = labels.to(logits.device)
              # Shift so that tokens < n predict n
              shift_logits = logits[..., :-1, :].contiguous()
              shift_labels = labels[..., 1:].contiguous()
              # Flatten the tokens
              loss_fct = CrossEntropyLoss()
              loss = loss_fct(
                  shift_logits.view(-1, self.config.vocab_size),
                  shift_labels.view(-1),
              )
          if not return_dict:
              output = (logits,) + outputs[1:]
              return (loss,) + output if loss is not None else output
  
          return CausalLMOutputWithPast(
              loss=loss,
              logits=logits,
              past_key_values=outputs.past_key_values,
              hidden_states=outputs.hidden_states,
              attentions=outputs.attentions,
          )
      else:
          hidden_states = outputs.get("hidden_states")
          return {"hidden_states": hidden_states}
  
file_path: colossalai/colossalai/shardformer/modeling/opt.py
imports: import random
  from typing import List, Optional, Tuple, Union
  import torch
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast
  from transformers.models.opt.modeling_opt import OPTForCausalLM, OPTForQuestionAnswering, OPTForSequenceClassification, OPTModel
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
label_desc: Use image node_id 4 for calling the OPTPipelineForwards obj's underlying member method code with example usage: obj.opt_for_causal_lm_forward(input_ids, attention_mask, head_mask, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, stage_manager, hidden_states, stage_index) and returns: CausalLMOutputWithPast, dict, unknown