南京大学LLM开发基础（四）MoE, LoRA, 数的精度 + MLP层实验

https://njudeepengine.github.io/llm-course-lecture/2025/lecture8.html#1

1. Mixture-of-experts (MoE)

1.1 优势

1.2 结构

1.3 训练

2. Low-rank adaptation (LoRA)

3. 数的精度 -- 混合精度 + 量化操作

Task1：DenseMLPWithLoRA

一、任务背景

二、任务要求

Task2：Sparse MLP

1. Mixture-of-experts (MoE)

通过gating model/network 选择最适合的专家。

Switch Transformers 将Transformers中的 FFN替换为MoE 结构。

1.1 优势

1. 推理阶段的效率提升：

每次只用一部分专家的参数；条件计算与动态资源分配

门控网络会根据输入样本的难度和类型，动态地选择最合适的专家。

处理简单任务时：模型可能会路由到一些相对“简单”或计算量较小的专家，从而更快地给出响应。

处理复杂任务时：模型会调用**更“专业”**的专家，虽然计算量稍大，但能保证输出质量。

2. 以恒定计算成本，实现模型参数规模的巨量提升。

传统模型（稠密模型）： 增加参数意味着每一层、每一个输入都要经过所有这些参数的计算。模型变大，计算成本（FLOPS）和时间几乎线性增长。
MoE 模型（稀疏模型）： 模型总参数量可以变得极其巨大（例如万亿参数），但对于每个具体的输入，只有一小部分“专家”（通常是1个或2个）被激活并参与计算。这意味着：
- 总参数量巨大：模型拥有海量的知识和能力储备。
- 激活参数量固定：无论总参数是千亿还是万亿，处理一个 token 实际使用的参数数量是基本不变的，因此计算成本得以保持在一个可接受的水平。

3. 预训练速度显著加快：

在相同的计算资源（如相同的GPU数量和训练时间）下，MoE 模型能获得比稠密模型更好的性能

1.2 结构

topk 得到前 k 名，对分数 softmax 得到每个export权重。

1class TopKGating(nn.Module):
2    def __init__(self, hidden_dim, num_experts, k=2):
3        super().__init__()
4        self.router = nn.Linear(hidden_dim, num_experts, bias=False)  # 路由网络
5        self.k = k  # 选择专家数量
6    def forward(self, x):
7        # 1. 计算专家得分
8        scores = self.router(x)  # [batch, experts]
9        # 输入x: [batch_size, hidden_dim]
10        # 输出scores: [batch_size, num_experts]
11        # 每个样本对每个专家得到一个得分
12        
13        # 2. 选择Top-K专家
14        topk_scores, topk_idx = torch.topk(scores, self.k, dim=-1)
15        # topk_scores: [batch_size, k] - 前k个最高得分
16        # topk_idx: [batch_size, k] - 对应的专家索引
17        
18        # 3. 计算选择概率
19        probs = torch.softmax(topk_scores, dim=-1)
20        # 对选中的k个专家得分做softmax，得到归一化的权重
21        
22        # 4. 构建门控矩阵
23        gate = torch.zeros_like(scores)  # 创建全零矩阵 [batch_size, num_experts]
24        gate.scatter_(-1, topk_idx, probs)  # 将概率值填充到对应位置
25        # 结果gate: [batch_size, num_experts]，每行只有k个非零值

1.3 训练

训练不稳定性问题：

强者恒强：某些专家训练得快，获得更高权重，被选择更多
弱者淘汰：其他专家得不到充分训练，逐渐被边缘化
模式坍塌：最终只有少数专家被使用

下面三个解决方案：

1. Auxiliary Loss 辅助损失：均衡各个专家的选择概率。

计算每个专家被选择的概率分布
计算批次中实际选择每个专家的频率分布
惩罚这两个分布之间的差异，鼓励均匀分布

2. Capacity Factor 容量因子

预先计算每个专家的"容量" 处理token数量上限。多余的强制给别的expert。

3. Shared Experts 共享专家机制

设计专家架构时（不同损失函数）就将专家分为"通用型"和"专用型"。

共享专家：处理通用模式和基础特征，保证基本利用率
专用专家：学习特定领域或复杂模式，允许一定程度的专业化

2. Low-rank adaptation (LoRA)

