Upload 6 files

0.9.3 版（新增：tokenize 結果磁碟快取，避免重複從頭訓練時的耗時重算）
# 14. 大規模語料（數億句、數百億字元等級）逐句 tokenize 往往要數小時，
# 但只要「資料集組合（含 extra_datasets）+ chunk_size + tokenizer
# 內容」三者不變，build_chunks() 的輸出結果必然相同。新增
# _chunks_cache_key()：以這三者算出 SHA-256 指紋（tokenizer 取其
# 序列化後的模型內容做 hash，而非僅看檔案路徑，確保 tokenizer 重新
# 訓練後即使路徑不變，舊快取也會因內容指紋不同而自動失效，不會誤讀
# 出用舊 vocab 編碼、卻對不上新 tokenizer 的錯誤資料）。build_chunks()
# 新增 cache_dir / use_cache 參數：若指紋對應的 .npy 快取檔已存在則
# 直接讀取回傳並跳過整段 tokenize；否則照常 tokenize，並在完成後將
# 結果寫入快取供下次使用。常見情境：同一份資料集反覆從頭訓練做除
# 錯、調整模型架構或超參數但資料不變時，可省去每次數小時的重複
# tokenize。MLX（train_model_mlx）與 PyTorch（train_model）兩條主
# 訓練路徑皆已接上此快取（快取目錄預設為 tokenizer_dir 下的
# chunks_cache 子目錄）；build_replay_chunks() 因本身已有獨立的
# replay_cache_dir 快取機制，呼叫 build_chunks() 時明確關閉內建快取
# （use_cache=False），避免同一份資料在兩處各存一份造成磁碟空間浪
# 費。注意：快取檔為未壓縮的 int32 .npy 陣列，規模與原始 token 總量
# 成正比（例如 593 萬筆 × 2049 tokens 的設定下，單一快取檔約 48.6
# GB），啟用前請確認磁碟空間充足。
# 0.9.4 版（新增：--cache-clean 參數，獨立控制 chunks 快取的清除）
# 15. 新增 CLI 參數 --cache-clean：訓練開始前清除 tokenizer_dir/
# chunks_cache 下所有 .npy 快取檔，強制下次 build_chunks() 重新
# tokenize。設計為獨立旗標、不綁在 --from-scratch 底下，因為
# 「要不要從頭訓練模型權重」與「要不要重用 tokenize 快取」是兩件
# 獨立的事：--from-scratch 控制的是模型權重，--cache-clean 控制的
# 是 chunks 快取，兩者可任意組合（同時用、各自單獨用、或都不用）。
# 典型情境是接續訓練（--resume）時用 --extra 新增了資料集——雖然
# 0.9.3 版的 cache key 已含資料集組合指紋，新組合本來就不會誤讀舊
# 快取，但若想徹底清空舊快取釋放磁碟空間，可加此參數手動觸發。清除
# 範圍僅限 chunks_cache 內的 .npy 檔，不觸碰 tokenizer.model/.vocab
# 本身、不觸碰已訓練的模型權重（model_dir），也不觸碰
# replay_cache_dir（Replay Buffer 為獨立快取機制，語意不同，使用者
# 未明確要求時不一併清除）。
# 0.9.5 版（效能修正：手寫 attention 改用各平台的 fused attention kernel，
# 解決記憶體佔用過高導致 batch_size 無法調大、訓練時間過長的問題）
# 16. 問題背景：0.9.2 版修正 epoch_mode 的 steps 計算後，1 epoch 真正需要
# 跑滿資料集對應的 micro-step 數（不再是先前 bug 版本因訓練量被低估
# 32 倍而「看似」很快跑完）。但 batch_size 維持原本的 4 太小，H100
# 93GB 顯存大部分閒置，GPU 利用率低，實測 1 epoch 預估需時約 215
# 小時。嘗試將 batch_size 提高到 16（accum_steps 對應降到 8，
# effective batch 維持 128 不變）以提升 GPU 利用率時，卻發生 CUDA
# OOM：錯誤訊息顯示 activation 記憶體佔用約 73.6GB，遠超模型權重+
# 梯度+optimizer state 的固定開銷（約 13GB）。
# 17. 根本原因：KuiXingAttention（PyTorch）與 KuiXingAttention_MLX（MLX）
# 皆為手寫實作 softmax(QK^T/√d)V，會把完整的 (B, H, L, L) attention
# weight matrix materialize 到記憶體裡，且為了 backward 必須在 12
# 層中各自保留一份。在 L=2048（chunk_size-1）下，這個矩陣的記憶體佔
# 用對 batch_size 是線性放大，batch_size 16 時 12 層加總估計即達
# 50GB 以上，這正是 OOM 的直接原因；也是 batch_size=4 時 H100 算力
# 嚴重未被利用、需要 215 小時才能跑完 1 epoch 的根本原因（micro-step
# 數量龐大、每個 micro-step 的 GPU 運算時間相對 Python/同步開銷占比
# 過低）。
# 18. 修法：兩個平台分別改用各自官方提供的 fused attention kernel，數學
# 上皆等價於原本的 softmax(QK^T/√d)V（已用獨立腳本以隨機 q/k/v 數值
# 交叉驗證，PyTorch 版輸出與原手寫版本最大絕對誤差僅 4.77e-07，屬浮
# 點運算順序差異的正常量級）：
# - PyTorch（KuiXingAttention）：改用
# torch.nn.functional.scaled_dot_product_attention，causal
# masking 改用 is_causal=True 由 kernel 處理（不再需要呼叫端
# 手動建構 (L,L) 的 -1e4 遮罩矩陣，KuiXingLM.forward() 中對應
# 移除該矩陣建構），attention-weight dropout 改用該函式的
# dropout_p 參數（訓練時為 config.dropout，eval 時自動為 0，
# 與原本 self.attn_drop 模組行為一致）。
# - MLX（KuiXingAttention_MLX）：改用
# mx.fast.scaled_dot_product_attention（Apple/Argmax 團隊實作
# 的 flash-attention-style fused kernel），causal masking 改用
# mask="causal" 字串模式由 kernel 處理。**與 PyTorch 版的差異**：
# 此函式官方參數簽名目前未提供 dropout_p，無法將
# attention-weight dropout 融入 kernel 內部計算，故 MLX 版移除
# 了原本的 self.attn_drop（其餘 embedding/residual/MLP 的
# dropout 不受影響、維持原樣）。若之後實測訓練出現過擬合傾向明
# 顯增加，可考慮以提高其他層的 dropout 或 weight decay 補償。
# 因記憶體佔用大幅降低，預期同樣 93GB 顯存下可支援的 batch_size 上限
# 將遠高於 16，可依實測顯存餘裕進一步調大、減少 micro-step 總數、
# 提升 GPU 利用率，藉此縮短訓練總時間。
# 注意：本次修改僅在 PyTorch（CUDA）環境以獨立數值腳本驗證過正確
# 性；MLX 版的程式邏輯雖已對照官方 API 文件與範例確認寫法正確，但
# 因本開發環境無 Apple Silicon / Metal，無法實際執行驗證，請在
# M3 Max 上實際跑訓練前幾步，確認無 runtime 錯誤、loss 數值正常
# （非 NaN 或異常飆高）後再正式長時間訓練。
# 0.9.6 版（除錯用：加入 GPU 記憶體峰值追蹤診斷，定位第二次 OOM 的真正原因）
# 19. 0.9.5 版修正 attention 記憶體問題後，batch_size=16 訓練時仍在
# loss.backward() 處發生 CUDA OOM（嘗試配置 13.38 GiB）。手動估算
# 顯示：vocab_size=109,568 搭配 batch_size=16、seq_len=2048 時，
# KuiXingLM.forward() 最終輸出的 logits 張量（(B,L,vocab) float32）
# 本身就約 14.36 GB，與錯誤訊息中的配置失敗量級相符，研判這次的瓶頸
# 很可能轉移到 lm_head 輸出的 logits 張量與 cross_entropy 計算，而
# 非先前已修正的 attention 部分。為了在動手修改架構（例如將
# cross-entropy 計算分塊、避免一次 materialize 完整 logits）之前
# 先取得確切證據，而非憑空臆測，新增暫時性的記憶體峰值診斷：在每個
# epoch 最前面 5 個 micro-step，分別於 forward 完成後、backward 完
# 成後呼叫 torch.cuda.reset_peak_memory_stats() /
# torch.cuda.max_memory_allocated()，印出 forward 前後即時用量、
# forward 峰值、logits 張量本身大小、backward 峰值，藉此確認記憶體
# 究竟在哪個階段、哪個張量達到尖峰，再據此決定後續修法（例如分塊
# cross entropy、降低 batch_size、或兩者並行）。此診斷不影響訓練邏
# 輯本身（不修改任何計算結果），純粹印出觀測值，待瓶頸定位後可移除
# 或保留作為日常監控。
# 0.9.7 版（修正：補上 chunked_lm_loss_backward 缺漏的 config 屬性，
# 並完成獨立數值驗證）
# 20. 0.9.6 版加入記憶體診斷後，實測證實 OOM 確切發生在 backward() 嘗試
# 配置 13.38GB（與 logits 梯度張量大小高度吻合）的當下，forward 結
# 束時已用掉 81.74GB（H100 93.12GB 總容量下僅剩約 11GB 餘裕），證實
# 瓶頸 100% 落在 lm_head 輸出的 (B,L,vocab) logits 張量。
# 21. 發現本檔案中已存在 forward_hidden()（停在 lm_head 之前，回傳遠小
# 於 logits 的 (B,L,d_model) hidden state）與
# chunked_lm_loss_backward()（在 batch 維度切 chunk，每個 chunk 各
# 自完成 lm_head 投影 → cross_entropy → 立即 backward，而非一次性
# materialize 完整 batch 的 logits）的實作，且訓練主迴圈已改為呼叫
# 此函式取代原本的 model(x) + F.cross_entropy + loss.backward() 三
# 段式寫法。然而 chunked_lm_loss_backward() 的呼叫端讀取
# config.loss_chunk_size，此屬性卻從未在 KuiXingConfig.__init__ 中
# 被定義，會導致訓練在第一個 micro-step 就因 AttributeError 直接崩
# 潰（甚至不會走到原本要解決的 OOM）。現補上此屬性，預設值
# loss_chunk_size=4：依實測 batch_size=16 時 logits 約 14.36GB（每
# 1 筆 batch 約 0.90GB）反推，chunk=4 時單一 chunk 的 logits 加上
# backward 所需梯度副本合計約 7.2GB，相較未分塊時約 28.7GB 大幅縮
# 減。注意：此設定僅降低 lm_head 輸出（logits）部分的記憶體，
# forward_hidden() 算出的 12 層 transformer 主體 activation 仍以完
# 整 batch_size 一次性計算、不受此設定影響，若該部�

