featured.svg

本文由 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,依序執行:

  1. Attention sublayerattn_norm → Q/K/V 投影 → RoPE → attention → attn_output 投影 → 殘差。
  2. FFN sublayer (SwiGLU)ffn_normffn_gateffn_up 兩個並列投影 → silu(gate) * upffn_down 投影 → 殘差。

attention 負責「跨 token 的資訊互換」,FFN 負責「在每個 token 內做非線性變換」。實證上這兩件事都很重要——你不能只有 attention(會少了 token 內部的資料豐富度),也不能只有 FFN(每個 token 各做各的、沒有上下文)。

n_ff 是什麼?

n_ff(也寫作 feed_forward_lengthintermediate_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 多一個投影,作者要讓總參數量持平)。直覺是:

  1. 更高維度提供更多「特徵組合」的空間。一層 FFN 等於是「在更寬的空間裡做一個 lookup table」,太窄就學不到複雜的模式。
  2. 大部分模型容量集中在 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_layern_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 級權重的數量。

理解了這個關係,你再看 LlamaConfign_layer = 22n_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——因為 TensorViewCopy,這只是一個 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
}

如果你不小心想把 modelmain 回傳出去:

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 來說,可以換成 ahashfxhash

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] 就行。但有兩個問題:

  1. endianness:GGUF 規定 little-endian,現代電腦也都是 LE,但 Rust 的 align_to 不保證安全做這個轉換。
  2. alignment:mmap 是 page-aligned 的,但 norm 張量的 offset 不一定 4-byte 對齊。

更安全的做法是 dequant_row_f32f32::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 的「樂高積木」。

系列文章: