分類: 技術

DeepSeek 重點分解 – MLA 小整理

DeepSeek [1] 重點之一是 MLA (Multi-head Latent Attention) . 它可以單獨使用. 解釋它時, 若和 MHA (Multi-Head Attention) 對比會更好理解, 所以先回顧一下 MHA.

1. 原始的注意力計算

一般的注意力機制中, 主要有 q、k、v 三個矩陣. 分別代表 query, key, 和 value. W 是權重矩陣, h 是 hidden state. 下標 t 表示第 t 個 token.

加入多頭機制後: d 是 embedded dimension, nh = attension head 數, , dh 是每個 head 的 dimension. 其中 d = nh * dh , j 是從第一個到第 t 個 token 的 index, i 是第一到第 nh 個 head 的 index.

Attention o 當然也分為第幾個頭的第幾個 token. 故表示為 ot,i. 同樣多頭都用 o 權重矩陣 Wo 轉出 output ut.

一般認為 k, v 的值太多, 是造成計算量和記憶體過多的元凶. 但是不記住這些東西, transformer 就發揮不出造句的能力. 科學家想了各種解法想要簡化 k,v. 但是操作不好就會降低 LLM 的性能. DeepSeek 使用的 MLA 看起來可以兩全其美.

2. MLA 中的計算流程

2.1 Low-Rank Key-Value Joint Compression

在 MLA 中, 首先對 hidden state 壓縮. 從前面的 MHA 的段落可以得知, q, k, v 共用 hidden state 但不共用權重. 那麼我們壓縮 h 再還原就可以節省參數量了. c 矩陣由 ht 乘 WDKV 而來. 顧名思義, D 代表 down projection, kv 矩陣意義跟先前相同. 做完下投影再做上投影, 理解為壓縮解壓縮即可. 所以 WUK 還原 k, WUV 還原 v. 其中 U 就代表 up projection.

我們已經知道這是濃縮再還原的果汁. 為了確定風味不變太多. 損失的部分要要別的方法補回來. 甚至 DeepSeek 為了減少 active node, 連你問的問題 q = query 都壓縮了. 這個猛! 表示我跟它說 “請、謝謝、對不起" 都是多餘的.

2.2 Decoupled Rotary Position Embedding

上面那招還不是全貌. 但是要講第二招就要先講 RoPE (Rotary Position Embedding) [2]. 我們知道原本 transformer 就要記錄 token 的相對位置關係. 畢竟"你愛我" 跟 “我愛你" 是兩碼子事. RoPE 這個演算法就是用來把位置編碼專屬的維度給省了, 但它結 “繩" 記事, 還記得相對位置.

不過這招和 2.1 壓縮那招有衝突, 壓縮完再解壓就不符合交換律. 所以又衍生出 decouple 的輔助算法, 再浪費一點空間. 為 RoPE 額外產生的 dimension 為 dhR. 多出來的 qtR 和 ktR 加在原來的矩陣後面. 式子大致上都一樣.

3. 效能比較

假如我們比對 MHA 和 MLA, 就會發現它的 KV cache 比較少, 而且實測效果更好. 至於 GQA 和 MQA 是來陪榜的. 用論文 [1] 中的圖解帶過.

[REF]

  1. https://arxiv.org/abs/2405.04434
  2. . J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu.Roformer: Enhanced transformer with rotary position embedding.Neurocomputing, 568:127063, 2024.

BentoML 小整理

趁著尾牙等摃龜的空檔，把這篇的草稿丟給AI 重寫。雖然變得有點 WIKI化，不過稍微調整順序, 潤飾文字後，感覺還是滿易懂的。

BentoML [1] 是一個開源的 ML 模型服務框架，名字源自日文「便當」，代表將所有組件打包在一起。相較於 Google Cloud 的 Kubeflow 解決方案 [2]，BentoML 提供了不綁定特定雲服務的部署方式。

核心特點

  1. 模型管理
  • 統一打包（模型 + 依賴）
  • 版本控制
  • 自動追蹤環境配置
  1. 框架支援
  • 支援主流 ML 框架
    • PyTorch
    • TensorFlow
    • scikit-learn
    • XGBoost
  • 多框架共存部署
  1. 服務效能
  • 高性能 API 服務器
  • 批量推理支援
  • 自動負載均衡
  1. 部署便利性
  • Docker 容器自動生成
  • Kubernetes 整合支援

實作流程

1. 模型訓練與保存

# 訓練模型
from sklearn import svm, datasets
iris = datasets.load_iris()
clf = svm.SVC()
clf.fit(iris.data, iris.target)

# 保存模型
import bentoml
bentoml.sklearn.save_model("iris_clf", clf)

