文件重點整理

run.c — 核心推理引擎

主要資料結構

1. Config：模型超參數（dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len）
2. TransformerWeights：所有權重指標（token_embedding_table, wq/wk/wv/wo, w1/w2/w3, rmsnorm 權重）
3. RunState：執行期間的激活值緩衝區（x, xb, q, k, v, att, logits, key_cache, value_cache）
4. Transformer：整合 Config + Weights + State + 檔案映射資訊

權重讀取方式

使用 mmap（記憶體映射）讀取 .bin 權重檔，而非傳統的 fread。優點：

• 作業系統按需載入頁面，節省記憶體
• 多個行程可共享同一映射
• memory_map_weights() 透過指標運算將權重指向正確位置

前向傳播流程 (forward())

token → embedding lookup → for each layer:
  ├─ RMSNorm → QKV matmul → RoPE → Attention → matmul(wo) → residual
  └─ RMSNorm → w1/w3 matmul → SwiGLU → w2 matmul → residual
→ final RMSNorm → classifier → logits

KV Cache 實作

以 (n_layers, seq_len, kv_dim) 的三維陣列儲存，每層每個位置保留一組 K/V。loff = l * seq_len * kv_dim 計算層偏移量，pos 參數決定寫入位置。

注意力機制

• 支援 GQA（Grouped Query Attention）：kv_mul = n_heads / n_kv_heads
• 每個 head 遍歷 0..pos 的所有時間步計算 score
• Score = dot(q, k) / sqrt(head_size)，再 softmax
• Weighted sum of values

RoPE 實作

在 C 中直接計算旋轉矩陣，而非預先計算 cos/sin 表格：

float freq = 1.0f / powf(10000.0f, head_dim / (float)head_size);
float val = pos * freq;
float fcr = cosf(val), fci = sinf(val);
vec[i]   = v0 * fcr - v1 * fci;
vec[i+1] = v0 * fci + v1 * fcr;

BPE Tokenizer

• 讀取 tokenizer.bin（由 tokenizer.py 匯出）
• encode()：字串 → token ID 序列（UTF-8 處理 + BPE 合併）
• decode()：token ID → 字串（處理 BOS、特殊位元組）
• str_lookup()：二分搜尋找到 token 對應的 ID

Sampler

三種採樣策略：

1. Argmax：溫度 = 0 時，取最高機率 token
2. Multinomial：依機率分布隨機採樣
3. Top-p (nucleus)：只從累積機率達 p 的最小 token 集合中採樣

RNG 實作使用 xorshift* 演算法。

---

model.py — PyTorch 模型定義

架構元件

元件	說明
RMSNorm	均方根層歸一化，無需減均值
Attention	多頭注意力 + RoPE + GQA + Flash Attention
FeedForward	SwiGLU MLP：w2(SiLU(w1(x)) * w3(x))
TransformerBlock	Attention + FFN + 殘差連接
Transformer	完整模型 + 權重綁定 + 預計算 RoPE

重要設計

• 權重綁定（Weight Tying）：tok_embeddings.weight = output.weight
• 特殊初始化：wo.weight 和 w3.weight 縮小初始化（除以 sqrt(2 * n_layers)）
• Flash Attention：PyTorch 2.0 以上自動使用 scaled_dot_product_attention
• MFU 估算：計算 Model Flops Utilization（相對於 A100 312 TFLOPS）

---

tokenizer.py

• 包裝 SentencePieceProcessor
• export() 方法將 tokenizer 寫為 .bin 格式（max_token_length + 每個 token 的分數與字串）

---

train.py — 訓練腳本

• 支援單 GPU debug 模式及多節點 DDP 訓練
• Cosine 學習率排程（含 warmup）
• 梯度累積（gradient_accumulation_steps）模擬更大 batch size
• 自動混合精度（bfloat16/float16）
• PyTorch 2.0 compile 支援
• Weights & Biases 整合
• 定期匯出 .bin 供 C 推理使用