code: def _pre_rms_test_helper(
      n_tokens: int,
      n_channels: int,
      dtype: torch.dtype,
      res_add: bool = False,
  ):
      config = DSNormConfig(
          max_tokens=2048,
          type="rms_norm",
          channels=n_channels,
          residual_dtype=dtype,
          input_dtype=dtype,
          output_dtype=dtype,
          eps=1e-5,
      )
      bundle = ConfigBundle(name="cuda_pre_rms", config=config)
  
      # Input vals
      if res_add:
          hidden_states = torch.randn(
              (n_tokens, n_channels),
              dtype=dtype,
              device=get_accelerator().current_device_name(),
          )
      else:
          hidden_states = None
  
      residual = torch.randn(
          (n_tokens, n_channels),
          dtype=dtype,
          device=get_accelerator().current_device_name(),
      )
      gamma = torch.randn(
          (n_channels),
          dtype=torch.float32,
          device=get_accelerator().current_device_name(),
      )
      epsilon = 1e-5
  
      # Reference output
      ref_residual, ref_output = reference_implementation(
          residual, hidden_states, gamma, epsilon
      )
  
      # New output
      pre_ln_module = DSPreNormRegistry.instantiate_config(bundle)
      gamma = pre_ln_module.transform_param(gamma)
  
      ds_residual, ds_output = pre_ln_module(
          residual, hidden_states, gamma
      )
  
      # Check
      assert allclose(ds_residual, ref_residual)
      assert allclose(ds_output, ref_output)
  
file_path: deepspeed/tests/unit/inference/v2/modules/test_pre_rms_module.py
imports: from typing import Optional, Tuple
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.modules import ConfigBundle
  from deepspeed.inference.v2.modules.configs import DSNormConfig
  from deepspeed.inference.v2.modules.interfaces import DSPreNormRegistry
  from ...v2.inference_test_utils import get_dtypes, allclose
label_desc: Use image node_id 2 for calling a global function with example usage: _pre_rms_test_helper(n_tokens, n_channels, dtype, res_add) without return types