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 — 三平台自動偵測，無需修改任何程式碼。
目錄
模型架構
| 項目 | 數值 |
|---|---|
| 總參數量 | ~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（推薦）
# 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+）
# 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：確認平台偵測
python train_llm.py --mode info
輸出範例（Apple M2 Pro）：
{
"device": "mps",
"device_name": "Apple M2 Pro",
"use_bf16": true,
"recommended_batch": 1,
"fused_adamw": false,
"dataloader_workers": 0
}
步驟 1：訓練分詞器（首次執行，只需一次）
python train_llm.py --mode tokenizer --data_dir ./data
從 FineWeb2 下載約 500 萬行語料訓練 SentencePiece BPE 分詞器（56K 詞彙，中英文及多語言 character coverage 99.95%）。
步驟 2：訓練模型
# 最簡單：使用預設 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：推論測試
python train_llm.py --mode demo
資料集設定
支援三種方式指定訓練資料，可任意混合多個來源。
方式 1：JSON 設定檔（最彈性）
python train_llm.py --mode train --dataset_config 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 快速指定
# 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（無需額外設定）
# 預設語言：簡體中文 + 繁體中文 + 英文
python train_llm.py --mode train
# 自訂語言過濾
python train_llm.py --mode train --langs "zho_Hans,eng_Latn"
多來源混合策略
# 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）
python train_llm.py --mode train \
--model_name kuixing-1.15b \
--max_steps 1000000 \
--seq_len 4096 \
--lr 2e-4
接續訓練（Continue）
從 Trainer Checkpoint 接續（訓練中斷後繼續，保留優化器狀態與步驟）：
python train_llm.py --mode train \
--train_mode continue \
--resume_from_checkpoint ./checkpoints/checkpoint-50000 \
--max_steps 1000000
從 Release Export 接續（換資料集繼續預訓練，步驟重設）：
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 重新打包（不需重新訓練）：
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
# 基本使用（讀取預設 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。
訓練資料 FineWeb2 由 HuggingFace 提供，請遵守其資料集授權。
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support
Model tree for jslin09/KuiXing
Unable to build the model tree, the base model loops to the model itself. Learn more.