docstr: Args:
      labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
          config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
          (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
  
  Returns:
  
  Example:
  
  ```python
  >>> from transformers import AutoTokenizer, LlamaForCausalLM
  
  >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
  >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
  
  >>> prompt = "Hey, are you conscious? Can you talk to me?"
  >>> inputs = tokenizer(prompt, return_tensors="pt")
  
  >>> # Generate
  >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
  >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
  "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
  ```
docstr_tok: ["Args",":","labels","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Labels","for","computing","the","masked","language","modeling","loss",".","Indices","should","either","be","in","`","[","0",",","...",",","config.vocab_size","]","`","or","-100","(","see","`","input_ids","`","docstring",")",".","Tokens","with","indices","set","to","`","-100","`","are","ignored","(","masked",")",",","the","loss","is","only","computed","for","the","tokens","with","labels","in","`","[","0",",","...",",","config.vocab_size","]","`",".","Returns",":","Example",":","``","`","python",">",">",">","from","transformers","import","AutoTokenizer",",","LlamaForCausalLM",">",">",">","model","=","LlamaForCausalLM.from_pretrained","(","PATH_TO_CONVERTED_WEIGHTS",")",">",">",">","tokenizer","=","AutoTokenizer.from_pretrained","(","PATH_TO_CONVERTED_TOKENIZER",")",">",">",">","prompt","=","``","Hey",",","are","you","conscious","?","Can","you","talk","to","me","?","''",">",">",">","inputs","=","tokenizer","(","prompt",",","return_tensors=","''","pt","''",")",">",">",">","#","Generate",">",">",">","generate_ids","=","model.generate","(","inputs.input_ids",",","max_length=30",")",">",">",">","tokenizer.batch_decode","(","generate_ids",",","skip_special_tokens=True",",","clean_up_tokenization_spaces=False",")","[","0","]","''","Hey",",","are","you","conscious","?","Can","you","talk","to","me","?","\\nI","'m","not","conscious",",","but","I","can","talk","to","you",".","''","``","`"]
code: def forward(
      self,
      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,
      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,
      chunk_head: Optional[bool] = True,
  ) -> Union[Tuple, CausalLMOutputWithPast]:
      r"""
      Args:
          labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
              config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
              (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
  
      Returns:
  
      Example:
  
      ```python
      >>> from transformers import AutoTokenizer, LlamaForCausalLM
  
      >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
      >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
  
      >>> prompt = "Hey, are you conscious? Can you talk to me?"
      >>> inputs = tokenizer(prompt, return_tensors="pt")
  
      >>> # Generate
      >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
      >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
      "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
      ```"""
      # reset moe loss
      MOE_MANAGER.reset_loss()
  
      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 = self.model(
          input_ids=input_ids,
          attention_mask=attention_mask,
          position_ids=position_ids,
          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,
      )
  
      hidden_states = outputs[0]
      if self.pretraining_tp > 1:
          lm_head_slices = self.lm_head.weight.split(
              self.vocab_size // self.pretraining_tp, dim=0
          )
          logits = [
              F.linear(hidden_states, lm_head_slices[i])
              for i in range(self.pretraining_tp)
          ]
          logits = torch.cat(logits, dim=-1)
  
      loss = None
      # if no training, just do forward
      if labels is None:
          logits = self.lm_head(hidden_states)
          logits = logits.float()
      # the vocab size for openmoe is 30w+
      # which causes great activation memory in training, up to 20G for one sequence
      # so we use chunk and checkpoint to reduce memory
      else:
          if chunk_head == True:
  
              def create_custom_forward(module):
                  def custom_forward(*inputs):
                      logits = module(inputs[0])
                      logits = logits.float()
                      # Shift so that tokens < n predict n
                      shift_logits = (
                          logits[..., :-1, :].contiguous().float()
                      )
                      shift_labels = inputs[1][..., 1:].contiguous()
                      # Flatten the tokens
                      loss = self._calculate_loss(
                          shift_logits, shift_labels
                      )
                      return loss
  
                  return custom_forward
  
              aux_loss, z_loss = self._calculate_router_loss()
              loss = aux_loss + z_loss
              for batch_idx in range(hidden_states.shape[0]):
                  loss = loss + torch.utils.checkpoint.checkpoint(
                      create_custom_forward(self.lm_head),
                      hidden_states[batch_idx : batch_idx + 1, :],
                      labels[batch_idx : batch_idx + 1, :],
                  )
              logits = None
          else:
              logits = self.lm_head(hidden_states)
              logits = logits.float()
              # Shift so that tokens < n predict n
              shift_logits = logits[..., :-1, :].contiguous()
              shift_labels = labels[..., 1:].contiguous()
              # Flatten the tokens
              aux_loss, z_loss = self._calculate_router_loss()
              loss = aux_loss + z_loss
              loss = loss + self._calculate_loss(
                  shift_logits, shift_labels
              )
  
      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,
      )
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch import nn
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.modeling_utils import PreTrainedModel
  from transformers.models.llama.modeling_llama import LlamaConfig, LlamaRMSNorm
  from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from colossalai.kernel.cuda_native.mha.flash_attn_2 import HAS_FLASH_ATTN
  from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
  from colossalai.moe.layers import SparseMLP
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.moe.utils import get_activation, set_moe_args
label_desc: Use image node_id 8 for calling the OpenMoeForCausalLM obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, chunk_head) and returns: CausalLMOutputWithPast, unknown, custom_forward, loss