phase 6+7+8: model loading, BPE tokenizer, GPT-2 inference (Milestone ①)

Phase 6 — Model Loading (xserv-model): - safetensors parser with single/sharded file support - ModelConfig with dual naming (GPT-2 n_embd/n_head + modern HF naming) - Weight loading flow: safetensors → mmap → CPU Tensor → GPU Phase 7 — BPE Tokenizer (xserv-tokenizer): - Full BPE encode/decode from tokenizer.json - GPT-2 byte-to-unicode mapping (printable ASCII identity + shifted bytes) - Pre-tokenization regex, special token handling - Chat template support structure Phase 8 — GPT-2 Complete Inference: - GPT-2 model definition: wte, wpe, 12 transformer blocks, ln_f - Forward pass: embedding → (LayerNorm → MHA → residual → LayerNorm → MLP → residual) × 12 → LN → logits - QKV split with correct [batch, heads, seq, dim] layout (fixed reshape bug) - Greedy sampling from last-position logits - Interactive CLI: xserv-cli <model-dir> [--max-tokens N] Verified: GPT-2 124M generates coherent English text on RTX 5090. "The future of AI is uncertain. The future of AI is uncertain..." "Once upon a time, the world was a place of great beauty..." Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-21 22:04:00 +08:00
parent 6035ffdc0b
commit e1e75fc7f6
13 changed files with 971 additions and 0 deletions
--- a/docs/06-model-loading.md
+++ b/docs/06-model-loading.md
@@ -0,0 +1,69 @@
+# Phase 6: Model Loading — Design Document
+
+## Goal
+
+从 HuggingFace safetensors 文件加载模型权重到 GPU Tensor。解析 config.json 获取模型结构参数。
+
+## Crate: `xserv-model`
+
+```
+crates/xserv-model/src/
+├── lib.rs
+├── config.rs       # ModelConfig from config.json
+├── loader.rs       # safetensors weight loading
+└── gpt2.rs         # (Phase 8) GPT-2 model definition
+```
+
+## Dependencies
+
+- `safetensors` crate: parse safetensors format
+- `serde` + `serde_json`: deserialize config.json
+- `memmap2`: mmap for zero-copy file access (safetensors uses this internally)
+
+## Weight Loading Flow
+
+```
+safetensors file (disk)
+  → safetensors crate parses header (tensor names, shapes, dtypes, offsets)
+  → mmap raw data
+  → for each tensor:
+      → read bytes at offset
+      → create CPU Tensor from raw bytes
+      → .to_device(Cuda(0)) → GPU Tensor
+  → return HashMap<String, Tensor>
+```
+
+## Config Parsing
+
+```rust
+#[derive(Deserialize)]
+pub struct ModelConfig {
+    pub architectures: Option<Vec<String>>,
+    pub model_type: Option<String>,
+    pub hidden_size: usize,
+    pub intermediate_size: Option<usize>,
+    pub num_attention_heads: usize,
+    pub num_key_value_heads: Option<usize>,
+    pub num_hidden_layers: usize,
+    pub vocab_size: usize,
+    pub max_position_embeddings: Option<usize>,
+    pub layer_norm_eps: Option<f64>,
+    pub rms_norm_eps: Option<f64>,
+    pub rope_theta: Option<f64>,
+    pub tie_word_embeddings: Option<bool>,
+}
+```
+
+## Test Plan
+
+- [x] Load GPT-2 124M: 160 tensors loaded successfully
+- [x] Parse GPT-2 config.json: hidden=768, layers=12, heads=12, vocab=50257
+- [x] Sharded loading path implemented (for larger models)
+
+## Takeaways
+
+1. **GPT-2 vs modern HF config naming**：GPT-2 uses `n_embd`/`n_head`/`n_layer`/`n_positions`，而不是 `hidden_size`/`num_attention_heads` 等。ModelConfig 需要支持两套命名并提供统一的 accessor methods（`hidden()`, `num_heads()` 等）。
+
+2. **safetensors 零拷贝读取**：`safetensors` crate 直接 mmap 文件，解析 header 得到 tensor 的 offset 和 shape，然后 zero-copy 读取 raw bytes。对于 GPT-2 的 500MB 权重文件，加载速度很快。
+
+3. **模型下载的网络问题**：HuggingFace 在中国网络下不可达。使用 modelscope.cn 或 hf-mirror.com 作为替代。大文件（>100MB）的 redirect 到 CDN 可能也会失败，modelscope 的 snapshot_download 更可靠。
--- a/docs/07-tokenizer.md
+++ b/docs/07-tokenizer.md
@@ -0,0 +1,57 @@
+# Phase 7: BPE Tokenizer — Design Document
+
+## Goal
+
+从零实现 Byte-Pair Encoding tokenizer，兼容 HuggingFace `tokenizer.json` 格式。支持 GPT-2 和 Qwen3。
+
+## Crate: `xserv-tokenizer`
+
+```
+crates/xserv-tokenizer/src/
+├── lib.rs
+├── bpe.rs          # BPE encode/decode core algorithm
+└── chat.rs         # Chat template formatting
+```
+
+## Dependencies
+
+- `serde` + `serde_json`: parse tokenizer.json
+- `regex`: pre-tokenization patterns
+
+## BPE Algorithm
+
+### Encode
+1. Pre-tokenize: split text by regex (GPT-2 pattern)
+2. Each word → byte sequence → initial token list (one token per byte)
+3. Repeatedly merge highest-priority pair until no more merges
+4. Map merged tokens to IDs via vocab
+
+### Decode
+Token IDs → lookup vocab → concatenate bytes → UTF-8 decode
+
+## Key Data Structures
+
+```rust
+pub struct Tokenizer {
+    vocab: HashMap<Vec<u8>, u32>,        // token bytes → ID
+    vocab_rev: Vec<Vec<u8>>,             // ID → token bytes
+    merges: Vec<(Vec<u8>, Vec<u8>)>,     // ordered merge rules
+    merge_ranks: HashMap<(u32, u32), usize>,  // (id_a, id_b) → priority
+    special_tokens: HashMap<String, u32>,
+    pre_tokenize_regex: Regex,
+}
+```
+
+## Test Plan
+
+- [x] Encode + decode roundtrip verified (GPT-2 tokenizer, English text)
+- [x] Special tokens handled (endoftext)
+- [x] Integrated into GPT-2 inference pipeline, generates coherent text
+
+## Takeaways
+
+1. **GPT-2 byte-to-unicode 映射**：GPT-2 的 vocab 中，每个 byte 都映射到一个 Unicode 字符。可打印 ASCII (0x21-0x7E) 映射到自身，其余字节（空格、控制字符等）映射到 U+0100 以上的 Unicode 码点。解码时需要反向映射。这个映射表是 BPE tokenizer 正确性的关键。
+
+2. **Rust regex 不支持 lookahead**：GPT-2 的 pre-tokenization regex 使用了 `(?!\S)` lookahead，Rust 的 `regex` crate 不支持。简化为去掉 lookahead 后功能等价（whitespace 仍然被正确分词）。如果需要精确匹配 Python 行为，需要 `fancy-regex` crate。
+
+3. **BPE merge 的 O(n²) 复杂度**：当前实现每次 merge 扫描整个 token 序列找最高优先级 pair，复杂度 O(n² × |merges|)。对于短文本够用，长文本需要 priority queue 优化。推理场景中 prompt 通常 < 10K tokens，暂时可接受。
--- a/docs/08-gpt2.md
+++ b/docs/08-gpt2.md
@@ -0,0 +1,71 @@
+# Phase 8: GPT-2 Complete Inference — Design Document (Milestone ①)
+
+## Goal
+
+Wire everything together: load GPT-2 124M, tokenize input, run forward pass, sample tokens, decode output. First time seeing the model "speak".
+
+## Model Architecture (GPT-2 124M)
+
+```
+hidden_size = 768
+num_heads = 12
+num_layers = 12
+vocab_size = 50257
+max_position_embeddings = 1024
+activation = GELU
+normalization = LayerNorm (pre-LN)
+tied embeddings (lm_head == wte)
+```
+
+## Forward Pass
+
+```
+tokens [S]
+  → wte[tokens] + wpe[0..S]           → [S, 768]
+  → for each layer:
+      residual = x
+      x = layernorm(x, ln_1)
+      x = attention(x)                 # Q,K,V from linear, MHA, output linear
+      x = x + residual
+      residual = x
+      x = layernorm(x, ln_2)
+      x = mlp(x)                       # linear→GELU→linear
+      x = x + residual
+  → layernorm(x, ln_f)
+  → logits = x @ wte.T                → [S, 50257]
+  → sample(logits[-1])                 → next token
+```
+
+## Sampling
+
+- Greedy: argmax
+- Temperature: logits / T → softmax → sample
+- Top-K: keep top-k logits, rest = -inf
+- Top-P: sorted by prob, cumsum ≤ p
+
+## CLI Binary
+
+```
+$ cargo run --release --bin xserv-cli -- --model path/to/gpt2
+
+xserv> The future of AI is
+GPT-2> ...generated text...
+```
+
+## Test Plan
+
+- [x] Greedy generation produces coherent English text
+- [x] Interactive CLI works (pipe and interactive mode)
+- [x] Multiple prompts verified: "The future of AI is", "Once upon a time"
+
+## Takeaways
+
+1. **QKV split + head reshape 的 layout 陷阱（最关键的 bug）**：GPT-2 的 `c_attn` 输出 `[S, 3H]` 需要 split 成 Q/K/V 再 reshape 成 `[1, num_heads, S, head_dim]`。关键错误：从 `[S, num_heads, head_dim]` 直接 `reshape` 到 `[1, num_heads, S, head_dim]` 不等于 transpose！Reshape 只是重新解释 flat data 的 shape，不会重排数据。必须手动按 `[batch, head, seq, dim]` 的目标 layout 写入数据。同理 merge_heads 也需要手动重排。
+
+2. **CPU round-trip 作为 correctness first 策略**：`add_tensors`、`add_bias`、`split_qkv`、`merge_heads` 都通过 CPU round-trip 实现。虽然慢（每次都有 GPU→CPU→GPU 拷贝），但确保了正确性。Phase 15 会写专门的 CUDA kernel 替换这些操作。
+
+3. **GPT-2 的 Conv1D 权重布局**：GPT-2 用 `Conv1D` 而非 `Linear`，权重存为 `[in, out]`（不是标准 Linear 的 `[out, in]`）。计算方式是 `x @ weight`（不需要转置）。这和 Qwen3/LLaMA 的 `[out, in]` 布局不同——Phase 10 需要注意。
+
+4. **Greedy decoding 的重复问题**：GPT-2 124M 在 greedy decoding 下极易陷入循环（"The world was a place of great danger, and..."）。这是已知行为，temperature + top-k/top-p sampling 可以缓解。当前实现只有 greedy，sampling 将在后续添加。
+
+5. **无 KV Cache 的性能代价**：每生成一个 token 都要重新跑完整 forward pass（O(S²) attention）。50 tokens 的生成需要 50 次 full forward，每次的 attention 复杂度还在增长。Phase 9 的 KV Cache 会将 decode 降到 O(S) per token。