wehub-resource-sync
591 lines
25 KiB
Python
591 lines
25 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
# Copyright (c) Microsoft Corporation.
|
|
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
# Copyright (C) 2024 THL A29 Limited, a Tencent company. 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 __future__ import annotations
|
|
|
|
from typing import Tuple
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
import paddle.nn as nn
|
|
import paddle.nn.functional as F
|
|
|
|
from ..utils.log import logger
|
|
|
|
|
|
class MoEGateMixin:
|
|
def gate_score_func(self, logits: paddle.Tensor) -> paddle.Tensor:
|
|
# [..., hidden_dim] -> [..., num_experts]
|
|
with paddle.amp.auto_cast(False):
|
|
scoring_func = getattr(self, "scoring_func", None)
|
|
if scoring_func == "softmax":
|
|
scores = F.softmax(logits.cast("float32"), axis=-1)
|
|
elif scoring_func == "sigmoid":
|
|
scores = F.sigmoid(logits.cast("float32"))
|
|
elif scoring_func == "tanh":
|
|
scores = F.tanh(logits.cast("float32"))
|
|
elif scoring_func == "relu":
|
|
scores = F.relu(logits.cast("float32"))
|
|
elif scoring_func == "gelu":
|
|
scores = F.gelu(logits.cast("float32"))
|
|
elif scoring_func == "leaky_relu":
|
|
scores = F.leaky_relu(logits.cast("float32"))
|
|
else:
|
|
logger.warning_once(
|
|
f"insupportable scoring function for MoE gating: {scoring_func}, use softmax instead"
|
|
)
|
|
scores = F.softmax(logits.cast("float32"), axis=-1)
|
|
return scores
|
|
|
|
def gumbel_rsample(self, logits: paddle.Tensor) -> paddle.Tensor:
|
|
gumbel = paddle.distribution.gumbel.Gumbel(0, 1)
|
|
return gumbel.rsample(logits.shape)
|
|
|
|
def uniform_sample(self, logits: paddle.Tensor) -> paddle.Tensor:
|
|
uniform = paddle.distribution.uniform.Uniform(0, 1)
|
|
return uniform.sample(logits.shape)
|
|
|
|
@paddle.no_grad()
|
|
def _one_hot_to_float(self, x, num_classes):
|
|
if x.dtype not in (paddle.int32, paddle.int64):
|
|
x = paddle.cast(x, paddle.int64)
|
|
return F.one_hot(x, num_classes=num_classes).cast(paddle.get_default_dtype())
|
|
|
|
@paddle.no_grad()
|
|
def _one_hot_to_int64(self, x, num_classes):
|
|
if x.dtype not in (paddle.int32, paddle.int64):
|
|
x = paddle.cast(x, paddle.int64)
|
|
return F.one_hot(x, num_classes=num_classes).cast(paddle.int64)
|
|
|
|
@paddle.no_grad()
|
|
def _capacity(
|
|
self,
|
|
gates: paddle.Tensor,
|
|
capacity_factor: float,
|
|
max_capacity: int,
|
|
min_capacity: int,
|
|
) -> paddle.Tensor:
|
|
"""Calculate the capacity for each expert based on the gates and capacity factor.
|
|
|
|
Args:
|
|
gates (paddle.Tensor): A tensor of shape [num_tokens, num_experts] representing the probability distribution
|
|
over experts for each token.
|
|
capacity_factor (float): A scalar float value representing the capacity factor for each expert.
|
|
min_capacity (int): A scalar integer value representing the minimum capacity for each expert.
|
|
|
|
Returns:
|
|
int: A tensor value representing the calculated capacity for each expert.
|
|
"""
|
|
assert gates.ndim == 2, f"gates should be 2D, but got {gates.ndim}, {gates.shape}"
|
|
# gates has shape of SE
|
|
num_tokens = gates.shape[0]
|
|
num_experts = gates.shape[1]
|
|
capacity = int((num_tokens // num_experts) * capacity_factor)
|
|
if capacity < min_capacity:
|
|
capacity = min_capacity
|
|
if capacity > max_capacity:
|
|
capacity = max_capacity
|
|
assert capacity > 0, f"requires capacity > 0, capacity_factor: {capacity_factor}, input_shape: {gates.shape}"
|
|
|
|
return capacity
|
|
|
|
def _cal_aux_loss(self, gates, mask):
|
|
"""
|
|
Calculate auxiliary loss
|
|
|
|
Args:
|
|
gates (paddle.Tensor): Represents the output probability of each expert. The shape is [batch_size, num_experts]
|
|
mask (paddle.Tensor): Represents whether each sample belongs to a certain expert. The shape is [batch_size, num_experts]
|
|
|
|
Returns:
|
|
paddle.Tensor: The value of auxiliary loss.
|
|
|
|
"""
|
|
# TODO: @DrownFish19 update aux_loss for Qwen2MoE and DeepSeekV2&V3
|
|
me = paddle.mean(gates, axis=0)
|
|
ce = paddle.mean(mask.cast("float32"), axis=0)
|
|
if self.global_aux_loss:
|
|
me_list, ce_list = [], []
|
|
dist.all_gather(me_list, me, group=self.group)
|
|
dist.all_gather(ce_list, ce, group=self.group)
|
|
|
|
me_list[self.rank] = me
|
|
ce_list[self.rank] = ce
|
|
me = paddle.stack(me_list).mean(0)
|
|
ce = paddle.stack(ce_list).mean(0)
|
|
aux_loss = paddle.sum(me * ce) * float(self.num_experts)
|
|
return aux_loss
|
|
|
|
def _cal_seq_aux_loss(self, gates, top_k, topk_idx) -> paddle.Tensor:
|
|
"""
|
|
Calculate sequence auxiliary loss.
|
|
|
|
Args:
|
|
logits (paddle.Tensor): Model output.
|
|
|
|
Returns:
|
|
paddle.Tensor: The value of sequence auxiliary loss.
|
|
"""
|
|
batch_size, seq_len, _ = gates.shape
|
|
ce = paddle.zeros([batch_size, self.num_experts])
|
|
topk_idx = topk_idx.reshape([batch_size, -1])
|
|
ce.put_along_axis_(indices=topk_idx, values=paddle.ones([batch_size, seq_len * top_k]), axis=1, reduce="add")
|
|
ce = ce / (seq_len * top_k / self.num_experts)
|
|
aux_loss = (ce * paddle.mean(gates, axis=1)).sum(axis=1).mean()
|
|
return aux_loss
|
|
|
|
def _cal_z_loss(self, logits) -> paddle.Tensor:
|
|
"""
|
|
Calculate the z loss.
|
|
|
|
Args:
|
|
logits (paddle.Tensor): Model output. The shape is [batch_size, num_experts].
|
|
|
|
Returns:
|
|
paddle.Tensor: The z loss value.
|
|
"""
|
|
l_zloss = paddle.logsumexp(logits, axis=1).square().mean()
|
|
return l_zloss
|
|
|
|
def _cal_orthogonal_loss(self) -> paddle.Tensor:
|
|
"""Gate weight orthogonal loss.
|
|
|
|
Returns:
|
|
Paddle.Tensor: orthogonal loss
|
|
"""
|
|
weight = F.normalize(self.weight, axis=0)
|
|
orthogonal_loss = paddle.mean(paddle.square(paddle.matmul(weight.T, weight) - paddle.eye(self.num_experts)))
|
|
return orthogonal_loss
|
|
|
|
|
|
class PretrainedMoEGate(nn.Layer, MoEGateMixin):
|
|
def __init__(self, config, num_experts, expert_hidden_size, **kwargs):
|
|
super(PretrainedMoEGate, self).__init__()
|
|
|
|
self.config = config
|
|
|
|
self.num_experts = num_experts
|
|
self.expert_hidden_size = expert_hidden_size
|
|
|
|
# force keep in float32 when using amp
|
|
self._cast_to_low_precision = False
|
|
|
|
self.capacity_factor = kwargs.pop("capacity_factor", 1.0)
|
|
self.eval_capacity_factor = kwargs.pop("eval_capacity_factor", 1.0)
|
|
self.min_capacity = kwargs.pop("min_capacity", 1.0)
|
|
self.max_capacity = kwargs.pop("max_capacity", pow(2, 32))
|
|
|
|
self.group = kwargs.pop("group", None)
|
|
self.global_aux_loss = kwargs.pop("global_aux_loss", False)
|
|
if self.global_aux_loss:
|
|
assert self.group is not None, "group is required when global_aux_loss is True"
|
|
self.rank = dist.get_rank(self.group)
|
|
|
|
self.expert_drop = kwargs.pop("expert_drop", False)
|
|
self.noisy_gate_policy = kwargs.pop("noisy_gate_policy", None)
|
|
self.drop_tokens = kwargs.pop("drop_tokens", True)
|
|
self.use_rts = kwargs.pop("use_rts", True)
|
|
self.top2_2nd_expert_sampling = kwargs.pop("top2_2nd_expert_sampling", True)
|
|
|
|
self.drop_policy = kwargs.pop("drop_policy", "probs")
|
|
# Qwen2MoE: greedy
|
|
# DeepSeekV2&V3: group_limited_greedy for training, and noaux_tc for inference
|
|
self.topk_method = kwargs.pop("topk_method", "greedy")
|
|
self.top_k = kwargs.pop("top_k", 2)
|
|
self.n_group = kwargs.pop("n_group", 1) # for group_limited_greedy
|
|
self.topk_group = kwargs.pop("topk_group", 1) # for group_limited_greedy
|
|
self.norm_topk_prob = kwargs.pop("norm_topk_prob", False)
|
|
self.routed_scaling_factor = kwargs.pop("routed_scaling_factor", 1.0)
|
|
|
|
def _priority(self, topk_idx: paddle.Tensor, capacity: int) -> paddle.Tensor:
|
|
"""_summary_
|
|
The priority is the cumulative sum of the expert indices.
|
|
|
|
This method is used in hunyuan model
|
|
Args:
|
|
topk_idx (paddle.Tensor): [batch_size * seq_len, topk]
|
|
|
|
Returns:
|
|
paddle.Tensor: cumsum locations
|
|
"""
|
|
_, k = topk_idx.shape
|
|
# Shape: [seq_len * k]
|
|
chosen_expert = topk_idx.reshape([-1])
|
|
# Shape: [seq_len * k, num_experts].
|
|
token_priority = F.one_hot(chosen_expert, self.num_experts).cast(paddle.int32)
|
|
token_priority = paddle.logical_and(token_priority > 0, token_priority.cumsum(axis=0) <= capacity)
|
|
# Shape: [seq_len, num_experts].
|
|
token_priority = token_priority.reshape([-1, k, self.num_experts]).sum(axis=1)
|
|
|
|
return (token_priority > 0.0).astype("float32")
|
|
|
|
def _topk_greedy(self, scores: paddle.Tensor, k: int) -> Tuple[paddle.Tensor, paddle.Tensor]:
|
|
"""_summary_
|
|
|
|
Args:
|
|
scores (paddle.Tensor): [bsz*seq_len, n_experts]
|
|
k (int): select the top k experts
|
|
|
|
Returns:
|
|
Tuple[paddle.Tensor, paddle.Tensor]: topk_weight, topk_idx
|
|
topk_weight: [bsz*seq_len, k]
|
|
topk_idx: [bsz*seq_len, k]
|
|
"""
|
|
topk_weight, topk_idx = paddle.topk(scores, k=k, axis=-1, sorted=True)
|
|
return topk_weight, topk_idx
|
|
|
|
def _topk_group_limited_greedy(
|
|
self, scores: paddle.Tensor, k: int, n_group: int, topk_group: int
|
|
) -> Tuple[paddle.Tensor, paddle.Tensor]:
|
|
"""_summary_
|
|
|
|
Args:
|
|
scores (paddle.Tensor): [bsz*seq_len, n_experts]
|
|
k (int): select the top k experts in each group
|
|
n_groups (int): the number of groups for all experts
|
|
topk_group (int): the number of groups selected
|
|
|
|
Returns:
|
|
Tuple[paddle.Tensor, paddle.Tensor]: topk_weight, topk_idx
|
|
topk_weight: [bsz*seq_len, k]
|
|
topk_idx: [bsz*seq_len, k]
|
|
|
|
Note: the group size is normal greater than the number of k
|
|
"""
|
|
bsz_seq_len, n_experts = scores.shape
|
|
assert n_experts % n_group == 0, "n_experts must be divisible by n_groups"
|
|
|
|
group_scores = scores.reshape([0, n_group, -1]).max(axis=-1) # [n, n_group]
|
|
group_idx = paddle.topk(group_scores, k=topk_group, axis=-1, sorted=True)[1] # [n, top_k_group]
|
|
group_mask = paddle.zeros_like(group_scores).put_along_axis(group_idx, paddle.to_tensor(1.0), axis=-1) # fmt:skip
|
|
score_mask = (
|
|
group_mask.unsqueeze(-1).expand([bsz_seq_len, n_group, n_experts // n_group]).reshape([bsz_seq_len, -1])
|
|
) # [n, e]
|
|
tmp_scores = scores * score_mask # [n, e]
|
|
topk_weight, topk_idx = paddle.topk(tmp_scores, k=k, axis=-1, sorted=True)
|
|
|
|
return topk_weight, topk_idx
|
|
|
|
def _topk_noaux_tc(
|
|
self, scores: paddle.Tensor, k: int, n_group: int, topk_group: int
|
|
) -> Tuple[paddle.Tensor, paddle.Tensor]:
|
|
"""_summary_
|
|
|
|
Args:
|
|
scores (paddle.Tensor): [bsz*seq_len, n_experts]
|
|
k (int): select the top k experts in each group
|
|
n_groups (int): the number of groups for all experts
|
|
topk_group (int): the number of groups selected
|
|
|
|
Returns:
|
|
Tuple[paddle.Tensor, paddle.Tensor]: topk_weight, topk_idx
|
|
topk_weight: [bsz*seq_len, k]
|
|
topk_idx: [bsz*seq_len, k]
|
|
|
|
Note: the group size is normal greater than the number of k
|
|
"""
|
|
bsz_seq_len, n_experts = scores.shape
|
|
assert n_experts % n_group == 0, "n_experts must be divisible by n_groups"
|
|
|
|
assert self.e_score_correction_bias is not None, "e_score_correction_bias is None"
|
|
scores_for_choice = scores.reshape([bsz_seq_len, -1]) + self.e_score_correction_bias.unsqueeze(0)
|
|
group_scores = (
|
|
scores_for_choice.reshape([bsz_seq_len, self.n_group, -1]).topk(2, axis=-1)[0].sum(axis=-1)
|
|
) # fmt:skip [n, n_group]
|
|
group_idx = paddle.topk(group_scores, k=topk_group, axis=-1, sorted=True)[1] # [n, top_k_group]
|
|
group_mask = paddle.zeros_like(group_scores).put_along_axis(group_idx, paddle.to_tensor(1.0, dtype="float32"), axis=-1) # fmt:skip
|
|
score_mask = (
|
|
group_mask.unsqueeze(-1).expand([bsz_seq_len, n_group, n_experts // n_group]).reshape([bsz_seq_len, -1])
|
|
) # [n, e]
|
|
tmp_scores = scores_for_choice * score_mask # [n, e]
|
|
topk_weight, topk_idx = paddle.topk(tmp_scores, k=k, axis=-1, sorted=True)
|
|
topk_weight = scores.take_along_axis(topk_idx, axis=1) if not self.training else topk_weight
|
|
|
|
return topk_weight, topk_idx
|
|
|
|
def top1gating(
|
|
self,
|
|
logits: paddle.Tensor,
|
|
used_token: paddle.Tensor = None,
|
|
) -> Tuple[int, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]:
|
|
"""Implements Top1Gating on logits."""
|
|
if self.noisy_gate_policy == "RSample":
|
|
logits += self.gumbel_rsample(logits.shape)
|
|
|
|
gates = self.gate_score_func(logits=logits)
|
|
|
|
capacity = self._capacity(gates, self.capacity_factor, self.max_capacity, self.min_capacity)
|
|
|
|
# Create a mask for 1st's expert per token
|
|
# noisy gating
|
|
# Only save the position of the maximum value
|
|
indices1_s = paddle.argmax(logits if self.noisy_gate_policy == "RSample" else gates, axis=1)
|
|
# Convert the position of the maximum value to a one-hot vector [s, e]
|
|
mask1 = self._one_hot_to_float(indices1_s, num_classes=self.num_experts)
|
|
|
|
# mask only used tokens
|
|
if used_token is not None:
|
|
mask1 = paddle.einsum(
|
|
"s,se->se", used_token, mask1
|
|
) # Element-wise multiply used_token with mask1 to obtain a new mask1
|
|
|
|
# gating decisions
|
|
exp_counts = paddle.sum(mask1, axis=0) # Calculate the number of tokens for each expert
|
|
|
|
# if we don't want to drop any tokens
|
|
if not self.drop_tokens:
|
|
new_capacity = paddle.max(exp_counts) # Calculate the number of tokens for each expert
|
|
# Communicate across expert processes to pick the maximum capacity.
|
|
if self.group is not None:
|
|
dist.all_reduce(
|
|
new_capacity, op=dist.ReduceOp.MAX, group=self.group
|
|
) # Calculate the maximum value among expert processes
|
|
# Make sure the capacity value does not exceed the number of tokens.
|
|
capacity = int(min(new_capacity, paddle.tensor(mask1.size(0))))
|
|
|
|
l_aux = self._cal_aux_loss(gates, mask1)
|
|
l_zloss = self._cal_z_loss(logits)
|
|
|
|
# Random Token Selection
|
|
if self.use_rts:
|
|
mask1_rand = mask1 * self.uniform_sample(mask1)
|
|
else:
|
|
mask1_rand = mask1
|
|
|
|
assert (
|
|
logits.shape[0] >= self.min_capacity
|
|
), "No. of tokens (batch-size) should be greater than min_capacity. Either set min_capacity to 0 or increase your batch size."
|
|
|
|
_, top_idx = paddle.topk(mask1_rand, k=capacity, axis=0) # Select top_capacity tokens
|
|
|
|
new_mask1 = mask1 * paddle.zeros_like(mask1).put_along_axis(
|
|
top_idx, paddle.to_tensor(1.0, dtype="float32"), axis=0
|
|
)
|
|
mask1 = new_mask1
|
|
|
|
# Compute locations in capacity buffer
|
|
locations1 = paddle.cumsum(mask1, axis=0) - 1 # Compute the position of each token in mask1
|
|
|
|
# Store the capacity location for each token
|
|
locations1_s = paddle.sum(locations1 * mask1, axis=1).cast(paddle.int64)
|
|
|
|
# Normalize gate probabilities
|
|
mask1_float = mask1.cast(paddle.float32)
|
|
gates = gates / gates * mask1_float
|
|
|
|
locations1_sc = self._one_hot_to_float(locations1_s, capacity)
|
|
combine_weights = paddle.einsum("se,sc->sec", gates, locations1_sc)
|
|
dispatch_mask = combine_weights.cast(paddle.bool).detach()
|
|
|
|
return capacity, combine_weights, dispatch_mask, exp_counts, l_aux, l_zloss
|
|
|
|
def top2gating(
|
|
self,
|
|
logits: paddle.Tensor,
|
|
) -> Tuple[int, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]:
|
|
# everything is in fp32 in this function
|
|
|
|
gates = self.gate_score_func(logits=logits)
|
|
|
|
# Create a mask for 1st's expert per token.
|
|
indices1_s = paddle.argmax(gates, axis=1) # [S, 1]
|
|
mask1 = self._one_hot_to_int64(indices1_s, self.num_experts) # [S, E]
|
|
|
|
if self.top2_2nd_expert_sampling:
|
|
# Create a mask for 2nd's expert per token using Gumbel-max trick.
|
|
# https://timvieira.github.io/blog/post/2014/07/31/gumbel-max-trick/
|
|
logits += self.gumbel_rsample(logits)
|
|
|
|
# Replace top-expert with min value
|
|
logits_except1 = logits.masked_fill(mask1.cast(paddle.bool), float("-inf")) # [S, E]
|
|
indices2_s = paddle.argmax(logits_except1, axis=1) # [S, 1]
|
|
mask2 = self._one_hot_to_int64(indices2_s, self.num_experts) # [S, E]
|
|
|
|
# Note: mask1 and mask2 can be combined to form a single mask.
|
|
# mask = paddle.concat([mask1, mask2], axis=0)
|
|
# locations = paddle.cumsum(mask, axis=0) - 1
|
|
# locations1, locations2 = locations.split(2, axis=0)
|
|
# Compute locations in capacity buffer.
|
|
locations1 = paddle.cumsum(mask1, axis=0) - 1 # [S, E]
|
|
locations2 = paddle.cumsum(mask2, axis=0) - 1 # [S, E]
|
|
# Update 2nd's location by accounting for locations of 1st.
|
|
locations2 += paddle.sum(mask1, axis=0, keepdim=True)
|
|
|
|
l_aux = self._cal_aux_loss(gates, mask1)
|
|
l_zloss = self._cal_z_loss(logits)
|
|
|
|
# gating decisions
|
|
exp_counts = paddle.sum(mask1 + mask2, axis=0)
|
|
if self.drop_tokens:
|
|
# Calculate configured capacity and remove locations outside capacity from mask
|
|
capacity = self._capacity(gates, self.capacity_factor, self.max_capacity, self.min_capacity)
|
|
# Remove locations outside capacity from mask.
|
|
mask1 *= (locations1 < capacity).cast(paddle.int64)
|
|
mask2 *= (locations2 < capacity).cast(paddle.int64)
|
|
else:
|
|
# Do not drop tokens - set capacity according to current expert assignments
|
|
new_capacity = paddle.max(exp_counts)
|
|
if self.group is not None:
|
|
dist.all_reduce(new_capacity, op=dist.ReduceOp.MAX, group=self.group)
|
|
capacity = int(new_capacity)
|
|
|
|
# Store the capacity location for each token.
|
|
locations1_s = paddle.sum(locations1 * mask1, axis=1)
|
|
locations2_s = paddle.sum(locations2 * mask2, axis=1)
|
|
|
|
# Normalize gate probabilities
|
|
mask1_float = mask1.cast(paddle.float32)
|
|
mask2_float = mask2.cast(paddle.float32)
|
|
gates1_s = paddle.einsum("se,se->s", gates, mask1_float)
|
|
gates2_s = paddle.einsum("se,se->s", gates, mask2_float)
|
|
denom_s = gates1_s + gates2_s
|
|
# Avoid divide-by-zero
|
|
denom_s = paddle.clip(denom_s, min=paddle.finfo(denom_s.dtype).eps)
|
|
gates1_s /= denom_s
|
|
gates2_s /= denom_s
|
|
|
|
# Calculate combine_weights and dispatch_mask
|
|
gates1 = paddle.einsum("s,se->se", gates1_s, mask1_float)
|
|
gates2 = paddle.einsum("s,se->se", gates2_s, mask2_float)
|
|
locations1_sc = self._one_hot_to_float(locations1_s, capacity)
|
|
locations2_sc = self._one_hot_to_float(locations2_s, capacity)
|
|
combine1_sec = paddle.einsum("se,sc->sec", gates1, locations1_sc)
|
|
combine2_sec = paddle.einsum("se,sc->sec", gates2, locations2_sc)
|
|
combine_weights = combine1_sec + combine2_sec
|
|
dispatch_mask = combine_weights.cast(paddle.bool)
|
|
|
|
return capacity, combine_weights, dispatch_mask, exp_counts, l_aux, l_zloss
|
|
|
|
def topkgating(
|
|
self,
|
|
gates: paddle.Tensor,
|
|
) -> Tuple[int, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]:
|
|
"""Implements TopKGating on logits."""
|
|
batch_size, seq_len, d_model = gates.shape
|
|
gates_ori = gates
|
|
gates = gates.reshape([-1, d_model])
|
|
|
|
l_zloss = self._cal_z_loss(gates)
|
|
|
|
# get topk gates
|
|
if self.topk_method == "greedy":
|
|
top_gate, top_idx = self._topk_greedy(gates, k=self.top_k)
|
|
elif self.topk_method == "group_limited_greedy":
|
|
top_gate, top_idx = self._topk_group_limited_greedy(
|
|
gates, k=self.top_k, n_group=self.n_group, topk_group=self.topk_group
|
|
)
|
|
elif self.topk_method == "noaux_tc":
|
|
top_gate, top_idx = self._topk_noaux_tc(
|
|
gates, k=self.top_k, n_group=self.n_group, topk_group=self.topk_group
|
|
)
|
|
# norm gate to sum 1
|
|
if self.top_k > 1 and self.norm_topk_prob:
|
|
denominator = top_gate.sum(axis=-1, keepdim=True) + 1e-20
|
|
top_gate = top_gate / denominator
|
|
top_gate = top_gate * self.routed_scaling_factor
|
|
|
|
# get topk mask
|
|
mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.to_tensor(1.0, dtype=gates.dtype), axis=1)
|
|
if hasattr(self.config, "seq_aux") and self.config.seq_aux:
|
|
l_aux = self._cal_seq_aux_loss(gates_ori, self.top_k, top_idx)
|
|
else:
|
|
l_aux = self._cal_aux_loss(gates, mask)
|
|
|
|
exp_counts = paddle.sum(mask.cast(paddle.int64), axis=0)
|
|
|
|
if self.drop_tokens:
|
|
# Calculate configured capacity and remove locations outside capacity from mask
|
|
capacity = self._capacity(
|
|
gates,
|
|
self.capacity_factor * self.top_k,
|
|
self.max_capacity,
|
|
self.min_capacity,
|
|
)
|
|
|
|
# update mask and locations by capacity
|
|
if self.drop_policy == "probs":
|
|
topk_masked_gates = paddle.zeros_like(gates).put_along_axis(top_idx, top_gate, axis=1)
|
|
capacity_probs, capacity_indices = paddle.topk(topk_masked_gates, k=capacity, axis=0, sorted=False)
|
|
token_priority = self._priority(capacity_indices, capacity)
|
|
|
|
elif self.drop_policy == "position":
|
|
token_priority = self._priority(top_idx, capacity)
|
|
else:
|
|
raise ValueError(f"Invalid drop_policy: {self.drop_policy}")
|
|
else:
|
|
# Do not drop tokens - set capacity according to current expert assignments
|
|
local_capacity = paddle.max(exp_counts)
|
|
if self.group is not None:
|
|
dist.all_reduce(local_capacity, op=dist.ReduceOp.MAX, group=self.group)
|
|
capacity = int(local_capacity)
|
|
token_priority = self._priority(top_idx, capacity)
|
|
|
|
# normalize gates
|
|
# gates_masked is equal to top_gate.
|
|
gates_masked = gates * mask
|
|
# if self.training:
|
|
gates_s = paddle.sum(gates_masked, axis=-1, keepdim=True)
|
|
denom_s = paddle.clip(gates_s, min=paddle.finfo(gates_masked.dtype).eps)
|
|
if self.norm_topk_prob:
|
|
gates_masked = gates_masked / denom_s
|
|
gates_masked *= self.routed_scaling_factor
|
|
|
|
return (
|
|
capacity,
|
|
gates_masked.take_along_axis(top_idx, axis=-1),
|
|
top_idx,
|
|
token_priority.take_along_axis(top_idx, axis=-1),
|
|
l_aux,
|
|
l_zloss,
|
|
)
|
|
|
|
def topkgating_nodrop(self, gates: paddle.Tensor):
|
|
"""Implements TopKGating on logits."""
|
|
batch_size, seq_len, d_model = gates.shape
|
|
gates_ori = gates
|
|
gates = gates.reshape([-1, d_model])
|
|
|
|
l_zloss = self._cal_z_loss(gates)
|
|
|
|
# get topk gates
|
|
if self.topk_method == "greedy":
|
|
top_gate, top_idx = self._topk_greedy(gates, k=self.top_k)
|
|
elif self.topk_method == "group_limited_greedy":
|
|
top_gate, top_idx = self._topk_group_limited_greedy(
|
|
gates, k=self.top_k, n_group=self.n_group, topk_group=self.topk_group
|
|
)
|
|
elif self.topk_method == "noaux_tc":
|
|
top_gate, top_idx = self._topk_noaux_tc(
|
|
gates, k=self.top_k, n_group=self.n_group, topk_group=self.topk_group
|
|
)
|
|
# norm gate to sum 1
|
|
if self.top_k > 1 and self.norm_topk_prob:
|
|
denominator = top_gate.sum(axis=-1, keepdim=True) + 1e-20
|
|
top_gate = top_gate / denominator
|
|
top_gate = top_gate * self.routed_scaling_factor
|
|
|
|
# get topk mask
|
|
mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.to_tensor(1.0), axis=1)
|
|
|
|
if hasattr(self.config, "seq_aux") and self.config.seq_aux:
|
|
l_aux = self._cal_seq_aux_loss(gates_ori, self.top_k, top_idx)
|
|
else:
|
|
l_aux = self._cal_aux_loss(gates, mask)
|
|
|
|
exp_counts = paddle.sum(mask.cast(paddle.int64), axis=0)
|
|
topk_masked_gates = paddle.zeros_like(gates).put_along_axis(top_idx, top_gate, axis=1)
|
|
return topk_masked_gates, mask, exp_counts, l_aux, l_zloss
|