import re import torch import torch.nn as nn import torch.distributed as dist from diffulex.layer.linear import ReplicatedLinear, divide, tp_all_reduce from diffulex.moe.layer.base import FusedMoE from diffulex.utils.checkpoint import LoadContext, ResolvedWeight from diffulex.distributed.parallel_state import fetch_parallel_state class TPFusedMoE(FusedMoE): """ Standard TP MoE without token dispatch. Every TP rank receives the full local token batch and owns a contiguous shard of experts. The router still selects global expert ids; each rank computes only the selected experts it owns, then TP all-reduce sums the partial expert contributions. """ def __init__( self, hidden_size: int, intermediate_size: int, num_experts: int, top_k: int, *, hidden_act: str = "silu", norm_topk_prob: bool = True, moe_gemm_impl: str = "triton", num_shared_experts: int = 0, shared_expert_intermediate_size: int | None = None, ) -> None: super().__init__( hidden_size, intermediate_size, num_experts, top_k, hidden_act=hidden_act, norm_topk_prob=norm_topk_prob, moe_gemm_impl=moe_gemm_impl, num_shared_experts=num_shared_experts, shared_expert_intermediate_size=shared_expert_intermediate_size, ) parallel_state = fetch_parallel_state() self.tp_rank = parallel_state.get_tp_rank() self.tp_size = parallel_state.get_tp_world_size() self.tp_group = parallel_state.get_tp_group() self.num_local_experts = divide(self.num_experts, self.tp_size) self.local_expert_start = self.tp_rank * self.num_local_experts self.local_expert_end = self.local_expert_start + self.num_local_experts self.gate = ReplicatedLinear(hidden_size, self.num_experts, bias=False) self.w13 = nn.Parameter( torch.empty(self.num_local_experts, self.intermediate_size * 2, hidden_size) ) self.w2 = nn.Parameter( torch.empty(self.num_local_experts, hidden_size, self.intermediate_size) ) def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: original_shape = hidden_states.shape flat_hidden_states = hidden_states.reshape(-1, original_shape[-1]) router_logits = self.gate(flat_hidden_states) topk_output = self.router(router_logits) topk_weights = topk_output.weights topk_ids = topk_output.ids final_hidden_states = self.expert_gemm( impl=self.moe_gemm_impl, hidden_states=flat_hidden_states, w13=self.w13, w2=self.w2, topk_ids=topk_ids, topk_weights=topk_weights, local_expert_start=self.local_expert_start, hidden_act=self.hidden_act, ) if self.tp_size > 1: final_hidden_states = tp_all_reduce(final_hidden_states, self.tp_group) final_hidden_states = final_hidden_states.reshape(original_shape) final_hidden_states = self.add_shared_experts(final_hidden_states, hidden_states) return final_hidden_states, router_logits def _local_expert_idx(self, expert_idx: int) -> int | None: if expert_idx < self.local_expert_start or expert_idx >= self.local_expert_end: return None return expert_idx - self.local_expert_start def load_w1(self, loaded_weight: torch.Tensor, expert_idx: int) -> None: local_expert_idx = self._local_expert_idx(expert_idx) if local_expert_idx is None: return # loaded_weight: [intermediate_size, hidden_size] self.w13.data[local_expert_idx, 0:self.intermediate_size].copy_(loaded_weight) def load_w3(self, loaded_weight: torch.Tensor, expert_idx: int) -> None: local_expert_idx = self._local_expert_idx(expert_idx) if local_expert_idx is None: return # loaded_weight: [intermediate_size, hidden_size] self.w13.data[local_expert_idx, self.intermediate_size:2*self.intermediate_size].copy_(loaded_weight) def load_w2(self, loaded_weight: torch.Tensor, expert_idx: int) -> None: local_expert_idx = self._local_expert_idx(expert_idx) if local_expert_idx is None: return target = self.w2.data[local_expert_idx] # [hidden_size, intermediate_size] if loaded_weight.shape == (self.hidden_size, self.intermediate_size): shard = loaded_weight elif loaded_weight.shape == (self.intermediate_size, self.hidden_size): shard = loaded_weight.transpose(0, 1).contiguous() else: raise ValueError( f"Unexpected down_proj weight shape: {loaded_weight.shape}, " f"target shape: {target.shape}" ) target.copy_(shard) def resolve_checkpoint_weight(self, suffix: str, ctx: LoadContext) -> ResolvedWeight | None: # Stacked format: experts.gate_proj.weight (all experts in one tensor) stacked_match = re.fullmatch(r"experts\.(gate_proj|up_proj|down_proj)\.weight", suffix) if stacked_match is not None: proj_name = stacked_match.group(1) return ResolvedWeight( loader=lambda loaded_weight, proj_name=proj_name: self._load_stacked_expert( loaded_weight, proj_name ) ) # Individual expert format: experts.0.gate_proj.weight match = re.fullmatch(r"experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight", suffix) if match is None: return None expert_idx = int(match.group(1)) proj_name = match.group(2) if proj_name == "gate_proj": return ResolvedWeight( loader=lambda loaded_weight, expert_idx=expert_idx: self.load_w1( loaded_weight, expert_idx, ) ) if proj_name == "up_proj": return ResolvedWeight( loader=lambda loaded_weight, expert_idx=expert_idx: self.load_w3( loaded_weight, expert_idx, ) ) if proj_name == "down_proj": return ResolvedWeight( loader=lambda loaded_weight, expert_idx=expert_idx: self.load_w2( loaded_weight, expert_idx, ) ) return None def _load_stacked_expert(self, loaded_weight: torch.Tensor, proj_name: str) -> None: """Load stacked expert weight [num_experts, ...] and slice own TP shard.""" # loaded_weight: [total_experts, ...] local_slice = loaded_weight[self.local_expert_start:self.local_expert_end] if proj_name == "gate_proj": self.w13.data[:, :self.intermediate_size].copy_(local_slice) elif proj_name == "up_proj": self.w13.data[:, self.intermediate_size:].copy_(local_slice) elif proj_name == "down_proj": # loaded_weight: [local_experts, hidden, intermediate] if local_slice.shape == self.w2.data.shape: self.w2.data.copy_(local_slice) else: # [local_experts, intermediate, hidden] → [local_experts, hidden, intermediate] self.w2.data.copy_(local_slice.transpose(1, 2).contiguous()) __all__ = ["TPFusedMoE"]