【PyTorch][chapter-38][MOE-load balancing】[3]
前言：
专家混合模型的基本思想是在深度学习时代之前提出的，可以追溯到 90 年代，当时罗伯特·雅各布斯 (Robert Jacobs) 与“人工智能教父”杰弗里·辛顿 (Geoffrey Hinton) 及其同事提出了“局部专家自适应混合模型”。他们提出了将神经网络划分为多个由门控网络管理的专业“专家”的想法。
随着深度学习的兴起，MoE 再次浮出水面。2017 年，Noam Shazeer 及其同事（再次包括 Geoffrey Hinton）提出了用于循环神经语言模型的稀疏门控混合专家层。
稀疏门控混合专家层由多位专家（前馈网络）和一个可训练的门控网络组成，后者选择专家组合来处理每个输入。门控机制支持条件计算，将处理引导到最适合输入文本各个部分的网络部分（专家）。
这样的 MoE 层可以集成到 LLM 中，取代 Transformer 块中的前馈层。其关键组件:
experts, the gating mechanism, and the load balancing.
目录：
- Expert
- Gating mechanism
- Load balancing
- Python 辅助损失函数代码
一 Expert
1: 传统的Transformer FFN 架构（ Dense Layers）
传统的Transformer 模型使用了Dense的架构,参数量很大,但是推理的时候大部分参数是没有激活的,为了加快推理时间，由此引入了MOE 架构
2 MoE Sparse Layer
MOE 把原来大的FFN 分割成一个个小的FFN，每个小的FFN 称为一个Expert
1.2 Dense Sparse 参数
MOE 里面有个概念是稀疏参数，这里指的是总的Expert Total Parames 参数量很大，但是实际推理过程中只会用到激活1-2个Expert,对应的Active Params 参数量很小。
二 Gating mechanism
1 What is the role of the gating mechanism in an MoE layer?
在上一节中，我们介绍了“专家”这一抽象概念，即模型参数的一个专门子集.这些参数应用于大语言模型（LLM）架构不同层次的输入高维表示。在训练过程中，这些子集会变得“擅长”处理特定类型的数据。门控机制在这一系统中起着关键作用。
门控机制在MoE层中的作用是什么？
当训练一个MoE（混合专家）LLM时，所有专家的参数都会被更新。门控机制学会将输入令牌分配给最合适的专家，而专家则会适应并优化处理经常分配给他们的输入类型。在推理过程中，只有与输入相关的专家会被激活。这使得一个具有专门部分的系统能够处理各种类型的输入。
用公司类比，门控机制就像是一个在团队中分配任务的经理。
门控组件是MoE层中的一个可训练网络。门控机制有几项职责：
基于输入为专家评分：对于N个专家，计算N个分数，这些分数对应于专家与输入token的相关性。
选择要激活的专家：根据专家的评分，选择一部分专家进行激活。这通常通过top-k选择来完成。
负载均衡：简单地选择top-k专家会导致令牌在专家之间的分布不平衡。一些专家可能因为只处理极小范围的输入而变得过于专门化，而其他专家则会过于泛化。在推理过程中，将大部分输入分配给一小部分专家会导致这些专家过载，而其他专家则未被充分利用。因此，门控机制必须在所有专家之间均匀分配负载。
2 How is gating implemented in MoE LLMs?
让我们考虑一个由n 个专家组成的 MoE 层，记为
，其中i =1,…,n，其输入词向量为x。门控层的输出计算为
其中是第 i位专家的得分，基于Softmax 函数建模。门控层的输出用作对专家输出进行平均以计算 MoE 层的最终输出时的权重。如果
MoE 门控机制的总体框架如下
一些具体的例子是：
- Top-1 门控：在仅选择得分最高的导出时，每个标记都会定向到单个专家。这用于Switch Transformer的 Switch 层。它在计算上很高效，但需要仔细平衡标记的负载，以便在专家之间均匀分布。
- Top-2 门控：每个 token 发送给两个专家。Mixtral 中就用到了这种方法。
- 噪声 top-k 门控：引入稀疏门控混合专家层，在应用Softmax之前添加噪声（标准正态性）以帮助实现负载平衡。GShard使用噪声 top-2 策略，添加了更先进的负载平衡技术。
三 Load balancing
3.1 router collapse (也称 token drop)
通过评分和选择前 k 名专家的直接门控可能会导致专家之间令牌分配不平衡。一些专家可能会超负荷工作，被分配处理更大比例的令牌，而其他专家则被选中的频率要低得多，并且利用率不足。这会导致路由“崩溃”，从两个方面损害 MoE 方法的有效性。
首先，经常选择的专家在训练过程中会不断更新，因此其表现比没有收到足够数据进行适当训练的专家更好。
其次，负载不平衡会导致内存和计算性能问题。当专家分布在不同的 GPU 和/或机器上时，专家选择的不平衡将转化为网络、内存和专家容量瓶颈。如果一个专家必须处理的令牌数量是另一个专家的十倍，这将增加总处理时间，因为后续计算将被阻止，直到所有专家完成其分配的负载处理。
3.2 Strategies for improving load balancing in MoE LLMs include:
1 Adding random noise in the scoring process helps redistribute tokens among experts.
2 Adding an auxiliary load-balancing loss to the overall model loss. It tries to minimize the fraction of the input routed to each expert. For example, in the Switch Transformer, for N experts and T tokens in batch B, the loss would be
3： 我们通常会调整 expert_capacity,限制每个FFN 的token 上限
Expert choice ，开源的模型很少用。
4 跳层
三 Load Balance(Auxiliary Loss)
为了在训练期间获得更均匀的专家分布，辅助损失（也称为负载平衡损失）被添加到网络的常规损失中。它增加了一个约束，迫使专家具有同等重要性。
通过评分和选择前 k 名专家的直接门控可能会导致专家之间令牌分配不平衡。一些专家可能会超负荷工作，被分配处理更大比例的令牌，而其他专家则被选中的频率要低得多，并且利用率不足。这会导致路由“崩溃”，从两个方面损害 MoE 方法的有效性。
首先，经常选择的专家在训练过程中会不断更新，因此其表现比没有收到足够数据进行适当训练的专家更好。
其次，负载不平衡会导致内存和计算性能问题。当专家分布在不同的 GPU 和/或机器上时，专家选择的不平衡将转化为网络、内存和专家容量瓶颈。如果一个专家必须处理的令牌数量是另一个专家的十倍，这将增加总处理时间，因为后续计算将被阻止，直到所有专家完成其分配的负载处理。
改善负载平衡的策略包括：
•1 在评分过程中添加随机噪声有助于在专家之间重新分配令牌。
2 在整体模型损失中添加辅助负载平衡损失。它试图最小化路由到每个专家的输入部分。例如，在Switch Transformer中，对于B批次中的N 个专家和T 个token ，损失将是
四 Python 辅助损失函数代码
# -*- coding: utf-8 -*-
"""
Created on Thu Mar 20 14:39:55 2025
@author: chengxf2
"""
import pandas as pd
def auxiliary_loss():
router= {
'Expert1':[0.51 ,0.51, 0.49,0.49],
'Expert2':[0.49, 0.49, 0.51,0.51 ]}
df = pd.DataFrame(router)
#f 是路由到专家i 的令牌的比例
f = pd.DataFrame([[0,0]], columns=['Expert1', 'Expert2'])
#p是分配给专家i 的路由器概率的比例。
p = pd.DataFrame([[0,0]], columns=['Expert1', 'Expert2'])
N = len(df)
for index, row in df.iterrows():
max_row_indices = row.idxmax()
#print("\n max_row_indices" ,max_row_indices)
f[max_row_indices]=f[max_row_indices]+1
p.loc[0, 'Expert1'] = df['Expert1'].sum()
p.loc[0, 'Expert2'] = df['Expert2'].sum()
f = f/N
print("\n 是路由到专家i 的令牌的比例",f)
p = p/N
print("\n 是分配给专家i 的路由器概率的比例",p)
loss = f*p
print("\n aux loss: %4.3f"%(loss.sum().sum()))
auxiliary_loss()
https://huggingface.co/blog/NormalUhr/moe-balance
https://www.youtube.com/watch?v=sOPDGQjFcuM&t=468s
https://www.youtube.com/watch?v=AgZuF7lsu-8
https://www.youtube.com/watch?v=AgZuF7lsu-8
A Visual Guide to Mixture of Experts (MoE)