目錄

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權重。

class TopKGating(nn.Module):
    def __init__(self, hidden_dim, num_experts, k=2):
        super().__init__()
        self.router = nn.Linear(hidden_dim, num_experts, bias=False)  # 路由網絡
        self.k = k  # 選擇專家數量
    def forward(self, x):
        # 1. 計算專家得分
        scores = self.router(x)  # [batch, experts]
        # 輸入x: [batch_size, hidden_dim]
        # 輸出scores: [batch_size, num_experts]
        # 每個樣本對每個專家得到一個得分
        
        # 2. 選擇Top-K專家
        topk_scores, topk_idx = torch.topk(scores, self.k, dim=-1)
        # topk_scores: [batch_size, k] - 前k個最高得分
        # topk_idx: [batch_size, k] - 對應的專家索引
        
        # 3. 計算選擇概率
        probs = torch.softmax(topk_scores, dim=-1)
        # 對選中的k個專家得分做softmax，得到歸一化的權重
        
        # 4. 構建門控矩陣
        gate = torch.zeros_like(scores)  # 創建全零矩陣 [batch_size, num_experts]
        gate.scatter_(-1, topk_idx, probs)  # 將概率值填充到對應位置
        # 結果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 ]

import torch.nn as nn

class LoRALayer(nn.Module):
    def __init__(self, in_dim, out_dim, rank, alpha):
        super().__init__()
        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
        self.A = nn.Parameter(torch.randn(in_dim, rank) * std_dev)
        self.B = nn.Parameter(torch.zeros(rank, out_dim))
        self.alpha = alpha

    def forward(self, x):
        x = self.alpha * (x @ self.A @ self.B)
        return x

線性層 + LoRA 層，參數更新

class LinearWithLoRA(nn.Module):

    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features, linear.out_features, rank, alpha
        )

    def forward(self, x):
        return self.linear(x) + self.lora(x)

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

浮點數 -- 雙精度: FP64    單精度: FP32     半精度: FP16

通過 exp和fraction 不同的範圍與精度。

混合精度是指在模型訓練過程中，同時使用不同數值精度的格式。

量化是指將模型中的浮點參數（如 FP32）轉換為低比特整數（如 INT8、INT4，甚至 INT2/INT1）或定點數，通過減少每個參數的存儲位數來壓縮模型。

對模型訓練的幫助

1. 減少內存佔用
2. 加速計算速度
3. 降低能源消耗

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，調用上一問。