chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,4 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum, is_gated
|
||||
from deepspeed.inference.v2.modules import ConfigBundle
|
||||
from deepspeed.inference.v2.modules.configs import DSLinearConfig
|
||||
from deepspeed.inference.v2.modules.interfaces import DSLinearRegistry
|
||||
from ...v2.inference_test_utils import allclose
|
||||
|
||||
|
||||
def reference_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
|
||||
act_type: ActivationType) -> torch.Tensor:
|
||||
dtype = hidden_states.dtype
|
||||
out_states = torch.nn.functional.linear(hidden_states, weight, bias)
|
||||
out_states.float()
|
||||
|
||||
if is_gated(act_type):
|
||||
act_func_map = {
|
||||
ActivationType.ReGLU: torch.nn.functional.relu,
|
||||
ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
|
||||
ActivationType.SiGLU: torch.nn.functional.silu,
|
||||
}
|
||||
|
||||
act_act = out_states[..., ::2]
|
||||
act_linear = out_states[..., 1::2]
|
||||
|
||||
act_act = act_func_map[act_type](act_act)
|
||||
out_states = act_act * act_linear
|
||||
else:
|
||||
act_func_map = {
|
||||
ActivationType.RELU: torch.nn.functional.relu,
|
||||
ActivationType.GELU: torch.nn.functional.gelu,
|
||||
ActivationType.SILU: torch.nn.functional.silu,
|
||||
ActivationType.IDENTITY: lambda x: x,
|
||||
}
|
||||
|
||||
out_states = act_func_map[act_type](out_states)
|
||||
return out_states.to(dtype)
|
||||
|
||||
|
||||
def _blas_linear_helper(tokens: int,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
dtype: DtypeEnum,
|
||||
act_fn: ActivationType,
|
||||
use_bias: bool = True) -> None:
|
||||
linear_config = DSLinearConfig(max_tokens=2048,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
activation=act_fn,
|
||||
input_dtype=dtype,
|
||||
output_dtype=dtype)
|
||||
|
||||
bundle = ConfigBundle(name='blas_fp_linear', config=linear_config)
|
||||
|
||||
module = DSLinearRegistry.instantiate_config(bundle)
|
||||
|
||||
# Input vals
|
||||
hidden_states = torch.randn(
|
||||
(tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
|
||||
|
||||
weight_out_channels = 2 * out_channels if is_gated(act_fn) else out_channels
|
||||
weight = torch.randn(
|
||||
(weight_out_channels, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
|
||||
if use_bias:
|
||||
bias = torch.randn(
|
||||
(weight_out_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
|
||||
else:
|
||||
bias = None
|
||||
|
||||
# Reference output
|
||||
ref_output = reference_implementation(hidden_states, weight, bias, act_fn)
|
||||
|
||||
# New output
|
||||
ds_output = module(hidden_states, weight, bias)
|
||||
|
||||
# Check
|
||||
assert allclose(ds_output, ref_output)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("tokens, in_channels, out_channels", [(1, 4608, 1728), (37, 8192, 4096), (1280, 3072, 6144)])
|
||||
def test_blas_linear_shapes(tokens: int, in_channels: int, out_channels: int) -> None:
|
||||
|
||||
_blas_linear_helper(tokens, in_channels, out_channels, DtypeEnum.fp16, ActivationType.IDENTITY)
|
||||
|
||||
|
||||
all_acts = [
|
||||
ActivationType.RELU,
|
||||
ActivationType.GELU,
|
||||
ActivationType.SILU,
|
||||
ActivationType.GEGLU,
|
||||
ActivationType.ReGLU,
|
||||
ActivationType.SiGLU,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("act_fn", all_acts)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_blas_linear_act_fn(act_fn: ActivationType, use_bias: bool) -> None:
|
||||
|
||||
_blas_linear_helper(283, 512, 4096, DtypeEnum.fp16, act_fn, use_bias=use_bias)
|
||||
@@ -0,0 +1,207 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import itertools
|
||||
|
||||
from typing import List, 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 DSSelfAttentionConfig, PositionalEmbeddingType, RotateHalfConfig
|
||||
from deepspeed.inference.v2.modules.interfaces import DSSelfAttentionRegistry, DSSelfAttentionBase
|
||||
|
||||
from ..kernels.ragged_ops.ragged_testing_utils import build_batch_and_manager
|
||||
from ...v2.inference_test_utils import allclose
|
||||
|
||||
try:
|
||||
from flash_attn.flash_attn_interface import flash_attn_varlen_func
|
||||
validate_accuracy = True
|
||||
except ImportError:
|
||||
validate_accuracy = False
|
||||
|
||||
|
||||
def _blocked_flash_testing_helper(head_size: int,
|
||||
n_heads_q: int,
|
||||
n_heads_kv: int,
|
||||
seq_params: List[Tuple[int, int]],
|
||||
trained_freqs: bool = None) -> None:
|
||||
"""
|
||||
Helper function for testing blocked flash attention. This implementation is based on
|
||||
the implemnentation in ``unit.inference.kernels.ragged_ops.test_blocked_flash`` but
|
||||
integrates functionality to validate the composability.
|
||||
"""
|
||||
if trained_freqs is None:
|
||||
embed_type = PositionalEmbeddingType.none
|
||||
embed_args = None
|
||||
else:
|
||||
embed_type = PositionalEmbeddingType.rotate_half
|
||||
embed_args = RotateHalfConfig(use_trained_freqs=trained_freqs)
|
||||
|
||||
attn_config = DSSelfAttentionConfig(max_tokens=2048,
|
||||
n_heads_q=n_heads_q,
|
||||
n_heads_kv=n_heads_kv,
|
||||
head_size=head_size,
|
||||
max_sequences=32,
|
||||
positional_embedding_type=embed_type,
|
||||
positional_embedding_config=embed_args)
|
||||
|
||||
config = ConfigBundle(name='dense_blocked_attention', config=attn_config)
|
||||
attn_module: DSSelfAttentionBase = DSSelfAttentionRegistry.instantiate_config(config)
|
||||
|
||||
kv_block_size = attn_module.kv_block_size
|
||||
|
||||
kvs = []
|
||||
for _, history_len in seq_params:
|
||||
if history_len > 0:
|
||||
kvs.append(
|
||||
torch.randn((history_len, 2 * n_heads_kv * head_size),
|
||||
device=get_accelerator().current_device(),
|
||||
dtype=torch.float16))
|
||||
else:
|
||||
kvs.append(None)
|
||||
|
||||
batch, state_manager, _ = build_batch_and_manager(seq_params, head_size, n_heads_kv, kv_block_size, kv_fill=kvs)
|
||||
|
||||
qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
|
||||
device=get_accelerator().current_device(),
|
||||
dtype=torch.float16)
|
||||
|
||||
kv_cache = state_manager.get_cache(0)
|
||||
|
||||
attn_module.build_atoms(batch)
|
||||
if not trained_freqs:
|
||||
out = attn_module(qkv, kv_cache, batch)
|
||||
else:
|
||||
inv_freqs = torch.randn((head_size // 2, ), device=get_accelerator().current_device(), dtype=torch.float16)
|
||||
out = attn_module(qkv, kv_cache, batch, inv_freqs)
|
||||
|
||||
if validate_accuracy and trained_freqs is None:
|
||||
cu_seqlens_q = torch.tensor([0] + list(itertools.accumulate([seq[0] for seq in seq_params])),
|
||||
dtype=torch.int32,
|
||||
device=get_accelerator().current_device())
|
||||
cu_seqlens_kv = torch.tensor([0] + list(itertools.accumulate([seq[1] + seq[0] for seq in seq_params])),
|
||||
dtype=torch.int32,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
inflight_kv = qkv[:, head_size * n_heads_q:]
|
||||
full_kvs = []
|
||||
for i, kv in enumerate(kvs):
|
||||
if kv is not None:
|
||||
full_kvs.append(torch.cat([kv, inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]]], dim=0))
|
||||
else:
|
||||
full_kvs.append(inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]])
|
||||
run_kvs = torch.cat(full_kvs, dim=0)
|
||||
k = run_kvs[:, :head_size * n_heads_kv]
|
||||
v = run_kvs[:, head_size * n_heads_kv:]
|
||||
|
||||
q = qkv[:, :head_size * n_heads_q]
|
||||
q_ref = q.reshape((batch.current_tokens, n_heads_q, head_size))
|
||||
k_ref = k.reshape((k.shape[0], n_heads_kv, head_size))
|
||||
v_ref = v.reshape((v.shape[0], n_heads_kv, head_size))
|
||||
|
||||
max_seqlen_q = max([seq[0] for seq in seq_params])
|
||||
max_seqlen_kv = max([seq[1] + seq[0] for seq in seq_params])
|
||||
|
||||
ref_o = flash_attn_varlen_func(q_ref,
|
||||
k_ref,
|
||||
v_ref,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv,
|
||||
max_seqlen_q,
|
||||
max_seqlen_kv,
|
||||
softmax_scale=1.0,
|
||||
causal=True)
|
||||
|
||||
ref_o = ref_o.reshape(batch.current_tokens, head_size * n_heads_q)
|
||||
|
||||
assert allclose(out, ref_o)
|
||||
|
||||
get_accelerator().synchronize()
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("n_tokens", [2, 33, 65, 128, 256, 2037])
|
||||
def test_single_prompt(n_tokens: int) -> None:
|
||||
head_size = 64
|
||||
n_heads_q = 16
|
||||
n_heads_kv = 16
|
||||
|
||||
seq_params = [(n_tokens, 0)]
|
||||
_blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("prompt_lengths", [(128, 128), (192, 38), (514, 713), (83, 312, 610)])
|
||||
def test_multiple_prompts(prompt_lengths: Tuple[int, int]) -> None:
|
||||
"""
|
||||
Test multiple prompts in a single batch.
|
||||
"""
|
||||
head_size = 64
|
||||
n_heads_q = 16
|
||||
n_heads_kv = 16
|
||||
|
||||
seq_params = [(prompt_lengths[i], 0) for i in range(len(prompt_lengths))]
|
||||
_blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("seq_params", [(1, 34), (43, 40), (1, 144), (64, 128), (332, 628)])
|
||||
def test_continuation(seq_params: Tuple[int, int]) -> None:
|
||||
"""
|
||||
Test continued generation/prompt processing.
|
||||
"""
|
||||
head_size = 64
|
||||
n_heads_q = 32
|
||||
n_heads_kv = 32
|
||||
|
||||
_blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, [seq_params])
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("head_size", [64, 128])
|
||||
def test_head_size(head_size: int) -> None:
|
||||
n_heads_q = 16
|
||||
n_heads_kv = 16
|
||||
seq_params = [(128, 128), (192, 38), (1, 814)]
|
||||
|
||||
_blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("head_config", [(32, 8), (64, 16), (40, 8)])
|
||||
def test_gqa(head_config: Tuple[int, int]) -> None:
|
||||
head_size = 128
|
||||
n_heads_q = head_config[0]
|
||||
n_heads_kv = head_config[1]
|
||||
|
||||
seq_params = [(128, 128), (192, 38), (1, 814)]
|
||||
|
||||
_blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
def test_fully_composed() -> None:
|
||||
head_size = 64
|
||||
n_heads_q = 16
|
||||
n_heads_kv = 16
|
||||
|
||||
seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
|
||||
|
||||
_blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("trained_freqs", [True, False])
|
||||
def test_rotary_emb(trained_freqs: bool) -> None:
|
||||
head_size = 64
|
||||
n_heads_q = 16
|
||||
n_heads_kv = 16
|
||||
|
||||
seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
|
||||
|
||||
_blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params, trained_freqs=trained_freqs)
|
||||
@@ -0,0 +1,88 @@
|
||||
# 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,
|
||||
beta: torch.Tensor, epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
dtype = residual.dtype
|
||||
|
||||
residual = residual.to(torch.float32)
|
||||
gamma = gamma.to(torch.float32)
|
||||
beta = beta.to(torch.float32)
|
||||
|
||||
if hidden_states is not None:
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
residual = residual + hidden_states
|
||||
hidden_states = torch.nn.functional.layer_norm(residual, (residual.size(-1), ),
|
||||
weight=gamma,
|
||||
bias=beta,
|
||||
eps=epsilon)
|
||||
return residual.to(dtype), hidden_states.to(dtype)
|
||||
|
||||
|
||||
def _pre_ln_test_helper(n_tokens: int, n_channels: int, dtype: torch.dtype, res_add: bool = False):
|
||||
config = DSNormConfig(max_tokens=2048,
|
||||
type="layer_norm",
|
||||
channels=n_channels,
|
||||
residual_dtype=dtype,
|
||||
input_dtype=dtype,
|
||||
output_dtype=dtype,
|
||||
eps=1e-5)
|
||||
bundle = ConfigBundle(name='cuda_pre_ln', 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())
|
||||
beta = torch.rand((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, beta, epsilon)
|
||||
|
||||
# New output
|
||||
pre_ln_module = DSPreNormRegistry.instantiate_config(bundle)
|
||||
gamma = pre_ln_module.transform_param(gamma)
|
||||
beta = pre_ln_module.transform_param(beta)
|
||||
|
||||
ds_residual, ds_output = pre_ln_module(residual, hidden_states, gamma, beta)
|
||||
|
||||
# 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_ln_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_ln_test_helper(733, 2560, dtype)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
def test_no_res_add():
|
||||
_pre_ln_test_helper(733, 2560, torch.float16, res_add=False)
|
||||
@@ -0,0 +1,76 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from deepspeed.inference.v2.modules import ConfigBundle
|
||||
from deepspeed.inference.v2.modules.interfaces import DSPostNormRegistry
|
||||
from deepspeed.inference.v2.modules.configs import DSNormConfig
|
||||
from deepspeed.inference.v2.modules.implementations import cuda_post_ln
|
||||
from ...v2.inference_test_utils import allclose
|
||||
|
||||
|
||||
def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
|
||||
beta: torch.Tensor, epsilon: float) -> torch.Tensor:
|
||||
residual_f = residual.to(torch.float32)
|
||||
hidden_states_f = hidden_states.to(torch.float32)
|
||||
gamma_f = gamma.to(torch.float32)
|
||||
beta_f = beta.to(torch.float32)
|
||||
return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
|
||||
weight=gamma_f,
|
||||
bias=beta_f,
|
||||
eps=epsilon).to(hidden_states.dtype)
|
||||
|
||||
|
||||
@DSPostNormRegistry.register_module
|
||||
class CustomPostLNModule(cuda_post_ln.DSPostLNCUDAModule):
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'custom_post_ln'
|
||||
|
||||
|
||||
"""
|
||||
Here, we explicitly register an LN implementation outside the core deepspeed repo. This should
|
||||
validate that the registry is working as expected and we can implement modules outside the core
|
||||
repo.
|
||||
"""
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
def test_custom_registration():
|
||||
channels = 4096
|
||||
dtype = torch.float16
|
||||
tokens = 1024
|
||||
|
||||
config = DSNormConfig(max_tokens=2048,
|
||||
type="layer_norm",
|
||||
channels=channels,
|
||||
residual_dtype=dtype,
|
||||
input_dtype=dtype,
|
||||
output_dtype=dtype,
|
||||
eps=1e-5)
|
||||
bundle = ConfigBundle(name='custom_post_ln', config=config)
|
||||
|
||||
# Input vals
|
||||
hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
|
||||
residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
|
||||
gamma = torch.randn((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
|
||||
beta = torch.rand((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
|
||||
epsilon = 1e-5
|
||||
|
||||
# Reference output
|
||||
ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
|
||||
|
||||
# New output
|
||||
post_ln_module = DSPostNormRegistry.instantiate_config(bundle)
|
||||
gamma = post_ln_module.transform_param(gamma)
|
||||
beta = post_ln_module.transform_param(beta)
|
||||
ds_output, _ = post_ln_module(residual, hidden_states, gamma, beta)
|
||||
|
||||
# Check
|
||||
assert allclose(ds_output, ref_output)
|
||||
@@ -0,0 +1,328 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
|
||||
from deepspeed.inference.v2.modules import ConfigBundle
|
||||
from deepspeed.inference.v2.modules.configs import DSMoEConfig
|
||||
from deepspeed.inference.v2.modules.interfaces import DSMoERegistry
|
||||
|
||||
from ..kernels.ragged_ops.ragged_testing_utils import build_simple_batch
|
||||
from ...v2.inference_test_utils import allclose, get_dtypes
|
||||
|
||||
|
||||
def _gating_reference(logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Reference gating code.
|
||||
"""
|
||||
logits = logits.float()
|
||||
probs = torch.nn.functional.softmax(logits, dim=1)
|
||||
|
||||
indices1_s = torch.argmax(probs, dim=-1)
|
||||
mask1 = torch.nn.functional.one_hot(indices1_s, num_classes=logits.shape[-1])
|
||||
indices_mask = mask1.sum(dim=1) * logits.shape[-1] - 1
|
||||
indices1_s = torch.min(indices1_s, indices_mask)
|
||||
|
||||
gates1_s = (probs * mask1).sum(dim=1)
|
||||
|
||||
sorted_indices = indices1_s.sort()[1]
|
||||
original_indices = sorted_indices.sort()[1]
|
||||
|
||||
exp_count = torch.bincount(indices1_s, minlength=logits.shape[-1]).long()
|
||||
exp_count_cumsum = exp_count.cumsum(dim=0)
|
||||
|
||||
return sorted_indices, original_indices, exp_count_cumsum, gates1_s
|
||||
|
||||
|
||||
def _reference_impl(hidden_states: torch.Tensor, gate_weight: torch.Tensor, mlp_1_w: torch.Tensor,
|
||||
mlp_2_w: torch.Tensor, mlp_1_b: torch.Tensor, mlp_2_b: torch.Tensor,
|
||||
act_fn: ActivationType) -> torch.Tensor:
|
||||
"""
|
||||
Reference implementation of the MoE module.
|
||||
"""
|
||||
|
||||
act_fn_dict = {
|
||||
ActivationType.GELU: torch.nn.functional.gelu,
|
||||
ActivationType.RELU: torch.nn.functional.relu,
|
||||
ActivationType.SILU: torch.nn.functional.silu,
|
||||
ActivationType.IDENTITY: lambda x: x,
|
||||
}
|
||||
|
||||
logits = torch.matmul(hidden_states, gate_weight.t())
|
||||
sorted_indices, original_indices, exp_count_cumsum, gate_scales = _gating_reference(logits)
|
||||
|
||||
moe_input = hidden_states[sorted_indices]
|
||||
|
||||
output_unordered = torch.empty_like(hidden_states)
|
||||
|
||||
for expert_idx in range(mlp_1_w.shape[0]):
|
||||
min_bound = 0 if expert_idx == 0 else exp_count_cumsum[expert_idx - 1]
|
||||
max_bound = exp_count_cumsum[expert_idx]
|
||||
|
||||
input_slice = moe_input[min_bound:max_bound]
|
||||
intermediate = torch.nn.functional.linear(input_slice, mlp_1_w[expert_idx], mlp_1_b[expert_idx])
|
||||
|
||||
intermediate = act_fn_dict[act_fn](intermediate)
|
||||
output_slice = torch.nn.functional.linear(intermediate, mlp_2_w[expert_idx], mlp_2_b[expert_idx])
|
||||
|
||||
output_unordered[min_bound:max_bound] = output_slice
|
||||
|
||||
output = output_unordered[original_indices]
|
||||
|
||||
output.mul_(gate_scales.unsqueeze(-1)).reshape(hidden_states.shape)
|
||||
return output
|
||||
|
||||
|
||||
def _cutlass_moe_testing_helper(tokens: int,
|
||||
in_channels: int,
|
||||
intermediate_dim: int,
|
||||
experts: int,
|
||||
dtype: int,
|
||||
activation_type: ActivationType = ActivationType.GELU,
|
||||
use_bias: bool = True,
|
||||
iters: int = 1) -> None:
|
||||
|
||||
config = DSMoEConfig(max_tokens=4096,
|
||||
model_dim=in_channels,
|
||||
intermediate_features=intermediate_dim,
|
||||
n_experts=experts,
|
||||
activation=activation_type,
|
||||
input_dtype=dtype,
|
||||
output_dtype=dtype)
|
||||
|
||||
implementation_config = {"weight_dtype": DtypeEnum(dtype)}
|
||||
|
||||
bundle = ConfigBundle(name='cutlass_multi_gemm_moe', config=config, implementation_config=implementation_config)
|
||||
moe_module = DSMoERegistry.instantiate_config(bundle)
|
||||
|
||||
batch = build_simple_batch([tokens])
|
||||
|
||||
# Parameters
|
||||
gate_weight = torch.randn(
|
||||
(experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
|
||||
mlp_1_w = torch.randn(
|
||||
(experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
mlp_2_w = torch.randn(
|
||||
(experts, in_channels, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
|
||||
if use_bias:
|
||||
mlp_1_b = torch.randn(
|
||||
(experts, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
mlp_2_b = torch.randn(
|
||||
(experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
else:
|
||||
mlp_1_b = None
|
||||
mlp_2_b = None
|
||||
|
||||
gate_ds = moe_module.transform_gate_param(gate_weight)
|
||||
mlp_1_w_ds = moe_module.transform_moe_mlp_1_param(mlp_1_w)
|
||||
mlp_1_b_ds = moe_module.transform_moe_mlp_1_param(mlp_1_b)
|
||||
mlp_2_w_ds = moe_module.transform_moe_mlp_2_param(mlp_2_w)
|
||||
mlp_2_b_ds = moe_module.transform_moe_mlp_2_param(mlp_2_b)
|
||||
|
||||
for _ in range(iters):
|
||||
# Input vals
|
||||
hidden_states = torch.randn(
|
||||
(tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
|
||||
# Reference implementation
|
||||
ref_output = _reference_impl(hidden_states, gate_weight, mlp_1_w, mlp_2_w, mlp_1_b, mlp_2_b, activation_type)
|
||||
|
||||
output = moe_module(hidden_states,
|
||||
batch,
|
||||
gate_ds,
|
||||
mlp_1_w_ds,
|
||||
mlp_2_w_ds,
|
||||
mlp_1_b=mlp_1_b_ds,
|
||||
mlp_2_b=mlp_2_b_ds)
|
||||
|
||||
# Increase the tolerance for larger meta ops since the error is additive
|
||||
assert allclose(output, ref_output, tolerances=(1e-2, 1e-2))
|
||||
|
||||
get_accelerator().synchronize()
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("experts", [2, 32, 64])
|
||||
def test_expert_variance(experts: int) -> None:
|
||||
_cutlass_moe_testing_helper(tokens=876,
|
||||
in_channels=4096,
|
||||
intermediate_dim=2048,
|
||||
experts=experts,
|
||||
dtype=DtypeEnum.fp16,
|
||||
activation_type=ActivationType.IDENTITY,
|
||||
use_bias=True)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
def test_successive_inputs():
|
||||
"""
|
||||
The CUTLASS MoE uses persistent state (expert counts) that is assumed to be cleared
|
||||
on each forward pass. This ensures that the module is clearing that metadata.
|
||||
"""
|
||||
_cutlass_moe_testing_helper(tokens=876,
|
||||
in_channels=4096,
|
||||
intermediate_dim=2048,
|
||||
experts=64,
|
||||
dtype=DtypeEnum.fp16,
|
||||
activation_type=ActivationType.IDENTITY,
|
||||
use_bias=True,
|
||||
iters=10)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
|
||||
def test_dtypes(dtype: torch.dtype) -> None:
|
||||
_cutlass_moe_testing_helper(tokens=876,
|
||||
in_channels=4096,
|
||||
intermediate_dim=2048,
|
||||
experts=64,
|
||||
dtype=DtypeEnum(dtype),
|
||||
activation_type=ActivationType.IDENTITY,
|
||||
use_bias=True)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("activation_type", [ActivationType.GELU, ActivationType.RELU, ActivationType.SILU])
|
||||
def test_activation_types(activation_type: ActivationType) -> None:
|
||||
_cutlass_moe_testing_helper(tokens=876,
|
||||
in_channels=4096,
|
||||
intermediate_dim=2048,
|
||||
experts=64,
|
||||
dtype=DtypeEnum.fp16,
|
||||
activation_type=activation_type,
|
||||
use_bias=True)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("in_channels, out_channels", [(4096, 2048), (2048, 8192), (6144, 3072)])
|
||||
def test_in_out_channels(in_channels: int, out_channels: int) -> None:
|
||||
_cutlass_moe_testing_helper(tokens=876,
|
||||
in_channels=in_channels,
|
||||
intermediate_dim=out_channels,
|
||||
experts=64,
|
||||
dtype=DtypeEnum.fp16,
|
||||
activation_type=ActivationType.IDENTITY,
|
||||
use_bias=True)
|
||||
|
||||
|
||||
def _mixtral_moe_baseline(hidden_states: torch.Tensor,
|
||||
gate_weight: torch.Tensor,
|
||||
mlp_w1: torch.Tensor,
|
||||
mlp_w2: torch.Tensor,
|
||||
mlp_w3: torch.Tensor,
|
||||
force_float: bool = False) -> torch.Tensor:
|
||||
"""
|
||||
Baseline implementation for mixtral MoE module.
|
||||
|
||||
Based on transformers implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py
|
||||
"""
|
||||
output_dtype = hidden_states.dtype
|
||||
if force_float:
|
||||
hidden_states = hidden_states.float()
|
||||
gate_weight = gate_weight.float()
|
||||
mlp_w1 = mlp_w1.float()
|
||||
mlp_w2 = mlp_w2.float()
|
||||
mlp_w3 = mlp_w3.float()
|
||||
|
||||
router_logits = torch.nn.functional.linear(hidden_states, gate_weight)
|
||||
routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
|
||||
routing_weights, selected_experts = routing_weights.topk(k=2, dim=-1)
|
||||
routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
|
||||
|
||||
# NOTE(cmikeh2): This is a difference implementation, ours will preserve the original scale
|
||||
# as float32 and perform in-kernel fused FP16->FP32->FP16 conversion.
|
||||
routing_weights = routing_weights.to(hidden_states.dtype)
|
||||
|
||||
final_hidden_states = torch.zeros_like(hidden_states)
|
||||
|
||||
expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=gate_weight.shape[0]).permute(2, 1, 0)
|
||||
get_accelerator().synchronize()
|
||||
|
||||
for expert_idx in range(gate_weight.shape[0]):
|
||||
exp_mlp_w1 = mlp_w1[expert_idx]
|
||||
exp_mlp_w2 = mlp_w2[expert_idx]
|
||||
exp_mlp_w3 = mlp_w3[expert_idx]
|
||||
|
||||
idx, top_x = torch.where(expert_mask[expert_idx])
|
||||
|
||||
if top_x.shape[0] == 0:
|
||||
continue
|
||||
|
||||
top_x_list = top_x.tolist()
|
||||
idx_list = idx.tolist()
|
||||
|
||||
current_state = hidden_states[top_x_list]
|
||||
|
||||
linear = torch.nn.functional.linear
|
||||
intermediate = torch.nn.functional.silu(linear(current_state, exp_mlp_w1)) * linear(current_state, exp_mlp_w3)
|
||||
output = linear(intermediate, exp_mlp_w2) * routing_weights[top_x_list, idx_list].unsqueeze(-1)
|
||||
final_hidden_states.index_add_(0, top_x, output.to(final_hidden_states.dtype))
|
||||
|
||||
return final_hidden_states.to(output_dtype)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
def test_mixtral_moe_config():
|
||||
|
||||
experts = 8
|
||||
n_top_k = 2
|
||||
in_channels = 4096
|
||||
intermediate_dim = 2048
|
||||
dtype = DtypeEnum.bf16
|
||||
|
||||
# Parameters
|
||||
gate_weight = torch.randn(
|
||||
(experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
|
||||
mlp_w1 = torch.randn(
|
||||
(experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
mlp_w3 = torch.randn(
|
||||
(experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
mlp_w2 = torch.randn(
|
||||
(experts, in_channels, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
|
||||
n_tokens = 256
|
||||
hidden_states = torch.randn(
|
||||
(n_tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
|
||||
|
||||
baseline = _mixtral_moe_baseline(hidden_states, gate_weight, mlp_w1, mlp_w2, mlp_w3)
|
||||
|
||||
mlp_w13_fused = torch.cat([mlp_w1, mlp_w3], dim=-1).reshape(experts, 2 * intermediate_dim, in_channels)
|
||||
|
||||
config = DSMoEConfig(max_tokens=4096,
|
||||
model_dim=in_channels,
|
||||
intermediate_features=intermediate_dim,
|
||||
n_experts=experts,
|
||||
activation=ActivationType.SiGLU,
|
||||
input_dtype=dtype,
|
||||
output_dtype=dtype,
|
||||
top_k=n_top_k,
|
||||
normalize_scores=True)
|
||||
|
||||
implementation_config = {"weight_dtype": DtypeEnum(dtype)}
|
||||
|
||||
bundle = ConfigBundle(name='cutlass_multi_gemm_moe', config=config, implementation_config=implementation_config)
|
||||
moe_module = DSMoERegistry.instantiate_config(bundle)
|
||||
|
||||
batch = build_simple_batch([n_tokens])
|
||||
|
||||
gate_ds = moe_module.transform_gate_param(gate_weight)
|
||||
mlp_w1_ds = moe_module.transform_moe_mlp_1_param(mlp_w13_fused)
|
||||
mlp_w2_ds = moe_module.transform_moe_mlp_2_param(mlp_w2)
|
||||
|
||||
output = moe_module(hidden_states, batch, gate_ds, mlp_w1_ds, mlp_w2_ds)
|
||||
|
||||
# NOTE(cmikeh2): These are higher than the other tests for reasons that aren't quite
|
||||
# clear to me. My best guess is that the SiGLU activation is causing larger numerical
|
||||
# divergence. The thresholds chosen here is based on the observed error between the
|
||||
# float and bfloat16 reference implementations.
|
||||
assert allclose(output, baseline.to(dtype.value), tolerances=(5e-2, 5e-2))
|
||||
@@ -0,0 +1,58 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
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 DSPostNormRegistry
|
||||
from ...v2.inference_test_utils import get_dtypes, allclose
|
||||
|
||||
|
||||
def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
|
||||
beta: torch.Tensor, epsilon: float) -> torch.Tensor:
|
||||
residual_f = residual.to(torch.float32)
|
||||
hidden_states_f = hidden_states.to(torch.float32)
|
||||
gamma_f = gamma.to(torch.float32)
|
||||
beta_f = beta.to(torch.float32)
|
||||
return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
|
||||
weight=gamma_f,
|
||||
bias=beta_f,
|
||||
eps=epsilon).to(hidden_states.dtype)
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
|
||||
@pytest.mark.parametrize("dtype", get_dtypes())
|
||||
def test_cuda_post_ln_module(tokens: int, channels: int, dtype: torch.dtype) -> None:
|
||||
config = DSNormConfig(max_tokens=2048,
|
||||
type="layer_norm",
|
||||
channels=channels,
|
||||
residual_dtype=dtype,
|
||||
input_dtype=dtype,
|
||||
output_dtype=dtype,
|
||||
eps=1e-5)
|
||||
bundle = ConfigBundle(name='cuda_post_ln', config=config)
|
||||
|
||||
# Input vals
|
||||
hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
|
||||
residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
|
||||
gamma = torch.randn((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
|
||||
beta = torch.rand((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
|
||||
epsilon = 1e-5
|
||||
|
||||
# Reference output
|
||||
ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
|
||||
|
||||
# New output
|
||||
post_ln_module = DSPostNormRegistry.instantiate_config(bundle)
|
||||
gamma = post_ln_module.transform_param(gamma)
|
||||
beta = post_ln_module.transform_param(beta)
|
||||
ds_output, _ = post_ln_module(residual, hidden_states, gamma, beta)
|
||||
|
||||
# Check
|
||||
assert allclose(ds_output, ref_output)
|
||||
@@ -0,0 +1,88 @@
|
||||
# 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)
|
||||
@@ -0,0 +1,183 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum, is_gated
|
||||
from deepspeed.inference.v2.modules import ConfigBundle
|
||||
from deepspeed.inference.v2.modules.configs import DSLinearConfig
|
||||
from deepspeed.inference.v2.modules.interfaces import DSLinearRegistry
|
||||
from ...v2.inference_test_utils import allclose
|
||||
|
||||
|
||||
def reference_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
|
||||
act_type: ActivationType) -> torch.Tensor:
|
||||
dtype = hidden_states.dtype
|
||||
out_states = torch.nn.functional.linear(hidden_states, weight, bias)
|
||||
out_states.float()
|
||||
|
||||
if is_gated(act_type):
|
||||
act_func_map = {
|
||||
ActivationType.ReGLU: torch.nn.functional.relu,
|
||||
ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
|
||||
ActivationType.SiGLU: torch.nn.functional.silu,
|
||||
}
|
||||
|
||||
act_act = out_states[..., ::2]
|
||||
act_linear = out_states[..., 1::2]
|
||||
|
||||
act_act = act_func_map[act_type](act_act)
|
||||
out_states = act_act * act_linear
|
||||
else:
|
||||
act_func_map = {
|
||||
ActivationType.RELU: torch.nn.functional.relu,
|
||||
ActivationType.GELU: torch.nn.functional.gelu,
|
||||
ActivationType.SILU: torch.nn.functional.silu,
|
||||
ActivationType.IDENTITY: lambda x: x,
|
||||
}
|
||||
|
||||
out_states = act_func_map[act_type](out_states)
|
||||
return out_states.to(dtype)
|
||||
|
||||
|
||||
def _fp6_quant_dequant_weights(weight: torch.Tensor) -> torch.Tensor:
|
||||
from deepspeed.inference.v2.modules.implementations.linear.quantized_linear import fp_quantize
|
||||
weight_quantized_fake_fp6, scales = fp_quantize(weight, num_bits=6, exp_bits=3)
|
||||
return weight_quantized_fake_fp6 * scales
|
||||
|
||||
|
||||
def quant_dequant_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
|
||||
act_type: ActivationType) -> torch.Tensor:
|
||||
dtype = hidden_states.dtype
|
||||
weight_dequantized = _fp6_quant_dequant_weights(weight)
|
||||
out_states = torch.nn.functional.linear(hidden_states, weight_dequantized, bias)
|
||||
out_states.float()
|
||||
|
||||
if is_gated(act_type):
|
||||
act_func_map = {
|
||||
ActivationType.ReGLU: torch.nn.functional.relu,
|
||||
ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
|
||||
ActivationType.SiGLU: torch.nn.functional.silu,
|
||||
}
|
||||
|
||||
act_act = out_states[..., ::2]
|
||||
act_linear = out_states[..., 1::2]
|
||||
|
||||
act_act = act_func_map[act_type](act_act)
|
||||
out_states = act_act * act_linear
|
||||
else:
|
||||
act_func_map = {
|
||||
ActivationType.RELU: torch.nn.functional.relu,
|
||||
ActivationType.GELU: torch.nn.functional.gelu,
|
||||
ActivationType.SILU: torch.nn.functional.silu,
|
||||
ActivationType.IDENTITY: lambda x: x,
|
||||
}
|
||||
|
||||
out_states = act_func_map[act_type](out_states)
|
||||
return out_states.to(dtype)
|
||||
|
||||
|
||||
def _fp6_quantized_linear_helper(tokens: int,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
dtype: DtypeEnum,
|
||||
act_fn: ActivationType,
|
||||
use_bias: bool = True,
|
||||
expect_failure: bool = False) -> None:
|
||||
# The current FP6 kernel only supports NVIDIA Ampere GPUs.
|
||||
if not 'cuda' in get_accelerator().current_device_name():
|
||||
return
|
||||
major, _ = torch.cuda.get_device_capability() #ignore-cuda
|
||||
if major != 8:
|
||||
return
|
||||
|
||||
# Input vals
|
||||
hidden_states = torch.randn(
|
||||
(tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
|
||||
|
||||
weight_out_channels = 2 * \
|
||||
out_channels if is_gated(act_fn) else out_channels
|
||||
weight = torch.randn(
|
||||
(weight_out_channels, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
|
||||
if use_bias:
|
||||
bias = torch.randn(
|
||||
(weight_out_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
|
||||
else:
|
||||
bias = None
|
||||
|
||||
# quantize and dequantize output
|
||||
ref_quant_dequant_output = quant_dequant_implementation(hidden_states, weight, bias, act_fn)
|
||||
|
||||
linear_config = DSLinearConfig(max_tokens=2048,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
activation=act_fn,
|
||||
input_dtype=dtype,
|
||||
output_dtype=dtype)
|
||||
bundle = ConfigBundle(name='quantized_wf6af16_linear', config=linear_config)
|
||||
fp6_linear_module = DSLinearRegistry.instantiate_config(bundle)
|
||||
weight_fp6 = fp6_linear_module.transform_param(weight.clone().cpu()).to(get_accelerator().current_device_name())
|
||||
|
||||
if expect_failure:
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
ds_output = fp6_linear_module(hidden_states, weight_fp6, bias)
|
||||
assert "The out and in channel should be multiple of 256 and 64 respectively." in str(excinfo.value)
|
||||
else:
|
||||
ds_output = fp6_linear_module(hidden_states, weight_fp6, bias)
|
||||
# The current FP6 kernel uses FP16 Tensor Core.
|
||||
tolerances = (3e-2, 2e-3) # tolerances for fp16
|
||||
|
||||
# Check DeepSpeed implementation
|
||||
assert allclose(ds_output, ref_quant_dequant_output, tolerances=tolerances)
|
||||
|
||||
|
||||
all_acts = [
|
||||
ActivationType.RELU,
|
||||
ActivationType.GELU,
|
||||
ActivationType.SILU,
|
||||
ActivationType.GEGLU,
|
||||
ActivationType.ReGLU,
|
||||
ActivationType.SiGLU,
|
||||
]
|
||||
all_tokens = [37]
|
||||
all_in_out_channels = [
|
||||
(4096, 4096),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("tokens", all_tokens)
|
||||
@pytest.mark.parametrize("in_channels, out_channels", all_in_out_channels)
|
||||
@pytest.mark.parametrize("act_fn", all_acts)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_fp6_quantized_linear_act_fn(tokens: int, in_channels: int, out_channels: int, act_fn: ActivationType,
|
||||
use_bias: bool) -> None:
|
||||
_fp6_quantized_linear_helper(tokens=tokens,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
dtype=DtypeEnum.fp16,
|
||||
act_fn=act_fn,
|
||||
use_bias=use_bias)
|
||||
|
||||
|
||||
# Other shapes, not supported by FP6 kernels. Will raise ValueError.
|
||||
@pytest.mark.inference_v2_ops
|
||||
@pytest.mark.parametrize("tokens", all_tokens)
|
||||
@pytest.mark.parametrize("in_channels, out_channels", [(4608, 1728)])
|
||||
@pytest.mark.parametrize("act_fn", all_acts)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_fp6_quantized_linear_act_fn_fail(tokens: int, in_channels: int, out_channels: int, act_fn: ActivationType,
|
||||
use_bias: bool) -> None:
|
||||
_fp6_quantized_linear_helper(tokens=tokens,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
dtype=DtypeEnum.fp16,
|
||||
act_fn=act_fn,
|
||||
use_bias=use_bias,
|
||||
expect_failure=True)
|
||||
Reference in New Issue
Block a user