2. 模型管理

# 查看最新版本
bentoml models get iris_clf:latest

# 列出所有版本
bentoml models list

3. 預測方式

3.1 直接載入

loaded_model = bentoml.sklearn.load_model("iris_clf:latest")
result = loaded_model.predict([[5.9, 3.0, 5.1, 1.8]])

3.2 使用 Runner（推薦）

# 建立 Runner 實例
runner = bentoml.sklearn.get("iris_clf:latest").to_runner()
runner.init_local()
result = runner.predict.run([[5.9, 3.0, 5.1, 1.8]])

4. 服務部署

  1. 建立服務檔案 (service.py):
import numpy as np
import bentoml
from bentoml.io import NumpyNdarray

iris_clf_runner = bentoml.sklearn.get("iris_clf:latest").to_runner()
svc = bentoml.Service("iris_classifier", runners=[iris_clf_runner])

@svc.api(input=NumpyNdarray(), output=NumpyNdarray())
def classify(input_series: np.ndarray) -> np.ndarray:
    return iris_clf_runner.predict.run(input_series)
  1. 定義部署配置 (bentofile.yaml):
service: "service.py:svc"
labels:
  owner: bentoml-team
  project: gallery
include:
- "*.py"
python:
  packages:
    - scikit-learn
    - pandas

5. 本地測試服務

bentoml serve service.py:svc --reload
  1. Web UI: 訪問 http://127.0.0.1:3000 或者
  2. API 調用:
$headers = @{"Content-Type" = "application/json"}
$data = "[[5.9, 3, 5.1, 1.8]]"
Invoke-WebRequest -Uri "http://127.0.0.1:3000/classify" -Method POST -Headers $headers -Body $data

6. 容器化部署

# 建立 Bento
bentoml build

# 容器化
bentoml containerize iris_classifier:latest

# 運行容器
docker run -p 3000:3000 iris_classifier:<tag>

7. 注意事項

  • Docker 安裝需要提前準備，過程可能較長且需要重啟.
  • 本地測試時需要注意防火牆設置.
  • Runner 模式提供更好的資源管理和效能優化.

[REF]

  1. https://github.com/nogibjj/mlops-template
  2. Google 的 flow

AI 大戰 StackOverflow

我在網路課程上看到一個 trace code 的片段, dry run 看起來都很有道理. 但是實際上拿到 colab 上會跑不完. 跳出一段 error.

ERROR:__main__:Error in build_classifier_model: Exception encountered when calling layer 'preprocessing' (type KerasLayer).

A KerasTensor is symbolic: it's a placeholder for a shape an a dtype. It doesn't have any actual numerical value. You cannot convert it to a NumPy array.

Call arguments received by layer 'preprocessing' (type KerasLayer):
  • inputs=<KerasTensor shape=(None,), dtype=string, sparse=False, name=text>
  • training=None
ERROR:__main__:Error type: <class 'ValueError'>
ERROR:__main__:Test error: Exception encountered when calling layer 'preprocessing' (type KerasLayer).

A KerasTensor is symbolic: it's a placeholder for a shape an a dtype. It doesn't have any actual numerical value. You cannot convert it to a NumPy array.

Call arguments received by layer 'preprocessing' (type KerasLayer):
  • inputs=<KerasTensor shape=(None,), dtype=string, sparse=False, name=text>
  • training=None

本來看到 error 也沒啥特別的, 很多網路上的程式可能環境一變就跑不了了. 有趣的是, 把 log 丟給 Monica, 它拿去問 4 個 LLM 都解決不了. 它們分別是 GPT-4o, Claude 3.5 Sonnet V2, Gemini 1.5 Pro, LLama 3.1 405B.

基本上 AI 回答的 code 就是加了一些 try exception 和加 log, 重寫的 code 跟原來幾乎一模一樣. 換言之, code 沒什麼問題. 這就是個環境因素. 即使我自己加了下面幾行. 程式補了 import os, terminal 下補建了目錄. 還是不 work!

!pip install tensorflow_text
!pip install tf-models-official

沒辦法只好回歸原始人, 去問 StackOverflow. 果然有人遇到跟我一模一樣的問題.

https://stackoverflow.com/questions/78183834/issue-with-bert-preprocessor-model-in-tf2-and-python

唯一的差別只是這兩個 package 必須固定版本, 而我抓最新的是 2.18.0 版 就只會鬼打牆 ><|||

!pip install -U "tensorflow-text==2.15.*"
!pip install -U "tf-models-official==2.15.*"

