class_docstr: This class serves as a micro library for forward function substitution of Llama models
  under pipeline setting.
class_docstr_tok: ["This","class","serves","as","a","micro","library","for","forward","function","substitution","of","Llama","models","under","pipeline","setting","."]
docstr: Args:
      input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
          Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
          provide it.
  
          Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and
          [`PreTrainedTokenizer.__call__`] for details.
  
          [What are input IDs?](../glossary#input-ids)
      attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
          Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
          - 1 for tokens that are **not masked**,
          - 0 for tokens that are **masked**.
  
          [What are attention masks?](../glossary#attention-mask)
      encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
          Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
          of the decoder.
      head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
          Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
  
          - 1 indicates the head is **not masked**,
          - 0 indicates the head is **masked**.
  
      cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
          Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention
          on hidden heads. Mask values selected in `[0, 1]`:
  
          - 1 indicates the head is **not masked**,
          - 0 indicates the head is **masked**.
  
      past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
          Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
          shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
          shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
  
          Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
          cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential 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)`. inputs_embeds (`torch.FloatTensor` of
          shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing
          `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more
          control over how to convert `input_ids` indices into associated vectors than the model's internal
          embedding lookup matrix.
      output_attentions (`bool`, *optional*):
          Whether or not to return the attentions tensors of all attention layers. See `attentions` under
          returned tensors for more detail.
      output_hidden_states (`bool`, *optional*):
          Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
          for more detail.
      return_dict (`bool`, *optional*):
          Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
docstr_tok: ["Args",":","input_ids","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",")",":","Indices","of","input","sequence","tokens","in","the","vocabulary",".","Padding","will","be","ignored","by","default","should","you","provide","it",".","Indices","can","be","obtained","using","[","`","WhisperTokenizer","`","]",".","See","[","`","PreTrainedTokenizer.encode","`","]","and","[","`","PreTrainedTokenizer.__call__","`","]","for","details",".","[","What","are","input","IDs","?","]","(","..","\/glossary","#","input-ids",")","attention_mask","(","`","torch.Tensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Mask","to","avoid","performing","attention","on","padding","token","indices",".","Mask","values","selected","in","`","[","0",",","1","]","`",":","-","1","for","tokens","that","are","*","*","not","masked","*","*",",","-","0","for","tokens","that","are","*","*","masked","*","*",".","[","What","are","attention","masks","?","]","(","..","\/glossary","#","attention-mask",")","encoder_hidden_states","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","encoder_sequence_length",",","hidden_size",")","`",",","*","optional","*",")",":","Sequence","of","hidden-states","at","the","output","of","the","last","layer","of","the","encoder",".","Used","in","the","cross-attention","of","the","decoder",".","head_mask","(","`","torch.Tensor","`","of","shape","`","(","decoder_layers",",","decoder_attention_heads",")","`",",","*","optional","*",")",":","Mask","to","nullify","selected","heads","of","the","attention","modules",".","Mask","values","selected","in","`","[","0",",","1","]","`",":","-","1","indicates","the","head","is","*","*","not","masked","*","*",",","-","0","indicates","the","head","is","*","*","masked","*","*",".","cross_attn_head_mask","(","`","torch.Tensor","`","of","shape","`","(","decoder_layers",",","decoder_attention_heads",")","`",",","*","optional","*",")",":","Mask","to","nullify","selected","heads","of","the","attention","modules","in","encoder","to","avoid","performing","cross-attention","on","hidden","heads",".","Mask","values","selected","in","`","[","0",",","1","]","`",":","-","1","indicates","the","head","is","*","*","not","masked","*","*",",","-","0","indicates","the","head","is","*","*","masked","*","*",".","past_key_values","(","`","tuple","(","tuple","(","torch.FloatTensor",")",")","`",",","*","optional","*",",","returned","when","`","use_cache=True","`","is","passed","or","when","`","config.use_cache=True","`",")",":","Tuple","of","`","tuple","(","torch.FloatTensor",")","`","of","length","`","config.n_layers","`",",","with","each","tuple","having","2","tensors","of","shape","`","(","batch_size",",","num_heads",",","sequence_length",",","embed_size_per_head",")","`",")","and","2","additional","tensors","of","shape","`","(","batch_size",",","num_heads",",","encoder_sequence_length",",","embed_size_per_head",")","`",".","Contains","pre-computed","hidden-states","(","key","and","values","in","the","self-attention","blocks","and","in","the","cross-attention","blocks",")","that","can","be","used","(","see","`","past_key_values","`","input",")","to","speed","up","sequential","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",")","`",".","inputs_embeds","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","sequence_length",",","hidden_size",")","`",",","*","optional","*",")",":","Optionally",",","instead","of","passing","`","input_ids","`","you","can","choose","to","directly","pass","an","embedded","representation",".","This","is","useful","if","you","want","more","control","over","how","to","convert","`","input_ids","`","indices","into","associated","vectors","than","the","model","'s","internal","embedding","lookup","matrix",".","output_attentions","(","`","bool","`",",","*","optional","*",")",":","Whether","or","not","to","return","the","attentions","tensors","of","all","attention","layers",".","See","`","attentions","`","under","returned","tensors","for","more","detail",".","output_hidden_states","(","`","bool","`",",","*","optional","*",")",":","Whether","or","not","to","return","the","hidden","states","of","all","layers",".","See","`","hidden_states","`","under","returned","tensors","for","more","detail",".","return_dict","(","`","bool","`",",","*","optional","*",")",":","Whether","or","not","to","return","a","[","`","~utils.ModelOutput","`","]","instead","of","a","plain","tuple","."]
code: def whisper_decoder_forward(
      self,
      input_ids=None,
      attention_mask=None,
      encoder_hidden_states=None,
      head_mask=None,
      cross_attn_head_mask=None,
      past_key_values=None,
      inputs_embeds=None,
      use_cache=None,
      output_attentions=None,
      output_hidden_states=None,
      return_dict=None,
      stage_manager: Optional[PipelineStageManager] = None,
      hidden_states: Optional[torch.FloatTensor] = None,
      stage_index: Optional[List[int]] = None,
      decoder_starting_stage: Optional[int] = None,
  ):
      r"""
      Args:
          input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
              Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
              provide it.
  
              Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and
              [`PreTrainedTokenizer.__call__`] for details.
  
              [What are input IDs?](../glossary#input-ids)
          attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
              Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
  
              [What are attention masks?](../glossary#attention-mask)
          encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
              Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
              of the decoder.
          head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
              Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
  
              - 1 indicates the head is **not masked**,
              - 0 indicates the head is **masked**.
  
          cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
              Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention
              on hidden heads. Mask values selected in `[0, 1]`:
  
              - 1 indicates the head is **not masked**,
              - 0 indicates the head is **masked**.
  
          past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
              Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
              shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
              shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
  
              Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
              cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential 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)`. inputs_embeds (`torch.FloatTensor` of
              shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing
              `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more
              control over how to convert `input_ids` indices into associated vectors than the model's internal
              embedding lookup matrix.
          output_attentions (`bool`, *optional*):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under
              returned tensors for more detail.
          output_hidden_states (`bool`, *optional*):
              Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
              for more detail.
          return_dict (`bool`, *optional*):
              Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
      """
      logger = logging.get_logger(__name__)
      stage = stage_manager.stage
      at_first_stage = stage == decoder_starting_stage
      at_last_stage = stage == stage_manager.num_stages - 1
  
      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
      )
  
      # decoder layers
      all_hidden_states = () if output_hidden_states else None
      all_self_attns = () if output_attentions else None
      all_cross_attentions = (
          ()
          if (output_attentions and encoder_hidden_states is not None)
          else None
      )
      next_decoder_cache = () if use_cache else None
  
      # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
      for attn_mask, mask_name in zip(
          [head_mask, cross_attn_head_mask],
          ["head_mask", "cross_attn_head_mask"],
      ):
          if attn_mask is not None:
              assert attn_mask.size()[0] == (len(self.layers)), (
                  f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
                  f" {head_mask.size()[0]}."
              )
  
      # past_key_values_length
      past_key_values_length = (
          past_key_values[0][0].shape[2]
          if past_key_values is not None
          else 0
      )
  
      if at_first_stage:
          # retrieve input_ids and inputs_embeds
          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:
              input_shape = input_ids.size()
              input_ids = input_ids.view(-1, input_shape[-1])
          elif inputs_embeds is not None:
              input_shape = inputs_embeds.size()[:-1]
          else:
              raise ValueError(
                  "You have to specify either decoder_input_ids or decoder_inputs_embeds"
              )
  
          if inputs_embeds is None:
              inputs_embeds = self.embed_tokens(input_ids)
  
          # embed positions
          if input_ids is not None:
              positions = self.embed_positions(
                  input_ids,
                  past_key_values_length=past_key_values_length,
              )
          else:
              positions = self.embed_positions(
                  inputs_embeds,
                  past_key_values_length=past_key_values_length,
              )
  
          attention_mask = self._prepare_decoder_attention_mask(
              attention_mask,
              input_shape,
              inputs_embeds,
              past_key_values_length,
          )
  
          hidden_states = inputs_embeds + positions
          hidden_states = nn.functional.dropout(
              hidden_states, p=self.dropout, training=self.training
          )
  
          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
  
      else:
          if hidden_states is None:
              raise ValueError(
                  "hidden_states shouldn't be None for stages other than the first stage of encoder/decoder."
              )
          input_shape = hidden_states.size()[:-1]
  
          attention_mask = self._prepare_decoder_attention_mask(
              attention_mask,
              input_shape,
              hidden_states,
              past_key_values_length,
          )
  
      start_idx, end_idx = stage_index[0], stage_index[1]
  
      for idx in range(start_idx, end_idx):
          # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
          decoder_layer = self.layers[idx]
  
          if output_hidden_states:
              all_hidden_states += (hidden_states,)
          dropout_probability = random.uniform(0, 1)
          if self.training and (dropout_probability < self.layerdrop):
              continue
  
          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, use_cache
                      )
  
                  return custom_forward
  
              layer_outputs = torch.utils.checkpoint.checkpoint(
                  create_custom_forward(decoder_layer),
                  hidden_states,
                  attention_mask,
                  encoder_hidden_states,
                  None,  # encoder attention mask
                  head_mask[idx] if head_mask is not None else None,
                  cross_attn_head_mask[idx]
                  if cross_attn_head_mask is not None
                  else None,
                  None,  # past_key_value
              )
          else:
              layer_outputs = decoder_layer(
                  hidden_states,
                  attention_mask=attention_mask,
                  encoder_hidden_states=encoder_hidden_states,
                  layer_head_mask=(
                      head_mask[idx] if head_mask is not None else None
                  ),
                  cross_attn_layer_head_mask=(
                      cross_attn_head_mask[idx]
                      if cross_attn_head_mask is not None
                      else None
                  ),
                  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[3 if output_attentions else 1],
              )
  
          if output_attentions:
              all_self_attns += (layer_outputs[1],)
  
              if encoder_hidden_states is not None:
                  all_cross_attentions += (layer_outputs[2],)
  
      if at_last_stage:
          hidden_states = self.layer_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
          if not return_dict:
              return tuple(
                  v
                  for v in [
                      hidden_states,
                      next_cache,
                      all_hidden_states,
                      all_self_attns,
                      all_cross_attentions,
                  ]
                  if v is not None
              )
          return BaseModelOutputWithPastAndCrossAttentions(
              last_hidden_state=hidden_states,
              past_key_values=next_cache,
              hidden_states=all_hidden_states,
              attentions=all_self_attns,
              cross_attentions=all_cross_attentions,
          )
  
      else:
          return {
              "head_mask": head_mask,
              "cross_attn_head_mask": cross_attn_head_mask,
              "hidden_states": hidden_states,
          }
  
file_path: colossalai/colossalai/shardformer/modeling/whisper.py
imports: import logging
  import random
  from typing import List, Optional, Tuple, Union
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss
  from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput, Seq2SeqModelOutput, SequenceClassifierOutput
  from transformers.models.whisper.modeling_whisper import WhisperEncoder, WhisperForAudioClassification, WhisperForConditionalGeneration, WhisperModel
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
label_desc: Use image node_id 2 for calling the WhisperPipelineForwards obj's underlying member method code with example usage: obj.whisper_decoder_forward(input_ids, attention_mask, encoder_hidden_states, head_mask, cross_attn_head_mask, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, stage_manager, hidden_states, stage_index, decoder_starting_stage) and returns: BaseModelOutputWithPastAndCrossAttentions, dict, tuple, custom_forward, module