Files
2026-07-13 13:18:33 +08:00

57 lines
1.3 KiB
Python

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from ...inference_utils import ActivationType, DtypeEnum
from ...modules.ds_module import DSModuleConfig
class DSMoEConfig(DSModuleConfig):
"""
Config class for DSMoEBase
"""
model_dim: int
"""
Size of input activation.
"""
intermediate_features: int
"""
Size of intermediate activation. Specifically, this is the number of input features
in the second linear layer. Depending on the activation function, the output of the first
linear layer may have increased dimensionality.
"""
n_experts: int
"""
Number of experts.
"""
top_k: int = 1
"""
top-k gating function (like top-1 or top-2)
"""
input_dtype: DtypeEnum = DtypeEnum.fp16
"""
Data type for the input activations.
"""
output_dtype: DtypeEnum = DtypeEnum.fp16
"""
Data type for the output activations.
"""
activation: ActivationType = ActivationType.IDENTITY
"""
Activation function of the first MLP1
"""
normalize_scores: bool = False
"""
Whether normalization is applied to the selected scores. If true, the module
should rescale the scores such that their sum is 1.0.
"""