code: def reference_rms_norm(
      vals: torch.Tensor, gamma: torch.Tensor, epsilon: float = 1e-5
  ) -> torch.Tensor:
      variance = vals.to(torch.float32).pow(2).mean(-1, keepdim=True)
      vals = vals * torch.rsqrt(variance + epsilon)
  
      if gamma.dtype in [torch.float16, torch.bfloat16]:
          vals = vals.to(gamma.dtype)
  
      return gamma * vals
  
file_path: deepspeed/tests/unit/inference/v2/kernels/core_ops/test_rms_norm.py
imports: import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.inference_utils import DtypeEnum
  from deepspeed.inference.v2.kernels.core_ops import CUDARMSNorm, CUDARMSPreNorm
  from ....v2.inference_test_utils import get_dtypes, allclose
label_desc: Use image node_id 1 for calling a global function with example usage: reference_rms_norm(vals, gamma, epsilon) and returns: unknown