深度学习基础理论————混合专家模型(MoE)/KV-cache

1、混合专家模型(MoE

参考HuggingFace中介绍:混合专家模型主要由两部分构成:
1、稀疏的MoE层:这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个),每个专家本身是一个独立的神经网络。在实际应用中,这些专家通常是前馈网络 (FFN),但它们也可以是更复杂的网络结构,甚至可以是 MoE 层本身,从而形成层级式的 MoE 结构。
2、门控网络/路由(Gate Layer/route Layer):这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如,在下图中,“More”这个令牌可能被发送到第二个专家,而“Parameters”这个令牌被发送到第一个专家。有时,一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点,因为路由器由学习的参数组成,并且与网络的其他部分一同进行预训练。

1

换言之也就是说:将原始的Transformer框架中的FFN Layer(全连接层)替换成一个由Gate Layer和若干的FFN Layer组成的结构,通过Gate来确定一个输入将会被那些FFN进行处理,而后对被FFN处理后的内容进行加权处理。


  • 1、稠密MoE稀疏MoE

1

检验而言:如上图所示,对于稠密的MoEDense MoE)而言(假设4个FFN)在通过Gate处理之后输入X要通过每一个FFN进行处理,而对于稀疏的MoESparse MoE)而言,通过Gate处理只去选择部分FFN进行处理

  • 2、MoE原理

1、Gate/route原理
输入数据x,通过一个线性层进行处理:

s=Wgatex+b

对于得到的score再通过Softmax函数处理,得到一个概率分布:

g(x)=Softmax(s)

对于稀疏的MoE而言还需要去选择部分专家进行激活:

h(x)=g(x)iTopkiTopkg(x)i

原理很简单,结合代码分析(以Deepseek-v3代码为例)

class Gate(nn.Module):
"""
Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.
"""
def __init__(self, args: ModelArgs):
"""
Initializes the Gate module.
Args:
args (ModelArgs): Model arguments containing gating parameters.
"""
super().__init__()
self.dim = args.dim
self.topk = args.n_activated_experts # 选择多少个专家进行使用
self.n_groups = args.n_expert_groups # Gate数量
self.topk_groups = args.n_limited_groups # 对于gate中分组数
self.score_func = args.score_func
self.route_scale = args.route_scale
self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))
self.bias = nn.Parameter(torch.empty(args.n_routed_experts)) if self.dim == 7168 else None
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass for the gating mechanism.
Args:
x (torch.Tensor): Input tensor.
Returns:
Tuple[torch.Tensor, torch.Tensor]: Routing weights and selected expert indices.
"""
scores = linear(x, self.weight) # 计算wx+b
# 归一化处理
if self.score_func == "softmax":
scores = scores.softmax(dim=-1, dtype=torch.float32)
else:
scores = scores.sigmoid()
original_scores = scores
if self.bias is not None:
scores = scores + self.bias
if self.n_groups > 1:
# 如果Gate数量>1
scores = scores.view(x.size(0), self.n_groups, -1)
if self.bias is None:
group_scores = scores.amax(dim=-1)
else:
group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)
indices = group_scores.topk(self.topk_groups, dim=-1)[1]
mask = torch.zeros_like(scores[..., 0]).scatter_(1, indices, True)
scores = (scores * mask.unsqueeze(-1)).flatten(1)
indices = torch.topk(scores, self.topk, dim=-1)[1]
weights = original_scores.gather(1, indices)
if self.score_func == "sigmoid":
weights /= weights.sum(dim=-1, keepdim=True)
weights *= self.route_scale
return weights.type_as(x), indices

整个过程分析:输入数据x(假设维度为:(bs, num_tokens, embed_dim)),通过线性(W:(n_routed_experts,embed_dim)计算得到:bs,num_tokens,n_routed_experts)和归一化处理:(bs,num_tokens, n_routed_experts)这样一来就可以知道每个tokens的一个分布概率(到底要走哪一个FFN)。如果n_expert_groups数量>1,也就是说如果Gate数量不止一个,那么也就只需要去选择topK即可然后用一个mask将其他的给mask掉即可(scores = (scores * mask.unsqueeze(-1)).flatten(1))然后再去对挑选得到的Gate选择topkindices = torch.topk(scores, self.topk, dim=-1)[1]

2、FFN Layer原理
这部分代码就比较简单(直接用线性层处理即可):

class Expert(nn.Module):
"""
Expert layer for Mixture-of-Experts (MoE) models.
Attributes:
w1 (nn.Module): Linear layer for input-to-hidden transformation.
w2 (nn.Module): Linear layer for hidden-to-output transformation.
w3 (nn.Module): Additional linear layer for feature transformation.
"""
def __init__(self, dim: int, inter_dim: int):
"""
Initializes the Expert layer.
Args:
dim (int): Input and output dimensionality.
inter_dim (int): Hidden layer dimensionality.
"""
super().__init__()
self.w1 = Linear(dim, inter_dim)
self.w2 = Linear(inter_dim, dim)
self.w3 = Linear(dim, inter_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass for the Expert layer.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output tensor after expert computation.
"""
return self.w2(F.silu(self.w1(x)) * self.w3(x))

