在上一章 后端缓冲区 (ggml_backend_buffer) 中,我们学会了如何为我们的"厨师"(后端)准备专用的"储物空间"(后端缓冲区),比如在 GPU 显存上预留一块内存。
这带来了一个巨大的挑战:对于一个复杂的计算图,它可能包含成百上千个中间张量。我们如何才能将所有这些张量,特别是那些临时的、仅在计算中途存在的张量,都塞进我们有限的缓冲区(比如仅有的 8GB 显存)里呢?手动计算每个张量应该放在哪里、什么时候可以被覆盖,简直是一场噩梦。
幸运的是,ggml
提供了一位专家来解决这个难题,他就是本章的主角------ggml_gallocr
,图分配器。
什么是图分配器?为什么需要它?
核心思想 :把它想象成一个极其精打细算的仓管员。在开始烹饪(计算)前,他会仔细研究菜谱(计算图),分析出每个步骤需要哪些临时碗碟和食材,并规划出一个最高效的碗碟使用和回收方案。
ggml_gallocr_t
负责为计算图中的所有中间张量分配内存,它会分析张量的生命周期,智能地复用内存,从而用最小的厨房空间(内存)完成整道菜。
让我们看一个简单的计算 d = (a*b) + (a*c)
。这个计算图如下:
一个朴素的内存分配方式是为 res1
, res2
, d
各自分配一块独立的内存。但一个聪明的仓管员会发现:一旦 d
被计算出来,res1
和 res2
就再也没用了!
ggml_gallocr
的高明之处在于它能分析出这一点。它可以规划出这样的内存复用策略:
- 分配一块内存给
res1
。 - 分配一块内存给
res2
。 - 计算
d = res1 + res2
。 - 关键 :
d
可以直接覆盖res1
或res2
使用过的内存空间,因为它们已经完成了历史使命。
通过这种方式,ggml_gallocr
可以大幅度降低计算一个复杂图所需的峰值内存。对于动辄有数千个节点的大型语言模型来说,这种优化是至关重要的,它直接决定了模型能否在你的硬件上运行。
动手实践:聘请仓管员来规划内存
让我们看看如何使用 ggml_gallocr
来为一个计算图自动规划和分配内存。
第 1 步:聘请仓管员(创建分配器)
首先,我们需要创建一个 ggml_gallocr
实例。创建时,需要告诉它我们打算使用哪种类型的后端缓冲区(比如是普通的 CPU 内存还是 GPU 显存)。
c
#include "ggml-alloc.h"
#include "ggml-backend.h"
// 假设我们使用 CPU 后端
ggml_backend_t cpu_backend = ggml_backend_cpu_init();
ggml_backend_buffer_type_t cpu_buft =
ggml_backend_get_default_buffer_type(cpu_backend);
// 创建一个图分配器,告诉它我们将使用 CPU 缓冲区类型
ggml_gallocr_t galloc = ggml_gallocr_new(cpu_buft);
if (!galloc) { /* 错误处理 */ }
这行代码实例化了我们的"仓管员",并告诉他之后规划出的所有空间都应该是 CPU 内存。
第 2 步:准备菜谱(构建计算图)
和之前一样,我们先定义计算的步骤,构建一个计算图。
c
// (此处省略创建 ctx 和 a, b, c 张量的代码)
struct ggml_tensor * res1 = ggml_mul(ctx, a, b);
struct ggml_tensor * res2 = ggml_mul(ctx, a, c);
struct ggml_tensor * d = ggml_add(ctx, res1, res2);
struct ggml_cgraph * graph = ggml_new_graph(ctx);
ggml_build_forward_expand(graph, d);
现在我们有了一份完整的"菜谱" (graph
),可以交给仓管员了。
第 3 步:规划与预留空间 (Reserve)
这是最关键的一步。我们调用 ggml_gallocr_reserve
,让分配器分析图并计算出所需的最小内存,并实际分配这块内存。
c
// 让分配器分析图,并预留足够的后端缓冲区
if (!ggml_gallocr_reserve(galloc, graph)) {
fprintf(stderr, "预留缓冲区失败\n");
// ... 清理并退出 ...
}
执行完这行代码后,galloc
内部发生了许多事情:
- 它遍历了
graph
中的每一个节点。 - 它分析了每个张量的生命周期(何时创建,何时不再需要)。
- 它运行了一个复杂的算法,计算出了一个最优的内存布局方案。
- 它根据这个方案,计算出所需的总内存大小,并创建了一个足够大的后端缓冲区。
我们可以查询这个缓冲区到底有多大:
c
// 获取 0 号缓冲区的大小 (因为我们只用了一种类型)
size_t buffer_size = ggml_gallocr_get_buffer_size(galloc, 0);
printf("分配器计算出所需的内存大小: %.2f MB\n", buffer_size / 1024.0 / 1024.0);
第 4 步:正式分配 (Allocate)
空间已经预留好了,最后一步是 ggml_gallocr_alloc_graph
。它会根据 reserve
阶段的规划,为图中每一个需要分配内存的张量设置其 data
指针,使其指向缓冲区中正确的位置。
c
// 将图中所有张量的指针指向预留好的缓冲区中的具体位置
if (!ggml_gallocr_alloc_graph(galloc, graph)) {
fprintf(stderr, "分配图内存失败\n");
// ... 清理并退出 ...
}
现在,graph
中所有的张量(res1
, res2
, d
)都已经"有家可归"了。它们的 data
指针都指向 galloc
所管理的那个大缓冲区里的某个偏移地址。
第 5 步:执行与清理
分配完成后,我们就可以像往常一样执行计算,最后释放所有资源。
c
// 使用 CPU 后端执行计算
ggml_backend_graph_compute(cpu_backend, graph);
// 获取结果 (此处省略)
// 释放分配器(它会同时释放内部创建的缓冲区)
ggml_gallocr_free(galloc);
// 释放后端
ggml_backend_free(cpu_backend);
// 释放上下文等...
深入幕后:生命周期分析的魔力
ggml_gallocr
的核心是 ggml_gallocr_alloc_graph_impl
函数 (位于 src/ggml-alloc.c
)。它的工作流程可以简化为以下步骤:
-
计数 :首先,遍历整个图,为每个张量计算出它被用作输入的次数(
n_children
)和被视图引用的次数(n_views
)。这就像仓管员统计每样半成品会被多少个后续步骤用到。 -
模拟执行:然后,它按照计算图的执行顺序,一个一个地"模拟"计算每个节点。
-
分配与释放循环:
- 处理当前节点
N
:当模拟到计算节点N
时,分配器认为N
的输出张量此刻诞生了。它会调用内部的动态分配器ggml_dyn_tallocr_alloc
,从内存池中为N
的输出张量找一块合适的空间。 - 更新父节点 :计算完
N
后,它的所有输入(父节点)张量的"待使用次数" (n_children
) 减一。 - 检查生命周期结束 :分配器检查每个父节点,如果它的
n_children
和n_views
都变成了 0,就意味着这个张量在未来的计算中再也用不到了。它就"死亡"了。 - 回收内存 :对于"死亡"的张量,分配器会调用
ggml_dyn_tallocr_free_tensor
将其占用的内存块归还到内存池中,以备后续的张量分配使用。
- 处理当前节点
下面的序列图展示了这个简化的流程:
内存复用可视化
让我们用一张图来直观感受内存是如何被复用的。假设缓冲区是一条长长的内存条。
- 时间 0-1 :计算
res1
,为res1
分配内存。 - 时间 1-2 :计算
res2
,为res2
分配另一块内存。此时res1
和res2
都必须在内存中。 - 时间 2-3 :计算
d
。一旦d
计算完成,res1
和res2
就无用了。ggml_gallocr
发现这一点后,可以立即将res1
或res2
的内存回收,并分配给d
。这样,峰值内存占用就是sizeof(res1) + sizeof(res2)
,而不是sizeof(res1) + sizeof(res2) + sizeof(d)
。
在ggml-alloc.c
中,你可以找到ggml_dyn_tallocr_alloc
和ggml_dyn_tallocr_free_tensor
这两个函数,它们实现了一个简单的动态内存管理器,维护一个空闲块列表,并支持合并相邻的空闲块,以减少内存碎片。
总结
在本章中,我们认识了 ggml
中强大的内存规划专家------图分配器 (ggml_gallocr
)。
- 它像一个精打细算的仓管员 ,通过分析计算图中张量的生命周期来工作。
- 它的核心能力是智能地复用内存,将已经"死亡"的临时张量的空间回收,给新生的张量使用。
- 这极大地降低了执行复杂计算图所需的峰值内存,让在有限资源(如消费级显卡)上运行大型模型成为可能。
- 我们学会了标准的使用流程:
ggml_gallocr_new
->ggml_gallocr_reserve
->ggml_gallocr_alloc_graph
。
我们已经走过了 ggml
的核心组件之旅:从数据(张量)、内存(上下文、缓冲区)到计算计算图、后端),最后到内存的自动规划(图分配器)。现在,所有的准备工作都已就绪。
但是,当一个模型非常大,一部分权重在 CPU 内存中,另一部分在 GPU 显存中时,我们该如何协调这两个"厨师"共同完成一份"菜谱"呢?