| --- |
| 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)。 |
|
|
