class_docstr: This class serves a micro library for bloom inference forwards.
  We intend to replace the forward methods for BloomForCausalLM, BloomModel, BloomBlock, and BloomAttention,
  as well as prepare_inputs_for_generation method for BloomForCausalLM.
  For future improvement, we might want to skip replacing methods for BloomForCausalLM,
  and call BloomModel.forward iteratively in TpInferEngine
class_docstr_tok: ["This","class","serves","a","micro","library","for","bloom","inference","forwards",".","We","intend","to","replace","the","forward","methods","for","BloomForCausalLM",",","BloomModel",",","BloomBlock",",","and","BloomAttention",",","as","well","as","prepare_inputs_for_generation","method","for","BloomForCausalLM",".","For","future","improvement",",","we","might","want","to","skip","replacing","methods","for","BloomForCausalLM",",","and","call","BloomModel.forward","iteratively","in","TpInferEngine"]
docstr: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
      Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
      `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
      are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
docstr_tok: ["labels","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Labels","for","language","modeling",".","Note","that","the","labels","*","*","are","shifted","*","*","inside","the","model",",","i.e",".","you","can","set","`","labels","=","input_ids","`","Indices","are","selected","in","`","[","-100",",","0",",","...",",","config.vocab_size","]","`","All","labels","set","to","`","-100","`","are","ignored","(","masked",")",",","the","loss","is","only","computed","for","labels","in","`","[","0",",","...",",","config.vocab_size","]","`"]
code: def bloom_for_causal_lm_forward(
      self: BloomForCausalLM,
      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.Tensor] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      labels: Optional[torch.Tensor] = None,
      use_cache: Optional[bool] = False,
      output_attentions: Optional[bool] = False,
      output_hidden_states: Optional[bool] = False,
      return_dict: Optional[bool] = False,
      infer_state: BatchInferState = None,
      stage_manager: Optional[PipelineStageManager] = None,
      hidden_states: Optional[torch.FloatTensor] = None,
      stage_index: Optional[List[int]] = None,
      tp_group: Optional[dist.ProcessGroup] = None,
      **deprecated_arguments,
  ):
      r"""
      labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
          `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
          are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
      """
      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,
          )
  
      # 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 is first stage and hidden_states is not None, go throught lm_head first
      if stage_manager.is_first_stage() and hidden_states is not None:
          lm_logits = self.lm_head(hidden_states)
          return {"logits": lm_logits}
  
      outputs = BloomInferenceForwards.bloom_model_forward(
          self.transformer,
          input_ids,
          past_key_values=past_key_values,
          attention_mask=attention_mask,
          head_mask=head_mask,
          inputs_embeds=inputs_embeds,
          use_cache=use_cache,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          infer_state=infer_state,
          stage_manager=stage_manager,
          hidden_states=hidden_states,
          stage_index=stage_index,
          tp_group=tp_group,
      )
  
      return outputs
  
file_path: colossalai/colossalai/inference/engine/modeling/bloom.py
imports: import math
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.distributed
  from torch.nn import functional
  from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomAttention, BloomBlock, BloomForCausalLM, BloomModel
  from transformers.utils import logging
  from colossalai.inference.kv_cache.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import bloom_context_attn_fwd, copy_kv_cache_to_dest, token_attention_fwd
  from colossalai.pipeline.stage_manager import PipelineStageManager
label_desc: Use image node_id 1 for calling the BloomInferenceForwards obj's underlying member method code with example usage: obj.bloom_for_causal_lm_forward(input_ids, past_key_values, attention_mask, head_mask, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, infer_state, stage_manager, hidden_states, stage_index, tp_group) and returns: outputs, dict