最后得到MoE代码如下:

class MoE(nn.Module):
"""
Mixture-of-Experts (MoE) module.
Attributes:
dim (int): Dimensionality of input features.
n_routed_experts (int): 模型中总的专家数量
n_local_experts (int): 在分布式系统中,当前设备(节点)负责的专家数量
n_activated_experts (int): 每个输入激活的专家数量(即每个输入会被路由到多少个专家)。
gate (nn.Module): Gating mechanism to route inputs to experts.
experts (nn.ModuleList): List of expert modules.
shared_experts (nn.Module): Shared experts applied to all inputs.
"""
def __init__(self, args: ModelArgs):
"""
Initializes the MoE module.
Args:
args (ModelArgs): Model arguments containing MoE parameters.
"""
super().__init__()
self.dim = args.dim
assert args.n_routed_experts % world_size == 0
self.n_routed_experts = args.n_routed_experts
self.n_local_experts = args.n_routed_experts // world_size
self.n_activated_experts = args.n_activated_experts
# 起始/结束(因为是涉及到分布训练,去寻找每个设备上的”专家“即可)
self.experts_start_idx = rank * self.n_local_experts
self.experts_end_idx = self.experts_start_idx + self.n_local_experts
self.gate = Gate(args)
self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None
for i in range(self.n_routed_experts)])
self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass for the MoE module.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output tensor after expert routing and computation.
"""
shape = x.size()
x = x.view(-1, self.dim)
weights, indices = self.gate(x)
y = torch.zeros_like(x)
counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
for i in range(self.experts_start_idx, self.experts_end_idx):
if counts[i] == 0:
continue
expert = self.experts[i]
idx, top = torch.where(indices == i)
y[idx] += expert(x[idx]) * weights[idx, top, None]
z = self.shared_experts(x)
if world_size > 1:
dist.all_reduce(y)
return (y + z).view(shape)

代码分析:
1、获取权重以及索引:weights, indices = self.gate(x)
2、计算每个专家索引:counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()通过torch.bincount返回数值个数,比如说:[1,2,2,2,3],那么返回:[0,1,3,1](0出现0次,1出现1次。。。。)相当于对每个专家都编号了,只需要根据出现次数,来索引出来
3、计算加权得分:

for i in range(self.experts_start_idx, self.experts_end_idx):
if counts[i] == 0:
continue
expert = self.experts[i]
idx, top = torch.where(indices == i)
y[idx] += expert(x[idx]) * weights[idx, top, None]

如果涉及到多个专家,很可能就会出现一种情况:只有部分的专家被使用进而导致负载不均衡。那么可以通过 负载均衡损失Load Balancing Loss)(如上图所示)计算主要为:
1、importance loss(重要性损失)
回到上面对于router函数(s=softmax(Wgatex+b))在论文中对于重要性损失定义如下:

Limportance(x)=(std(x)mean(x))2

也就是说,通过最小化上式子的最小值(均值大,方差小),进而保证所有的专家平均而言拥有相似的权重。

2、Load Loss:重要性损失旨在保证所有专家平均而言具有相似的路由权重。但是不难想到这些看上去有着总体趋于平衡的权重的路由配置,仍然有一小部分专家获得了所有分配,如下图:

首先,对于专家选择阈值:thresholdk(x):=maxk(Wx+ϕ),其中ϕ代表前向传播过程中的采样噪声,maxk代表选择第k个最大的值,对于专家的负载:指的是在整个批次中的使用情况:loadi(X)=xXpi(x)对于这个概率的计算:

pi(x):=P(ϕnewthresholdk(x)(Wx)i)

其中ϕnewN(0,σ2),σ=1/E最后对于load loss定义为:

Lload(X)=(std(load(X))mean(load(X)))2

最后损失函数为:L=12Limportance(X)+12Lload(X)

同时也可以采用另外的方法:
1、GShard设定一个 专家容量:设定一个阈值,定义一个专家能处理多少令牌。如果两个专家的容量都达到上限,令牌就会溢出,并通过残差连接传递到下一层,或在某些情况下被完全丢弃。对于专家容量可以借鉴下面处理方式:tokens per batchnumber of experts×capacity factorcapacity factor可以选择1-1.25之间

KV cache

Transformer模型推理过程为:
image.png

目前主流的LLM框架主要都是使用的decoder-only(也就是说只用Transformer中的decoder结构)
image.png

对于LLM任务(通常采用自回归过程)可以简单认为是一种“完形填空”的过程,在输入前面i-1个词然后推测第i个词

回归上面的推理过程(翻译输出:i am a student):模型中在输出'a'的时候会将'i am'都输入到模型中。理解这个过程(假设就是直接输出文本:i am a student):

参考:https://zhuanlan.zhihu.com/p/662498827

step1: in: Q=K<S> || out: i
Attention1:Q1K1T
step2: in: Q=K=<s>,i || out: i am
Attention1:Q1K1TAttention2:Q2K1T,Q2K2T
step3: in: Q=K=<s>, i, am || out: i am a
Attention1:Q1K1TAttention2:Q2K1T,Q2K2TAttention3:Q3K1T,Q3K2T,Q3K3T
step4: in: Q=K=<s>, i, am, a || out: i am a student
Attention1:Q1K1TAttention2:Q2K1T,Q2K2TAttention3:Q3K1T,Q3K2T,Q3K3TAttention4:Q4K1T,Q4K2T,Q4K3T,Q4K4T

不过上面操作过程中会有问题:
计算有很大冗余(每次生成新的词,都需要回归一下之前生成的词),并且每次计算Attentioni只与Qi相关
对于后面一点理解(以step2为例):
我目前已经有两个QQ1:<s>Q2:i 并且还有K和V(这两个也是有两个值),我会初始化一个Q3对于下一个值我就用Q3进行表示,然后我就需要去计算注意力得分(只用Q,K,V这三个值计算过程举例):

QKT=(bs,3,embeddim)(bs,embeddim,2)=(bs,3,2),接下来计算QKTV=(bs,3,2)(bs,2,embeddim)=(bs,3,embeddim)那么在这个过程中就会有一个有意思问题:Q会有重复的(dim=3,前面两个都是前面已经计算过的)(观察上面Attention计算可以发现:每次计算Attentioni只与Qi相关)。因此就有KV-cache理论:既然每次都是Q在变化,但是K和V都是用的之前的,那我之前每次就只用新的Q去和旧的KV计算即可(将KV存储起来),KV-cache一种典型的用内存换速度的方法

image

简易Demo:

import torch
class KVCache:
def __init__(self):
self.k = None
self.v = None
def update(self, k, v):
if self.k is None:
self.k = k
self.v = v
else:
self.k = torch.cat([self.k, k], dim=1) # 在序列维度上拼接
self.v = torch.cat([self.v, v], dim=1)
def get(self):
return self.k, self.v
class Decoder(torch.nn.Module):
def __init__(self, embed_dim, hidden_dim, vocab_size, num_heads=8):
super().__init__()
self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
self.attention = torch.nn.MultiheadAttention(embed_dim, num_heads)
self.linear = torch.nn.Linear(embed_dim, vocab_size)
self.kv_cache = KVCache()
def forward(self, input_ids):
x = self.embedding(input_ids) # (batch_size, seq_len, embed_dim)
# 获取 KV-cache
k, v = self.kv_cache.get()
# 计算 Attention
if k is not None and v is not None:
# 使用 KV-cache
attn_output, _ = self.attention(x, k, v) # (batch_size, seq_len, embed_dim)
else:
# 初始状态,没有 KV-cache
attn_output, _ = self.attention(x, x, x) # (batch_size, seq_len, embed_dim)
# 更新 KV-cache
self.kv_cache.update(x, x)
# 残差连接
x = x + attn_output
# 线性变换
logits = self.linear(x) # (batch_size, seq_len, vocab_size)
return logits
batch_size = 2
seq_len = 4
embed_dim = 64
hidden_dim = 256
vocab_size = 10000 # 假设词汇表大小为 10000
decoder = Decoder(embed_dim, hidden_dim, vocab_size)
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len)) # (batch_size, seq_len)
logits = decoder(input_ids) # (batch_size, seq_len, vocab_size)
print("Logits shape:", logits.shape)

使用Huggingfacetransformers框架代码:https://huggingface.co/docs/transformers/main/en/kv_cache。只需要类似下面操作:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda:0")
inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)
# 具体参数:https://huggingface.co/docs/transformers/en/main_classes/text_generation
out = model.generate(**inputs, do_sample=False, max_new_tokens=23, use_cache=True)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
out = model.generate(**inputs, do_sample=False, max_new_tokens=23)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])

Transformers中不同cache方式:

缓存类型 描述 适用场景 优点 缺点
StaticCache 静态缓存,缓存所有的 K 和 V,不更新。 短序列生成、内存充足的场景 实现简单,快速 不适合长序列生成,内存消耗较大
OffloadedStaticCache 静态缓存,但将缓存内容卸载到外部存储。 内存受限的环境,长序列生成 减少显存占用,适合大规模生成 存取速度较慢,可能影响生成速度
SlidingWindowCache 滑动窗口缓存,缓存一个固定大小的窗口。 长序列生成、内存有限的场景 限制内存消耗,适合长序列生成 窗口太小可能丢失上下文信息,影响生成效果
HybridCache 混合缓存,结合静态缓存和滑动窗口缓存。 长序列生成,要求平衡内存和上下文 平衡内存消耗和上下文保留 比静态缓存更复杂,可能需要更多内存管理和计算资源
MambaCache 高效的缓存实现,针对推理速度和内存占用进行了优化。 高性能计算环境、高并发推理任务 高度优化,适合大规模并行推理 可能需要特定硬件支持,复杂度较高
QuantizedCache 量化缓存,减少存储需求。 内存受限的设备、需要减少内存占用的场景 大幅度减少内存占用,适合嵌入式设备 量化可能导致精度损失,影响生成质量

争对上面描述其实KV-cahce是一种用存储换速度的方法,因此,对于KV存储进行优化就十分有必要了!

参考

1、https://arxiv.org/pdf/2101.03961
2、混合专家模型 (MoE) 详解
3、https://arxiv.org/pdf/2407.06204
4、https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py
5、https://arxiv.org/pdf/2106.05974
6、https://arxiv.org/pdf/2006.16668
7、https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf
8、https://jalammar.github.io/illustrated-transformer/
9、https://zhuanlan.zhihu.com/p/662498827

posted @   Big-Yellow-J  阅读(145)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· DeepSeek “源神”启动!「GitHub 热点速览」
· 我与微信审核的“相爱相杀”看个人小程序副业
· 微软正式发布.NET 10 Preview 1:开启下一代开发框架新篇章
· 如何使用 Uni-app 实现视频聊天(源码,支持安卓、iOS)
· C# 集成 DeepSeek 模型实现 AI 私有化(本地部署与 API 调用教程)
点击右上角即可分享
微信分享提示