谈谈DeepSeek MoE模型优化和未来演进以及字节Ultra-Sparse Memory相关的工作

大概一个月前发了一篇 《一个关于MoE的猜想》 ,  最近优化DeepSeek MoE推理的时候, 发现这是一个值得去解决的问题, 从算法上, 这是DeepSeek-V3/R1 MoE的一个不足之处, 值得进一步去分析和提高效率. 后文顺便也分析了一下字节的Ultra-sparse Memory Network的工作.

1. DeepSeek-MoE Profiling

其实任何算法的效率提升, 无非就是在计算/通信/存储上取得平衡. 我们以DeepSeek MoE为例. 它的一个Expert结构如下:

import torch
from torch import nn
import torch.nn.functional as F

class Expert(nn.Module):
    def __init__(self, dim: int, inter_dim: int):
        super().__init__()
        self.w1 = nn.Linear(dim, inter_dim)
        self.w2 = nn.Linear(inter_dim, dim)
        self.w3 = nn.Linear(dim, inter_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

按照DeepSeek-R1定义的dim = 7168, moe-inter-dim=2048, 以及论文中提到的对256个Tokens Batch处理, 进行算力评估

dim = 7168
inter_dim = 2048
tokens = 256
e = Expert(dim, inter_dim)

from ptflops import get_model_complexity_info
input_tokens = (1,tokens,dim)
flops, params = get_model_complexity_info(e, input_tokens, as_strings=True,print_per_layer_stat=True)

Expert(
  44.05 M, 100.000% Params, 11.28 GMac, 99.995% MACs, 
  (w1): Linear(14.68 M, 33.329% Params, 3.76 GMac, 33.328% MACs, in_features=7168, out_features=2048, bias=True)
  (w2): Linear(14.69 M, 33.341% Params, 3.76 GMac, 33.340% MACs, in_features=2048, out_features=7168, bias=True)
  (w3): Linear(14.68 M, 33.329% Params, 3.76 GMac, 33.328% MACs, in_features=7168, out_features=2048, bias=True)
)

单个专家要加载44.05MB的参数, 整个模型的专家参数为(256+1) x 60层 x 44.05M = 663B, 几乎占到了98%的参数量. 对于训练而言, 每个Batch大, 加载数据和加载参数的比值影响比较小. 但是对于推理单个Batch, 加载44.05MB参数仅处理256个Token数据为256 x 7168B = 1.8MB数据,开销非常大, 而实际运算为11.28GMAC, 开销非常小, 可以说细粒度MoE的算法和FP8本质上先解决了国内算力受限的问题.

但是访问内存和通信的问题还需要进一步解决, 因此DeepSeek-V3才提及了需要做Prefill Decode分离,以及在Decoding集群需要按照EP320并行, 甚至是从袁进辉老师的一段话可以知道, 梁总为啥要推荐性能最好需要80台, 通过EP并行获得更好的Data Locality.

2. MoE演进

因此, 渣B一直在提一个问题就是, 能否进一步优化来降低MoE阶段的数据/参数访存比? 所以才会提到2层Gate Function的做法

所以期望的方式是构造2级的Routing Gate, 使得本来Attention里面携带的信息通过两个Gate找到矩阵中(x,y)对应的某个Expert,或者多个expert. 例如我们继续按照Hidden-dim切分成16片, 然后也实现一个256选8的MoE算法, 这样每个专家的inter-dim维度应该是可以继续降低的.同时整个模型的参数量还可以进一步提升.

其实某种意义上来看, 这就是最早来自于PKM 《Large Memory Layers with Product Keys》 ^[1] 相关的工作, 也是字节Ultra-Sparse Memory Network一开始讲的一个图

其实你会看到这本质上和渣B提到的通过两个Gate找到矩阵中(x,y)对应的某个Expert是一致的. 而字节论文的描述里, 谈到的其实更适合于用渣B另外一篇 《谈谈大模型架构的演进之路, The Art of memory.》 中的描述

通过这样的视角来看, 将MoE作为内存层, 然后两级页表的方式作为Gating函数, 从这个视角上看, 大模型本身就可以构造一个 自己能够产生代码运行的通用计算机架构 .