code: def get_triton_rmsnorm_forward():
      if HAS_TRITON_RMSNORM:
  
          def _triton_rmsnorm_forward(
              self: LlamaRMSNorm, hidden_states: torch.Tensor
          ):
              return lightllm_rmsnorm_forward(
                  hidden_states, self.weight.data, self.variance_epsilon
              )
  
          return _triton_rmsnorm_forward
      else:
          return None
  
file_path: ColossalAI/colossalai/legacy/inference/tensor_parallel/policies/llama.py
imports: from functools import partial
  import torch
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel, LlamaRMSNorm
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, SubModuleReplacementDescription
  from colossalai.shardformer.policies.llama import LlamaForCausalLMPolicy
  from ..modeling._utils import init_to_get_rotary
  from ..modeling.llama import LlamaInferenceForwards
label_desc: Use image node_id 1 for calling a global function with example usage: get_triton_rmsnorm_forward() and returns: _triton_rmsnorm_forward, None, lightllm_rmsnorm_forward