code: def rms_norm_forward(
      self: LlamaRMSNorm, hidden_states: torch.Tensor
  ) -> torch.Tensor:
      """
      Formard function for RMS Norm
      """
      return rms_norm(
          x=hidden_states,
          weight=self.weight,
          epsilon=self.variance_epsilon,
      )
  
file_path: ColossalAI/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
imports: from types import MethodType
  from typing import Optional, Tuple
  import torch
  import torch.nn.functional
  from einops import rearrange
  from flash_attn.bert_padding import pad_input, unpad_input
  from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_varlen_kvpacked_func
  from flash_attn.ops.rms_norm import rms_norm
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaForCausalLM, LlamaModel, LlamaRMSNorm, apply_rotary_pos_emb, repeat_kv
  from colossalai.logging import get_dist_logger
label_desc: Use image node_id 3 for calling a global function with example usage: rms_norm_forward(hidden_states) and returns: rms_norm