目录
- 前言
- 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推理