https://arxiv.org/abs/2106.09685

随着预训练的模型规模不断扩大，对所有模型参数进行全量微调（full fine-tuning）不可行。

冻结预训练模型的权重，并在 Transformer 架构的每一层中注入可训练的低秩分解矩阵，

从而大幅减少下游任务所需的可训练参数数量。

LoRALayer 求更新的变化值。 A = [ in, rank] B = [ rank, out ]

1import torch.nn as nn
2
3class LoRALayer(nn.Module):
4    def __init__(self, in_dim, out_dim, rank, alpha):
5        super().__init__()
6        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
7        self.A = nn.Parameter(torch.randn(in_dim, rank) * std_dev)
8        self.B = nn.Parameter(torch.zeros(rank, out_dim))
9        self.alpha = alpha
10
11    def forward(self, x):
12        x = self.alpha * (x @ self.A @ self.B)
13        return x

线性层 + LoRA 层，参数更新

1class LinearWithLoRA(nn.Module):
2
3    def __init__(self, linear, rank, alpha):
4        super().__init__()
5        self.linear = linear
6        self.lora = LoRALayer(
7            linear.in_features, linear.out_features, rank, alpha
8        )
9
10    def forward(self, x):
11        return self.linear(x) + self.lora(x)

3. 数的精度 -- 混合精度 + 量化操作

浮点数 -- 双精度: FP64 单精度: FP32 半精度: FP16

通过 exp和fraction 不同的范围与精度。

混合精度是指在模型训练过程中，同时使用不同数值精度的格式。

量化是指将模型中的浮点参数（如 FP32）转换为低比特整数（如 INT8、INT4，甚至 INT2/INT1）或定点数**，通过减少每个参数的存储位数来压缩模型。**

对模型训练的帮助

减少内存占用
加速计算速度
降低能源消耗

https://njudeepengine.github.io/LLM-Blog/2025/06/29/A3-modeling-mlp/

Task1：DenseMLPWithLoRA

Dense 的 MLP 模块：先将 hidden_states 从 h 维上投影（up-project）到更高的 ffh 维，

再通过 gating 机制下投影（down-project）回原始维度。

一、任务背景

MLP 模块中的参数通常占据了 LLMs 中超过 90% 的可训练参数，

因此采用全量线性参数监督微调（supervised fine-tuning, SFT）的方式效率非常低。

LoRA 的核心基本假设是： δ W 是一个低秩且稀疏的矩阵，其可以通过低秩分解表示为：

二、任务要求

对三个投影矩阵 W 进行初始化，初始化方式为从正态分布中采样。

如果 activation_type 使用 MLPActivationType.SIGMOID 或 MLPActivationType.BILINEAR，则使用 Xavier Initialization；
否则，对于其余 ReLU-family 的激活函数，则使用 Kaiming Initialization；
标准差 std = sqrt ( 2/ fan_in )

LoRA 低秩分解矩阵 A、B初始化：均匀分布 seed = lora_init_base_seed +1 / +2

对 lora_gate\up\down_A\B 的 seed 分别设置 offset 偏移量。

1. 若 lora_rank = 0，你应跳过任何与 LoRA 有关的逻辑

2. 参数 lora_alpha 是一个正缩放因子，默认为 None 时，表示应设置 = lora_rank（不缩放）

3. Dropout 层有 lora_dropout_seed 和 lora_dropout_rate 参数

Task2：Sparse MLP

类似多头注意力，将投影矩阵的 ffh 维度划分为 ne 个大小相等的 shard（分片）对应一个专家。

hidden states 的每个 token 通过一个机制，映射到对应的 k 个 experts。

最终，每个 token 的最终输出是来自这 k 个 experts 子输出的加权和。

从而：降低高维计算开销；通过并行学习，提高潜在模式的多样性。

下文实现方式 Mixtral 方法：

1. 训练一个 (h, ne) 的网络 G；把 h 维的输入，转换为 ne 个专家分别的概率。

2. 只取前 k 名的 P，用 P / SUM 分配权重。将这几个 experts 的输出结果加权。

3. 每一个小专家相当于一个 DenseMLP，调用上一问。

《南京大学LLM开发基础（四）MoE, LoRA, 数的精度 + MLP层实验》是转载文章，点击查看原文。