Update modeling_motif.py

modeling_motif.py

@@ -464,263 +464,6 @@ class MorehMoeFusedMLP(nn.Module):
         return output
 
 
-class MoEGate(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.top_k = config.num_experts_per_tok
-        self.n_routed_experts = config.n_routed_experts
-        self.routed_scaling_factor = config.routed_scaling_factor
-        self.scoring_func = config.scoring_func
-        self.seq_aux = config.seq_aux
-        self.topk_method = config.topk_method
-        self.n_group = config.n_group
-        self.topk_group = config.topk_group
-
-        # topk selection algorithm
-        self.norm_topk_prob = config.norm_topk_prob
-        self.gating_dim = config.hidden_size
-        self.weight = nn.Parameter(
-            torch.empty((self.n_routed_experts, self.gating_dim)))
-        if self.topk_method == "noaux_tc":
-            self.e_score_correction_bias = nn.Parameter(
-                torch.empty((self.n_routed_experts)))
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        import torch.nn.init as init
-
-        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-
-    def forward(self, hidden_states):
-        bsz, seq_len, h = hidden_states.shape
-        ### compute gating score
-        hidden_states = hidden_states.view(-1, h)
-        logits = F.linear(hidden_states.type(torch.float32),
-                          self.weight.type(torch.float32), None)
-        if self.scoring_func == "sigmoid":
-            scores = logits.sigmoid()
-        else:
-            raise NotImplementedError(
-                f"insupportable scoring function for MoE gating: {self.scoring_func}"
-            )
-
-        ### select top-k experts
-        if self.topk_method == "greedy":
-            topk_weight, topk_idx = torch.topk(scores,
-                                               k=self.top_k,
-                                               dim=-1,
-                                               sorted=False)
-        elif self.topk_method == "group_limited_greedy":
-            group_scores = (scores.view(bsz * seq_len, self.n_group,
-                                        -1).max(dim=-1).values)  # [n, n_group]
-            group_idx = torch.topk(group_scores,
-                                   k=self.topk_group,
-                                   dim=-1,
-                                   sorted=False)[1]  # [n, top_k_group]
-            group_mask = torch.zeros_like(group_scores)  # [n, n_group]
-            group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
-            score_mask = (group_mask.unsqueeze(-1).expand(
-                bsz * seq_len, self.n_group,
-                self.n_routed_experts // self.n_group).reshape(
-                    bsz * seq_len, -1))  # [n, e]
-            tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0)  # [n, e]
-            topk_weight, topk_idx = torch.topk(tmp_scores,
-                                               k=self.top_k,
-                                               dim=-1,
-                                               sorted=False)
-        elif self.topk_method == "noaux_tc":
-            ###  will be used. ###
-            scores_for_choice = scores.view(
-                bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)
-            group_scores = (scores_for_choice.view(
-                bsz * seq_len, self.n_group,
-                -1).topk(2, dim=-1)[0].sum(dim=-1))  # [n, n_group]
-            group_idx = torch.topk(group_scores,
-                                   k=self.topk_group,
-                                   dim=-1,
-                                   sorted=False)[1]  # [n, top_k_group]
-            group_mask = torch.zeros_like(group_scores)  # [n, n_group]
-            group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
-            score_mask = (group_mask.unsqueeze(-1).expand(
-                bsz * seq_len, self.n_group,
-                self.n_routed_experts // self.n_group).reshape(
-                    bsz * seq_len, -1))  # [n, e]
-            tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(),
-                                                       0.0)  # [n, e]
-            _, topk_idx = torch.topk(tmp_scores,
-                                     k=self.top_k,
-                                     dim=-1,
-                                     sorted=False)
-            topk_weight = scores.gather(1, topk_idx)
-        else:
-            raise NotImplementedError(
-                f"insupportable TopK function for MoE gating: {self.topk_method}"
-            )
-
-        ### norm gate to sum 1
-        if self.top_k > 1 and self.norm_topk_prob:
-            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
-            topk_weight = topk_weight / denominator
-        topk_weight = topk_weight * self.routed_scaling_factor  # must multiply the scaling factor
-
-        return topk_idx, topk_weight
-
-
-class MotifMoE(nn.Module):
-    """
-    A mixed expert module containing shared experts.
-    """
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.use_moreh_moe = config.use_moreh_moe
-        self.use_fused_mlp = config.use_fused_mlp
-
-        if hasattr(config, "ep_size") and config.ep_size > 1:
-            assert config.ep_size == dist.get_world_size()
-            assert not config.use_moreh_moe
-            self.ep_size = config.ep_size
-            self.experts_per_rank = config.n_routed_experts // config.ep_size
-            self.ep_rank = dist.get_rank()
-            self.experts = nn.ModuleList([
-                (DeepseekV3MLP(config,
-                               intermediate_size=config.moe_intermediate_size)
-                 if i >= self.ep_rank * self.experts_per_rank and i <
-                 (self.ep_rank + 1) * self.experts_per_rank else None)
-                for i in range(config.n_routed_experts)
-            ])
-        else:
-            self.ep_size = 1
-            self.experts_per_rank = config.n_routed_experts
-            self.ep_rank = 0
-            if self.use_moreh_moe:
-                if not self.use_fused_mlp:
-                    self.experts = MorehMoeMLP(
-                        ffn_dim=config.moe_intermediate_size,
-                        hidden_dim=config.hidden_size,
-                        hidden_act_moe=config.hidden_act_moe,
-                        num_experts=config.n_routed_experts,
-                        device=None)
-                else:
-                    ## group expert.
-                    self.experts = MorehMoeFusedMLP(
-                        ffn_dim=config.moe_intermediate_size,
-                        hidden_dim=config.hidden_size,
-                        hidden_act_moe=config.hidden_act_moe,
-                        num_experts=config.n_routed_experts,
-                        num_groups=config.n_group,
-                        device=None,
-                        continual_training=config.continual_training,
-                        )
-            else:
-                self.experts = nn.ModuleList([
-                    DeepseekV3MLP(
-                        config, intermediate_size=config.moe_intermediate_size)
-                    for i in range(config.n_routed_experts)
-                ])
-
-        self.gate = MoEGate(config)
-
-    def forward(self, hidden_states):
-        identity = hidden_states
-        orig_shape = hidden_states.shape
-        topk_idx, topk_weight = self.gate(hidden_states)
-        if self.use_moreh_moe:
-            y = self.experts(hidden_states, topk_idx.view(*orig_shape[:-1], -1),
-                             topk_weight.view(*orig_shape[:-1], -1))
-            y = y.type(hidden_states.dtype)
-        else:
-            hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-            flat_topk_idx = topk_idx.view(-1)
-            if self.training:
-                hidden_states = hidden_states.repeat_interleave(
-                    self.num_experts_per_tok, dim=0)
-                y = torch.empty_like(hidden_states)
-                for i, expert in enumerate(self.experts):
-                    y[flat_topk_idx == i] = expert(
-                        hidden_states[flat_topk_idx == i])
-                y = (y.view(*topk_weight.shape, -1) *
-                     topk_weight.unsqueeze(-1)).sum(dim=1)
-                y = y.type(hidden_states.dtype)
-                y = y.view(*orig_shape)
-                # y = AddAuxiliaryLoss.apply(y, aux_loss)
-            else:
-                y = self.moe_infer(hidden_states, topk_idx,
-                                   topk_weight).view(*orig_shape)
-        return y, identity
-
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens = x[idxs // topk_ids.shape[1]]
-        sorted_tokens_shape = sorted_tokens.shape
-        if self.ep_size > 1:
-            tokens_per_ep_rank = tokens_per_expert.view(self.ep_size,
-                                                        -1).sum(dim=1)
-            tokens_per_expert_group = tokens_per_expert.new_empty(
-                tokens_per_expert.shape[0])
-            dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)
-            output_splits = (tokens_per_expert_group.view(
-                self.ep_size, -1).sum(1).cpu().numpy().tolist())
-            gathered_tokens = sorted_tokens.new_empty(
-                tokens_per_expert_group.sum(dim=0).cpu().item(),
-                sorted_tokens.shape[1])
-            input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()
-            dist.all_to_all(
-                list(gathered_tokens.split(output_splits)),
-                list(sorted_tokens.split(input_split_sizes)),
-            )
-            tokens_per_expert_post_gather = tokens_per_expert_group.view(
-                self.ep_size, self.experts_per_rank).sum(dim=0)
-            gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],),
-                                    dtype=np.int32)
-            s = 0
-            for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
-                gatherd_idxs[s:s + k] = i % self.experts_per_rank
-                s += k
-            gatherd_idxs = gatherd_idxs.argsort()
-            sorted_tokens = gathered_tokens[gatherd_idxs]
-            tokens_per_expert = tokens_per_expert_post_gather
-        tokens_per_expert = tokens_per_expert.cpu().numpy()
-
-        outputs = []
-        start_idx = 0
-        for i, num_tokens in enumerate(tokens_per_expert):
-            end_idx = start_idx + num_tokens
-            if num_tokens == 0:
-                continue
-            expert = self.experts[i + self.ep_rank * self.experts_per_rank]
-            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
-            expert_out = expert(tokens_for_this_expert)
-            outputs.append(expert_out)
-            start_idx = end_idx
-
-        outs = torch.cat(outputs,
-                         dim=0) if len(outputs) else sorted_tokens.new_empty(0)
-        if self.ep_size > 1:
-            new_x = torch.empty_like(outs)
-            new_x[gatherd_idxs] = outs
-            gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
-            dist.all_to_all(
-                list(gathered_tokens.split(input_split_sizes)),
-                list(new_x.split(output_splits)),
-            )
-            outs = gathered_tokens
-
-        new_x = torch.empty_like(outs)
-        new_x[idxs] = outs
-        final_out = (new_x.view(
-            *topk_ids.shape, -1).type(topk_weight.dtype).mul_(
-                topk_weight.unsqueeze(dim=-1)).sum(dim=1).type(new_x.dtype))
-        return final_out
-
-
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: