目錄
- 前言
- 1. 簡述
- 2. 加載模型超參數
- 3. 加載詞匯表
- 4. 初始化計算上下文
- 5. 初始化計算后端
- 6. 創建模型張量
- 7. 分配緩沖區
- 8. 加載模型權重
- 結語
- 下載鏈接
- 參考
前言
學習 UP 主 比飛鳥貴重的多_HKL 的 GGML源碼逐行調試 視頻,記錄下個人學習筆記,僅供自己參考😄
refer1:【大模型部署】GGML源碼逐行調試
refer2:llama.cpp源碼解讀–ggml框架學習
refer3:https://github.com/ggml-org/ggml
refer4:https://chatgpt.com/
1. 簡述
我們接著 上篇文章 來講,在上篇文章中我們梳理了 ggml 推理 gpt-2 的總體流程,這篇文章我們就來具體地看看每個流程具體是如何實現的
由于篇幅原因,這里我們只來分析 gpt2_model_load()
函數,看 ggml 框架是如何加載模型并創建詞匯表的
gpt2_model_load()
函數主要負責從本地加載 GPT-2 模型文件并初始化模型相關的結構。該函數通過加載權重、詞匯表、超參數等信息,并為模型的計算和推理階段分配必要的內存和資源。下面是對該函數的整體工作流程和關鍵步驟的梳理
2. 加載模型超參數
首先我們以二進制格式打開指定路徑的模型文件(gguf 格式),接著會讀取并校驗 gguf 模型文件中的魔術值(magic number),然后讀取模型的超參數,包括:
n_vocab
:詞匯表大小(50257)n_ctx
:上下文窗口的大小(1024)n_embd
:嵌入維度(768)n_heads
:注意力頭數(12)n_layers
:模型的層數(12)ftype
:數據的類型(FP32)
這些超參數在后續的模型加載過程中會被用來確定模型的結構和計算方式
3. 加載詞匯表
具體代碼分析如下:(from ChatGPT)
1. 讀取詞匯表大小并驗證一致性
int32_t n_vocab = 0;
fin.read((char *) &n_vocab, sizeof(n_vocab));if (n_vocab != model.hparams.n_vocab) {fprintf(stderr, "%s: invalid model file '%s' (bad vocab size %d != %d)\n",__func__, fname.c_str(), n_vocab, model.hparams.n_vocab);return false;
}
首先從模型中讀取詞匯表大小 n_vocab
,并驗證其是否與模型超參數中記錄的大小一致。
2. 準備臨時存儲空間
std::string word;
std::vector<char> buf(128);
接著創建一個 string
類型的對象 word
用于存儲單個 token 的內容,同時創建一個字符向量 buf
作為臨時緩沖區,初始容量為 128 個字符
3. 循環讀取每個 token
for (int i = 0; i < n_vocab; i++) {uint32_t len;fin.read((char *) &len, sizeof(len));
循環進行 n_vocab
次,每次處理一個 token
3.1 調整緩存和讀取 token 數據
buf.resize(len);fin.read((char *) buf.data(), len);
通過 buf.resize()
將緩沖區調整為合適的大小,以容納長度為 len
的 token 字符串。接著使用 fin.read(...)
讀取 len
個字節,將數據存放在緩沖區 buf
中
3.2 構建 token 字符串并建立映射
word.assign(buf.data(), len);vocab.token_to_id[word] = i;vocab.id_to_token[i] = word;
}
使用 word.assign(...)
將緩沖區內的數據轉換為一個 string
類型的 token 字符串。隨后,將該 token 字符串與當前的索引 i
建立雙向映射:
- 插入到
vocab.token_to_id
中,形成 token 到 id 的映射 - 插入到
vocab.id_to_token
中,形成 id 到 token 的映射
這段代碼的工作流程大致如下:
- 讀取詞匯表的總數:從文件中讀取包含在詞匯表中的 token 數量,并與模型超參數驗證一致性
- 動態讀取每個 token:逐個 token 讀取其長度及內容,通過調整緩存大小來適應每個 token 的長度
- 構建雙向映射:將讀取的 token 與對應的序號建立映射,構成完整的詞匯表
4. 初始化計算上下文
具體代碼分析如下:(from ChatGPT)
1. 上下文變量與內存大小估算
auto & ctx = model.ctx_w;
通過引用獲取 model.ctx_w
,后續將其設置為初始化好的 ggml 計算上下文(ggml_context
),該上下文負責管理后續的張量分配和計算圖
接下來,通過計算 tensor 數量來估算需要為模型分配的內存大小:
size_t n_tensors = 2 + 6 + 12 * model.hparams.n_layer;
struct ggml_init_params params = {/*.mem_size =*/ ggml_tensor_overhead() * n_tensors,/*.mem_buffer =*/ NULL,/*.no_alloc =*/ true,
};
n_tensors
計算:
- 公式為
2 + 6 + 12 * model.hparams.n_layer
- 固定分配的 2 個 tensor 可能是歸一化相關的張量,如
ln_f_g
和ln_f_b
- 再加上 6 個固定的張量,可能是嵌入向量(
wte
和wpe
)和模型頭部(lm_head
)相關的張量 - 加上每一層需要創建 12 個張量,模型有 12 層,共
12 * n_layers
個即 144 個張量
- 固定分配的 2 個 tensor 可能是歸一化相關的張量,如
- 該計算目的是預估模型中將會創建的總張量數量,用于確定內存池大小
內存大小計算:
- 調用
ggml_tensor_overhead()
函數得到單個 tensor 的 “額外開銷” 大小,再乘以n_tensors
得到整個模型最小內存池需要的字節數 - Note:
ggml_tensor_overhead()
函數返回的是GGML_OBJECT_SIZE + GGML_TENSOR_SIZE
,也就是ggml_object
結構體大小(32 bytes)加ggml_tensor
結構體大小(336 bytes)。在上篇文章中我們也提到過ggml_context
的結構,對于其中的 Tensor 數據類型,它就是一個 object 再加一個 tensor 的形式,如下圖所示:
初始化參數:
mem_buffer
設為NULL
:如果沒有預先提供內存池,會在ggml_init()
內部調用ggml_aligned_malloc()
申請一塊內存no_alloc
設為true
:表示在創建 tensor 時不自動分配數據空間,僅僅保留元數據管理的內存空間
2. 調用 ggml_init()
創建計算上下文
ctx = ggml_init(params);
if (!ctx) {fprintf(stderr, "%s: ggml_init() failed\n", __func__);return false;
}
ggml_init()
函數會根據傳入的參數創建一個 ggml_context
對象,初始化模型中用于張量管理和內存分配的核心上下文。若創建失敗,則打印錯誤信息并返回 false
3. ggml_init()
函數核心步驟
struct ggml_context * ggml_init(struct ggml_init_params params) {static bool is_first_call = true;ggml_critical_section_start();if (is_first_call) {// 初始化時間系統(例如 Windows 環境)以及 FP16 轉換表ggml_time_init();for (int i = 0; i < (1 << 16); ++i) {union {uint16_t u16;ggml_fp16_t fp16;} u = {i};ggml_table_f32_f16[i] = GGML_COMPUTE_FP16_TO_FP32(u.fp16);}is_first_call = false;}ggml_critical_section_end();struct ggml_context * ctx = GGML_MALLOC(sizeof(struct ggml_context));// 若申請的內存大小為 0,則設置為默認對齊大小if (params.mem_size == 0) {params.mem_size = GGML_MEM_ALIGN;}// 若未提供 mem_buffer 則對 mem_size 按對齊標準進行補齊const size_t mem_size = params.mem_buffer ? params.mem_size : GGML_PAD(params.mem_size, GGML_MEM_ALIGN);*ctx = (struct ggml_context) {/*.mem_size =*/ mem_size,/*.mem_buffer =*/ params.mem_buffer ? params.mem_buffer : ggml_aligned_malloc(mem_size),/*.mem_buffer_owned =*/ params.mem_buffer ? false : true,/*.no_alloc =*/ params.no_alloc,/*.n_objects =*/ 0,/*.objects_begin =*/ NULL,/*.objects_end =*/ NULL,};GGML_ASSERT(ctx->mem_buffer != NULL);GGML_ASSERT_ALIGNED(ctx->mem_buffer);GGML_PRINT_DEBUG("%s: context initialized\n", __func__);return ctx;
}
3.1 線程安全與初始化一次性工作
臨界區處理:
- 使用
ggml_critical_section_start()
與ggml_critical_section_end()
鎖定臨界區,確保ggml_init()
中的一次性初始化工作只在第一次調用時進行
一次性初始化:
- 如果第一次調用(
is_first_call
為 true),調用ggml_time_init()
初始化時間計時系統 - 填充
ggml_table_f32_f16
表,它提供 FP16 到 FP32 的快速轉換映射 - 將
is_first_call
標記為 false,避免后續重復初始化
3.2 分配 ggml_context
內存
內存分配:
- 使用
GGML_MALLOC(sizeof(struct ggml_context))
分配上下文結構體的內存 - 如果未提供預先的內存池(
params.mem_buffer
為 NULL),則使用ggml_aligned_malloc()
申請一塊對齊的內存
內存大小對齊:
- 如果
params.mem_size
為 0,則將其設置為默認的對齊大小(GGML_MEM_ALIGN
),并通過GGML_PAD
宏按對齊要求調整申請的內存大小
上下文初始化:
- 將上下文結構體
ggml_context
的各字段賦值:mem_size
:已對齊的內存大小mem_buffer
:如果用戶沒提供,則為內部申請的內存區域mem_buffer_owned
:標記是否由上下文內部管理內存釋放no_alloc
:傳遞原始參數,控制后續 tensor 創建時是否申請數據存儲空間n_objects
、objects_begin
、objects_end
:初始化為 0 或 NULL,用于管理上下文中創建的對象鏈表
斷言檢查:
- 通過
GGML_ASSERT
系列宏確保mem_buffer
非空且內存對齊正確
最終返回:
- 打印調試信息,返回初始化好的
ggml_context
指針供后續模型張量和計算圖構建使用
整個初始化計算上下文的過程主要包括:
- 1. 預估內存需求:根據模型所需張量數量和單個 tensor 的開銷,計算出所需內存池大小
- 2. 構建初始化參數:通過
ggml_init_params
傳遞內存大小、緩沖區指針(NULL 表示內部自動分配)和是否分配實際數據空間(no_alloc
)的標志 - 3. 調用
ggml_init()
:- 內部用臨界區確保一次性初始化(如時間系統、FP16 轉換表)
- 根據是否提供外部內存池對內存進行對齊并分配
- 初始化
ggml_context
對象,其后續管理所有模型張量的元數據及對象鏈表
- 4. 結果驗證:斷言和調試信息確保上下文正確創建,返回一個有效的計算上下文供后續模型權重和計算圖構建使用
5. 初始化計算后端
具體代碼分析如下:(from ChatGPT)
1. 后端初始化的總體邏輯
首先,由于我們在編譯時開啟了 CUDA 后端(條件編譯宏 GGML_USE_CUDA
被定義)且 n_gpu_layers > 0
,程序執行以下邏輯:
#ifdef GGML_USE_CUDAif (n_gpu_layers > 0) {fprintf(stderr, "%s: using CUDA backend\n", __func__);model.backend = ggml_backend_cuda_init(0);if (!model.backend) {fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__);}}
#endif
通過調用 ggml_backend_cuda_init(0)
初始化 CUDA 后端,這里設置的 device 默認為 0
2. ggml_backend_cuda_init()
的實現
2.1 參數檢查
在 ggml_backend_cuda_init()
函數的一開始,會檢查傳入的 device 編號:
if (device < 0 || device >= ggml_backend_cuda_get_device_count()) {GGML_LOG_ERROR("%s: invalid device %d\n", __func__, device);return nullptr;
}
驗證 device 編號是否合法,如果 device 編號超出系統可用 CUDA 設備數量,則打印錯誤并返回空指針
2.2 上下文構造
接下來通過如下代碼為 CUDA 后端構建一個上下文對象:
ggml_backend_cuda_context * ctx = new ggml_backend_cuda_context(device);
if (ctx == nullptr) {GGML_LOG_ERROR("%s: failed to allocate context\n", __func__);return nullptr;
}
分配一個 ggml_backend_cuda_context
對象,并傳入 device 編號,在構造函數中,它會保存當前 device 編號,并根據 device 編號生成一個名稱
ggml_backend_cuda_context
結構體內容如下:
struct ggml_backend_cuda_context {int device;std::string name;cudaEvent_t copy_event = nullptr;cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = { { nullptr } };cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES] = {nullptr};std::unique_ptr<ggml_cuda_graph> cuda_graph;explicit ggml_backend_cuda_context(int device) :device(device),name(GGML_CUDA_NAME + std::to_string(device)) {}~ggml_backend_cuda_context() {if (copy_event != nullptr) {CUDA_CHECK(cudaEventDestroy(copy_event));}for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) {for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) {if (streams[i][j] != nullptr) {CUDA_CHECK(cudaStreamDestroy(streams[i][j]));}}if (cublas_handles[i] != nullptr) {CUBLAS_CHECK(cublasDestroy(cublas_handles[i]));}}}cudaStream_t stream(int device, int stream) {if (streams[device][stream] == nullptr) {ggml_cuda_set_device(device);CUDA_CHECK(cudaStreamCreateWithFlags(&streams[device][stream], cudaStreamNonBlocking));}return streams[device][stream];}cudaStream_t stream() {return stream(device, 0);}cublasHandle_t cublas_handle(int device) {if (cublas_handles[device] == nullptr) {ggml_cuda_set_device(device);CUBLAS_CHECK(cublasCreate(&cublas_handles[device]));CUBLAS_CHECK(cublasSetMathMode(cublas_handles[device], CUBLAS_TF32_TENSOR_OP_MATH));}return cublas_handles[device];}cublasHandle_t cublas_handle() {return cublas_handle(device);}// poolstd::unique_ptr<ggml_cuda_pool> pools[GGML_CUDA_MAX_DEVICES];static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device);ggml_cuda_pool & pool(int device) {if (pools[device] == nullptr) {pools[device] = new_pool_for_device(device);}return *pools[device];}ggml_cuda_pool & pool() {return pool(device);}
};
這個結構體封裝了 CUDA 后端需要維護的所有資源和輔助函數,它的主要字段包括:
- 設備與名稱:
device
:記錄使用的 device 編號name
:device 名稱,由預定義名稱與 device 編號組合生成
- CUDA Event:
copy_event
:用于記錄數據拷貝過程中的 CUDA 事件,便于在異步操作中進行同步
- CUDA Stream:
streams
:二維數組管理每個 device 上多條 CUDA 流。函數stream(int device, int stream)
檢查對應的流是否已創建,如果沒有,則創建帶有非阻塞標記的 CUDA 流
- cuBLAS 句柄
cublas_handles
:用于在 CUDA 后端調用 cuBLAS 庫進行矩陣計算。調用cublas_handle(int device)
檢查是否已創建對應的句柄,如果未創建則初始化,并設置運算模式為 TF32
- CUDA 圖與內存池:
cuda_graph
:用于管理 CUDA Graph 計算的資源pools
:每個 device 上各自的內存池管理器,通過new_pool_for_device(int device)
實例化,負責分配和管理 GPU 上的臨時內存資源
- 析構函數:
- 在析構函數中,會依次銷毀已創建的 CUDA Event、各個 CUDA Stream 以及 cuBLAS 句柄,保證資源正確釋放,防止內存泄漏和 CUDA 資源未釋放問題
- 輔助函數:
stream()
:用于獲取當前設備默認的 CUDA 流cublas_handle()
:獲取當前設備默認的 cuBLAS 句柄,并確保其已正確初始化pool()
:獲取當前設備對應的內存池
2.3 構造 ggml_backend
對象
一旦上下文對象構建完成,接下來構造后端對象:
ggml_backend_t cuda_backend = new ggml_backend {/* .guid = */ ggml_backend_cuda_guid(),/* .interface = */ ggml_backend_cuda_interface,/* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device),/* .context = */ ctx,
};
其中:
gpuid
:通過ggml_backend_cuda_guid()
獲得一個全局唯一標識符,便于識別該后端類型interface
:指向ggml_backend_cuda_interface
,這是一個結構體,包含了該后端實現的各個函數接口,例如設置異步傳輸、同步、圖計算、事件記錄與等待等device
:調用ggml_backend_reg_dev_get(...)
獲得當前設備相關的注冊信息。設備信息包括設備名稱、描述等context
:上面分配的 CUDA 上下文對象,包含了設備號、CUDA 流、cuBLAS 句柄、內存池等資源
該函數最終返回一個 ggml_backend
對象指針,用于后續張量數據的異步傳輸和計算調度
ggml_backend_cuda_interface
結構體包含了 CUDA 后端實現的具體函數指針:
static const ggml_backend_i ggml_backend_cuda_interface = {/* .get_name = */ ggml_backend_cuda_get_name,/* .free = */ ggml_backend_cuda_free,/* .set_tensor_async = */ ggml_backend_cuda_set_tensor_async,/* .get_tensor_async = */ ggml_backend_cuda_get_tensor_async,/* .cpy_tensor_async = */ ggml_backend_cuda_cpy_tensor_async,/* .synchronize = */ ggml_backend_cuda_synchronize,/* .graph_compute = */ ggml_backend_cuda_graph_compute,/* .event_record = */ ggml_backend_cuda_event_record,/* .event_wait = */ ggml_backend_cuda_event_wait,
};
這些接口函數定義了 CUDA 后端支持的操作,如異步傳輸、同步以及圖計算等。調用者可以通過統一的后端接口調用這些函數,而無需關心底層 CUDA 細節。
3. CPU 后端回退
如果上述 CUDA 或 Metal 后端初始化失敗(model.backend
為 nullptr
),邏輯會回退至 CPU 后端:
if (!model.backend) {// fallback to CPU backendfprintf(stderr, "%s: using CPU backend\n", __func__);model.backend = ggml_backend_cpu_init();
}
如果 CPU 后端初始化也失敗,則打印錯誤并返回 false,終止模型加載流程:
if (!model.backend) {fprintf(stderr, "%s: ggml_backend_cpu_init() failed\n", __func__);return false;
}
整個初始化計算后端的過程(以 CUDA 后端為例)主要包括:
- CUDA 后端的初始化:
- 驗證設備:首先檢查所傳入設備編號是否合法,確保設備數量足夠
- 構建 CUDA 上下文:通過
new ggml_backend_cuda_context(device)
創建包含 CUDA 流、cuBLAS 句柄、事件和內存池等資源的上下文 - 構造后端對象:將全局唯一標識符(
guid
)、接口函數集、設備信息和上下文整合到一個ggml_backend
對象中返回
- 回退機制:如果 CUDA 后端初始化失敗,則嘗試使用 Metal(針對 Apple 平臺)或直接回退到 CPU 后端,保證模型加載流程的健壯性
6. 創建模型張量
具體代碼分析如下:(from ChatGPT)
1. 模型張量創建與映射
1.1 上下文與超參數
首先獲取模型超參數:
const auto & hparams = model.hparams;
const int n_embd = hparams.n_embd;
const int n_layer = hparams.n_layer;
const int n_ctx = hparams.n_ctx;
const int n_vocab = hparams.n_vocab;
這里將模型所需的嵌入維度、層數、上下文大小(序列長度)和詞匯表大小保存為局部變量,以便后續創建各個張量時使用
1.2 創建模型整體的固定張量
隨后為模型創建整體共有的張量,這里包括:
- 歸一化參數:
ln_f_g
和ln_f_b
,用于最終的 layer norm 處理。
model.ln_f_g = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
model.ln_f_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
- 嵌入與投影層:
wte
:token embedding,大小為 n_embd x n_vocabwpe
:position embedding,大小為 n_embd x n_ctxlm_head
:語言模型輸出頭,同樣維度為 n_embd x n_vocab
model.wte = ggml_new_tensor_2d(ctx, wtype, n_embd, n_vocab);
model.wpe = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_ctx);
model.lm_head = ggml_new_tensor_2d(ctx, wtype, n_embd, n_vocab);
- 映射到查找表:為了后續在加載權重時能夠通過名字訪問對應的張量,將上述張量按名稱保存到
model.tensors
的 map 容器中:
model.tensors["model/ln_f/g"] = model.ln_f_g;
model.tensors["model/ln_f/b"] = model.ln_f_b;
model.tensors["model/wte"] = model.wte;
model.tensors["model/wpe"] = model.wpe;
model.tensors["model/lm_head"] = model.lm_head;
1.3 創建每一層的張量
接下來,針對 gpt-2 模型的每一層(共 n_layer
個 Transformer block 堆疊),代碼遍歷循環進行以下工作:
- 為當前層創建歸一化層參數:
layer.ln_1_g = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
layer.ln_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
layer.ln_2_g = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
layer.ln_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
- 為注意力模塊創建參數:
- 注意力部分中的權重矩陣(
c_attn_attn_w
)和偏置(c_attn_attn_b
),矩陣尺寸為 n_embd x (3*n_embd)(用于 Q、K、V 的拼接) - 注意力投影權重(
c_attn_proj_w
)和偏置(c_attn_proj_b
),尺寸為 n_embd x n_embd
- 注意力部分中的權重矩陣(
layer.c_attn_attn_w = ggml_new_tensor_2d(ctx, wtype, n_embd, 3*n_embd);
layer.c_attn_attn_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 3*n_embd);layer.c_attn_proj_w = ggml_new_tensor_2d(ctx, wtype, n_embd, n_embd);
layer.c_attn_proj_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
- 為 MLP 模塊創建參數:
- 第一層全連接(
c_mlp_fc_w
、c_mlp_fc_b
),尺寸為 n_embd x (4*n_embd) - 第二層投影(
c_mlp_proj_w
、c_mlp_proj_b
),尺寸為 (4*n_embd) x n_embd
- 第一層全連接(
layer.c_mlp_fc_w = ggml_new_tensor_2d(ctx, wtype, n_embd, 4*n_embd);
layer.c_mlp_fc_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 4*n_embd);layer.c_mlp_proj_w = ggml_new_tensor_2d(ctx, wtype, 4*n_embd, n_embd);
layer.c_mlp_proj_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);
所有這些張量都由相應的調用 ggml_new_tensor_1d
(用于一維張量,如偏置或歸一化參數)或 ggml_new_tensor_2d
(用于二維張量,如權重矩陣)生成
- 映射:每個生成的張量都會通過拼接字符串(例如
"model/h" + std::to_string(i) + "/ln_1/g"
)的方式注冊到model.tensors
中,以便后續查找和加載對應的權重數據
// map by name
model.tensors["model/h" + std::to_string(i) + "/ln_1/g"] = layer.ln_1_g;
model.tensors["model/h" + std::to_string(i) + "/ln_1/b"] = layer.ln_1_b;model.tensors["model/h" + std::to_string(i) + "/ln_2/g"] = layer.ln_2_g;
model.tensors["model/h" + std::to_string(i) + "/ln_2/b"] = layer.ln_2_b;model.tensors["model/h" + std::to_string(i) + "/attn/c_attn/w"] = layer.c_attn_attn_w;
model.tensors["model/h" + std::to_string(i) + "/attn/c_attn/b"] = layer.c_attn_attn_b;model.tensors["model/h" + std::to_string(i) + "/attn/c_proj/w"] = layer.c_attn_proj_w;
model.tensors["model/h" + std::to_string(i) + "/attn/c_proj/b"] = layer.c_attn_proj_b;model.tensors["model/h" + std::to_string(i) + "/mlp/c_fc/w"] = layer.c_mlp_fc_w;
model.tensors["model/h" + std::to_string(i) + "/mlp/c_fc/b"] = layer.c_mlp_fc_b;model.tensors["model/h" + std::to_string(i) + "/mlp/c_proj/w"] = layer.c_mlp_proj_w;
model.tensors["model/h" + std::to_string(i) + "/mlp/c_proj/b"] = layer.c_mlp_proj_b;
2. 底層張量創建函數實現
張量創建的核心是 ggml_new_tensor_1d
與 ggml_new_tensor_2d
,它們都最終調用了 ggml_new_tensor()
,并由其調用 ggml_new_tensor_impl()
2.1 ggml_new_tensor_1d
和 ggml_new_tensor_2d
ggml_new_tensor_1d(ctx, type, ne0)
:內部調用ggml_new_tensor(ctx, type, 1, &ne0)
,創建 1 維張量,其元素數量為ne0
struct ggml_tensor * ggml_new_tensor(struct ggml_context * ctx,enum ggml_type type,int n_dims,const int64_t * ne) {return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, 0);
}struct ggml_tensor * ggml_new_tensor_1d(struct ggml_context * ctx,enum ggml_type type,int64_t ne0) {return ggml_new_tensor(ctx, type, 1, &ne0);
}
ggml_new_tensor_2d(ctx, type, ne0, ne1)
:構造一個包含兩個維度的數組ne[2] = {ne0, ne1}
,然后調用ggml_new_tensor(ctx, type, 2, ne)
創建二維張量
struct ggml_tensor * ggml_new_tensor_2d(struct ggml_context * ctx,enum ggml_type type,int64_t ne0,int64_t ne1) {const int64_t ne[2] = { ne0, ne1 };return ggml_new_tensor(ctx, type, 2, ne);
}
2.2 ggml_new_tensor_impl
這是最核心的張量創建函數,其主要流程如下:
1. 檢查與處理 view 信息
static struct ggml_tensor * ggml_new_tensor_impl(struct ggml_context * ctx,enum ggml_type type,int n_dims,const int64_t * ne,struct ggml_tensor * view_src,size_t view_offs) {GGML_ASSERT(type >= 0 && type < GGML_TYPE_COUNT);GGML_ASSERT(n_dims >= 1 && n_dims <= GGML_MAX_DIMS);// find the base tensor and absolute offsetif (view_src != NULL && view_src->view_src != NULL) {view_offs += view_src->view_offs;view_src = view_src->view_src;}
如果傳入了 view tensor(用于創建張量視圖),則調整 view_src
和偏移量,否則 view_src
為 NULL
2. 計算數據大小
size_t data_size = ggml_row_size(type, ne[0]);
for (int i = 1; i < n_dims; i++) {data_size *= ne[i];
}GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src));
使用 ggml_row_size(type, ne[0])
計算第一維每行的數據大小,再乘以后續各維度的大小,得到整個張量的數據所需字節數
同時驗證如果 view_src
不為空,數據大小加上偏移量不超過原始張量的總字節數
3. 分配數據空間
void * data = view_src != NULL ? view_src->data : NULL;
if (data != NULL) {data = (char *) data + view_offs;
}size_t obj_alloc_size = 0;if (view_src == NULL && !ctx->no_alloc) {// allocate tensor data in the context's memory poolobj_alloc_size = data_size;
}
如果 view_src
為 NULL 且上下文允許分配(ctx->no_alloc
為 false),則將 obj_alloc_size
設置為 data_size
,表示需要為張量數據申請空間;否則,張量的數據指針將直接指向 view_src
的數據或置為空
4. 分配對象內存
struct ggml_object * const obj_new = ggml_new_object(ctx, GGML_OBJECT_TYPE_TENSOR, GGML_TENSOR_SIZE + obj_alloc_size);
GGML_ASSERT(obj_new);
調用 ggml_new_object(...)
在內存池中分配一個新對象,該函數通過:
- 計算當前對象結束偏移:從內存池中當前對象的末尾得到下一對象插入位置
- 對齊分配:根據
GGML_MEM_ALIGN
要求對齊分配大小 - 檢查內存池空間:檢查內存池空間:確保新的對象可以在當前上下文的內存池中放下,否則發出警告并終止
返回的新對象包含元數據存儲區,其大小等于 GGML_TENSOR_SIZE + obj_alloc_size
ggml_new_object
具體實現如下:
static struct ggml_object * ggml_new_object(struct ggml_context * ctx, enum ggml_object_type type, size_t size) {// always insert objects at the end of the context's memory poolstruct ggml_object * obj_cur = ctx->objects_end;const size_t cur_offs = obj_cur == NULL ? 0 : obj_cur->offs;const size_t cur_size = obj_cur == NULL ? 0 : obj_cur->size;const size_t cur_end = cur_offs + cur_size;// align to GGML_MEM_ALIGNsize_t size_needed = GGML_PAD(size, GGML_MEM_ALIGN);char * const mem_buffer = ctx->mem_buffer;struct ggml_object * const obj_new = (struct ggml_object *)(mem_buffer + cur_end);if (cur_end + size_needed + GGML_OBJECT_SIZE > ctx->mem_size) {GGML_LOG_WARN("%s: not enough space in the context's memory pool (needed %zu, available %zu)\n",__func__, cur_end + size_needed + GGML_OBJECT_SIZE, ctx->mem_size);
#ifndef NDEBUGGGML_ABORT("not enough space in the context's memory pool");
#endifreturn NULL;}*obj_new = (struct ggml_object) {.offs = cur_end + GGML_OBJECT_SIZE,.size = size_needed,.next = NULL,.type = type,};GGML_ASSERT_ALIGNED(mem_buffer + obj_new->offs);if (obj_cur != NULL) {obj_cur->next = obj_new;} else {// this is the first object in this contextctx->objects_begin = obj_new;}ctx->objects_end = obj_new;//printf("%s: inserted new object at %zu, size = %zu\n", __func__, cur_end, obj_new->size);return obj_new;
}
5. 初始化張量結構
struct ggml_tensor * const result = (struct ggml_tensor *)((char *)ctx->mem_buffer + obj_new->offs);*result = (struct ggml_tensor) {/*.type =*/ type,/*.buffer =*/ NULL,/*.ne =*/ { 1, 1, 1, 1 },/*.nb =*/ { 0, 0, 0, 0 },/*.op =*/ GGML_OP_NONE,/*.op_params =*/ { 0 },/*.flags =*/ 0,/*.src =*/ { NULL },/*.view_src =*/ view_src,/*.view_offs =*/ view_offs,/*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data,/*.name =*/ { 0 },/*.extra =*/ NULL,/*.padding =*/ { 0 },
};// TODO: this should not be needed as long as we don't rely on aligned SIMD loads
//GGML_ASSERT_ALIGNED(result->data);for (int i = 0; i < n_dims; i++) {result->ne[i] = ne[i];
}result->nb[0] = ggml_type_size(type);
result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type));
for (int i = 2; i < GGML_MAX_DIMS; i++) {result->nb[i] = result->nb[i - 1]*result->ne[i - 1];
}ctx->n_objects++;
獲得 ggml_object
對象后,張量的內存區域位于對象元數據后的內存空間中,設置張量的各項字段:
- 類型和 view 信息:存入張量的
type
、view_src
、view_offs
- 數據指針:如果需要分配數據,則指針指向對象后面分配的區域
- 維度信息:將傳入的各維度值拷貝到
ne
數組中 - 步長計算:
nb[0]
設為每個元素的字節大小nb[1]
設為nb[0]
乘以每行的元素數- 后續維度的 stride 依次遞推計算
- 對象計數:最后上下文中的對象計數
ctx->n_objects
加 1,標記新對象的存在
6. 返回結果
return result;
返回新創建的 ggml_tensor
指針,該指針既包含張量元數據,也包含數據存儲區(如果有分配)
總的來說創建模型張量主要分為以下幾個部分:
- 高層邏輯:模型加載過程中會依次創建模型的整體參數張量(如嵌入、歸一化參數、輸出頭)和每一層的各個子模塊張量。創建后將它們以固定名稱存儲在
model.tensors
內,便于后續加載對應權重數據。 - 底層實現:張量創建調用鏈為:
ggml_new_tensor_1d
或ggml_new_tensor_2d
? 調用ggml_new_tensor
? 進入ggml_new_tensor_impl
ggml_new_tensor_impl
內部通過ggml_new_object
在上下文的內存池分配內存,并設置張量的各項屬性(類型、維度、步長、數據指針等)
- 內存管理:整個過程充分考慮內存對齊、動態內存分配和對象管理,每個張量對象除了數據之外,還有額外的元數據開銷。上下文(
ggml_context
)負責整體內存池的管理,并通過鏈表組織所有張量對象
7. 分配緩沖區
具體代碼分析如下:(from ChatGPT)
在 GPT-2 推理過程中,模型緩沖區分配代碼分為兩個主要部分:
- 模型權重(models tensors)的分配:調用
ggml_backend_alloc_ctx_tensors(ctx, model.backend)
分配內存用于存放模型的張量(權重、激活值等) - KV 內存(key/value memory)的分配:專門為 kv cache 創建單獨的 ggml 上下文(
ggml_context
),并在這個上下文中通過相同方式分配內存,用于存放 key 和 value 張量
對于這兩部分,代碼都采用了后端緩沖區的方式,調用統一的接口將上下文中所有尚未分配內存的張量分配到后端特定的緩沖區中
1. 模型權重部分
// allocate the model tensors in a backend buffer
model.buffer_w = ggml_backend_alloc_ctx_tensors(ctx, model.backend);printf("%s: ggml tensor size = %d bytes\n", __func__, (int) sizeof(ggml_tensor));
printf("%s: backend buffer size = %6.2f MB\n", __func__, ggml_backend_buffer_get_size(model.buffer_w)/(1024.0*1024.0));
這個部分調用 ggml_backend_alloc_ctx_tensors
為整個模型的張量創建創建后端緩存區,函數實現如下:
ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend) {return ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_get_default_buffer_type(backend));
}
這個函數是入口,會從傳入的 backend
中獲取默認的緩沖區類型(buft
),然后交給 ggml_backend_alloc_ctx_tensors_from_buft
進行實際分配,函數實現如下:
ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) {GGML_ASSERT(ggml_get_no_alloc(ctx) == true);size_t alignment = ggml_backend_buft_get_alignment(buft);size_t max_size = ggml_backend_buft_get_max_size(buft);ggml_backend_buffer_t * buffers = NULL;size_t n_buffers = 0;size_t cur_buf_size = 0;struct ggml_tensor * first = ggml_get_first_tensor(ctx);for (struct ggml_tensor * t = first; t != NULL; t = ggml_get_next_tensor(ctx, t)) {size_t this_size = 0;if (t->data == NULL && t->view_src == NULL) {// 對于未分配的 tensor,根據該后端分配要求獲得 tensor 所需大小并按對齊要求進行填充this_size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, t), alignment);}if (cur_buf_size > 0 && (cur_buf_size + this_size) > max_size) {// 當前緩沖區已不能容納當前 tensor 時,調用 alloc_tensor_range 分配當前這段范圍if (!alloc_tensor_range(ctx, first, t, buft, cur_buf_size, &buffers, &n_buffers)) {return NULL;}// 重置 first 為當前 tensor,并重置計數first = t;cur_buf_size = this_size;} else {cur_buf_size += this_size;}}// 最后分配剩余未分配的張量if (cur_buf_size > 0) {if (!alloc_tensor_range(ctx, first, NULL, buft, cur_buf_size, &buffers, &n_buffers)) {return NULL;}}if (n_buffers == 0) {
#ifndef NDEBUGGGML_LOG_DEBUG("%s: all tensors in the context are already allocated\n", __func__);
#endifreturn NULL;}ggml_backend_buffer_t buffer;if (n_buffers == 1) {buffer = buffers[0];} else {// 如果分配了多個緩沖區,則調用后端函數合并為一個邏輯上的緩沖區buffer = ggml_backend_multi_buffer_alloc_buffer(buffers, n_buffers);}free(buffers);return buffer;
}
流程如下:
- 斷言與基礎設置
- 斷言當前上下文
no_alloc
為true
,確保在調用此函數前沒有自動進行其它內存分配 - 從
buft
中獲取內存對齊要求(alignment
)以及單個緩沖區的最大容量(max_size
)
- 斷言當前上下文
- 遍歷 ggml 上下文中的張量
- 使用
ggml_get_first_tensor
得到上下文中第一個 tensor,然后不斷遍歷(用ggml_get_next_tensor
) - 對于每個 tensor,如果該 tensor 目前還沒有綁定內存(即
t->data == NULL
且t->view_src == NULL
),就調用ggml_backend_buft_get_alloc_size
得到張量所需內存大小,再利用GGML_PAD
函數按對齊數做填充,得到正確的分配大小
- 使用
- 分割緩沖區的邏輯
- 用
cur_buf_size
累加一段連續 tensor 的內存需求。如果累加后超過當前緩沖區最大容量max_size
,則通過alloc_tensor_range
函數把從first
到當前 tensor 前一個的這段張量分配到一個緩沖區中 - 分配完后,重置
fitst
和cur_buf_size
,開始為下一段張量分配內存
- 用
- 最后的分配
- 遍歷完成后,如果還有剩余的
cur_buf_size
,則進行最后一段內存分配
- 遍歷完成后,如果還有剩余的
- 合并與返回
- 如果只分配了一個緩沖區,直接返回它
- 如果分配了多個緩存區,則調用
ggml_backend_multi_buffer_alloc_buffer
把多個緩沖區邏輯上合并為一個緩沖區
通過這種方式,ggml 可以高效地管理內存,將所有未分配的張量集中放入一個或多個連續的后端內存區域,同時保證內存對齊和不超過設備最大分配量
2. kv cache 部分
{auto * ctx = model.ctx_kv;// create the ggml context{size_t n_tensors = 2;struct ggml_init_params params = {/*.mem_size =*/ ggml_tensor_overhead() * n_tensors,/*.mem_buffer =*/ NULL,/*.no_alloc =*/ true,};ctx = ggml_init(params);if (!ctx) {fprintf(stderr, "%s: ggml_init() failed\n", __func__);return false;}}const auto & hparams = model.hparams;const int n_embd = hparams.n_embd;const int n_layer = hparams.n_layer;const int n_ctx = hparams.n_ctx;const int n_mem = n_layer * n_ctx;const int n_elements = n_embd * n_mem;model.memory_k = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_elements);model.memory_v = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_elements);// allocate the KV memory in a backend buffermodel.buffer_kv = ggml_backend_alloc_ctx_tensors(ctx, model.backend);const size_t memory_size = ggml_backend_buffer_get_size(model.buffer_kv);printf("%s: memory size = %8.2f MB, n_mem = %d\n", __func__, memory_size/1024.0/1024.0, n_mem);
}
首先,為 kv cache(key 和 value)新建一個 ggml 上下文 ctx
,構建時 no_alloc
參數為 true,意味著在初始化時不進行自動內存分配
接著根據模型超參數(embeddding 大小、層數、上下文長度),計算出需要存儲多少個元素(n_elements
),并通過 ggml_new_tensor_1d
創建兩個一維張量分別用于存儲 key 和 value
最后,再次調用 ggml_backend_alloc_ctx_tensors
給 kv cache 的所有張量分配后端緩沖區
8. 加載模型權重
具體代碼分析如下:(from ChatGPT)
1. 讀取和循環處理模型文件中的每個張量
while (true) {int32_t n_dims;int32_t length;int32_t ttype;fin.read(reinterpret_cast<char *>(&n_dims), sizeof(n_dims));fin.read(reinterpret_cast<char *>(&length), sizeof(length));fin.read(reinterpret_cast<char *>(&ttype), sizeof(ttype));if (fin.eof()) {break;}// ...
}
循環從模型文件中不斷讀取每個 tensor 的元數據,直到遇到文件結束
2. 讀取每個張量的元數據
int32_t nelements = 1;
int32_t ne[2] = { 1, 1 };
for (int i = 0; i < n_dims; ++i) {fin.read(reinterpret_cast<char *>(&ne[i]), sizeof(ne[i]));nelements *= ne[i];
}std::string name(length, 0);
fin.read(&name[0], length);
讀取各個維度的大小并累乘得到總的元素數 nelements
,此外將 tensor 的名稱讀取為 name
字符串
3. 根據名稱查找對應的 tensor 并進行校驗
if (model.tensors.find(name) == model.tensors.end()) {fprintf(stderr, "%s: unknown tensor '%s' in model file\n", __func__, name.c_str());return false;
}auto tensor = model.tensors[name];
ggml_set_name(tensor, name.c_str());
if (ggml_nelements(tensor) != nelements) {fprintf(stderr, "%s: tensor '%s' has wrong size in model file\n", __func__, name.c_str());return false;
}if (tensor->ne[0] != ne[0] || tensor->ne[1] != ne[1]) {fprintf(stderr, "%s: tensor '%s' has wrong shape in model file: got [%d, %d], expected [%d, %d]\n",__func__, name.c_str(), (int) tensor->ne[0], (int) tensor->ne[1], ne[0], ne[1]);return false;
}
利用 tensor 的名稱在預先構建的 model.tensors
哈希表中查詢該 tensor,并使用 ggml_set_name
設置 tensor 的名字,用 ggml_nelements
檢查元素數量是否匹配,通過 tensor 內部 ne
數組校驗 shape
4. 計算 tensor 內存大小與驗證
const size_t bpe = ggml_type_size(ggml_type(ttype));if ((nelements * bpe) / ggml_blck_size(tensor->type) != ggml_nbytes(tensor)) {fprintf(stderr, "%s: tensor '%s' has wrong size in model file: got %zu, expected %zu\n",__func__, name.c_str(), ggml_nbytes(tensor), nelements * bpe);return false;
}
使用 ggml_type_size
得到每個元素的字節數,調整后通過 ggml_blck_size
檢查文件中計算出的總字節數是否與 tensor 實際內存計算結果 ggml_nbytes
一致
5. 讀取權重數據并寫入 tensor 內存
if (ggml_backend_buffer_is_host(model.buffer_w)) {// 對于 CPU 或 Metal 后端,數據直接讀入系統內存(tensor->data)fin.read(reinterpret_cast<char *>(tensor->data), ggml_nbytes(tensor));
} else {// 對于設備內存,先讀入臨時緩存,再復制到設備內存read_buf.resize(ggml_nbytes(tensor));fin.read(read_buf.data(), ggml_nbytes(tensor));ggml_backend_tensor_set(tensor, read_buf.data(), 0, ggml_nbytes(tensor));
}
利用 ggml_backend_buffer_is_host
判斷當前后端緩沖區是否為 host 內存(在我們的例子中不是),如果是系統內存,直接讀取到 tensor->data
中,如果不是,需要先存到臨時緩存 read_buf
,再調用 ggml_backend_tensor_set
將數據寫入(適用于例如 CUDA 后端)
6. 處理特殊的 LM head 邏輯
// GPT-2 models share the WTE tensor as the LM head
if (name == "model/wte" && has_lm_head == false) {//ggml_backend_tensor_copy(tensor, model.lm_head);model.lm_head = tensor;
}if (name == "model/lm_head") {has_lm_head = true;
}
針對 gpt-2 模型中常見的共享權重場景:
- 如果當前讀取的是
"model/wte"
且未設置 lm_head,則將model.lm_head
指向該 tensor - 如果讀取到
"model/lm_head"
,則標記has_lm_head
為 true
這部分代碼通過文件流方式逐個加載模型中的 tensor 數據,并進行嚴格的 shape 匹配和數據有效性檢查
至此,我們完成了模型加載函數的全部代碼分析,下面我們來看看其他部分的處理
結語
這篇文章我們利用 ChatGPT 過了一遍 ggml 框架加載 gpt-2 模型的整體過程,包括文件驗證、詞匯表加載、上下文和后端初始化、模型張量的構建、kv cache 準備以及權重數據的載入等各個主要步驟
下篇文章我們來看 gpt-2 推理的其它流程,敬請期待🤗
下載鏈接
- ggml源碼下載鏈接【提取碼:1234】
- gpt-2-117M模型下載【提取碼:1234】
參考
- 【大模型部署】GGML源碼逐行調試
- llama.cpp源碼解讀–ggml框架學習
- https://github.com/ggml-org/ggml
- https://chatgpt.com/
- 理解llama.cpp如何進行LLM推理