Initialize the Moe layer with given parameters. Args: fd_config (FDConfig): Arguments related to inference, containing attributes such as weight_dtype, act_dtype, mp_size, hidden_size, head_dim, num_attention_heads, and ffn_hidden_size.
(
self,
fd_config,
reduce_results: bool = True,
renormalize: bool = False,
moe_intermediate_size: int = -1,
num_experts: int = -1,
expert_id_offset: int = 0,
top_k: int = -1,
topk_method: str = "",
topk_group: int = -1,
n_group: int = -1,
routed_scaling_factor: float = 1.0,
layer_idx: int = -1,
moe_tag: str = "",
gate_correction_bias=None,
redundant_table_manger: RedundantExpertManger = None,
weight_key_map: dict = {},
with_bias: bool = False,
activation="swiglu",
model_format: Optional[str] = None,
)
| 128 | """ |
| 129 | |
| 130 | def __init__( |
| 131 | self, |
| 132 | fd_config, |
| 133 | reduce_results: bool = True, |
| 134 | renormalize: bool = False, |
| 135 | moe_intermediate_size: int = -1, |
| 136 | num_experts: int = -1, |
| 137 | expert_id_offset: int = 0, |
| 138 | top_k: int = -1, |
| 139 | topk_method: str = "", |
| 140 | topk_group: int = -1, |
| 141 | n_group: int = -1, |
| 142 | routed_scaling_factor: float = 1.0, |
| 143 | layer_idx: int = -1, |
| 144 | moe_tag: str = "", |
| 145 | gate_correction_bias=None, |
| 146 | redundant_table_manger: RedundantExpertManger = None, |
| 147 | weight_key_map: dict = {}, |
| 148 | with_bias: bool = False, |
| 149 | activation="swiglu", |
| 150 | model_format: Optional[str] = None, |
| 151 | ): |
| 152 | """ |
| 153 | Initialize the Moe layer with given parameters. |
| 154 | Args: |
| 155 | fd_config (FDConfig): Arguments related to inference, containing |
| 156 | attributes such as weight_dtype, act_dtype, mp_size, hidden_size, head_dim, |
| 157 | num_attention_heads, and ffn_hidden_size. |
| 158 | """ |
| 159 | super().__init__() |
| 160 | |
| 161 | self.fd_config = fd_config |
| 162 | self.layer_idx = layer_idx |
| 163 | self.reduce_results = reduce_results |
| 164 | self.renormalize = renormalize |
| 165 | self.tp_rank = fd_config.parallel_config.tensor_parallel_rank |
| 166 | self.tp_size = fd_config.parallel_config.tensor_parallel_size |
| 167 | self.ep_size = fd_config.parallel_config.expert_parallel_size |
| 168 | self.ep_rank = fd_config.parallel_config.expert_parallel_rank |
| 169 | self.tp_group = fd_config.parallel_config.tp_group |
| 170 | # NOTE(Zhenyu Li): just supports tp_size = 1 when ep_size > 1 in MOE now. |
| 171 | if self.ep_size > 1: |
| 172 | self.tp_size = 1 |
| 173 | self.tp_rank = 0 |
| 174 | |
| 175 | self.attn_tp_size = fd_config.parallel_config.tensor_parallel_size |
| 176 | self.attn_tp_rank = fd_config.parallel_config.tensor_parallel_rank |
| 177 | |
| 178 | assert (self.tp_size >= 1 and self.ep_size == 1) or ( |
| 179 | self.tp_size == 1 and self.ep_size > 1 |
| 180 | ), "MoE only support parallelism on TP or EP dimension." |
| 181 | |
| 182 | self.hidden_size = fd_config.model_config.hidden_size |
| 183 | self.num_experts = num_experts |
| 184 | |
| 185 | self.num_local_experts = self.num_experts // self.ep_size |
| 186 | |
| 187 | self.moe_intermediate_size = moe_intermediate_size // self.tp_size |
nothing calls this directly
no test coverage detected