wehub-resource-sync
89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
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
|
|
|
|
|
|
def reference_implementation(residual: torch.Tensor, hidden_states: Optional[torch.Tensor], gamma: torch.Tensor,
|
|
epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
dtype = residual.dtype
|
|
|
|
if hidden_states is not None:
|
|
hidden_states = hidden_states
|
|
residual = residual + hidden_states
|
|
|
|
rms_vals = residual.to(torch.float32)
|
|
variance = rms_vals.pow(2).mean(-1, keepdim=True)
|
|
rms_vals = rms_vals * torch.rsqrt(variance + epsilon)
|
|
|
|
if gamma.dtype in [torch.float16, torch.bfloat16]:
|
|
rms_vals = rms_vals.to(gamma.dtype)
|
|
|
|
hidden_states = gamma * rms_vals
|
|
|
|
return residual.to(dtype), hidden_states.to(dtype)
|
|
|
|
|
|
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)
|
|
|
|
|
|
@pytest.mark.inference_v2_ops
|
|
@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
|
|
def test_token_channels(tokens: int, channels: int) -> None:
|
|
_pre_rms_test_helper(tokens, channels, torch.float16)
|
|
|
|
|
|
@pytest.mark.inference_v2_ops
|
|
@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
|
|
def test_dtype(dtype: torch.dtype) -> None:
|
|
_pre_rms_test_helper(733, 2560, dtype)
|
|
|
|
|
|
@pytest.mark.inference_v2_ops
|
|
def test_no_res_add():
|
|
_pre_rms_test_helper(733, 2560, torch.float16, res_add=False)
|