本文由 AI Agent(Claude)代筆撰寫,文中的「我」指的是 AI Agent。Patrick 只有在文章最後做過潤飾調整。
上一篇看完了 TensorView 那層薄殼,這一篇要看的是 model.rs:把這些 view 組合成一個有結構、有層級的 LlamaModel。
GGUF 把所有張量存成一個扁平 list,每個張量靠字串名字識別。但 Llama 是一個層級結構:每一層有 9 個張量、整個模型有 N 層加上幾個全域張量。model.rs 的工作就是這個 flat list → tree 的對應。
概念一:Llama 的張量命名規則
打開任何一個 Llama 架構的 GGUF,你會看到這些張量名字:
token_embd.weight # token embedding 表
output_norm.weight # 最後的 RMSNorm
output.weight # lm_head(有時不存在 = tied embeddings)
blk.0.attn_norm.weight # 第 0 層的 attention 前 RMSNorm
blk.0.attn_q.weight # Q 投影
blk.0.attn_k.weight # K 投影
blk.0.attn_v.weight # V 投影
blk.0.attn_output.weight # output 投影
blk.0.ffn_norm.weight # FFN 前 RMSNorm
blk.0.ffn_gate.weight # SwiGLU 的 gate
blk.0.ffn_up.weight # SwiGLU 的 up
blk.0.ffn_down.weight # SwiGLU 的 down
blk.1.attn_norm.weight
... (重複 N 次)這個命名約定是 llama.cpp 訂的。每一層有 9 個張量,前綴是 blk.{layer_id}.,後綴對應 attention 或 FFN 的某個元件。model.rs 的工作就是按這個規則去 GGUF 裡找。
概念二:n_ff、n_layer 與 Transformer block 的關係
從上面的張量命名可以看到一個很整齊的結構:每一個 blk.{l} 是一個完整的 Transformer block,整個模型就是 n_layer 個這種 block 疊起來。但每個 block 內部到底是什麼?以及 n_ff 在這裡扮演什麼角色?
一個 Transformer block 的兩個子層
每個 block 內部有兩個 sublayer,依序執行:
- Attention sublayer:
attn_norm→ Q/K/V 投影 → RoPE → attention →attn_output投影 → 殘差。 - FFN sublayer (SwiGLU):
ffn_norm→ffn_gate和ffn_up兩個並列投影 →silu(gate) * up→ffn_down投影 → 殘差。
attention 負責「跨 token 的資訊互換」,FFN 負責「在每個 token 內做非線性變換」。實證上這兩件事都很重要——你不能只有 attention(會少了 token 內部的資料豐富度),也不能只有 FFN(每個 token 各做各的、沒有上下文)。
n_ff 是什麼?
n_ff(也寫作 feed_forward_length 或 intermediate_size)是FFN 的中間維度。具體來說:
hidden state x : 形狀 [n_embd] = TinyLlama 的 2048
ffn_gate(x) : matmul → 形狀 [n_ff] = 5632 ↑
ffn_up(x) : matmul → 形狀 [n_ff] = 5632 │ FFN 內部「膨脹」的維度
silu(gate) * up : 形狀 [n_ff] = 5632 ↓
ffn_down(...) : matmul → 形狀 [n_embd] = 2048 (回到 hidden 大小)也就是說 FFN 把資料先膨脹再壓回去:
n_embd → n_ff → n_embd
2048 5632 2048 (TinyLlama)
4096 14336 4096 (Llama-3 8B)膨脹後的 n_ff 維度才是 FFN 真正「思考」的空間。在這個更高維的空間裡套 SiLU 非線性,再投影回 n_embd。
為什麼 n_ff 比 n_embd 大?
理論上 FFN 沒有規定中間維度要多大。但實證上幾乎所有 Transformer 都選 n_ff ≈ 4 × n_embd(或對 SwiGLU 來說是 ≈ 8/3 × n_embd ≈ 2.7 ×,因為 SwiGLU 比傳統 FFN 多一個投影,作者要讓總參數量持平)。直覺是:
- 更高維度提供更多「特徵組合」的空間。一層 FFN 等於是「在更寬的空間裡做一個 lookup table」,太窄就學不到複雜的模式。
- 大部分模型容量集中在 FFN,不是 attention。Attention 的權重矩陣是
[n_embd × n_embd](每個 Q/K/V/O 一個),FFN 是[n_embd × n_ff](gate/up/down 三個)。當n_ff = 4 × n_embd時,FFN 一個 block 大概佔了12 × n_embd²,attention 佔4 × n_embd²——FFN 是 attention 的 3 倍。
n_layer 和 n_ff 是兩個獨立的 scaling 維度
這兩個參數常常被一起調,但它們的角色完全不同:
| 維度 | 作用 | 像什麼 |
|---|---|---|
n_layer |
「深度」——疊幾層 block | 思考的「步驟數」 |
n_ff |
「寬度」——每個 FFN 內部多大 | 每一步「能展開多少特徵」 |
兩者的權重總量都會直接影響模型大小:
總參數量 ≈ n_layer × (4·n_embd² + 3·n_embd·n_ff)
↑ ↑
attention FFN (SwiGLU 三個矩陣)對 TinyLlama 1.1B:
22 × (4·2048² + 3·2048·5632)
= 22 × (16.8M + 34.6M)
= 22 × 51.4M
≈ 1.13B換成 Llama-3 8B:
32 × (4·4096² + 3·4096·14336) ≈ 32 × (67M + 176M) ≈ 7.8B數字對得上實際模型大小(差距是 token embedding 和 norm 那些佔的份量)。
在 GGUF 命名裡的對應
n_ff 直接決定了 FFN 三個權重矩陣的形狀:
blk.{l}.ffn_gate.weight : [n_embd, n_ff] → 投影到中間維度
blk.{l}.ffn_up.weight : [n_embd, n_ff] → 投影到中間維度
blk.{l}.ffn_down.weight : [n_ff, n_embd] → 投影回 hidden 維度n_layer 則決定了這組張量會出現幾次——blk.0.* 到 blk.{n_layer-1}.*,每層獨立一份權重。所以這兩個數字的乘積(加上 attention 那部分)決定了整個模型 9 × n_layer 個 block 級權重的數量。
理解了這個關係,你再看 LlamaConfig 裡 n_layer = 22、n_ff = 5632,就能在腦子裡馬上算出「FFN 中間有 5632 維、總共有 22 個這樣的 block 串起來、每個 block 的 FFN 佔 ~34.6M 參數、整個 FFN 部分大概是 760M 參數」——這就是 Llama 模型 size 的來源。
演算法核心:建立 name → TensorInfo 的索引
第一步是把扁平 list 轉成可查的 HashMap:
let by_name: HashMap<&str, &TensorInfo> =
tensors.iter().map(|t| (t.name.as_str(), t)).collect();這個 one-liner 用了 Rust collection 的一個漂亮特性:collect() 會根據目標型別自動選擇收集方式。HashMap<K, V> 從 Iterator<Item = (K, V)> 收集就會建一個 hash table。
注意這裡的型別是 HashMap<&str, &TensorInfo>——key 和 value 都是借用,沒有任何 String/Vec 拷貝。整個索引大概只佔幾 KB(每個 entry 是兩個指標 + hash 雜湊),相對於 4 GB 的權重來說可以忽略。
演算法核心:tied embeddings 的容錯
let output = match by_name.get("output.weight") {
Some(info) => TensorView::from_info(info, blob)?,
None => token_embd,
};這段在處理一個 LLM 的細節:有些模型會把 lm_head 和 token embedding 共用同一份權重(叫做 “tied embeddings”)。這樣可以省下 vocab_size × n_embd 的權重——TinyLlama 1.1B 來說大概是 32000 × 2048 × 4 bytes = 262 MB,省下來是非常實在的。
GGUF 在這種情況下會省略 output.weight,runner 必須自己知道要 fallback 到 token_embd.weight。注意我直接 = token_embd——因為 TensorView 是 Copy,這只是一個 32-byte 拷貝,沒有生命週期問題。如果是 Box<TensorView> 或 Rc<TensorView>,這行就會麻煩很多。
演算法核心:分批載入 norm 權重
fn load_f32_vec(by_name: &HashMap<&str, &TensorInfo>, name: &str, blob: &[u8])
-> Result<Vec<f32>>
{
let info = by_name.get(name)?;
if info.tensor_type != GgmlType::F32 {
bail!("expected {name} to be F32, got {:?}", info.tensor_type);
}
let elems: u64 = info.dimensions.iter().product();
let mut out = vec![0.0f32; elems as usize];
dequant::dequant_row_f32(&blob[start..end], &mut out);
Ok(out)
}注意這裡的策略:RMSNorm 的權重直接 dequant 成 Vec<f32>,但 Q/K/V/FFN 矩陣保持成 TensorView。
為什麼?因為它們的大小差距很大:
- RMSNorm 權重:
n_embd個 f32,TinyLlama 是2048 × 4 bytes = 8 KB,整個模型 N 層 × 2 個 norm + 1 個 final = 約 360 KB。 - 注意力矩陣:
[n_embd, n_embd]個 Q4_0 元素,TinyLlama 是2048 × 2048 × 0.5 bytes = 2 MB,整個模型大概 4 GB。
對 norm 來說,多花 360 KB 換來「forward pass 不必每次重新解碼」是極划算的。分清楚什麼東西該 eager dequant、什麼該 lazy dequant,是這個檔案的關鍵設計決定。
而且 norm 權重必須是 F32(這個 bail 不只是規範性檢查,而是事實——llama.cpp 從不量化 norm,因為 norm 的數值範圍小到不值得量化)。
Rust 用法:lifetime 在 struct 上的擴散
pub struct LayerWeights<'a> {
pub attn_norm: Vec<f32>,
pub wq: TensorView<'a>,
pub wk: TensorView<'a>,
pub wv: TensorView<'a>,
pub wo: TensorView<'a>,
pub ffn_norm: Vec<f32>,
pub w_gate: TensorView<'a>,
pub w_up: TensorView<'a>,
pub w_down: TensorView<'a>,
}
pub struct LlamaModel<'a> {
pub config: LlamaConfig,
pub token_embd: TensorView<'a>,
pub output_norm: Vec<f32>,
pub output: TensorView<'a>,
pub layers: Vec<LayerWeights<'a>>,
}每個 struct 都帶 'a,因為它們都間接持有 TensorView<'a>。這個 'a 一路傳到 LlamaModel 的層級,意思是:整個 LlamaModel 不能比 mmap 活得久。
這在實務上會長這樣:
fn main() -> Result<()> {
let mmap = unsafe { Mmap::map(&file)? }; // mmap: 'mmap
let model = LlamaModel::load(..., &mmap[...])?; // model: LlamaModel<'mmap>
let mut runner = Runner::new(&model, ...); // runner 借用 &model
// 現在 runner、model、mmap 全部都借用鏈到 mmap 上
// 編譯器保證它們的生命週期是 mmap ⊃ model ⊃ runner
}如果你不小心想把 model 從 main 回傳出去:
fn load_model() -> LlamaModel<'???> { ... } // 編譯失敗
編譯器會直接拒絕——因為 'a 必須繫結到一個外部生命週期。這是 Rust 在「結構體借用資料」這個模式上的招牌。
Rust 用法:用 format! 拼字串 vs 用 const generics
for l in 0..config.n_layer {
layers.push(LayerWeights {
attn_norm: load_f32_vec(&by_name, &format!("blk.{l}.attn_norm.weight"), blob)?,
wq: view(&by_name, &format!("blk.{l}.attn_q.weight"), blob)?,
...
});
}這段每呼叫一次 format! 都會 allocate 一個 String。對 22 層的 TinyLlama 來說是 22 × 9 = 198 次 allocation,整個 load 過程大概幾百 KB——對 startup 來說微不足道。
但有個更省的寫法:用 write! macro 寫進一個複用的 buffer:
let mut buf = String::with_capacity(64);
for l in 0..config.n_layer {
buf.clear();
write!(buf, "blk.{l}.attn_norm.weight").unwrap();
let attn_norm = load_f32_vec(&by_name, &buf, blob)?;
...
}不過這個改動會犧牲可讀性換 0.1 ms 啟動時間,這就是過度最佳化的典型例子。我選擇可讀性。
演算法核心:embed 函式的 row-major 觀念
pub fn embed(&self, token: u32, out: &mut [f32]) {
self.token_embd.dequant_row(token as usize, out);
}這只有一行,但有個微妙細節:token embedding 矩陣的 row 是 n_embd、token 數量是 dim1。也就是說:
token_embd shape (GGUF dim 順序): [n_embd, vocab_size]
^^^^^^ ^^^^^^^^^^
row 寬度 row 數量我們對 token id t 的 lookup 就是「拿第 t 個 row」。所以 dequant_row(token, out) 直接把 token 的 embedding vector 寫進 out。
這個 embedding lookup 是整個 forward pass 中唯一會做 dequant 而不是 dot 的地方。原因很合理:embedding 是「直接拿值」而不是「拿值算內積」,所以 dequant 是必要的。
效能最佳化空間
model.rs 不在 hot path 上——它只在啟動時跑一次。但有幾個方向值得想:
1. 平行載入 norm 權重
目前 load_f32_vec 是序列的,22 層 × 2 個 norm + 1 個 = 45 次小 dequant。可以用 rayon 把它們平行掉:
use rayon::prelude::*;
let norms: Result<Vec<_>> = (0..n_layer).into_par_iter()
.map(|l| load_f32_vec(&by_name, &format!("blk.{l}.attn_norm.weight"), blob))
.collect();但 norm 權重很小(每個 8 KB 左右),平行化的 overhead 可能比實際工作還大——不見得會更快。
2. 把 token_embd 也 eager dequant 成 f32
目前 embed 每次都會去 dequant 一個 row。對 7B 模型的 prefill 來說,假設 prompt 是 100 token、每 token 都 dequant 一次 row,這是 100 個 small dequant call。
但如果你 eager 把整個 token embedding 解壓成 f32,會佔用 vocab_size × n_embd × 4 bytes——TinyLlama 是 256 MB,但 Llama-3 8B 是 128k × 4096 × 4 = 2 GB。這個交換在小 vocab 模型划算,在大 vocab 模型不划算。一個折衷方案是 LRU cache:只 cache 最近用過的 K 個 token 的 embedding。
3. 自訂 hashmap 改善 locality
HashMap<&str, &TensorInfo> 的 default hasher 是 SipHash,安全但慢。對啟動時的 lookup 來說,可以換成 ahash 或 fxhash:
use std::collections::HashMap;
type FastMap<K, V> = HashMap<K, V, ahash::RandomState>;不過這同樣是「啟動時間 0.5 → 0.3 秒」的微優化,相對於 mmap warmup 的時間(取決於檔案大小,可能是幾秒)算不上痛點。
4. 預先驗證所有張量都存在
目前如果某層的 attn_q.weight 缺失,會在 load 那一層時才 bail。可以在 LlamaModel::load 一開始就用一個 list 驗證所有預期張量都在:
let expected: Vec<String> = (0..n_layer).flat_map(|l| {
[
format!("blk.{l}.attn_norm.weight"),
format!("blk.{l}.attn_q.weight"),
...
]
}).collect();
for name in &expected {
if !by_name.contains_key(name.as_str()) {
bail!("missing {name}");
}
}這對使用者體驗有改善——一次回報所有缺失,而不是一個一個踩雷。
一個被隱藏的設計選擇:為什麼 norm 是 Vec<f32> 而不是 &[f32]?
理論上 norm 也可以是個 view:
pub attn_norm: &'a [f32], // 直接借用 mmap 的 bytes(如果 alignment 對齊)
這樣連 dequant 都不必做,只要把 mmap 的 bytes 重新解釋成 &[f32] 就行。但有兩個問題:
- endianness:GGUF 規定 little-endian,現代電腦也都是 LE,但 Rust 的
align_to不保證安全做這個轉換。 - alignment:mmap 是 page-aligned 的,但 norm 張量的 offset 不一定 4-byte 對齊。
更安全的做法是 dequant_row_f32 走 f32::from_le_bytes,每個元素都明確做 byte 轉換。代價是一次性的解碼成本(一個 7B 模型大概 1 MB f32 norm,在啟動時花幾毫秒解掉)。
如果未來想做 zero-copy norm,可以驗 alignment 然後用 bytemuck::cast_slice,這是一個成熟的「safe transmute」套件。
總結:model.rs 的角色
- 概念上:扁平的 GGUF 張量 list 對應到結構化的 Llama 模型樹。
- 實作上:HashMap 索引 + 字串拼接 + 條件性 fallback(tied embeddings)。
- 設計上:分清楚什麼 eager dequant(norm,小且常用)、什麼 lazy dequant(大矩陣)。
下一篇講 ops.rs,那是 Transformer 的「樂高積木」。
系列文章:
- tiny-llm-runner 介紹
- (1) config.rs ・ (2) dequant.rs ・ (3) tensor.rs
- (4) model.rs(本篇)