chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Imports for registering ops
|
||||
from .attention import *
|
||||
from .linear import *
|
||||
from .post_norm import *
|
||||
from .pre_norm import *
|
||||
from .embedding import *
|
||||
from .unembed import *
|
||||
from .moe import *
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .dense_blocked_attention import DSDenseBlockedAttention
|
||||
@@ -0,0 +1,180 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ....allocator import empty_from
|
||||
from ....inference_utils import DtypeEnum
|
||||
from ....kernels.ragged_ops import (
|
||||
AtomBuilder,
|
||||
BlockedFlashAttn,
|
||||
BlockedRotaryEmbeddings,
|
||||
BlockedTrainedRotaryEmbeddings,
|
||||
get_q_block_size,
|
||||
get_kv_block_size,
|
||||
LinearBlockedKVCopy,
|
||||
)
|
||||
from ....ragged import RaggedBatchWrapper, split_kv
|
||||
from deepspeed.ops.op_builder import RaggedUtilsBuilder
|
||||
|
||||
from ...interfaces import DSSelfAttentionBase, DSSelfAttentionRegistry
|
||||
from ...configs import DSSelfAttentionConfig, PositionalEmbeddingType, MaskingType
|
||||
|
||||
try:
|
||||
from functools import cached_property
|
||||
except ImportError:
|
||||
|
||||
def cached_property(func):
|
||||
return property(func)
|
||||
|
||||
|
||||
@DSSelfAttentionRegistry.register_module
|
||||
class DSDenseBlockedAttention(DSSelfAttentionBase):
|
||||
"""
|
||||
Self attention implementation for dense, blocked self attention.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def name() -> str:
|
||||
return 'dense_blocked_attention'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSSelfAttentionConfig) -> bool:
|
||||
|
||||
if config.input_dtype != config.output_dtype:
|
||||
return False
|
||||
|
||||
if DtypeEnum(config.input_dtype) not in (DtypeEnum.fp16, DtypeEnum.bf16):
|
||||
return False
|
||||
|
||||
if PositionalEmbeddingType(config.positional_embedding_type) not in [
|
||||
PositionalEmbeddingType.none, PositionalEmbeddingType.rotate_half
|
||||
]:
|
||||
return False
|
||||
|
||||
if MaskingType(config.masking_type) != MaskingType.causal:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSSelfAttentionConfig, implementation_config: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Create the Attention DSModule.
|
||||
|
||||
Args:
|
||||
config (DSSelfAttentionConfig): The self attention config for all attention DSModules.
|
||||
implementation_config (Dict[str, Any]):
|
||||
There are two (dependent) potential components in the implementtion config.
|
||||
|
||||
1. `trained_freqs` - If the embedding weights for RoPE are trained, the implementation
|
||||
config should contain {'trained_freqs': True}. This will mean the implementation will
|
||||
expect a `trained_freqs` tensor in the `forward` method and will not synthesize the
|
||||
values internally.
|
||||
|
||||
2. `theta_base` - The base value for synthesized frequencies in the rotary embeddings.
|
||||
This will only be used if `trained_freqs` is False or not present in the `implementation_config`. If this is not included, the default value of 10000.0 will be used.
|
||||
"""
|
||||
super().__init__(config, implementation_config)
|
||||
|
||||
embed_type = PositionalEmbeddingType(config.positional_embedding_type)
|
||||
if embed_type == PositionalEmbeddingType.none:
|
||||
self._kv_copy = LinearBlockedKVCopy(self._config.head_size, self._config.n_heads_q,
|
||||
self._config.n_heads_kv, self._config.input_dtype)
|
||||
elif embed_type == PositionalEmbeddingType.rotate_half:
|
||||
rotary_config = config.positional_embedding_config
|
||||
assert rotary_config is not None, "Rotary config must be provided if using rotate_half as Positional Embedding Type."
|
||||
|
||||
if rotary_config.use_trained_freqs:
|
||||
# Theta and rotary dim are effectively embedded into either the values (theta) or the shape (rotary_dim)
|
||||
# of the trained_freqs tensor.
|
||||
self._kv_copy = BlockedTrainedRotaryEmbeddings(self._config.head_size, self._config.n_heads_q,
|
||||
self._config.n_heads_kv, self._config.input_dtype)
|
||||
else:
|
||||
theta_base = rotary_config.theta_base
|
||||
rotary_dim = rotary_config.rotate_dim if rotary_config.rotate_dim is not None else self._config.head_size
|
||||
self._kv_copy = BlockedRotaryEmbeddings(self._config.head_size, self._config.n_heads_q,
|
||||
self._config.n_heads_kv, self._config.input_dtype, rotary_dim,
|
||||
theta_base)
|
||||
|
||||
self._softmax_scale = self._config.scale_factor
|
||||
|
||||
# TODO(cmikeh2): Attention kernel gets created here.
|
||||
self._attn_kernel = BlockedFlashAttn(self._config.head_size, self._config.input_dtype)
|
||||
self._atom_builder = AtomBuilder()
|
||||
|
||||
self.model_dim = self._config.head_size * self._config.n_heads_q
|
||||
self._output = torch.empty((self._config.max_tokens, self._config.head_size * self._config.n_heads_q),
|
||||
dtype=self._config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
# TODO(cmikeh2): Pre-allocate storage buffer for the attention atoms.
|
||||
self._max_atoms = self._config.max_sequences
|
||||
self._atoms = torch.empty((self._max_atoms, 8), dtype=torch.int32, device=get_accelerator().current_device())
|
||||
|
||||
alloc_func = RaggedUtilsBuilder().load().allocate_fast_host_buffer
|
||||
self._atoms_shadow = alloc_func(self._atoms)
|
||||
self._cur_atoms = 0
|
||||
|
||||
@cached_property
|
||||
def kv_block_size(self) -> int:
|
||||
"""
|
||||
Return preferred granulatity for blocked KV-cache implementation.
|
||||
"""
|
||||
return get_kv_block_size(self._config.head_size)
|
||||
|
||||
@cached_property
|
||||
def q_block_size(self) -> int:
|
||||
"""
|
||||
Property to calculate blocking granularity for the query dimension.
|
||||
This has no impact on the KV-cache structure, but will affect the
|
||||
number of attention atoms associated with a batch.
|
||||
"""
|
||||
return get_q_block_size(self._config.head_size)
|
||||
|
||||
def build_atoms(self, ragged_batch: RaggedBatchWrapper) -> None:
|
||||
"""
|
||||
Build the atoms for the attention kernel.
|
||||
|
||||
Args:
|
||||
ragged_batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata.
|
||||
"""
|
||||
host_atoms, n_atoms = self._atom_builder(self._atoms_shadow, ragged_batch, self.q_block_size,
|
||||
self.kv_block_size)
|
||||
|
||||
self._cur_atoms = n_atoms
|
||||
self._atoms[:n_atoms].copy_(host_atoms[:n_atoms], non_blocking=True)
|
||||
|
||||
def forward(self,
|
||||
q_k_v: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
batch: RaggedBatchWrapper,
|
||||
inv_freqs: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
"""
|
||||
Forward implementation.
|
||||
|
||||
Args:
|
||||
q_k_v (torch.Tensor): Query/Key/Value projection Tensor of shape
|
||||
[n_heads, (n_heads_q + 2 * n_heads_kv) * head_size].
|
||||
kv_cache (torch.Tensor): Blocked persistent cache of shape
|
||||
[2, batch, block_size, n_heads_kv, head_size].
|
||||
batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata.
|
||||
inv_freqs (Optional[torch.Tensor]): The inverse frequencies for the rotary embeddings if they
|
||||
have been modified from synthesizable values.
|
||||
"""
|
||||
if inv_freqs is not None:
|
||||
self._kv_copy(kv_cache, q_k_v, batch, inv_freqs)
|
||||
else:
|
||||
self._kv_copy(kv_cache, q_k_v, batch)
|
||||
|
||||
q = q_k_v[:, :self._config.head_size * self._config.n_heads_q]
|
||||
output = empty_from(self._output, q.shape)
|
||||
k_cache, v_cache = split_kv(kv_cache)
|
||||
|
||||
self._attn_kernel(output, q, k_cache, v_cache, self._atoms[:self._cur_atoms], self._softmax_scale)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .ragged_embedding import DSRaggedEmbedding
|
||||
@@ -0,0 +1,77 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ....allocator import empty_from
|
||||
from ....inference_utils import DtypeEnum
|
||||
from ....kernels.ragged_ops import RaggedEmbeddingKernel
|
||||
from ....ragged import RaggedBatchWrapper
|
||||
from ...interfaces import DSEmbeddingBase, DSEmbeddingRegistry
|
||||
from ...configs import DSEmbeddingsConfig
|
||||
|
||||
|
||||
@DSEmbeddingRegistry.register_module
|
||||
class DSRaggedEmbedding(DSEmbeddingBase):
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'ragged_embedding'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSEmbeddingsConfig) -> bool:
|
||||
|
||||
if DtypeEnum(config.residual_dtype) not in [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]:
|
||||
return False
|
||||
|
||||
if config.use_token_type:
|
||||
return False
|
||||
|
||||
if config.output_normalization is not None:
|
||||
return False
|
||||
|
||||
try:
|
||||
_ = RaggedEmbeddingKernel(config.residual_dtype, torch.int32, config.embedding_dim)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSEmbeddingsConfig, implementation_config: Dict[str, Any]) -> None:
|
||||
super().__init__(config, implementation_config)
|
||||
|
||||
self.embed_offset = self._config.positional_offset
|
||||
|
||||
# TODO(cmikeh2): How do we want to avoid the int32 vs int64 issue?
|
||||
self._ragged_embed = RaggedEmbeddingKernel(self._config.residual_dtype, torch.int32,
|
||||
self._config.embedding_dim)
|
||||
|
||||
self._output = torch.empty((self._config.max_tokens, self._config.embedding_dim),
|
||||
dtype=self._config.residual_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
@property
|
||||
def output(self) -> torch.Tensor:
|
||||
return self._output
|
||||
|
||||
def forward(self,
|
||||
ragged_batch: RaggedBatchWrapper,
|
||||
word_embeddings: torch.Tensor,
|
||||
position_embeddings: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
"""
|
||||
Parameters:
|
||||
ragged_batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata.
|
||||
word_embeddings (torch.Tensor): The word embedding table
|
||||
"""
|
||||
output = empty_from(self._output, (ragged_batch.tensor_toks, self._config.embedding_dim))
|
||||
self._ragged_embed(output,
|
||||
ragged_batch,
|
||||
word_embeddings,
|
||||
position_embed_weight=position_embeddings,
|
||||
position_embed_offset=self.embed_offset)
|
||||
return output
|
||||
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .blas_fp_linear import BlasFPLinear
|
||||
from .quantized_linear import QuantizedWf6Af16Linear, fp_quantize
|
||||
@@ -0,0 +1,103 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ....allocator import empty_from
|
||||
from ....inference_utils import is_gated
|
||||
from ....kernels.core_ops import (
|
||||
BlasLibLinear,
|
||||
CUDABiasActivation,
|
||||
CUDAGatedActivation,
|
||||
)
|
||||
|
||||
from ...interfaces import DSLinearBase, DSLinearRegistry
|
||||
from ...configs import DSLinearConfig
|
||||
from ....inference_parameter import InferenceParameter
|
||||
|
||||
|
||||
@DSLinearRegistry.register_module
|
||||
class BlasFPLinear(DSLinearBase):
|
||||
"""
|
||||
Linear DSModule based on BLAS library and standalone bias + activation kernel implementation.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'blas_fp_linear'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSLinearConfig) -> bool:
|
||||
if config.input_dtype != config.output_dtype:
|
||||
return False
|
||||
|
||||
if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16:
|
||||
return False
|
||||
|
||||
if is_gated(config.activation):
|
||||
try:
|
||||
_ = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
except ValueError:
|
||||
return False
|
||||
else:
|
||||
try:
|
||||
_ = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None:
|
||||
super().__init__(config, implementation_config)
|
||||
|
||||
self._linear_impl = BlasLibLinear(self._config.input_dtype)
|
||||
|
||||
if is_gated(config.activation):
|
||||
self._is_gated = True
|
||||
self._act_fn = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
self._double_buffer = torch.empty((config.max_tokens, config.out_channels * 2),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
else:
|
||||
self._is_gated = False
|
||||
self._act_fn = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
|
||||
self._output = torch.empty((config.max_tokens, config.out_channels),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
def transform_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Converts param to same data type as input and output.
|
||||
|
||||
Parameters:
|
||||
param (torch.Tensor): Weight or bias tensor.
|
||||
"""
|
||||
param = param.to(self._config.output_dtype)
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
|
||||
output = empty_from(self._output, (hidden_states.shape[0], self._config.out_channels))
|
||||
|
||||
if self._is_gated:
|
||||
staging_output = empty_from(self._double_buffer, (hidden_states.shape[0], self._config.out_channels * 2))
|
||||
self._linear_impl(staging_output, hidden_states, w)
|
||||
self._act_fn(output, staging_output, b)
|
||||
else:
|
||||
self._linear_impl(output, hidden_states, w)
|
||||
self._act_fn(output, b)
|
||||
|
||||
return output
|
||||
|
||||
@property
|
||||
def output(self) -> torch.Tensor:
|
||||
"""
|
||||
Return the padded, pre-allocated output Tensor.
|
||||
"""
|
||||
return self._output
|
||||
@@ -0,0 +1,205 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from deepspeed.ops.op_builder import InferenceCoreBuilder
|
||||
from ....allocator import empty_from
|
||||
from ....inference_utils import is_gated
|
||||
from ....kernels.core_ops import (
|
||||
CUDAWf6Af16Linear,
|
||||
CUDABiasActivation,
|
||||
CUDAGatedActivation,
|
||||
)
|
||||
|
||||
from ...interfaces import DSLinearBase, DSLinearRegistry
|
||||
from ...configs import DSLinearConfig
|
||||
from ....inference_parameter import InferenceParameter
|
||||
|
||||
|
||||
def fp_quantize(input: torch.FloatTensor,
|
||||
num_bits: int = 6,
|
||||
exp_bits: int = 3,
|
||||
min_value: torch.FloatTensor = None,
|
||||
max_value: torch.FloatTensor = None,
|
||||
group_size: int = -1):
|
||||
"""
|
||||
Args:
|
||||
inputs (`torch.FloatTensor`)
|
||||
The input which needs to be quantized
|
||||
num_bits (int, >=4)
|
||||
Number of bits to use for quantization
|
||||
exp_bits:
|
||||
fp exp_bits
|
||||
min_value/max_vlue (torch.FloatTensor)
|
||||
Used for static activation quantization
|
||||
group_size (int) N
|
||||
The quantization block size, each N numbers has its own scaling
|
||||
factor and off-site. -1 means use the last dim as the group_size
|
||||
Returns:
|
||||
quantized_fake_fp6
|
||||
The quantized weights, in fp16 format and contains fp6 value.
|
||||
scales
|
||||
Quantization scales
|
||||
"""
|
||||
|
||||
try:
|
||||
from qtorch.quant import float_quantize
|
||||
except ImportError:
|
||||
raise ImportError("Please install qtorch to use this function")
|
||||
|
||||
assert (min_value is None and max_value is None) or (min_value is not None and max_value is not None)
|
||||
|
||||
assert input.dtype == torch.float16
|
||||
|
||||
orig_device = input.device
|
||||
input = input.to(torch.float32).to(get_accelerator().current_device())
|
||||
if num_bits == 6 and exp_bits == 3: # this is default
|
||||
q_range = 28
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
man_bits = num_bits - exp_bits - 1
|
||||
input_shape = input.shape
|
||||
|
||||
if group_size == -1:
|
||||
group_size = input_shape[-1]
|
||||
else:
|
||||
# Only support per-channel quantization
|
||||
raise NotImplementedError
|
||||
num_groups = input.numel() // group_size
|
||||
input = input.reshape(num_groups, -1)
|
||||
|
||||
if min_value is None:
|
||||
max_input = torch.amax(torch.abs(input), dim=-1).view(num_groups, -1)
|
||||
else:
|
||||
max_input = torch.max(min_value.abs(), max_value) # .view(-1)
|
||||
scales = max_input / q_range # q_range + 1
|
||||
scales[scales == 0] = 1 # avoid zero scales
|
||||
scaled_input = input / scales
|
||||
|
||||
quantized_fake_fp6 = float_quantize(scaled_input, exp_bits, man_bits, rounding="nearest")
|
||||
|
||||
quantized_fake_fp6 = quantized_fake_fp6.reshape(input_shape).contiguous().to(torch.float16).to(orig_device)
|
||||
scales = scales.to(torch.float16).to(orig_device)
|
||||
# Now the dequantized value is quantized_fake_fp6 * scales
|
||||
|
||||
return quantized_fake_fp6, scales
|
||||
|
||||
|
||||
@DSLinearRegistry.register_module
|
||||
class QuantizedWf6Af16Linear(DSLinearBase):
|
||||
"""
|
||||
Linear DSModule for FP6 weight-only quantization kernel, where weight is FP6
|
||||
and activation is FP16.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'quantized_wf6af16_linear'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSLinearConfig) -> bool:
|
||||
if config.input_dtype != config.output_dtype:
|
||||
return False
|
||||
|
||||
# As for fp6 data items, they are packed and stored in a set of fp16
|
||||
# tensors. E.g., 8 fp6 data items are stored in 3 fp16 tensor.
|
||||
if config.input_dtype != torch.float16:
|
||||
return False
|
||||
|
||||
if is_gated(config.activation):
|
||||
try:
|
||||
_ = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
except ValueError:
|
||||
return False
|
||||
else:
|
||||
try:
|
||||
_ = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None:
|
||||
super().__init__(config, implementation_config)
|
||||
|
||||
self._linear_impl = CUDAWf6Af16Linear()
|
||||
|
||||
if is_gated(config.activation):
|
||||
# In the FP6 kernel implementation, the MatMul is W * A, where W is
|
||||
# the weight and A is activation. M is the output channel size.
|
||||
self.out_channels = self._config.out_channels * 2
|
||||
self.in_channels = self._config.in_channels
|
||||
self._is_gated = True
|
||||
self._act_fn = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
self._double_buffer = torch.empty((config.max_tokens, config.out_channels * 2),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
else:
|
||||
self.out_channels = self._config.out_channels
|
||||
self.in_channels = self._config.in_channels
|
||||
self._is_gated = False
|
||||
self._act_fn = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
|
||||
|
||||
self._output = torch.empty((config.max_tokens, config.out_channels),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
self.inf_module = InferenceCoreBuilder().load()
|
||||
self.inf_module.create_handle()
|
||||
self.preprocess_weight = self.inf_module.preprocess_weight
|
||||
|
||||
self.quantizer = fp_quantize
|
||||
|
||||
def transform_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Converts param to same data type as input and output.
|
||||
|
||||
Parameters:
|
||||
param (torch.Tensor): Weight or bias tensor.
|
||||
"""
|
||||
# It expects that the quantization scales are store in the attribute `scales`.
|
||||
|
||||
if param.ndim == 1: # bias, do nothing
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
quantized_fake_fp6, scales = self.quantizer(param, num_bits=6, exp_bits=3)
|
||||
|
||||
# This is for debugging, will delete before release.
|
||||
assert (quantized_fake_fp6.dtype == torch.float16)
|
||||
assert quantized_fake_fp6.shape[0] == self.out_channels
|
||||
assert scales.numel() == self.out_channels
|
||||
|
||||
weights_2bit, weights_4bit = self.preprocess_weight(quantized_fake_fp6)
|
||||
|
||||
return InferenceParameter.initialize(weights_2bit, weights_4bit=weights_4bit, scales=scales)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
weights_2bit = w
|
||||
weights_4bit = w.weights_4bit
|
||||
scales = w.scales
|
||||
output = empty_from(self._output, (hidden_states.shape[0], self._config.out_channels))
|
||||
if self._is_gated:
|
||||
staging_output = empty_from(self._double_buffer, (hidden_states.shape[0], self.out_channels))
|
||||
self._linear_impl(staging_output, hidden_states, weights_2bit, weights_4bit, scales, self.out_channels,
|
||||
hidden_states.shape[0], self.in_channels)
|
||||
self._act_fn(output, staging_output, b)
|
||||
else:
|
||||
self._linear_impl(output, hidden_states, weights_2bit, weights_4bit, scales, self.out_channels,
|
||||
hidden_states.shape[0], self.in_channels)
|
||||
self._act_fn(output, b)
|
||||
|
||||
return output
|
||||
|
||||
@property
|
||||
def output(self) -> torch.Tensor:
|
||||
"""
|
||||
Return the padded, pre-allocated output Tensor.
|
||||
"""
|
||||
return self._output
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .cutlass_multi_gemm import DSMultiGemmMoE
|
||||
@@ -0,0 +1,249 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ....allocator import empty_from
|
||||
from ....inference_utils import ActivationType, is_gated
|
||||
from ....kernels.core_ops import BlasLibLinear, CUDAGatedActivation
|
||||
from ....kernels.ragged_ops import (
|
||||
MoEGather,
|
||||
MoEScatter,
|
||||
RaggedTopKGating,
|
||||
)
|
||||
from ....ragged import RaggedBatchWrapper
|
||||
|
||||
from ...interfaces import DSMoEBase, DSMoERegistry
|
||||
from ...configs import DSMoEConfig
|
||||
from ....kernels.cutlass_ops import MoEGEMM
|
||||
from ....inference_parameter import InferenceParameter
|
||||
|
||||
|
||||
@DSMoERegistry.register_module
|
||||
class DSMultiGemmMoE(DSMoEBase):
|
||||
"""
|
||||
MoE implementation based on the CUTLASS multi-GEMM.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'cutlass_multi_gemm_moe'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSMoEConfig) -> bool:
|
||||
if config.input_dtype != config.output_dtype:
|
||||
return False
|
||||
|
||||
if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16:
|
||||
return False
|
||||
|
||||
if config.top_k != 1 and config.top_k != 2 and config.top_k != 4 and config.top_k != 8:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSMoEConfig, implementation_config: Dict[str, Any]) -> None:
|
||||
super().__init__(config, implementation_config)
|
||||
|
||||
# Convenience variables for frequently accessed items.
|
||||
self.max_tokens = self._config.max_tokens
|
||||
self.n_experts = self._config.n_experts
|
||||
self.n_top_k = self._config.top_k
|
||||
self.intermediate_dim = self._config.intermediate_features
|
||||
|
||||
moe_op_act_fn = ActivationType.IDENTITY if is_gated(self._config.activation) else self._config.activation
|
||||
|
||||
self._mlp_1 = MoEGEMM(fp_dtype=implementation_config['weight_dtype'], act_fn=moe_op_act_fn)
|
||||
self._mlp_2 = MoEGEMM(fp_dtype=implementation_config['weight_dtype'], act_fn=ActivationType.IDENTITY)
|
||||
|
||||
if is_gated(self._config.activation):
|
||||
self._activation = CUDAGatedActivation(self._config.model_dim, self._config.input_dtype,
|
||||
self._config.activation)
|
||||
else:
|
||||
self._activation = None
|
||||
|
||||
self._gate_proj = BlasLibLinear(self._config.input_dtype)
|
||||
self._top_1_gate = RaggedTopKGating(config.input_dtype)
|
||||
self._moe_scatter = MoEScatter(config.input_dtype, config.model_dim)
|
||||
self._moe_gather = MoEGather(config.input_dtype, config.model_dim, config.normalize_scores)
|
||||
|
||||
self._create_buffers()
|
||||
|
||||
def _create_buffers(self):
|
||||
|
||||
# Gating buffers
|
||||
self._logits = torch.empty((self._config.max_tokens, self.n_experts),
|
||||
dtype=self._config.input_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
self._expert_counts = torch.empty((self.n_experts, ),
|
||||
dtype=torch.int32,
|
||||
device=get_accelerator().current_device())
|
||||
self._scores = torch.empty((self._config.max_tokens, self.n_top_k),
|
||||
dtype=torch.float32,
|
||||
device=get_accelerator().current_device())
|
||||
self._assignments = torch.empty((self._config.max_tokens, self.n_top_k),
|
||||
dtype=torch.int32,
|
||||
device=get_accelerator().current_device())
|
||||
self._offsets = torch.empty((self._config.max_tokens, self.n_top_k),
|
||||
dtype=torch.int32,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
# Scatter buffers
|
||||
self._moe_input = torch.empty((self._config.max_tokens * self.n_top_k, self._config.model_dim),
|
||||
dtype=self._config.input_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
self._expert_cumsum = torch.empty((self._config.n_experts, ),
|
||||
dtype=torch.int64,
|
||||
device=get_accelerator().current_device())
|
||||
self._mapped_slots = torch.empty((self._config.max_tokens, self.n_top_k),
|
||||
dtype=torch.int32,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
# GEMM Buffers
|
||||
self._intermediate = torch.empty((self._config.max_tokens * self.n_top_k, self._config.intermediate_features),
|
||||
dtype=self._config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
if self._activation is not None:
|
||||
self._gated_intermediate = torch.empty(
|
||||
(self._config.max_tokens * self.n_top_k, self._config.intermediate_features * 2),
|
||||
dtype=self._config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
self._output_unordered = torch.empty((self._config.max_tokens * self.n_top_k, self._config.model_dim),
|
||||
dtype=self._config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
# Gather buffer
|
||||
self._output = torch.empty((self._config.max_tokens, self._config.model_dim),
|
||||
dtype=self._config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
def transform_gate_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Ensures gate param is going to match the activation data type.
|
||||
"""
|
||||
param = param.to(self._config.input_dtype)
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Converts param to same data type as input and output.
|
||||
|
||||
Parameters:
|
||||
param (torch.Tensor): Weight or bias tensor.
|
||||
"""
|
||||
param = param.to(self._config.input_dtype)
|
||||
|
||||
if len(param.shape) == 3:
|
||||
param = param.permute(0, 2, 1).contiguous()
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
def transform_moe_mlp_2_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
"""
|
||||
Converts param to same data type as input and output.
|
||||
|
||||
Parameters:
|
||||
param (torch.Tensor): Weight or bias tensor.
|
||||
"""
|
||||
param = param.to(self._config.input_dtype)
|
||||
|
||||
if len(param.shape) == 3:
|
||||
param = param.permute(0, 2, 1).contiguous()
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
@property
|
||||
def output(self) -> torch.Tensor:
|
||||
return self._output
|
||||
|
||||
def _gate(self, hidden_states: torch.Tensor, batch_metadata: RaggedBatchWrapper,
|
||||
gate_w: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Helper function to isolate the logit for gating. This will take the hidden states
|
||||
and produce the metadata + tensors for the CUTLASS ragged GEMMs. If the input has
|
||||
been padded for CG, this will strip the padding for MoE.
|
||||
|
||||
Parameters:
|
||||
hidden_states (torch.Tensor): Hidden states tensor. Expected shape is [n_tokens, model_dim].
|
||||
batch_metadata (RaggedBatchWrapper): Batch metadata for the hidden states.
|
||||
gate_w (torch.Tensor): Gate weight tensor. Expected shape is [num_experts, model_dim].
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: The MoE input, the cumsum of the offsets (for the MoE kernels themselves), the scores, and the mapped slots (to recover the original order of the tokens)
|
||||
"""
|
||||
|
||||
# Get views on the buffers for gating
|
||||
logits = empty_from(self._logits, (hidden_states.shape[0], self._logits.shape[-1]))
|
||||
scores = empty_from(self._scores, (hidden_states.shape[0], self.n_top_k))
|
||||
assignments = empty_from(self._assignments, (hidden_states.shape[0], self.n_top_k))
|
||||
offsets = empty_from(self._offsets, (hidden_states.shape[0], self.n_top_k))
|
||||
mapped_slots = empty_from(self._mapped_slots, (hidden_states.shape[0], self.n_top_k))
|
||||
moe_input = empty_from(self._moe_input, (hidden_states.shape[0] * self.n_top_k, self._moe_input.shape[-1]))
|
||||
|
||||
self._gate_proj(logits, hidden_states, gate_w)
|
||||
self._expert_counts.zero_()
|
||||
self._top_1_gate(self._expert_counts, scores, assignments, offsets, logits, batch_metadata)
|
||||
self._moe_scatter(moe_input, self._expert_cumsum, mapped_slots, hidden_states, self._expert_counts,
|
||||
assignments, offsets)
|
||||
|
||||
return moe_input, self._expert_cumsum, scores, mapped_slots
|
||||
|
||||
def forward(self,
|
||||
hidden_states: torch.Tensor,
|
||||
batch_metadata: RaggedBatchWrapper,
|
||||
gate_w: torch.Tensor,
|
||||
mlp_1_w: torch.Tensor,
|
||||
mlp_2_w: torch.Tensor,
|
||||
mlp_1_b: Optional[torch.Tensor] = None,
|
||||
mlp_2_b: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
"""
|
||||
MoE forward pass built on top of CUTLASS multi-GEMM.
|
||||
|
||||
Parameters:
|
||||
hidden_states (torch.Tensor): Hidden states tensor. Expected shape is [batch, seq_len, model_dim].
|
||||
gate_w (torch.Tensor): Gate weight tensor. Expected shape is [num_experts, model_dim].
|
||||
"""
|
||||
|
||||
moe_input, expert_cumsum, scores, mapped_slots = self._gate(hidden_states, batch_metadata, gate_w)
|
||||
|
||||
# Get views on the buffers for GEMM
|
||||
intermediate = empty_from(self._intermediate,
|
||||
(hidden_states.shape[0] * self.n_top_k, self._intermediate.shape[-1]))
|
||||
output_unordered = empty_from(self._output_unordered,
|
||||
(hidden_states.shape[0] * self.n_top_k, self._output_unordered.shape[-1]))
|
||||
output = empty_from(self._output, (hidden_states.shape[0], self._output.shape[-1]))
|
||||
|
||||
if self._activation is not None:
|
||||
gated_intermediate = empty_from(
|
||||
self._gated_intermediate, (hidden_states.shape[0] * self.n_top_k, self._gated_intermediate.shape[-1]))
|
||||
self._mlp_1(
|
||||
gated_intermediate,
|
||||
moe_input,
|
||||
mlp_1_w,
|
||||
expert_cumsum,
|
||||
mlp_1_b,
|
||||
)
|
||||
self._activation(intermediate, gated_intermediate)
|
||||
else:
|
||||
self._mlp_1(
|
||||
intermediate,
|
||||
moe_input,
|
||||
mlp_1_w,
|
||||
expert_cumsum,
|
||||
mlp_1_b,
|
||||
)
|
||||
|
||||
self._mlp_2(
|
||||
output_unordered,
|
||||
intermediate,
|
||||
mlp_2_w,
|
||||
expert_cumsum,
|
||||
mlp_2_b,
|
||||
)
|
||||
|
||||
self._moe_gather(output, output_unordered, scores, mapped_slots, self._expert_counts)
|
||||
return output
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .cuda_post_ln import DSPostLNCUDAModule
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ...interfaces import DSPostNormBase, DSPostNormRegistry
|
||||
from ...configs import DSNormConfig
|
||||
from ....kernels.core_ops.cuda_layer_norm.cuda_post_ln import CUDAFPPostLN
|
||||
from ....allocator import empty_from
|
||||
from ....inference_parameter import InferenceParameter
|
||||
|
||||
|
||||
@DSPostNormRegistry.register_module
|
||||
class DSPostLNCUDAModule(DSPostNormBase):
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'cuda_post_ln'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSNormConfig):
|
||||
if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1:
|
||||
return False
|
||||
|
||||
try:
|
||||
_ = CUDAFPPostLN(config.channels, config.residual_dtype)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
|
||||
super().__init__(config, implementation_config)
|
||||
self._fp_post_ln = CUDAFPPostLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
|
||||
|
||||
self._output = torch.empty((config.max_tokens, config.channels),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
def transform_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
param = param.to(self._config.input_dtype)
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
def forward(self, residual: torch.Tensor, hidden_in: torch.Tensor, gamma: torch.Tensor,
|
||||
beta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Since the CUDA FP only supports all data types being the same, we will alias the residual
|
||||
with our output.
|
||||
"""
|
||||
self._residual_output = empty_from(self._output, residual.shape)
|
||||
self._fp_post_ln(residual, residual, hidden_in, gamma, beta)
|
||||
return residual, residual
|
||||
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .cuda_pre_ln import DSPreLNCUDAModule
|
||||
from .cuda_pre_rms import DSPreRMSCUDAModule
|
||||
@@ -0,0 +1,69 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ...interfaces import DSPreNormBase, DSPreNormRegistry
|
||||
from ...configs import DSNormConfig, NormTypeEnum
|
||||
from ....kernels.core_ops.cuda_layer_norm.cuda_pre_ln import CUDAFPPreLN
|
||||
from ....kernels.core_ops.cuda_layer_norm.cuda_ln import CUDAFPLN
|
||||
from ....allocator import empty_from
|
||||
from ....inference_parameter import InferenceParameter
|
||||
|
||||
|
||||
@DSPreNormRegistry.register_module
|
||||
class DSPreLNCUDAModule(DSPreNormBase):
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'cuda_pre_ln'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSNormConfig):
|
||||
type = NormTypeEnum(config.type)
|
||||
if type != NormTypeEnum.LayerNorm:
|
||||
return False
|
||||
|
||||
if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1:
|
||||
return False
|
||||
|
||||
try:
|
||||
_ = CUDAFPPreLN(config.channels, config.residual_dtype)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
|
||||
super().__init__(config, implementation_config)
|
||||
self._fp_pre_ln = CUDAFPPreLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
|
||||
self._fp_ln = CUDAFPLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
|
||||
|
||||
# Buffers for the hidden output (residual is updated in-place)
|
||||
self._hidden_output = torch.empty((config.max_tokens, config.channels),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
def transform_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
param = param.to(self._config.input_dtype)
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
def forward(self, residual: torch.Tensor, hidden_in: Optional[torch.Tensor], gamma: torch.Tensor,
|
||||
beta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Since the CUDA FP only supports all data types being the same, we will alias the residual
|
||||
with our output.
|
||||
|
||||
If hidden_in is None, that means we do not need to perform the residual add and will
|
||||
only return the hidden output modified.
|
||||
"""
|
||||
hidden_out = empty_from(self._hidden_output, residual.shape)
|
||||
if hidden_in is None:
|
||||
self._fp_ln(hidden_out, residual, gamma, beta)
|
||||
else:
|
||||
self._fp_pre_ln(residual, hidden_out, residual, hidden_in, gamma, beta)
|
||||
return residual, hidden_out
|
||||
@@ -0,0 +1,79 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ...interfaces import DSPreNormBase, DSPreNormRegistry
|
||||
from ...configs import DSNormConfig, NormTypeEnum
|
||||
from ....kernels.core_ops import CUDARMSNorm, CUDARMSPreNorm
|
||||
from ....allocator import empty_from
|
||||
from ....inference_parameter import InferenceParameter
|
||||
|
||||
|
||||
@DSPreNormRegistry.register_module
|
||||
class DSPreRMSCUDAModule(DSPreNormBase):
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'cuda_pre_rms'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSNormConfig):
|
||||
type = NormTypeEnum(config.type)
|
||||
if type != NormTypeEnum.RMSNorm:
|
||||
return False
|
||||
|
||||
if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1:
|
||||
return False
|
||||
|
||||
try:
|
||||
# Only need to check one since the support matrix for the two rms kernels is the same
|
||||
_ = CUDARMSPreNorm(config.channels, config.residual_dtype)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
|
||||
super().__init__(config, implementation_config)
|
||||
self._fp_rms = CUDARMSNorm(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
|
||||
self._fp_rms_pre = CUDARMSPreNorm(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
|
||||
|
||||
# Buffers for both the hidden and residual outputs
|
||||
self._hidden_output = torch.empty((config.max_tokens, config.channels),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
self._residual_output = torch.empty((config.max_tokens, config.channels),
|
||||
dtype=config.output_dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
def transform_param(self, param: torch.Tensor) -> InferenceParameter:
|
||||
param = param.to(self._config.input_dtype)
|
||||
return InferenceParameter.initialize(param)
|
||||
|
||||
def forward(self,
|
||||
residual: torch.Tensor,
|
||||
hidden_in: Optional[torch.Tensor],
|
||||
gamma: torch.Tensor,
|
||||
beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Since the CUDA FP only supports all data types being the same, we will alias the residual
|
||||
with our output.
|
||||
|
||||
If hidden_in is None, that means we do not need to perform the residual add and will
|
||||
only return the hidden output modified.
|
||||
"""
|
||||
assert beta is None, "Beta is not supported for RMSNorm"
|
||||
|
||||
hidden_out = empty_from(self._hidden_output, residual.shape)
|
||||
if hidden_in is None:
|
||||
self._fp_rms(hidden_out, residual, gamma)
|
||||
residual_out = residual
|
||||
else:
|
||||
residual_out = empty_from(self._residual_output, residual.shape)
|
||||
self._fp_rms_pre(residual_out, hidden_out, residual, hidden_in, gamma)
|
||||
return residual_out, hidden_out
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .ragged_unembed import DSRaggedUnembed
|
||||
@@ -0,0 +1,120 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ....allocator import empty_from
|
||||
from ....inference_utils import DtypeEnum, ActivationType
|
||||
from ....kernels.core_ops import CUDAFPLN, BlasLibLinear, CUDARMSNorm, CUDABiasActivation
|
||||
from ....kernels.ragged_ops import RaggedLogitsGather
|
||||
from ....ragged import RaggedBatchWrapper
|
||||
from ...interfaces import DSUnembedBase, DSUnembedRegistry
|
||||
from ...configs import DSUnembedConfig
|
||||
|
||||
|
||||
@DSUnembedRegistry.register_module
|
||||
class DSRaggedUnembed(DSUnembedBase):
|
||||
"""
|
||||
Ragged unembedding implementation. This implementation will gather only the last token
|
||||
of each sequence in the ragged inflight batch and calculate the logits only for those rows.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def name():
|
||||
return 'ragged_unembed'
|
||||
|
||||
@staticmethod
|
||||
def supports_config(config: DSUnembedConfig):
|
||||
|
||||
if DtypeEnum(config.dtype) not in [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]:
|
||||
return False
|
||||
|
||||
try:
|
||||
_ = RaggedLogitsGather(config.model_dim, config.dtype)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
if config.norm_type == 'rms_norm':
|
||||
try:
|
||||
_ = CUDARMSNorm(config.model_dim, config.dtype)
|
||||
except ValueError:
|
||||
return False
|
||||
elif config.norm_type == 'layer_norm':
|
||||
try:
|
||||
_ = CUDAFPLN(config.model_dim, config.dtype)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def __init__(self, config: DSUnembedConfig, implementation_config: Dict[str, Any]) -> None:
|
||||
super().__init__(config, implementation_config)
|
||||
|
||||
self._logits_gather = RaggedLogitsGather(config.model_dim, self._config.dtype)
|
||||
|
||||
if self._config.norm_type == 'layer_norm':
|
||||
self._norm = CUDAFPLN(self._config.model_dim, self._config.dtype)
|
||||
elif self._config.norm_type == 'rms_norm':
|
||||
self._norm = CUDARMSNorm(self._config.model_dim, self._config.dtype)
|
||||
else:
|
||||
self._norm = None
|
||||
|
||||
self._linear = BlasLibLinear(self._config.dtype)
|
||||
# Here the activation kernel is being used to apply bias, hence the identity activation type!
|
||||
self._act_fn = CUDABiasActivation(self._config.vocab_size, self._config.dtype, ActivationType.IDENTITY)
|
||||
|
||||
self._intermediate = torch.empty((self._config.max_sequences, self._config.model_dim),
|
||||
dtype=self._config.dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
self._output = torch.empty((self._config.max_sequences, self._config.vocab_size),
|
||||
dtype=self._config.dtype,
|
||||
device=get_accelerator().current_device())
|
||||
|
||||
@property
|
||||
def output(self) -> torch.Tensor:
|
||||
return self._output
|
||||
|
||||
def forward(self,
|
||||
hidden_states: torch.Tensor,
|
||||
vocab_embedding: torch.Tensor,
|
||||
ragged_metadata: RaggedBatchWrapper,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
gamma: Optional[torch.Tensor] = None,
|
||||
beta: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
"""
|
||||
Return final model logits.
|
||||
|
||||
Args:
|
||||
hidden_states (torch.Tensor): The hidden states from the model. This is the output of the
|
||||
final layer of the model.
|
||||
vocab_embedding (torch.Tensor): The vocab embedding table.
|
||||
raged_metadata (RaggedBatchWrapper): The ragged batch metadata.
|
||||
gamma (Optional[torch.Tensor]): The gamma tensor for normalization.
|
||||
beta (Optional[torch.Tensor]): The beta tensor for normalization.
|
||||
"""
|
||||
|
||||
cut_down_hidden_states = empty_from(self._intermediate,
|
||||
(ragged_metadata.current_sequences, self._config.model_dim))
|
||||
self._logits_gather(cut_down_hidden_states, hidden_states, ragged_metadata)
|
||||
|
||||
if self._config.norm_type == 'rms_norm':
|
||||
if gamma is None:
|
||||
raise ValueError('RMS Normalization enabled but gamma not provided.')
|
||||
self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma)
|
||||
elif self._config.norm_type == 'layer_norm':
|
||||
if gamma is None or beta is None:
|
||||
raise ValueError('Normalization enabled but gamma and/or beta not provided.')
|
||||
self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma, beta)
|
||||
|
||||
output = empty_from(self._output, (ragged_metadata.current_sequences, self._config.vocab_size))
|
||||
self._linear(output, cut_down_hidden_states, vocab_embedding)
|
||||
if bias is not None:
|
||||
self._act_fn(output, bias)
|
||||
|
||||
return output
|
||||
Reference in New Issue
Block a user