模型加载与 GGUF
本章对应 s2_model.cpp 的加载部分,回答三个问题:GGUF 是怎么打开的、元数据怎么变成 hparams_、权重张量最终放在哪块后端 buffer 上。
1. “只取元数据”的 GGUF 打开方式
入口是 SlowARModel::load_shared(),由 Pipeline 调用。它和普通 ggml 示例最大的不同是:GGUF 被打开两次,数据只读一次。
struct gguf_init_params params = { true, &weights_.ctx_w }; // s2_model.cpp:212
gguf_context * local_gguf = gguf_init_from_file(gguf_path, params);
gguf_free(local_gguf); // :219
- 第一个参数
true表示用ggml表示所有张量,于是weights_.ctx_w成为一个持有全部张量”形状/名字/类型”的 ggml 上下文,但它此时并不装实际权重数据; - 调用者(Pipeline)手里还保留着一个纯粹的
gguf_context(ctx_gguf),用于读取键值元数据和后续的统一数据读取; - 本地这个 gguf 上下文读完就释放。
这样模型和 codec 可以共享同一个 GGUF 文件句柄和同一次磁盘读取(统一读取在 read_all_tensor_data(),第 10 章详述)。加载第一步永远先初始化 CPU 后端(:206),因为无论是否用 GPU,总有张量要落在 CPU。
2. 元数据 → hparams_
:223-243 定义了三个小工具 get_u32 / get_f32 / get_bool:键缺失时不报错,而是打一条 warning 并使用默认值。这意味着引擎对元数据不全的旧 GGUF 有向前兼容性。
2.1 架构字符串决定一切
int arch_id = gguf_find_key(ctx_gguf, "general.architecture");
std::string arch = gguf_get_val_str(...); // "fish-speech"
hparams_.has_fast_decoder = (arch == "fish-speech"); // :253
arch_prefix = arch + "."; // 后续键都加此前缀
这是整个引擎里最关键的一个判断:只有架构为 fish-speech 的新模型才带 Fast-AR 解码器。老的单 AR 权重仍然能加载,只是 fast_* 路径全部跳过。
2.2 两组超参数
Slow-AR 超参数在 :258-274,默认值即第 02 章那张维度表(32768 / 155776 / 2560 / 9728 / 36 / 32 / 8 / RoPE 1e6 / eps 1e-6)。其中码本相关参数不加架构前缀,固定在 fish_speech.* 命名空间下:codebook_size=4096、num_codebooks=10、semantic_begin/end=151678/155773、tie_word_embeddings=true、attention_qk_norm、scale_codebook_embeddings(:268-274)。
当 has_fast_decoder 为真时,再读一组 fast_* 参数(:276-288):窗口 11、4 层、head_dim 128、32/8 头、可选的 fast_project_in 投影。
3. 张量名绑定
元数据就绪后,:374-424 通过 req_t()(:366,找不到张量直接抛异常)把 GGUF 张量绑到 weights_ 结构上。命名严格对应 Qwen3 风格:
| 组件 | 张量名 |
|---|---|
| 全局 | embeddings.weight、codebook_embeddings.weight、norm.weight |
| 每层 i | layers.i.attention_norm.weight、ffn_norm.weight |
layers.i.attention.wqkv.weight(融合 QKV)、attention.wo.weight | |
layers.i.feed_forward.w1/w2/w3.weight(SwiGLU 三件套) | |
| 可选 | layers.i.attention.q_norm.weight、k_norm.weight |
| Fast-AR | fast_project_in、fast_embeddings、fast_norm、fast_output + fast_layers.i.* |
绑定结束后,:444-479 做一件重要的 bookkeeping:把所有真正需要读数据的张量指针收进 weight_tensor_set_(用 push_unique_tensor 去重——tied 权重、共享张量只读一次)。后续磁盘读取只服务这个集合,其他张量即使存在于 GGUF 也被忽略。
4. 权重放置:full offload 与混合卸载
放置策略由两个布尔/谓词决定。
const bool full_model_offload =
backend_gpu_ != nullptr &&
backend_type != BackendType::CUDA && // CUDA 永远不走全卸载
n_gpu_layers_ == hparams_.block_count; // :481-484
随后的 get_weight_backend(name) 谓词(:486 起)按名字和层号指派后端,规则是:
- 全卸载(Vulkan/Metal,36 层全上):一切权重都在 GPU;
- 混合路径:
embeddings、codebook_embeddings、fast_embeddings、norm等非层张量 → CPU;layers.i.*:i < n_gpu_layers的层在 GPU,其余在 CPU——这就是--gpu-layers/-ngl的语义;- Fast-AR 权重:只要有任意 GPU 层就放 GPU(它是每帧都要跑 9 次的热路径);
- CUDA 特殊处理:嵌入表无条件留 CPU。:645-650 还会在嵌入确实是量化类型时明确日志说明——K 量化嵌入的
get_rows在 CUDA 上不被支持,强推上去曾导致克隆 prefill 不稳定。
n_gpu_layers = -1(默认自动)在前面被展开成 block_count,即”能卸多少卸多少”。GPU 设备初始化(CUDA/Vulkan/Metal,:316-358)失败时不致命:若请求了 GPU 但设备不可用且 n_gpu_layers>0,直接返回 false(:360-364);Pipeline 层面对 codec 之外的 GPU 失败则由上层处理。
5. 多 buffer 切分:绕开驱动的单块分配上限
这是 s2.cpp 最有工程含量的加载细节。很多 GPU 驱动(尤其老卡、某些 Vulkan 实现)对单个 buffer 对象有大小上限,而 S2 Pro 单层/全部权重远超这个数。若按 ggml 常见做法申请一块连续大 buffer,分配会直接失败。
allocate_weight_buffers() 的解法:
- 查询后端默认 buffer type 的对齐值和
ggml_backend_buft_get_max_size;返回 0 视为无上限(:117-121); - 逐张量累加
get_alloc_size(按对齐补齐),一旦当前 chunk 再放一个张量就超限,立刻封箱开新 chunk(:138-150); - 对每个 chunk 调
ggml_backend_buft_alloc_buffer单独申请,标记为GGML_BACKEND_BUFFER_USAGE_WEIGHTS(:168); - 用
ggml_tallocr(buffer 内的线性分配器)把该 chunk 的张量依次摆放到 buffer 里(:171-180)。
GPU 张量和 CPU 张量分别收集、分别走一遍这个流程,产出 weights_.model_bufs_gpu / model_bufs_cpu 两组 buffer(:512-581)。GPU 分配失败时还会给出可操作的建议——提示用户把 --gpu-layers 减半(:554-555)。
6. 调度器:跨后端分片执行
权重分属两个后端后,图计算需要一个能跨后端搬运中间结果的调度器:
backends[0] = backend_gpu_; backends[1] = backend_cpu_; // 有 GPU 时
sched_ = ggml_backend_sched_new(backends, NULL, n_backends, 32768, false, true);
fast_sched_ = ggml_backend_sched_new(backends, NULL, n_backends, 16384, false, true);
见 :586-607。两个调度器图容量不同:Slow-AR 要容纳 prefill 的大图(32768),Fast-AR 窗口小(16384)。调度器会自动在后端边界插入跨 buffer 的拷贝,于是”前几层在 GPU、后几层在 CPU”的混合执行对建图代码完全透明。
7. 真正的数据读取
注意 load_shared() 结束时权重数据还不在 buffer 里——之前绑定的只是形状。真正的字节读取在 read_tensor_data():
- 只处理
weight_tensor_set_里登记过的张量; - 通过
fseeko定位到张量在 GGUF 中的偏移,fread读到临时缓冲,再ggml_backend_tensor_set写进对应后端 buffer(CPU/GPU 均可,因为张量已经由后端 buffer 支持); - 实际部署中 Pipeline 不分别调模型和 codec 的读取,而是用 read_all_tensor_data() 做单遍扫描分发两类张量,读完后在 Linux 上还会
posix_fadvise(DONTNEED)释放页缓存(:377-385)。
独立的 load() 是”不共享 GGUF”的便捷路径:自己打开文件 → load_shared → read_tensor_data,供测试或 C API 的单独模型初始化使用(InitializeS2ModelWithGpuLayers,s2_export_api.cpp:453)。
8. 小结
加载阶段的核心设计可以归纳为:元数据与数据分离、按架构开功能、按名字指派后端、按驱动上限切 buffer、由调度器抹平后端边界。权重各就各位之后,第 05 章进入模型计算本身——Slow-AR 的 36 层计算图是如何搭建的,KV cache 又如何组织。
关键文件
| 位置 | 职责 |
|---|---|
| s2_model.cpp:204 | load_shared 入口 |
| s2_model.cpp:258 | Slow-AR 默认超参 |
| s2_model.cpp:374 | 张量名绑定 |
| s2_model.cpp:481 | 全卸载判定与后端谓词 |
| s2_model.cpp:101 | 多 buffer 切分 |
| s2_model.cpp:595 | 两个后端调度器 |
| s2_model.cpp:655 | 实际权重字节读取 |