Architecture Improvement: Feedforward network part

[LoserCheems]

Shared Expert Isolation of DeepeSeek has show that:

• Improving the granularity of routing experts can improve the effectiveness
• Having a shared expert with resident activation reduces redundancy

However, the routing of classical loop traversal strategy will greatly increase the invalid runtime due to the increase of granularity.

The parameter efficient expert composed of Embedding can greatly increase the expert granularity without causing too many invalid computations.

CDMoE combines the characteristics of both shared experts and parametrically efficient experts, but currently it has not been verified for large-scale tensor parallelism, perhaps we need to develop an Embedding parallel strategy.

paper: https://arxiv.org/pdf/2412.11834

[LoserCheems]

Hi guys, after my tests, I found that the multi-head combination version of CDMoE doesn't seem to be suitable for very large scale parameters, although it is quite effect.

So here I provide a version of CDMoE that uses matrix multiplication, which can reduce the memory footprint by ~ 50% for large-scale parameter training, and the variable names are changed to be consistent with deepseek, which is easier to understand.

class DogeCDMoE(DogeMLP):
    def __init__(self, config: DogeConfig):
        super().__init__(config)
        self.hidden_dim = config.hidden_size
        self.act_fn = ACT2FN[config.hidden_act]

        self.expert_retrieval_dim = config.expert_retrieval_size
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_tok
        self.num_keys = int(math.sqrt(self.num_experts))

        # queries and keys for retrieval experts
        self.queries_proj = nn.Linear(self.hidden_dim, self.expert_retrieval_dim, bias=False)
        self.keys = nn.Parameter(torch.zeros(2, self.expert_retrieval_dim // 2, self.num_keys))

        # experts
        self.down_embed = nn.Embedding(self.num_experts, self.hidden_dim)
        self.up_embed = nn.Embedding(self.num_experts, self.hidden_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        bsz, seq_len, _ = hidden_states.shape

        # get routing weights with queries and keys
        queries = self.queries_proj(hidden_states).view(2, bsz * seq_len, -1)
        routing_weights = torch.matmul(queries, self.keys)

        # get experts with the highest routing weights
        (scores_x, scores_y), (indices_x, indices_y) = routing_weights.topk(self.top_k, dim=-1)
        all_scores = scores_x.unsqueeze(-1) + scores_y.unsqueeze(-2)
        all_scores = all_scores.view(*scores_x.shape[:-1], -1)
        all_indices = (indices_x.unsqueeze(-1) * self.num_keys) + indices_y.unsqueeze(-2)
        all_indices = all_indices.view(*indices_x.shape[:-1], -1)
        scores, pk_indices = all_scores.topk(self.top_k, dim=-1)
        indices = all_indices.gather(-1, pk_indices)
        down_embed = self.down_embed(indices).transpose(1, 2)
        up_embed = self.up_embed(indices)

        # mix experts states with cross domain states
        experts_weights = torch.matmul(hidden_states.view(bsz * seq_len, 1, -1), down_embed).view(bsz * seq_len, -1)
        experts_weights = self.act_fn(experts_weights) * scores.softmax(dim=-1)
        experts_states = torch.matmul(experts_weights.view(bsz * seq_len, 1, -1), up_embed).view(bsz, seq_len, -1)
        hidden_states = self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
        hidden_states = hidden_states + experts_states
        return hidden_states