出发点

本来想延续之前的博客风格（从输入到输出），但是这里Deepseek-V3属实太大了，我在服务器上跑不起来了，所以这里挑出一些和之前的代码不同的进行解析，本篇从核心之一MOE出发解析loss free和topkrouter以及最后的MOE

auxiliary loss free

AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY FOR MIXTURE-OF-EXPERTS
MoE环游记：3、换个思路来分配

这里使用的MOE和之前不太一样，这里引入了一个偏置项b（训练的时候更新，预测的时候直接使用）来帮助专家系统更均衡
根据伪代码写出了它的更新函数

from torch import nn
import torch.nn.functional as F
import torch
import numpy as np
class DeepseekV3TopkRouter(nn.Module):
    def __init__(self):
        super().__init__()
        self.n_routed_experts = 256
        self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts)))
        self.lr = 0.1
    def update_bias(self, expert_counts):  # 每个batch后调用
        """
        expert_counts: [n_experts] 统计当前批次各专家的token处理数
        """
        # 计算 the load violation error
        e = expert_counts.float().mean() - expert_counts  # 
        self.e_score_correction_bias += self.lr * np.sign(e.float())  # 
        # 数值稳定性约束
        self.e_score_correction_bias.clamp_(min=-5.0, max=5.0)  # 防止梯度爆炸
    def collect_expert_counts(self, topk_indices):
	    """
	    topk_indices: [batch_size * seq_len, top_k] 每个token对应的Top-K专家索引
	    """
	    # 维度展开与类型转换
	    expert_indices = topk_indices.view(-1).long()  # [batch_size*seq_len*top_k]
	    # 统计各专家出现的总次数
	    counts = torch.bincount(expert_indices, minlength=self.n_routed_experts)
	    return counts  # shape: [n_routed_experts]

DeepseekV3TopkRouter

DeepSeek-V3 Technical Report

from torch import nn
import torch.nn.functional as F
import torch
import numpy as np
class DeepseekV3TopkRouter(nn.Module):
    def __init__(self):
        super().__init__()
        self.top_k = 8
        self.n_routed_experts = 256
        self.routed_scaling_factor = 2.5
        self.n_group = 8
        self.topk_group = 4
        self.norm_topk_prob = True
        self.weight = torch.randn(self.n_routed_experts, 768)
        self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts)))
        self.lr = 0.1# 论文里指定是0.001，这里调的比较大，方便看变化过程
    def update_bias(self, expert_counts):  # 每个batch后调用
        """
        expert_counts: [n_experts] 统计当前批次各专家的token处理数
        """
        # 计算 the load violation error
        e = expert_counts.float().mean() - expert_counts  # 
        self.e_score_correction_bias += self.lr * np.sign(e.float())  # 
        # 数值稳定性约束
        self.e_score_correction_bias.clamp_(min=-5.0, max=5.0)  # 防止梯度爆炸
    def collect_expert_counts(self, topk_indices):
	    """
	    topk_indices: [batch_size * seq_len, top_k] 每个token对应的Top-K专家索引
	    """
	    # 维度展开与类型转换
	    expert_indices = topk_indices.view(-1).long()  # [batch_size*seq_len*top_k]
	    # 统计各专家出现的总次数
	    counts = torch.bincount(expert_indices, minlength=self.n_routed_experts)
	    return counts  # shape: [n_routed_experts]
    @torch.no_grad()
    def get_topk_indices(self, scores):
        scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0)
        group_scores = (
            scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group)
            .topk(2, dim=-1)[0]
            .sum(dim=-1)
        )
        # scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) 4， 8， 32
        # .topk(2, dim=-1)[0] 4,8,2
        # .sum(dim=-1) 4,8
        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
        group_mask = torch.zeros_like(group_scores)
        group_mask.scatter_(1, group_idx, 1)
        score_mask = (
            group_mask.unsqueeze(-1).expand(-1, self.n_group, self.n_routed_experts // self.n_group)
            .reshape(-1, self.n_routed_experts)
            )
        scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)
        topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
        return topk_indices
    def forward(self, hidden_states):
        hidden_states = hidden_states.view(-1, 768)# 4, 768
        router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))# 4, 256
        scores = router_logits.sigmoid()# 4, 256 计算的是每个值的sigmoid值
        topk_indices = self.get_topk_indices(scores)
        # counts = self.collect_expert_counts(topk_indices)
        # self.update_bias(counts)
        # print (self.e_score_correction_bias[:20])
        topk_weights = scores.gather(1, topk_indices)
        if self.norm_topk_prob:
            denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
            topk_weights /= denominator
        topk_weights = topk_weights * self.routed_scaling_factor
        return topk_indices, topk_weights

流程图如下：

这里针对b做了一下测试，其他参数不变，一直调用DeepseekV3TopkRouter，可以简单看一下b的变化（会一直较均匀的增长）

hidden_states = torch.randn(4, 8, 768)
router = DeepseekV3TopkRouter()
for i in range(20):
    topk_indices, topk_weights = router(hidden_states)

DeepseekV3MoE

class DeepseekV3MoE(nn.Module):
    """
    A mixed expert module containing shared experts.
    """

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList(
            [
                DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size)
                for _ in range(config.n_routed_experts)# MLP和之前的是一样的
            ]
        )
        self.gate = DeepseekV3TopkRouter(config)# 上面已经介绍过了
        self.shared_experts = DeepseekV3MLP(
            config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
        )

    def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor):
    # 这里求专家头的方式和Qwen2MoeSparseMoeBlock是一样的
    # https://blog.csdn.net/zpp13hao1/article/details/146153335?spm=1011.2415.3001.5331
    # 不一样的地方是qwen2针对共享专家多了一个shared_expert_gate，即多了一个权重（sigmoid），这里没有了
        r"""
        CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused
        to not have to do a loop here (deepseek has 256 experts soooo yeah).
        # 我也很好奇有没有取消loop的方式
        """
        final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype)
        expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts))
        expert_mask = expert_mask.permute(2, 0, 1)

        for expert_idx in range(len(self.experts)):
            expert = self.experts[expert_idx]
            mask = expert_mask[expert_idx]
            token_indices, weight_indices = torch.where(mask)

            if token_indices.numel() > 0:
                expert_weights = topk_weights[token_indices, weight_indices]
                expert_input = hidden_states[token_indices]
                expert_output = expert(expert_input)
                weighted_output = expert_output * expert_weights.unsqueeze(-1)
                final_hidden_states.index_add_(0, token_indices, weighted_output)

        # in original deepseek, the output of the experts are gathered once we leave this module
        # thus the moe module is itelsf an IsolatedParallel module
        # and all expert are "local" meaning we shard but we don't gather
        return final_hidden_states.type(hidden_states.dtype)

    def forward(self, hidden_states):
        residuals = hidden_states
        orig_shape = hidden_states.shape
        topk_indices, topk_weights = self.gate(hidden_states)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape)
        hidden_states = hidden_states + self.shared_experts(residuals)
        return hidden_states