config.json

@@ -7,17 +7,16 @@
     "AutoConfig": "modeling_kuixing.KuiXingHFConfig",
     "AutoModelForCausalLM": "modeling_kuixing.KuiXingForCausalLM"
   },
-  "vocab_size": 99384,
+  "vocab_size": 109568,
   "hidden_size": 2400,
-  "num_hidden_layers": 12,
+  "num_hidden_layers": 16,
   "num_attention_heads": 32,
   "intermediate_size": 9600,
-  "max_position_embeddings": 2048,
-  "hidden_act": "gelu",
+  "max_position_embeddings": 1024,
+  "hidden_act": "silu",
   "dropout": 0.1,
   "pad_token_id": 0,
   "bos_token_id": 2,
   "eos_token_id": 3,
-  "torch_dtype": "float32",
-  "transformers_version": "5.9.0"
+  "torch_dtype": "float32"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:744326669bf372fb164a19f25692dd5abd76198997ca33241d4c94d6b8ab8397
-size 4291757984
+oid sha256:cc60ca46ed58edc2644a2b66a931dd28ced77b4df94b9fea62369cc9bbd46572
+size 6537547232

modeling_kuixing.py

@@ -19,12 +19,12 @@ class KuiXingHFConfig(PretrainedConfig):
     model_type = "kuixing"
     def __init__(
         self,
-        vocab_size=99384,
+        vocab_size=109568,
         hidden_size=2400,
-        num_hidden_layers=12,
+        num_hidden_layers=16,
         num_attention_heads=32,
         intermediate_size=9600,
-        max_position_embeddings=2048,
+        max_position_embeddings=1024,
         dropout=0.1,
         pad_token_id=0, bos_token_id=2, eos_token_id=3,
         **kwargs,
@@ -40,6 +40,7 @@ class KuiXingHFConfig(PretrainedConfig):
         self.max_position_embeddings = max_position_embeddings
         self.dropout                 = dropout
 
+
 class _Attention(nn.Module):
     def __init__(self, cfg):
         super().__init__()
@@ -68,7 +69,7 @@ class _MLP(nn.Module):
     def __init__(self, cfg):
         super().__init__()
         self.fc1  = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
-        self.act  = nn.GELU()
+        self.act  = nn.SiLU()
         self.fc2  = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
     def forward(self, x):
         return self.fc2(self.act(self.fc1(x)))
@@ -142,10 +143,11 @@ class KuiXingForCausalLM(PreTrainedModel, GenerationMixin):
         import json
         from huggingface_hub import snapshot_download
         import os
+
         # 如果不是本地目錄，先下載到快取
         if not os.path.isdir(model_path):
             model_path = snapshot_download(model_path)
-            
+
         with open(os.path.join(model_path, "config.json")) as f:
             cfg_dict = json.load(f)
         valid = set(KuiXingHFConfig.__init__.__code__.co_varnames)
@@ -158,7 +160,7 @@ class KuiXingForCausalLM(PreTrainedModel, GenerationMixin):
         # lm_head.weight 是刻意省略的（weight tying），從 missing 中排除再判斷
         missing = [k for k in missing if k != "model.lm_head.weight"]
         if not missing and not unexpected:
-            print("✅ 所有權重 key 完整對映，無缺漏。\n如需以 bfloat16 推理：model = model.to(torch.bfloat16).eval()")
+            print("✅ 所有權重 key 完整對映，無缺漏。")
         else:
             if missing:    print(f"⚠️  缺少 key（{len(missing)}）：{missing[:5]}")
             if unexpected: print(f"⚠️  多餘 key（{len(unexpected)}）：{unexpected[:5]}")

tokenizer.model

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ea55192fcb4081708ba5f50ea10a73974665eb4ed33f6ede0705ce3dc85b682d
-size 1484540
+oid sha256:ac0318ec811b0697f8ddffec0ef7b4603483f0c4fced21ba1c2a9b79e42c5662
+size 1621927

tokenizer_config.json

@@ -1,6 +1,6 @@
 {
   "tokenizer_class": "PreTrainedTokenizerFast",
-  "model_max_length": 2048,
+  "model_max_length": 1024,
   "padding_side": "right",
   "bos_token": "<s>",
   "eos_token": "</s>",