diff --git a/mindformers/pynative/transformers/moe/router.py b/mindformers/pynative/transformers/moe/router.py index 4c2fe818a4e579bd0553d30bfdc3dc1ffa7798ac..7f6b3dc716c8aa56f1fd39ffca740757422321b9 100644 --- a/mindformers/pynative/transformers/moe/router.py +++ b/mindformers/pynative/transformers/moe/router.py @@ -340,6 +340,8 @@ class TopKRouterWithLoadBalancing(nn.Cell): self.reshape = mint.reshape self.cat = mint.cat self.zeros = mint.zeros + self.stack = ops.stack + self.logical_and = ops.logical_and def update_expert_bias(self, num_tokens_per_expert: Tensor, total_tokens: int): """ @@ -383,8 +385,15 @@ class TopKRouterWithLoadBalancing(nn.Cell): def compute_load_balancing_loss(self, router_probs: Tensor, expert_indices: Tensor) -> Tensor: """ - Grouped load balancing loss: + Grouped load balancing loss per LongCat-Flash (arxiv.org/abs/2509.01322) Sec 2.1.2: L_LB = alpha * sum_{j=1}^{D+1} (f_j * P_j) + where: + P_j = (1/T) * sum_{i in Group_j} sum_t R(x_t)_i + f_j = (D/(K_e*T)) * sum_t I(token t selects Group_j) for FFN groups j=1..D + f_j = (1/((K-K_e)*T)) * sum_t I(token t selects zero-comp experts) for j=D+1 + + Note: I(token t selects Group_j) = 1 iff token t selects AT LEAST ONE expert in Group j. + The count is over TOKENS, not over (token, slot) expert selections. """ if self.aux_loss_coeff <= 0: return self.zeros((), dtype=router_probs.dtype) @@ -393,33 +402,44 @@ class TopKRouterWithLoadBalancing(nn.Cell): if bs_slen == 0: return self.zeros((), dtype=router_probs.dtype) - # 1) f_j: frequency of selection per group - selected_flat = self.reshape(expert_indices, (-1,)) - expert_counts = self.histc( - selected_flat, - bins=self.num_total_experts, - min=0, - max=self.num_total_experts, - ) - expert_counts = self.cast(expert_counts, router_probs.dtype) - + T = bs_slen + K = self.top_k + K_e = self.expected_ffn_k + D = self.num_ffn_groups experts_per_ffn_group = self.num_ffn_experts // self.num_ffn_groups - ffn_counts = expert_counts[: self.num_ffn_experts] - ffn_counts_reshaped = self.reshape( - ffn_counts, (self.num_ffn_groups, experts_per_ffn_group) - ) - group_counts_ffn = self.sum(ffn_counts_reshaped, dim=1) + + # 1) f_j: frequency of TOKENS selecting each group (paper Eq. 5) + # f_j = (D/(K_e*T)) * sum_t I(token t selects Group_j) for FFN groups + # I(token t selects Group_j) = 1 iff token t selects AT LEAST ONE expert in Group j + # NOT the count of expert selections (which would overcount by K per token) + f_j_ffn_list = [] + for j in range(D): + low = j * experts_per_ffn_group + high = (j + 1) * experts_per_ffn_group + in_group = self.logical_and( + expert_indices >= low, + expert_indices < high, + ) + in_group_count = self.sum(self.cast(in_group, router_probs.dtype), dim=1) + token_selects_group_j = in_group_count > 0 + num_tokens_selecting_group = self.sum(self.cast(token_selects_group_j, router_probs.dtype)) + coeff = D / max(K_e * T, 1e-8) + f_j_ffn_list.append(self.mul(num_tokens_selecting_group, coeff)) + + f_j_ffn = self.stack(f_j_ffn_list) if self.num_copy_experts > 0: - copy_counts = expert_counts[self.num_ffn_experts :] - group_count_copy = self.sum(copy_counts).unsqueeze(0) - all_group_counts = self.cat((group_counts_ffn, group_count_copy)) + in_copy = expert_indices >= self.num_ffn_experts + in_copy_count = self.sum(self.cast(in_copy, router_probs.dtype), dim=1) + token_selects_copy = in_copy_count > 0 + num_tokens_selecting_copy = self.sum(self.cast(token_selects_copy, router_probs.dtype)) + coeff_copy = 1.0 / max((K - K_e) * T, 1e-8) + f_j_copy = self.mul(num_tokens_selecting_copy, coeff_copy).unsqueeze(0) + f_j = self.cat((f_j_ffn, f_j_copy)) else: - all_group_counts = group_counts_ffn - - f_j = all_group_counts / (bs_slen * self.top_k) + f_j = f_j_ffn - # 2) P_j: probability mass per group + # 2) P_j: probability mass per group (paper Eq. 4) expert_prob_sum = self.sum(router_probs, dim=0) ffn_probs = expert_prob_sum[: self.num_ffn_experts] ffn_probs_reshaped = self.reshape( @@ -434,10 +454,10 @@ class TopKRouterWithLoadBalancing(nn.Cell): else: all_group_probs = group_probs_ffn - P_j = all_group_probs / bs_slen + P_j = self.div(all_group_probs, T) - num_groups = self.num_ffn_groups + (1 if self.num_copy_experts > 0 else 0) - loss = self.aux_loss_coeff * num_groups * self.sum(f_j * P_j) + # L_LB = alpha * sum_j (f_j * P_j) (paper Eq. 3) + loss = self.aux_loss_coeff * self.sum(self.mul(f_j, P_j)) return loss def construct( diff --git a/mindformers/pynative/transformers/nsa.py b/mindformers/pynative/transformers/nsa.py index f302d7db105acedc2aa0d0c32c12364545b71b2f..9ba9d0282b9fe6d22756d89ea6fceda4337b3e81 100644 --- a/mindformers/pynative/transformers/nsa.py +++ b/mindformers/pynative/transformers/nsa.py @@ -297,9 +297,13 @@ class NSABlockIndexer(nn.Cell): def _bool_to_score_mask(mask: Tensor, dtype: ms.dtype) -> Tensor: - """Convert boolean mask to additive score mask.""" - # Use a large negative value to avoid NaNs when all positions are masked. - neg_inf = mint.full(mask.shape, float(-1e9), dtype=dtype) + """Convert boolean mask to additive score mask. + + Uses float('-inf') for masked positions, consistent with Triton's + parallel_nsa implementation. LSE-based softmax handles -inf correctly + (exp(-inf)=0) while preserving full gradient flow for valid positions. + """ + neg_inf = mint.full(mask.shape, float('-inf'), dtype=dtype) zero = mint.zeros_like(neg_inf) return mint.where(mask, zero, neg_inf) @@ -320,16 +324,24 @@ def _normalize_attention_mask(attention_mask: Tensor, dtype: ms.dtype) -> Tensor def safe_softmax(x: Tensor, dim: int = -1) -> Tensor: """ - Numerically stable softmax that prevents overflow/underflow. - - Uses the identity: softmax(x) = softmax(x - max(x)) - This prevents exp() overflow while maintaining mathematical equivalence. + Numerically stable softmax using full LSE (Log-Sum-Exp) trick. + + Uses softmax(x) = exp(x - max(x)) / sum(exp(x - max(x))), which: + - Prevents exp() overflow (subtracting max keeps values <= 0) + - Handles -inf mask correctly (exp(-inf)=0, masked positions get 0 weight) + - Preserves complete gradient information (no clamping truncation) + - Handles all-masked rows (replace -inf max with 0 to avoid NaN from -inf-(-inf)) + + Reference: Triton parallel_nsa kernels in native-sparse-attention. """ x_max = mint.max(x, dim=dim, keepdim=True)[0] - x_shifted = x - x_max + # When all positions are masked (-inf), x_max=-inf causes x_shifted = x-(-inf) = NaN. + # Replace -inf max with 0 so x_shifted = -inf for masked, exp(-inf)=0, sum=0 -> uniform. + is_neg_inf = ops.logical_and(x_max.isinf(), x_max < 0) + x_max_safe = mint.where(is_neg_inf, mint.zeros_like(x_max), x_max) + x_shifted = x - x_max_safe exp_x = mint.exp(x_shifted) sum_exp = mint.sum(exp_x, dim=dim, keepdim=True) - # Avoid division by zero sum_exp = mint.where(sum_exp == 0.0, mint.ones_like(sum_exp), sum_exp) return exp_x / sum_exp @@ -566,9 +578,8 @@ class NSAAttention(nn.Cell): scores = scores * self.softmax_scale if mask is not None: scores = scores + mask - # Clip scores to prevent overflow/underflow in softmax - # This is critical for training stability - scores = mint.clamp(scores, min=-50.0, max=50.0) + # Full LSE softmax: no clamping. safe_softmax uses x - max(x) for numerical + # stability, preserving gradient flow while avoiding overflow/underflow. attn = safe_softmax(scores, dim=-1) # Replace NaN values with uniform distribution if they occur # This prevents cascading NaN issues during training