Co-design of a TinyLLM using Programmable Logic and Software on an FPGA

### **1. Framework (框架)**
這篇論文提出了一個 **TinyLLM 設計框架**，專門針對 **FPGA (Field Programmable Gate Array) 與軟體 (Software) 的協同設計 (Co-design)**，目標是實現在低功耗、低計算資源環境下執行 LLM（Large Language Model）。這個框架的核心概念包含 **硬體/軟體 (PS/PL) 分割、量化 (Quantization)、優化數據傳輸** 等關鍵技術。

#### **(1) 硬體與軟體的劃分 (Hardware/Software Partitioning)**
- **硬體部分 (Programmable Logic, PL)**：
- **加速矩陣乘法 (Matrix Multiplication)**，因為矩陣計算在 PS (Processing System) 上效率低，故移到 PL 來加速。
- **使用 DSP (Digital Signal Processing) Slice 進行 MAC (Multiply-Accumulate) 運算**，透過流水線 (Pipeline) 增加吞吐量。
- **透過 Direct Memory Access (DMA) 在 PS/PL 之間傳輸資料**，減少延遲。

- **軟體部分 (Programmable Software, PS)**：
- **負責 LayerNorm、Softmax、GELU 等函數計算**，因為這些操作需要浮點數計算，而 PL 的資源有限。
- **處理 Tokenization 與文本處理**，使用 Hugging Face Tokenizer 將輸入轉換為模型可讀的格式。
- **數據管理與調度**，決定哪些部分要交給 PL 計算，並控制數據流。

#### **(2) 設計空間探索 (Design Space Exploration)**
- 選擇 **TinyStories LLM** 作為基礎模型，因為該模型比典型 LLM 更小巧，適合 FPGA。
- 使用 **SmoothQuant 量化技術**，將部分層轉換為 8-bit 來降低計算成本。
- 在 FPGA 內部進行軟體分析，發現 **95% 的計算時間來自矩陣乘法 (matmuls)**，因此選擇將這部分放到 PL。

#### **(3) PL 內部架構**
- 設計了一個 **AXI-Stream DMA-based 資料傳輸架構**：
- **權重 (Weights) 存在 PS 端**，使用 AXI-Stream 送到 PL 。
- **矩陣運算結果 (Matrix Multiplication Output) 透過 DMA 回傳 PS 記憶體**。
- **利用 4 個 DSP 單元進行矩陣運算**，並透過流水線來加速計算。

---

### **2. 模型 (Model)**
TinyLLM 的模型基於 **TinyStories LLM**，其架構如下：

#### **(1) TinyStories LLM 特色**
- **較小的參數數量 (約 200MB)**，適合嵌入式設備。
- **仍能生成可讀文本**，但效能比 GPT-2、OPT-175B 低。
- **原始模型使用 32-bit 浮點數運算**，但為了適應 FPGA 硬體，本文使用了 **SmoothQuant 量化**，將部分層轉為 **8-bit 整數運算**。

#### **(2) 模型架構**
- **輸入處理 (Tokenization)**：
- 使用 Hugging Face Tokenizer 分割文本，並將其轉換為 Token ID。
- **Transformer Block**：
- **Self-Attention**
- **LayerNorm**
- **Feedforward Fully Connected (FC) Layers**
- **輸出層 (Logits Layer)**：
- 預測下一個 token。

---

### **3. 模型訓練 (Training)**
本研究的 TinyLLM 主要是 **基於預訓練模型 (Pre-trained Model) 進行推論 (Inference)**，而非從頭開始訓練，但它仍然使用了一些訓練技巧來優化其適應性。

#### **(1) 量化 (Quantization)**
- 使用 **SmoothQuant** 技術：
- 透過 **a priori scaling** 來減少量化時的誤差，提升 8-bit 運算的準確度。
- 在 **7 層 Transformer Block** 上應用量化，使得矩陣計算能夠在 FPGA PL 上執行。

#### **(2) 訓練策略**
- **先在雲端 x86 CPU 上測試** 量化 TinyStories LLM，確保其運作正常後再部署至 FPGA。
- **轉換 PyTorch 為 NumPy**，以適應 FPGA 上的 ARM PS 核心，因為 PyTorch 目前沒有適用於 ARMv7 PS 的官方支持。
- **透過 Profiling 分析瓶頸**，優化計算部分，如：
- **矩陣乘法移至 PL**
- **LayerNorm、Softmax 留在 PS**

#### **(3) 訓練數據**
- 使用 **TinyStories Corpus**，這是一個專門為小型 LLM 設計的語料庫，確保模型可以在有限資源下生成合適的文本。

---

### **4. 實驗結果 (Results)**
- **效能提升**：
- **FPGA Hybrid 設計比純 PS 版本快 37 倍**（0.26 秒 vs 9.7 秒）。
- **Fully Connected 層的加速高達 64 倍**。
- **整體 Token 生成速度為 3.78 tokens/s**。

- **功耗 (Power Efficiency)**：
- **整體功耗為 1.47W**，其中 **PS 消耗 1.27W**，PL 消耗極少。
- 比傳統 GPU 更省電，適合嵌入式應用。

- **資源使用率**：
- **僅使用 9.29% 的 LUTs**，顯示還有進一步擴展空間。
- **DSP 使用率僅 1.8%**，可增加更多並行運算來提高速度。

---

### **5. 總結**
| **類別** | **內容** |
|----------|---------|
| **Framework (框架)** | FPGA + ARM 雙架構，PS 負責控制 & 軟體處理，PL 負責矩陣計算 |
| **模型 (Model)** | 基於 TinyStories LLM，應用量化技術 (SmoothQuant) |
| **訓練 (Training)** | 先在 x86 測試，轉換為 NumPy，移動矩陣計算至 PL |
| **優化方法** | 量化、硬體/軟體劃分、DMA 傳輸加速 |
| **效能提升** | 37 倍加速、功耗僅 1.47W、Token 生成速度 3.78 tokens/s |

這個設計展示了一種 **適用於低功耗 FPGA 平台的 LLM 部署方法**，透過量化和 FPGA 加速技術，使得 LLM 可以在嵌入式環境中運行，而無需依賴雲端運算。這對於 IoT、邊緣 AI、智慧設備等應用具有重要意義。🚀