""" # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ from functools import partial from typing import Callable, Optional import paddle from paddle import nn from paddleformers.utils.log import logger from fastdeploy import envs from fastdeploy.distributed.communication import ( tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce_custom, ) from fastdeploy.model_executor.forward_meta import ForwardMeta from fastdeploy.model_executor.layers.moe.routing_indices_cache import ( save_routing_to_buffer, ) from fastdeploy.model_executor.layers.utils import get_tensor from fastdeploy.model_executor.utils import h2d_copy, slice_fn from fastdeploy.platforms import current_platform from fastdeploy.worker.experts_manager import RedundantExpertManger try: from fastdeploy.model_executor.ops.gpu import noaux_tc, noaux_tc_redundant except: logger.warning("import noaux_tc Failed!") import numpy as np def get_moe_method(layer=None): """ return moe method based on device platform """ if current_platform.is_cuda() or current_platform.is_iluvatar(): from .fused_moe_cutlass_backend import CutlassMoEMethod return CutlassMoEMethod(None) elif current_platform.is_xpu(): from fastdeploy.model_executor.layers.backends import XPUMoEMethod return XPUMoEMethod(None, layer) elif current_platform.is_gcu(): from fastdeploy.model_executor.layers.backends import GCUFusedMoeMethod return GCUFusedMoeMethod(None) elif current_platform.is_intel_hpu(): from fastdeploy.model_executor.layers.backends import HpuMoEMethod return HpuMoEMethod(None) elif current_platform.is_maca(): from fastdeploy.model_executor.layers.backends import ( MetaxCutlassUnquantizedFusedMoEMethod, ) return MetaxCutlassUnquantizedFusedMoEMethod(None) return None def get_moe_scores( gating_output: paddle.Tensor, n_group, topk_group, top_k, routed_scaling_factor, e_score_correction_bias, renormalize: bool = False, expert_id_to_ep_rank_array: paddle.Tensor = None, expert_in_rank_num_list: paddle.Tensor = None, tokens_per_expert_stats_list: paddle.Tensor = None, redundant_ep_rank_num_plus_one: int = 1, ) -> paddle.Tensor: """ compute moe scores using e_score_correction_bias. """ scores = paddle.nn.functional.sigmoid(gating_output) assert e_score_correction_bias is not None, "e_score_correction_bias is none!" scores_with_bias = scores + e_score_correction_bias if expert_id_to_ep_rank_array is None: scores, topk_values, topk_idx = noaux_tc( scores, scores_with_bias, n_group if n_group > 0 else 1, topk_group if topk_group > 0 else 1, top_k, renormalize, routed_scaling_factor, ) else: scores, topk_values, topk_idx = noaux_tc_redundant( scores, scores_with_bias, expert_id_to_ep_rank_array, expert_in_rank_num_list, tokens_per_expert_stats_list, n_group if n_group > 0 else 1, topk_group if topk_group > 0 else 1, top_k, renormalize, routed_scaling_factor, redundant_ep_rank_num_plus_one, ) return scores, topk_values, topk_idx class FusedMoE(nn.Layer): """ FusedMoE is a layer that performs MoE (Mixture of Experts) computation. """ def __init__( 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, ): """ 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. """ super().__init__() self.fd_config = fd_config self.layer_idx = layer_idx self.reduce_results = reduce_results self.renormalize = renormalize self.tp_rank = fd_config.parallel_config.tensor_parallel_rank self.tp_size = fd_config.parallel_config.tensor_parallel_size self.ep_size = fd_config.parallel_config.expert_parallel_size self.ep_rank = fd_config.parallel_config.expert_parallel_rank self.tp_group = fd_config.parallel_config.tp_group # NOTE(Zhenyu Li): just supports tp_size = 1 when ep_size > 1 in MOE now. if self.ep_size > 1: self.tp_size = 1 self.tp_rank = 0 self.attn_tp_size = fd_config.parallel_config.tensor_parallel_size self.attn_tp_rank = fd_config.parallel_config.tensor_parallel_rank assert (self.tp_size >= 1 and self.ep_size == 1) or ( self.tp_size == 1 and self.ep_size > 1 ), "MoE only support parallelism on TP or EP dimension." self.hidden_size = fd_config.model_config.hidden_size self.num_experts = num_experts self.num_local_experts = self.num_experts // self.ep_size self.moe_intermediate_size = moe_intermediate_size // self.tp_size self.top_k = top_k self.weight_key_map = weight_key_map self.use_method = envs.FD_MOE_BACKEND.lower() self.moe_tag = moe_tag self.with_bias = with_bias self.activation = activation if self.ep_size > 1: expert_id_offset = expert_id_offset + self.ep_rank * self.num_local_experts self.expert_id_offset = expert_id_offset self.gate_correction_bias_key = self.weight_key_map.get("gate_correction_bias_key", None) if self.gate_correction_bias_key is not None: self.moe_use_gate_correction_bias = True else: self.moe_use_gate_correction_bias = False # used for deepseek_v3 self.topk_method = topk_method self.topk_group = topk_group self.n_group = n_group self.routed_scaling_factor = routed_scaling_factor self._dtype = self._helper.get_default_dtype() self.weight_dtype = self._dtype self.is_moe_quantized = getattr(self.fd_config.model_config, "is_moe_quantized", False) self.is_quantized = self.is_moe_quantized or ( fd_config.model_config.is_quantized and not (fd_config.quant_config.name() == "mix_quant" and fd_config.quant_config.moe_quant_type is None) ) moe_quant_config = fd_config.quant_config self.moe_quant_config = moe_quant_config self.moe_quant_type = None if moe_quant_config and moe_quant_config.get_quant_method(self): self.quant_method = moe_quant_config.get_quant_method(self) self.moe_quant_type = moe_quant_config.name() else: # unquantized quant_method self.quant_method = get_moe_method(self) assert self.quant_method is not None, "self.quant_method should not be None" self.redundant_table_manger = redundant_table_manger self.is_rearrange = False if self.ep_size > 1: self.quant_method.init_ep(self) self.enable_routing_replay = fd_config.routing_replay_config.enable_routing_replay # Merge normal and RL build model if gate_correction_bias is not None: self.gate_correction_bias = gate_correction_bias else: self.gate_correction_bias = None self.quant_method.create_weights( self, weight_loader=self.weight_loader, model_format=fd_config.model_config.model_format if model_format is None else model_format, num_experts=self.num_local_experts if self.ep_size > 1 else self.num_experts, hidden_size=self.hidden_size, moe_intermediate_size=self.moe_intermediate_size, ) logger.info( f"{moe_tag}MoE config is {num_experts=}[{expert_id_offset}, {expert_id_offset + self.num_local_experts}), \ {top_k=}, hidden_size={self.hidden_size}, {moe_intermediate_size=}, \ , ep_size={self.ep_size}, \ tp_size={self.tp_size}." ) def weight_loader( self, param, loaded_weight, expert_id, shard_id: Optional[str] = None, source: Optional[str] = None ): """ source:Avoid redundant transpose of fused weights when weight_loader is called iteratively """ if expert_id is None and shard_id is None: # MoE experts has been fused in disk self._load_fused_experts_weight(param, loaded_weight) return if hasattr(param, "SHARD_ID_TO_SHARDED_DIM"): SHARD_ID_TO_SHARDED_DIM = param.SHARD_ID_TO_SHARDED_DIM elif current_platform.is_cuda() or current_platform.is_iluvatar() or current_platform.is_maca(): SHARD_ID_TO_SHARDED_DIM = {"gate": 1, "down": 0, "up": 1} else: SHARD_ID_TO_SHARDED_DIM = {"gate": 0, "down": 1, "up": 0} if not (expert_id - self.expert_id_offset >= 0 and expert_id - self.expert_id_offset < self.num_local_experts): return if not param._is_initialized(): param.initialize() weight_need_transpose = getattr(param, "weight_need_transpose", False) if self.ep_size > 1 or weight_need_transpose: loaded_weight = get_tensor(loaded_weight) if shard_id is None: # 1.gate up fused in disk if weight_need_transpose: loaded_weight = loaded_weight.transpose([1, 0]) output_size = param[expert_id - self.expert_id_offset].shape[SHARD_ID_TO_SHARDED_DIM["gate"]] shard_offsets = [ # (shard_id, shard_offset, shard_size) ("gate", 0, output_size // 2 * self.tp_size), ("up", output_size // 2 * self.tp_size, output_size // 2 * self.tp_size), ] for shard_id, shard_offset, shard_size in shard_offsets: loaded_weight_shard = slice_fn( loaded_weight, SHARD_ID_TO_SHARDED_DIM[shard_id], shard_offset, shard_offset + shard_size ) self.weight_loader(param, loaded_weight_shard, expert_id, shard_id, "fused") else: if weight_need_transpose and source != "fused": loaded_weight = loaded_weight.transpose([1, 0]) # 2.gate up splited in disk assert shard_id in ["gate", "down", "up"] self._load_expert_weight( param=param, expert_id=expert_id, loaded_weight=loaded_weight, shard_id=shard_id, shard_dim=SHARD_ID_TO_SHARDED_DIM[shard_id], ) def _load_gate_up_weight(self, param, expert_id, loaded_weight, shard_id, shard_dim=None, is_sharded=False): if self.tp_size > 1 and not is_sharded: tp_shard_dim = shard_dim weight_dim = -1 if tp_shard_dim else 0 size = loaded_weight.shape[weight_dim] block_size = size // self.tp_size shard_offset = self.tp_rank * block_size shard_size = (self.tp_rank + 1) * block_size loaded_weight = slice_fn(loaded_weight, tp_shard_dim, shard_offset, shard_size) expert_param = param[expert_id - self.expert_id_offset] dim = -1 if shard_dim else 0 param_shard_size = expert_param.shape[dim] // 2 if shard_id == "gate": param_shard_offset = 0 else: # shard_id == "up": param_shard_offset = param_shard_size expert_param = slice_fn( expert_param, shard_dim, start=param_shard_offset, end=param_shard_offset + param_shard_size ) if hasattr(param, "tensor_track"): # for dyn quant param.tensor_track.mark( start=param_shard_offset, end=param_shard_offset + param_shard_size, batch_id=expert_id - self.expert_id_offset, ) # To ensure compatibility across backends, apply an extra transpose for GCU and XPU if expert_param.shape != loaded_weight.shape: loaded_weight = loaded_weight.transpose([1, 0]) assert expert_param.shape == loaded_weight.shape, ( f"Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({expert_param.shape})" ) if expert_param.dtype != loaded_weight.dtype: if loaded_weight.dtype == paddle.int8 and expert_param.dtype == paddle.float8_e4m3fn: loaded_weight = loaded_weight.view(expert_param.dtype) else: loaded_weight = loaded_weight.cast(expert_param.dtype) h2d_copy(dst=expert_param, src=loaded_weight) def _load_down_weight(self, param, expert_id, loaded_weight, shard_id, shard_dim=None): if self.tp_size > 1 and shard_dim is not None: tp_shard_dim = shard_dim dim = -1 if tp_shard_dim else 0 size = loaded_weight.shape[dim] block_size = size // self.tp_size shard_offset = self.tp_rank * block_size shard_size = (self.tp_rank + 1) * block_size loaded_weight = slice_fn(loaded_weight, tp_shard_dim, shard_offset, shard_size) expert_param = param[expert_id - self.expert_id_offset] if hasattr(param, "tensor_track"): # for dyn quant param.tensor_track.mark(start=0, batch_id=expert_id - self.expert_id_offset) # To ensure compatibility across backends, apply an extra transpose for GCU and XPU and opensource weight if expert_param.shape != loaded_weight.shape: loaded_weight = loaded_weight.transpose([1, 0]) assert expert_param.shape == loaded_weight.shape, ( f"Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({expert_param.shape})" ) if expert_param.dtype != loaded_weight.dtype: if loaded_weight.dtype == paddle.int8 and expert_param.dtype == paddle.float8_e4m3fn: loaded_weight = loaded_weight.view(expert_param.dtype) else: loaded_weight = loaded_weight.cast(expert_param.dtype) h2d_copy(dst=expert_param, src=loaded_weight) def _load_fused_experts_weight(self, param, loaded_weight): if self.tp_size > 1: dim = -1 if isinstance(loaded_weight, (np.ndarray, paddle.Tensor)): size = loaded_weight.shape[dim] else: size = loaded_weight.get_shape()[dim] block_size = size // self.tp_size shard_offset = self.tp_rank * block_size shard_size = (self.tp_rank + 1) * block_size loaded_weight = slice_fn(loaded_weight, dim, shard_offset, shard_size) assert param.shape == loaded_weight.shape, ( f"Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({param.shape})" ) h2d_copy(dst=param, src=loaded_weight) if hasattr(param, "tensor_track"): for i in range(self.num_local_experts): param.tensor_track.mark(start=0, batch_id=i) def _load_expert_weight( self, param, expert_id, loaded_weight, shard_id, shard_dim=None, ): if shard_id == "down": self._load_down_weight(param, expert_id, loaded_weight, shard_id, shard_dim) elif shard_id in ["gate", "up"]: self._load_gate_up_weight(param, expert_id, loaded_weight, shard_id, shard_dim) @classmethod def make_expert_params_mapping( cls, num_experts: int, ckpt_gate_proj_name: Optional[str] = None, ckpt_up_proj_name: Optional[str] = None, ckpt_down_proj_name: Optional[str] = None, ckpt_gate_up_proj_name: Optional[str] = None, param_gate_up_proj_name: Optional[str] = None, param_down_proj_name: Optional[str] = None, ckpt_expert_key_name: str = "experts", experts_offset: int = 0, num_experts_start_offset: int = 0, ) -> list[tuple[str, str, int, str]]: param_name_maping = [] if ckpt_gate_up_proj_name: param_name_maping.append((None, ckpt_gate_up_proj_name)) if ckpt_gate_proj_name: param_name_maping.append(("gate", ckpt_gate_proj_name)) if ckpt_down_proj_name: param_name_maping.append(("down", ckpt_down_proj_name)) if ckpt_up_proj_name: param_name_maping.append(("up", ckpt_up_proj_name)) return [ # (param_name, weight_name, expert_id, shard_id) ( ( param_gate_up_proj_name if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name, ckpt_gate_up_proj_name] else param_down_proj_name ), f"{ckpt_expert_key_name}.{expert_id}.{weight_name}.", expert_id, shard_id, ) for expert_id in range( experts_offset + num_experts_start_offset, experts_offset + num_experts_start_offset + num_experts ) for shard_id, weight_name in param_name_maping ] def load_experts_weight( self, state_dict: dict, up_gate_proj_expert_weight_key: str, down_proj_expert_weight_key: str, is_rearrange: bool = False, ): """ Load experts weight from state_dict. Args: state_dict (dict): The state_dict of model. up_gate_proj_expert_weight_key (str): The key of up_gate_proj expert weight. down_proj_expert_weight_key (str): The key of down_proj expert weight. """ logical_expert_ids = [ i for i in range( self.expert_id_offset, self.expert_id_offset + self.num_local_experts, ) ] ep_rank_to_expert_id_list = [i for i in range(self.num_experts)] if self.redundant_table_manger is not None and is_rearrange is True: ( ep_rank_to_expert_id_list, expert_id_to_ep_rank_array, expert_in_rank_num_list, tokens_per_expert_stats_list, ) = self.redundant_table_manger.get_ep_rank_to_expert_id_list_by_layer(self.layer_idx) logical_expert_ids = ep_rank_to_expert_id_list[ self.expert_id_offset : self.expert_id_offset + self.num_local_experts ] up_gate_proj_weights = [] down_proj_weights = [] if isinstance(state_dict, list): state_dict = dict(state_dict) is_ffn_merged = ( up_gate_proj_expert_weight_key.format(logical_expert_ids[0] if is_rearrange else self.expert_id_offset) in state_dict ) if is_ffn_merged: for expert_idx in logical_expert_ids: down_proj_expert_weight_key_name = down_proj_expert_weight_key.format(expert_idx) up_gate_proj_expert_weight_key_name = up_gate_proj_expert_weight_key.format(expert_idx) up_gate_proj_weights.append( get_tensor( ( state_dict.pop(up_gate_proj_expert_weight_key_name) if up_gate_proj_expert_weight_key_name in state_dict else up_gate_proj_expert_weight_key_name ), self.fd_config.model_config.model, ) ) down_proj_weights.append( get_tensor( ( state_dict.pop(down_proj_expert_weight_key_name) if down_proj_expert_weight_key_name in state_dict else down_proj_expert_weight_key_name ), self.fd_config.model_config.model, ) ) else: gate_expert_weight_key = up_gate_proj_expert_weight_key.replace("up_gate_proj", "gate_proj") up_expert_weight_key = up_gate_proj_expert_weight_key.replace("up_gate_proj", "up_proj") for expert_idx in logical_expert_ids: gate_expert_weight_key_name = gate_expert_weight_key.format(expert_idx) up_expert_weight_key_name = up_expert_weight_key.format(expert_idx) down_proj_expert_weight_key_name = down_proj_expert_weight_key.format(expert_idx) gate = get_tensor( ( state_dict.pop(gate_expert_weight_key_name) if gate_expert_weight_key_name in state_dict else gate_expert_weight_key_name ), self.fd_config.model_config.model, ) up = get_tensor( ( state_dict.pop(up_expert_weight_key_name) if up_expert_weight_key_name in state_dict else up_expert_weight_key_name ), self.fd_config.model_config.model, ) up_gate_proj_weights.append(paddle.concat([gate, up], axis=-1)) down_proj_weights.append( get_tensor( ( state_dict.pop(down_proj_expert_weight_key_name) if down_proj_expert_weight_key_name in state_dict else down_proj_expert_weight_key_name ), self.fd_config.model_config.model, ) ) return up_gate_proj_weights, down_proj_weights, logical_expert_ids, ep_rank_to_expert_id_list def extract_moe_ffn_weights(self, state_dict: dict): """ Extract MoE FFN weights from state dict based on weight key mapping. Args: state_dict (dict): Model state dictionary containing the weights. Returns: tuple: A tuple containing two lists: - up_gate_proj_weights: List of tensors for first FFN layer weights - down_proj_weights: List of tensors for second FFN layer weights Raises: AssertionError: If required weight keys are missing or number of weights doesn't match number of local experts. """ up_gate_proj_expert_weight_key = self.weight_key_map.get("up_gate_proj_expert_weight_key", None) down_proj_expert_weight_key = self.weight_key_map.get("down_proj_expert_weight_key", None) assert up_gate_proj_expert_weight_key is not None, "up_gate_proj_expert_weight_key should not be none." assert down_proj_expert_weight_key is not None, "down_proj_expert_weight_key should not be none." up_gate_proj_weights, down_proj_weights, logical_expert_ids, ep_rank_to_expert_id_list = ( self.load_experts_weight( state_dict, up_gate_proj_expert_weight_key, down_proj_expert_weight_key, ) ) assert ( len(up_gate_proj_weights) == self.num_local_experts ), "up_gate_proj_weights length should be equal to num_local_experts." assert ( len(down_proj_weights) == self.num_local_experts ), "down_proj_weights length should be equal to num_local_experts." return up_gate_proj_weights, down_proj_weights, logical_expert_ids, ep_rank_to_expert_id_list def extract_gate_correction_bias(self, gate_correction_bias_key, state_dict): """ extract_gate_correction_bias function. """ gate_correction_bias_tensor = get_tensor(state_dict.pop(gate_correction_bias_key)).astype("float32") return gate_correction_bias_tensor def load_state_dict(self, state_dict, is_rearrange: bool = False): """ load_state_dict function. """ if self.is_quantized or self.fd_config.model_config.is_moe_quantized: if getattr(self.fd_config.quant_config, "is_permuted", True): self.quant_method.process_prequanted_weights(self, state_dict, is_rearrange) else: self.quant_method.process_loaded_weights(self, state_dict) else: self.quant_method.process_loaded_weights(self, state_dict) def forward_split_allgather(self, x: paddle.Tensor, gate: nn.Layer, topk_ids_hookfunc: Callable = None): """ Forward split allgather function. """ token_num = x.shape[0] token_num_per_rank = (token_num + self.attn_tp_size - 1) // self.attn_tp_size # AllGather will hang when the data shapes on multi-ranks are different! part_x = paddle.zeros(shape=[token_num_per_rank, x.shape[1]], dtype=x.dtype) start_offset = self.attn_tp_rank * token_num_per_rank end_offset = (self.attn_tp_rank + 1) * token_num_per_rank if start_offset >= token_num: start_offset = token_num if end_offset > token_num: end_offset = token_num part_x[: (end_offset - start_offset), :] = x[start_offset:end_offset, :] out = self.quant_method.apply(self, part_x, gate, topk_ids_hookfunc=topk_ids_hookfunc) multi_outs = paddle.zeros([token_num_per_rank * self.attn_tp_size, x.shape[1]], dtype=x.dtype) paddle.distributed.all_gather(multi_outs, out, self.tp_group) out = multi_outs[:token_num, :] return out def forward(self, x: paddle.Tensor, gate: nn.Layer, forward_meta: ForwardMeta = None): """ Defines the forward computation of the moe layer. Args: x (Tensor): Input tensor to the moe layer. Returns: Tensor: Output tensor.s """ topk_ids_hookfunc = None if self.enable_routing_replay: if forward_meta is not None: # forward_meta is None when execute empty_input_forward topk_ids_hookfunc = partial( save_routing_to_buffer, routing_replay_table=forward_meta.routing_replay_table, batch_id_per_token=forward_meta.batch_id_per_token, seq_lens_decoder=forward_meta.seq_lens_decoder, cu_seqlens_q=forward_meta.cu_seqlens_q, layer_idx=self.layer_idx, tp_size=self.fd_config.parallel_config.tensor_parallel_size, ep_size=self.fd_config.parallel_config.expert_parallel_size, tp_group=self.fd_config.parallel_config.tp_group, ) token_num = x.shape[0] if ( self.ep_size > 1 and self.attn_tp_size > 1 and (not self.fd_config.parallel_config.use_sequence_parallel_moe) and token_num >= self.attn_tp_size ): out = self.forward_split_allgather(x, gate, topk_ids_hookfunc=topk_ids_hookfunc) elif self.fd_config.parallel_config.use_ep and self.fd_config.parallel_config.enable_chunked_moe: out = self.forward_chunked_moe( x, gate, forward_meta, topk_ids_hookfunc=topk_ids_hookfunc, ) else: out = self.forward_normal(x, gate, forward_meta, topk_ids_hookfunc=topk_ids_hookfunc) if self.reduce_results and self.tp_size > 1: if current_platform.is_intel_hpu(): tensor_model_parallel_all_reduce_custom(out) else: out = tensor_model_parallel_all_reduce(out, self.tp_group) return out def forward_chunked_moe( self, x: paddle.Tensor, gate: nn.Layer, forward_meta: ForwardMeta, topk_ids_hookfunc: Callable = None ): """ Split input to multi chunk to reduce the memory usage of moe. Args: x (Tensor): Input tensor to the moe layer. Returns: Tensor: Output tensor.s """ chunk_size = self.fd_config.parallel_config.chunked_moe_size token_num = x.shape[0] fake_x = paddle.empty( shape=[0, self.fd_config.model_config.hidden_size], dtype=paddle.get_default_dtype(), ) # input size that are less than a chunk, less than the max size data or empty input # need to be repeated until the max chunk data infer MOE finished. if token_num > chunk_size: # chunked moe x_split_list = paddle.tensor_split(x, forward_meta.moe_num_chunk, axis=0) out_split_list = [None] * forward_meta.moe_num_chunk for i in range(forward_meta.max_moe_num_chunk): if i < forward_meta.moe_num_chunk: out_split_list[i] = self.quant_method.apply( self, x_split_list[i], gate, topk_ids_hookfunc=topk_ids_hookfunc ) else: # just need to use real data to infer max_moe_num_chunk times. self.quant_method.apply(self, fake_x, gate, topk_ids_hookfunc=topk_ids_hookfunc) out = paddle.concat(out_split_list, axis=0) else: # when only one chunk, just need to use real data to infer once. out = self.quant_method.apply(self, x, gate, topk_ids_hookfunc=topk_ids_hookfunc) for i in range(forward_meta.max_moe_num_chunk - 1): self.quant_method.apply(self, fake_x, gate, topk_ids_hookfunc=topk_ids_hookfunc) return out def forward_normal( self, x: paddle.Tensor, gate: nn.Layer, forward_meta: ForwardMeta, topk_ids_hookfunc: Callable = None ): """ Normal mode of forward. Args: x (Tensor): Input tensor to the moe layer. Returns: Tensor: Output tensor.s """ out = self.quant_method.apply(self, x, gate, topk_ids_hookfunc=topk_ids_hookfunc) return out