wehub-resource-sync
250 lines
11 KiB
Python
250 lines
11 KiB
Python
# 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
|