README.md

---
license: apache-2.0
datasets:
- HuggingFaceFW/fineweb-2
language:
- en
- zh
base_model:
- jslin09/KuiXing
---
# KuiXing(魁星)-1.15B — 繁中 / 英 MoE 大型語言模型訓練框架

**KuiXing(魁星)-1.15B** 是一個從零開始訓練的 ~1.15B 參數稀疏 Mixture-of-Experts (MoE) 語言模型，支援繁體中文、簡體中文與英文，最大 context window 為 **1M tokens**（訓練 128K + YaRN 外推 ×8）。詞彙量擴充至 **56K**，對中文字符與多語言符號有更完整的覆蓋。

> **平台**：NVIDIA CUDA · Apple Silicon (MPS) · CPU — 三平台自動偵測，無需修改任何程式碼。

---

## 目錄

- [模型架構](#模型架構)
- [環境安裝](#環境安裝)
- [快速開始](#快速開始)
- [資料集設定](#資料集設定)
- [訓練模式](#訓練模式)
- [CLI 完整參數](#cli-完整參數)
- [發布存檔格式](#發布存檔格式)
- [專案結構](#專案結構)
- [常見問題](#常見問題)
- [授權](#授權)

---

## 模型架構

| 項目 | 數值 |
|------|------|
| 總參數量 | ~1.15B |
| 激活參數量（推理時）| ~460M |
| 隱藏層總數 | 24 |
| 其中 Dense 層 | 16（每 3 層中的前 2 層）|
| 其中 MoE 層 | 8（每 3 層中的第 3 層）|
| 隱藏維度 | 2048 |
| Attention 機制 | GQA — Q heads: 16 / KV heads: 4 |
| Head 維度 | 128 |
| Dense FFN 中間維度 | 5632（SwiGLU）|
| MoE 專家數 | 16（top-2 稀疏激活）|
| 每個 Expert 中間維度 | 2048 |
| 位置編碼 | YaRN RoPE（θ=500000, factor=8）|
| 訓練 context 長度 | 128K tokens |
| 推理 context 長度 | 最大 **1M tokens**（YaRN 外推）|
| 注意力策略 | 偶數層：全注意力；奇數層：Sliding Window (4096)|
| Normalization | RMSNorm（ε=1e-5，float32 計算）|
| 詞彙量 | 56,000（SentencePiece BPE）|
| MoE 輔助損失 | Load Balancing Loss + Router Z-Loss |

---

## 環境安裝

### NVIDIA CUDA（推薦）

```bash
# PyTorch with CUDA 12.1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers datasets sentencepiece tensorboard accelerate safetensors
```

### Apple Silicon（macOS 12.3+）

```bash
# PyTorch with MPS（BF16 需 PyTorch >= 2.3）
pip install torch torchvision torchaudio
pip install transformers datasets sentencepiece tensorboard accelerate safetensors
```

> **注意**：macOS 上的 OpenMP 衝突問題已由程式自動處理，不需要手動設定環境變數。

### 最低 Python 版本

Python **3.10+**（使用了 walrus operator `:=`）

---

## 快速開始

### 步驟 0：確認平台偵測

```bash
python train_llm.py --mode info
```

輸出範例（Apple M2 Pro）：
```json
{
  "device": "mps",
  "device_name": "Apple M2 Pro",
  "use_bf16": true,
  "recommended_batch": 1,
  "fused_adamw": false,
  "dataloader_workers": 0
}
```

### 步驟 1：訓練分詞器（首次執行，只需一次）

```bash
python train_llm.py --mode tokenizer --data_dir ./data
```

從 FineWeb2 下載約 500 萬行語料訓練 SentencePiece BPE 分詞器（56K 詞彙，中英文及多語言 character coverage 99.95%）。

### 步驟 2：訓練模型

```bash
# 最簡單：使用預設 FineWeb2，所有超參數自動偵測
python train_llm.py --mode train --model_name kuixing-1.15b

# 使用自訂資料集設定
python train_llm.py --mode train \
    --dataset_config dataset_config.json \
    --mix_strategy weighted \
    --model_name kuixing-1.15b
```

### 步驟 3：推論測試

```bash
python train_llm.py --mode demo
```

---

## 資料集設定

支援三種方式指定訓練資料，可任意混合多個來源。

### 方式 1：JSON 設定檔（最彈性）

```bash
python train_llm.py --mode train --dataset_config dataset_config.json
```

參見 [`dataset_config.json`](dataset_config.json) 範例（FineWeb2 + Wikipedia 混合）。

`dataset_config_local.json` 示範本地 JSONL 檔案的用法。

**DatasetSource 完整欄位：**

| 欄位 | 型別 | 預設 | 說明 |
|------|------|------|------|
| `source` | str | `"huggingface"` | `huggingface` / `local_files` / `local_dir` |
| `path` | str | — | HF dataset id 或本地路徑（多檔用逗號分隔）|
| `name` | str | `""` | HF subset name，如 `"20231101.zh"` |
| `split` | str | `"train"` | HF split |
| `text_field` | str | `"text"` | 要讀取的欄位名稱 |
| `filter_field` | str | `""` | 過濾欄位（空=不過濾）|
| `filter_values` | list | `[]` | 允許通過的值清單 |
| `streaming` | bool | `true` | 串流載入（省記憶體）|
| `shuffle` | bool | `true` | 是否打亂 |
| `buffer_size` | int | `10000` | shuffle 緩衝大小 |
| `min_length` | int | `50` | 最短文字長度（字元）|
| `max_samples` | int | `0` | 最多取用樣本數（0=不限）|
| `weight` | float | `1.0` | weighted 混合時的取樣比例 |
| `file_format` | str | `"txt"` | 本地格式：`txt` / `jsonl` / `csv` |
| `glob_pattern` | str | `"**/*.txt"` | local_dir 模式的 glob |
| `csv_delimiter` | str | `","` | CSV 分隔符 |
| `seed` | int | `42` | 隨機種子 |
| `label` | str | _(path basename)_ | 日誌顯示標籤 |

### 方式 2：CLI 快速指定

```bash
# HuggingFace 資料集
python train_llm.py --mode train \
    --dataset_path "wikimedia/wikipedia" \
    --dataset_name "20231101.zh" \
    --text_field "text"

# 本地 JSONL 檔案
python train_llm.py --mode train \
    --dataset_path "/data/corpus.jsonl" \
    --source_type local_files \
    --file_format jsonl \
    --text_field "content"

# 本地目錄（掃描所有 .txt）
python train_llm.py --mode train \
    --dataset_path "/data/articles/" \
    --source_type local_dir \
    --file_format txt
```

### 方式 3：預設 FineWeb2（無需額外設定）

```bash
# 預設語言：簡體中文 + 繁體中文 + 英文
python train_llm.py --mode train

# 自訂語言過濾
python train_llm.py --mode train --langs "zho_Hans,eng_Latn"
```

### 多來源混合策略

```bash
# sequential（預設）：依序消耗每個來源
python train_llm.py --mode train --dataset_config dataset_config.json

# weighted：依 weight 比例同時交錯取樣（推薦多語料混訓）
python train_llm.py --mode train \
    --dataset_config dataset_config.json \
    --mix_strategy weighted
```

---

## 訓練模式

### 從頭訓練（Pretrain）

```bash
python train_llm.py --mode train \
    --model_name kuixing-1.15b \
    --max_steps 1000000 \
    --seq_len 4096 \
    --lr 2e-4
```

### 接續訓練（Continue）

**從 Trainer Checkpoint 接續**（訓練中斷後繼續，保留優化器狀態與步驟）：

```bash
python train_llm.py --mode train \
    --train_mode continue \
    --resume_from_checkpoint ./checkpoints/checkpoint-50000 \
    --max_steps 1000000
```

**從 Release Export 接續**（換資料集繼續預訓練，步驟重設）：

```bash
python train_llm.py --mode train \
    --train_mode continue \
    --resume_from_checkpoint ./checkpoints/release/kuixing-1.15b-20250101_120000/model \
    --dataset_path "wikimedia/wikipedia" \
    --dataset_name "20231101.zh" \
    --lr 1e-4 \
    --max_steps 200000
```

### 重新發布存檔

對已完成訓練的 checkpoint 重新打包（不需重新訓練）：

```bash
python train_llm.py --mode export \
    --output_dir ./checkpoints \
    --tokenizer_model ./data/spm_tokenizer.model \
    --model_name kuixing-1.15b
```

---

## CLI 完整參數

```
--mode              train | tokenizer | demo | info | export
--model_name        發布名稱前綴（預設: kuixing-1.15b）
--data_dir          分詞器語料目錄（預設: ./data）
--output_dir        訓練輸出目錄（預設: ./checkpoints）
--tokenizer_model   SPM 模型路徑（預設: ./data/spm_tokenizer.model）

訓練模式：
--train_mode        pretrain（預設）| continue
--resume_from_checkpoint  接續訓練的 checkpoint 路徑

資料集（三選一）：
--dataset_config    JSON 設定檔路徑（最高優先）
--dataset_path      單一資料集路徑（HF id 或本地路徑）
--dataset_name      HF subset name
--dataset_split     HF split（預設: train）
--text_field        文字欄位名稱（預設: text）
--filter_field      過濾欄位名稱
--filter_values     過濾值，逗號分隔
--source_type       huggingface | local_files | local_dir
--file_format       txt | jsonl | csv
--langs             FineWeb2 語言清單，逗號分隔（預設行為）
--mix_strategy      sequential（預設）| weighted
--no_streaming      停用串流，完整下載後載入

超參數：
--batch_size        per-device batch size（-1=自動）
--grad_accum        gradient accumulation steps（-1=自動）
--lr                學習率（預設: 2e-4）
--max_steps         總訓練步數（串流資料集用；預設 1,000,000）
                    與 --num_epochs 同時指定時 epoch 模式優先
--num_epochs        訓練 epoch 數（有限資料集用；-1=停用，改用 --max_steps）
                    適合本地資料集或固定大小的 HuggingFace 資料集
--seq_len           訓練序列長度（預設: 4096）
--warmup_steps      warmup 步數（預設: 4000）

Checkpoint 儲存（--save_steps 與 --save_total_limit 搭配使用）：
--save_steps        每幾步儲存一個 checkpoint（預設: 5000）
--save_total_limit  最多同時保留幾份 checkpoint（預設: 3）
                      0 = 無限制，保留所有 checkpoint
                      例: --save_steps 2000 --save_total_limit 10

精度：
--bf16              BF16：-1=自動，0=關，1=開
--fp16              FP16：-1=自動，0=關，1=開

Loss 記錄：
--loss_log_file     Training loss CSV 路徑
                      空字串 = 自動使用 {output_dir}/training_loss.csv
                      接續訓練時自動 append，不覆蓋已有記錄

其他：
--no_grad_ckpt      停用 gradient checkpointing
--workers           DataLoader workers（-1=自動）
```

---

## Training Loss 記錄與繪圖

### CSV 格式

訓練過程中自動產生 `{output_dir}/training_loss.csv`（或 `--loss_log_file` 指定路徑）：

```
step,epoch,loss,learning_rate,grad_norm,samples_seen,elapsed_sec
100,0.0,8.312451,0.0002,1.2341,3200,45.2
200,0.0,7.891234,0.00019,1.1892,6400,89.7
...
1000000,0.0,2.341200,2e-05,0.8123,3200000,18420.0
1000000,0.0,END,,,,18421.1
```

- **接續訓練**：自動 append，`END` 行標記每段訓練結束，可區分多次訓練
- **即時 flush**：每個 log step 寫入後立即 flush，中斷也不丟失記錄

### 繪圖工具 `plot_loss.py`

```bash
# 基本使用（讀取預設 CSV，輸出 PNG）
python plot_loss.py

# 指定路徑與輸出
python plot_loss.py \
    --csv ./checkpoints/training_loss.csv \
    --out ./loss_curve.png

# 印出訓練摘要統計（最終 loss、最低點、收斂步數）
python plot_loss.py --summary

# 比較 pretrain + continue 兩段訓練
python plot_loss.py \
    --csv ./run1/training_loss.csv ./run2/training_loss.csv \
    --labels "Pretrain" "Continue (Wikipedia)" \
    --out compare.png

# 以訓練時間為 X 軸，顯示互動視窗
python plot_loss.py --x_axis elapsed_sec --show

# 過濾初期 spike，調整平滑視窗
python plot_loss.py --max_loss 10.0 --smooth 100
```

圖表包含三個面板：

| 面板 | 內容 |
|------|------|
| Loss 曲線 | 原始 loss（半透明）+ 滾動平均平滑，自動標記最低點 |
| Learning Rate | Cosine decay + warmup 排程曲線 |
| Gradient Norm | L2 norm 趨勢（反映訓練穩定性）|

---

## 發布存檔格式

訓練完成後自動產生（也可用 `--mode export` 手動觸發）：

```
checkpoints/release/{model_name}-{timestamp}/
├── model/
│   ├── model.safetensors      # 模型權重（SafeTensors，推薦）
│   ├── pytorch_model.bin      # 模型權重（PyTorch bin，相容備用）
│   ├── config.json            # 架構與超參數設定
│   └── generation_config.json # 預設生成參數
├── tokenizer/
│   ├── spm_tokenizer.model    # SentencePiece 模型
│   ├── spm_tokenizer.vocab    # 詞彙表（piece + BPE score）
│   ├── tokenizer_config.json  # HuggingFace tokenizer 設定
│   └── special_tokens_map.json
├── model_card.md              # HuggingFace Hub 模型說明卡
├── manifest.json              # 所有檔案 SHA-256 + 大小清單
└── release_info.json          # 訓練環境、超參數完整快照
```

---

## 建議硬體配置

| 硬體 | batch | grad_accum | seq_len | 精度 |
|------|-------|------------|---------|------|
| A100 80G | 4 | 8 | 4096 | BF16 |
| A100 40G | 2 | 16 | 4096 | BF16 |
| RTX 4090 24G | 1 | 32 | 2048 | BF16 |
| RTX 3090 24G | 1 | 32 | 2048 | FP16 |
| M3 Max 128G | 2 | 16 | 4096 | BF16 |
| M2 Ultra 192G | 2 | 16 | 4096 | BF16 |
| M2 Max 96G | 1 | 32 | 2048 | BF16 |
| M1/M2 16-24G | 1 | 32 | 1024 | FP32 |

> **1M context 推理**需搭配 Flash Attention 2（CUDA A100/H100）或足夠的 Apple Silicon Unified Memory。訓練時 `seq_len=4096` 即可；長上下文外推由 YaRN 在推理時自動完成。

---

## 專案結構

```
.
├── train_llm.py                 # 主程式（分詞器 / 訓練 / 推論 / 存檔）
├── plot_loss.py                 # Training loss 曲線繪圖工具
├── dataset_config.json          # 多來源混合範例（FineWeb2 + Wikipedia）
├── dataset_config_local.json    # 本地 JSONL 資料集範例
├── requirements.txt             # Python 套件需求
├── README.md                    # 本文件
├── CHANGELOG.md                 # 版本變更記錄
├── LICENSE                      # Apache 2.0 授權
└── .gitignore                   # Git 排除規則
```

---

## 常見問題

**Q: macOS 出現 `OMP: Error #15: Initializing libomp.dylib` 然後 abort？**

A: 這是 macOS 上 PyTorch、sentencepiece 等套件各自靜態連結不同版本 libomp 所導致的衝突。本程式已在啟動時自動設定 `KMP_DUPLICATE_LIB_OK=TRUE` 等環境變數，理論上不會出現此問題。若仍發生，請確認您使用的是 `python train_llm.py` 而非直接 import 本模組。

**Q: MPS 上訓練比 CPU 還慢？**

A: 部分運算（如 MoE router 的 scatter/gather）在 MPS 上會 fallback 至 CPU，導致額外的資料搬移開銷。可嘗試減小 `--seq_len` 或 `--batch_size` 以提高 Metal GPU 利用率。

**Q: 1M context 真的能用嗎？**

A: 訓練時固定使用 `--seq_len 4096`（或自訂），推理時 YaRN 外推讓模型能處理更長序列。實際最大長度受限於可用記憶體：80GB A100 約可推理 128K–512K tokens（使用 Flash Attention + KV cache 量化）。

**Q: `safetensors` 套件未安裝時怎麼辦？**

A: 程式會自動偵測。若未安裝，則跳過 SafeTensors 格式，仍輸出 `pytorch_model.bin`。建議安裝：`pip install safetensors`。

**Q: 如何只訓練分詞器而不訓練模型？**

A: `python train_llm.py --mode tokenizer --data_dir ./data`

---

## 授權

本專案採用 **Apache License 2.0**。詳見 [LICENSE](LICENSE)。

訓練資料 FineWeb2 由 HuggingFace 提供，請遵守其[資料集授權](https://huggingface.co/datasets/HuggingFaceFW/fineweb-2)。