File size: 3,050 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | from diffulex.moe.config import (
get_mlp_only_layers,
get_moe_intermediate_size,
get_num_experts,
get_num_experts_per_tok,
get_norm_topk_prob,
is_moe_layer,
)
from diffulex.distributed.parallel_state import fetch_parallel_state
def build_mlp_or_moe(config, layer_idx: int, dense_factory):
"""Build a dense MLP or MoE block according to the config."""
if is_moe_layer(config, layer_idx):
from diffulex.moe.layer import build_moe_block
parallel_state = fetch_parallel_state()
if parallel_state.is_ep_enabled():
return build_moe_block("ep", config)
if parallel_state.is_tp_enabled():
return build_moe_block("tp", config)
return build_moe_block("naive", config)
return dense_factory()
def __getattr__(name: str):
if name in ("FusedMoE", "SparseMoEBlock"):
from diffulex.moe.layer import FusedMoE
return FusedMoE
if name == "build_moe_block":
from diffulex.moe.layer import build_moe_block
return build_moe_block
if name in ("DispatcherOutput", "TokenDispatcher", "build_token_dispatcher"):
from diffulex.moe.dispatcher import DispatcherOutput, TokenDispatcher, build_token_dispatcher
return {
"DispatcherOutput": DispatcherOutput,
"TokenDispatcher": TokenDispatcher,
"build_token_dispatcher": build_token_dispatcher,
}[name]
if name == "DeepEPDispatcher":
from diffulex.moe.dispatcher.deepep_dispatcher import DeepEPDispatcher
return DeepEPDispatcher
if name == "NaiveA2ADispatcher":
from diffulex.moe.dispatcher.naive_dispatcher import NaiveA2ADispatcher
return NaiveA2ADispatcher
if name == "DeepEPMode":
from diffulex.moe.mode import DeepEPMode
return DeepEPMode
if name in ("DispatchMetadata", "DeepEPDispatchMetadata", "DispatcherStage", "ExpertExecutionMetadata", "RouterMetadata"):
from diffulex.moe.metadata import (
DeepEPDispatchMetadata,
DispatchMetadata,
DispatcherStage,
ExpertExecutionMetadata,
RouterMetadata,
)
return {
"DeepEPDispatchMetadata": DeepEPDispatchMetadata,
"DispatchMetadata": DispatchMetadata,
"DispatcherStage": DispatcherStage,
"ExpertExecutionMetadata": ExpertExecutionMetadata,
"RouterMetadata": RouterMetadata,
}[name]
raise AttributeError(name)
__all__ = [
"FusedMoE",
"SparseMoEBlock",
"DeepEPDispatchMetadata",
"DispatchMetadata",
"DispatcherStage",
"DeepEPDispatcher",
"DeepEPMode",
"DispatcherOutput",
"ExpertExecutionMetadata",
"NaiveA2ADispatcher",
"RouterMetadata",
"TokenDispatcher",
"build_token_dispatcher",
"build_mlp_or_moe",
"get_mlp_only_layers",
"get_moe_intermediate_size",
"get_num_experts",
"get_num_experts_per_tok",
"get_norm_topk_prob",
"is_moe_layer",
]
|