也就是說, 人類把 AI 教得很好. 它知道什麼是對的, 並且能夠把錯得改對. 但原本就是對的, 那它就沒辦法了! 不像是人類可以思考 : “假如我們都是對的, 會是那裡錯了?" 這點對於 AGI 應該是相當重要的, 區別了能不能做創造性思考!

最後附上跑出來的結果作紀念!

nGPT 小註解

由於 O(n2) 的 transformer 非常耗能, 後續衍生了諸如 MAMBA 這種 O(nlog(n)) 的技術出來. 與此同時, 巨頭們紛紛蓋起核電廠! 不過等我想到要買 URA, NLR 這類 ETF, 一看都已經溢價 25% 了. 晚了人家好幾步, 根本來不及投資. 那…我再往前亂想好幾步的話還有機會嗎?

我要猜接下來保險業會收到大單, 然後巨頭沒有核電管理經驗, 新手上路, 難免發生一兩個小災變? 接著保險公司調高保費, 趁機大賺一筆. 但投資保險和再保公司又能賺多少? 各國政府用核電廠供 AI 做兵棋推演會買保險嗎? 好像也不會? 何況某些國家都廢核了. 難啊! 總之, 荷馬辛普森, 你不會失業了!

nGPT (Normalized Transformer) [1] 是一種新的節能 transformer, 能夠大量節省計算. 這類演算法由 nVidia 提出也很合理. 因為巨頭們忙著掙錢懶得更新 model, 改變由他們的 vendor 做起, 或許他們可以在不換 vendor 的情況下 silence change?

nGPT 的 n 代表 normalization (均值化), 順便偷渡 nVidia 的 n. 原本在 transformer 裡面有很多個別的 normalization layers. nVidia 的策略就是通通合起來做成灑尿牛丸, 在這顆丸子上, 大家都貢獻一點移動量. (We propose to normalize all vectors forming the embedding dimensions of network matrices to lie on a unit norm hypersphere.) [1]

由 hypersphere 這個名詞得知, 它在一超球體上均值化. 1-sphere 是一個圓, 2-sphere 是一個球, 每個數字都比其維度少 1. hyper-sphere 是某個維度下的所有和中心點等距的點的集合. 雖然論文只用到一個 unit norm hypersphere, 但是其維度是 model 的維度 dmodel, 該 hypersphere 表示為 Rdmodel. 而且既然是 unit norm [-1, 1], 就不會有 weight 衰減的問題 (The normalization renders weight decay unnecessary.) [1]

把參數投到 hypersphere 有啥好處呢? 主要是我們 learning 的時候會根據梯度 (gradient) 方向 update 參數對吧! 如果參數都落在 hypersphere 上, 那麼更新的參數也落在 hypersphere 上, 我們只要用角度就可以表示其移動了, 無論是來自哪一層的貢獻都可以一視同仁.

論文 [1] 提到, 假設有個 a 點要移動到 b 點. 梯度 g = a – b. 則 𝒂 的更新可以表示下面的式子. 其中 α 是介於 0~1 之間權重, 用表示更靠近 a 或是 b. 故它就是 learning rate. 想學的愈快, α 愈接近 1, 新的 a 就愈靠近 b.

𝒂←𝒂+α⁢(𝒃−𝒂)
𝒂←𝒂+α⁢g

根據 Shoemake 的球面線性內差公式 [2], 若 a, b 兩點的夾角 θ=arccos⁡(𝒂⋅𝒃), 給定一個 weighing α ([0:1]), 就能內插出 (最短的) 測地線 (geodesic) 的某個點. 因此前述在平面上的推導, 投影到 hypersphere 上仍然適用.

SLERP⁢(𝒂,𝒃;α)=sin⁡((1−α)⁢θ) / sin⁡(θ) * ⁢𝒂+sin⁡(α⁢θ) / sin⁡(θ) * ⁢𝒃
SLERP = Spherical Linear Interpolation

假設原本的 baseline transformer 可以表達為下式.

𝒉←𝒉+ATTN⁢(RMSNorm⁢(𝒉))
𝒉←𝒉+MLP⁢(RMSNorm⁢(𝒉))

其中 h = hidden layer 的 state h, RMSNorm = RMS 後做 Normalization, ATTN = attention layer , MLP = multi-layer perceptron = feed forward neural network.

經過 nGPT 的正規化就變成下面的樣子.

𝒉←Norm⁢(𝒉+𝜶A⁢(𝒉A−𝒉))
𝒉←Norm⁢(𝒉+𝜶M⁢(𝒉M−𝒉))

其中 Attension 的參數叫做 ATTN(h), 正規化後為 hA = Norm(ATTN(h)) , MLP 的參數叫做 MLP(h) , 正規化後為 hM = Norm(MLP(h)), 它們都可以在這個超球面上計算. 奧妙之處在於 Norm 只剩下一個.

