Switch Transformer由谷歌2021年提出,是一种基于Transformer架构的改进模型,其通过引入MoE(Mixture of Experts)机制提升模型的效率与扩展性。相关论文可以参考《Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity》。
在标准Transformer中,每个输入的token都会经过相同的前馈网络(FFN)层进行处理,而Switch Transformer将其中的FFN替换为多个独立的专家子网络。模型通过一个路由机制动态地将每个token分配给最合适的专家进行处理,这种设计允许模型在不显著增加单个样本推理计算量的情况下扩展参数规模,从而提升处理能力。
需要注意的是,MoE中的专家并不会都全被激活,一般只会激活1-2个专家。而其训练过程需要使用数据并行、模型并行及专家并行方法,从而可以让模型支持万亿参数规模。
数据并行过程将训练数据分片到多个设备,每个设备独立计算损失和梯度,随后通过All-Reduce操作聚合梯度并同步模型参数。
专家并行将MoE结构中的专家子网络分布到不同设备,每个设备仅存储部分专家的参数。通过稀疏路由机制(Switch Routing)确保每个token仅由一个专家进行处理,通过All-to-All通信将token路由到对应设备,实现专家数量的线性扩展并减少通信开销。
参考文章:
https://zhuanlan.zhihu.com/p/705443412 https://developer.aliyun.com/article/1690017
如果喜欢这篇文章或对您有帮助，可以：[☕] 请我喝杯咖啡 | [💓] 小额赞助