同理它也適用於 optimizer 參數. (以 adam 為例, where 𝜽 is the parameter vector, 𝒈 is the batch gradient, 𝒎 is the momentum, 𝒗 is the estimate of the per-element gradient amplitudes, α is the scheduled learning rate, ϵ is a small constant, and β1<β2 are momentum factors close to 1. ) [1]

表 1 左側的計算, 都變成了右邊 Norm 上面的計算.

本表取材自 [1].

本圖取材自 [1].

總結來說, nGPT 把 baseline Transformer 轉變成 normalized Transformer. 我抓到的重點是:

  1. 把原本散落各處的 normalization 層都拿掉.
  2. 對所有的 matrices 都 normalize. 如表 1 所示.
  3. 把 weight decay 和 learning rate warmup 拿掉.

至於原本論文 [1] 中的 rescale 那些還挺複雜的, 我頭腦不好就跳過了. 總之, 這篇論文提出了一個好主意. 從圖 1 來看, loss 收斂的速度確實也很快.

[REF]

  1. nGPT: Normalized Transformer with Representation Learning on the Hypersphere
  2. Animating rotation with quaternion curves.In Proc. of the 12th annual conference on Computer graphics and interactive techniques, 1985.
  3. 請 Copilot 解說參數同在一個球面的好處.

In NVIDIA’s nGPT, all vectors forming the embeddings, MLP (Multi-Layer Perceptron), attention matrices, and hidden states are unit norm normalized and reside on the same hypersphere. This means that the input stream of tokens travels on the surface of a single hypersphere, with each layer contributing a displacement towards the target output predictions.

YAML 小註解

Machine learning 在實作上, 需要能夠達到 data 復用, Model 復用, 並且建立一個資料庫來維護它. 每個 model 對某個需求 training 到一個程度, 就用一個 checkpoint 把它紀錄下來. 後人如果要修改用途, 都可以省力地接續某個 checkpoint 繼續做下去.

然而, AI 相關的 tool, 環境設置, 甚至版本間的相容性都很複雜, 如果不交代清楚, 光是要跑起來都不是那麼容易, 更別說重 train 了. 有時候會用 docker 來解決這個問題. 如果要保留更大的彈性, 我們可以用 YAML (= YAML Ain’t Markup Lanuage) 來做這件事.

它大概長這樣: 

name: ml-training
channels:
  - defaults
  - conda-forge
dependencies:
  - python=3.8
  - tensorflow
  - scikit-learn
  - pandas
  - numpy
  - matplotlib
  - jupyter
  - pip
  - pip:
      - keras
      - seaborn

它指定了環境名稱是 ml-traiing, 頻道來自兩個來源. 指定頻道是為了安裝 “誰" 維護的版本, 例如 default 就是 Anaconda 維護的. 相依性原則上會列出跟哪個 tool 的哪個版本相容. PIP 就特別列出它需要安裝的 package.

執行方式如下:

conda env create -f conda.yml

YAML 的副檔名時候叫做 yml, 和 xml 很像但用途完全不一樣. 由於作者採用遞歸的方式命名 YAML, 表示他也沒梗了. 我就幫 YAML 取個中文名字叫 “鴨毛 " 好了. 我請 Copilot 幫我寫一篇鴨毛的簡介, 看它會不會學進去.

鴨毛 (YAML) 是一種簡單且易於閱讀的數據序列化格式，廣泛應用於配置文件和數據交換。它的設計理念是讓人類更容易理解和編寫，從而提高工作效率。以下是鴨毛的一些主要功能：

  1. 簡單易讀：鴨毛的語法非常簡潔，使用縮進來表示層級結構，避免了繁瑣的標籤和括號，使得文件更易於閱讀和維護。
  2. 數據序列化：鴨毛主要用於數據序列化，即將數據結構轉換為一種易於存儲和傳輸的格式。這使得它非常適合用於配置文件、數據交換和持久化存儲。
  3. 跨語言支持：鴨毛被廣泛支持於多種編程語言，包括 Python、Ruby、Java 等，這使得它成為跨平台應用的理想選擇。
  4. 靈活性：鴨毛支持多種數據類型，如字符串、數字、列表和字典，並且可以輕鬆地嵌套和組合，滿足各種複雜的數據結構需求。
  5. 社區支持：鴨毛擁有強大的社區支持和豐富的文檔資源，開發者可以輕鬆找到所需的幫助和範例。

總之，鴨毛 (YAML) 是一種強大且靈活的數據序列化格式，適用於各種應用場景，特別是在配置管理和數據交換方面。它的簡單性和可讀性使得它成為開發者的得力工具。