ggml介绍 (8) 图分配器 (ggml_gallocr)

在上一章 后端缓冲区 (ggml_backend_buffer) 中,我们学会了如何为我们的"厨师"(后端)准备专用的"储物空间"(后端缓冲区),比如在 GPU 显存上预留一块内存。

这带来了一个巨大的挑战:对于一个复杂的计算图,它可能包含成百上千个中间张量。我们如何才能将所有这些张量,特别是那些临时的、仅在计算中途存在的张量,都塞进我们有限的缓冲区(比如仅有的 8GB 显存)里呢?手动计算每个张量应该放在哪里、什么时候可以被覆盖,简直是一场噩梦。

幸运的是,ggml 提供了一位专家来解决这个难题,他就是本章的主角------ggml_gallocr,图分配器。

什么是图分配器?为什么需要它?

核心思想 :把它想象成一个极其精打细算的仓管员。在开始烹饪(计算)前,他会仔细研究菜谱(计算图),分析出每个步骤需要哪些临时碗碟和食材,并规划出一个最高效的碗碟使用和回收方案。ggml_gallocr_t 负责为计算图中的所有中间张量分配内存,它会分析张量的生命周期,智能地复用内存,从而用最小的厨房空间(内存)完成整道菜。

让我们看一个简单的计算 d = (a*b) + (a*c)。这个计算图如下:

graph TD a --> mul1("乘法 1") b --> mul1 mul1 --> res1("中间结果 res1 = a*b") a --> mul2("乘法 2") c --> mul2 mul2 --> res2("中间结果 res2 = a*c") res1 --> add("加法") res2 --> add add --> d("最终结果 d")

一个朴素的内存分配方式是为 res1, res2, d 各自分配一块独立的内存。但一个聪明的仓管员会发现:一旦 d 被计算出来,res1res2 就再也没用了!

ggml_gallocr 的高明之处在于它能分析出这一点。它可以规划出这样的内存复用策略:

  1. 分配一块内存给 res1
  2. 分配一块内存给 res2
  3. 计算 d = res1 + res2
  4. 关键d 可以直接覆盖 res1res2 使用过的内存空间,因为它们已经完成了历史使命。

通过这种方式,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 内部发生了许多事情:

  1. 它遍历了 graph 中的每一个节点。
  2. 它分析了每个张量的生命周期(何时创建,何时不再需要)。
  3. 它运行了一个复杂的算法,计算出了一个最优的内存布局方案。
  4. 它根据这个方案,计算出所需的总内存大小,并创建了一个足够大的后端缓冲区。

我们可以查询这个缓冲区到底有多大:

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)。它的工作流程可以简化为以下步骤:

  1. 计数 :首先,遍历整个图,为每个张量计算出它被用作输入的次数(n_children)和被视图引用的次数(n_views)。这就像仓管员统计每样半成品会被多少个后续步骤用到。

  2. 模拟执行:然后,它按照计算图的执行顺序,一个一个地"模拟"计算每个节点。

  3. 分配与释放循环

    • 处理当前节点 N :当模拟到计算节点 N 时,分配器认为 N 的输出张量此刻诞生了。它会调用内部的动态分配器 ggml_dyn_tallocr_alloc,从内存池中为 N 的输出张量找一块合适的空间。
    • 更新父节点 :计算完 N 后,它的所有输入(父节点)张量的"待使用次数" (n_children) 减一。
    • 检查生命周期结束 :分配器检查每个父节点,如果它的 n_childrenn_views 都变成了 0,就意味着这个张量在未来的计算中再也用不到了。它就"死亡"了。
    • 回收内存 :对于"死亡"的张量,分配器会调用 ggml_dyn_tallocr_free_tensor 将其占用的内存块归还到内存池中,以备后续的张量分配使用。

下面的序列图展示了这个简化的流程:

sequenceDiagram participant User as 用户代码 participant Galloc as ggml_gallocr_reserve participant DynAlloc as 动态内存池 participant Graph as 计算图 User->>Galloc: 调用 ggml_gallocr_reserve(galloc, graph) Galloc->>Graph: 遍历图, 统计每个张量的"孩子"数量 Galloc->>Graph: 开始按顺序模拟执行节点 loop 对图中的每个节点 node Galloc->>DynAlloc: 为 node 的输出张量分配内存 DynAlloc-->>Galloc: 返回内存地址 Note over Galloc: 模拟计算 node... loop 对 node 的每个输入(父)张量 parent Galloc->>Galloc: parent 的"孩子"数量减 1 alt parent 的"孩子"和"视图"数量都为 0 Galloc->>DynAlloc: 释放 parent 占用的内存 DynAlloc-->>Galloc: 内存已回收 end end end Note over Galloc: 记录整个过程中的峰值内存使用量 Galloc->>Galloc: 根据峰值内存大小, 分配最终的缓冲区 Galloc-->>User: 规划和预留完成

内存复用可视化

让我们用一张图来直观感受内存是如何被复用的。假设缓冲区是一条长长的内存条。

gantt title 内存使用时间线 dateFormat X axisFormat %s section 任务 计算 res1=a*b: 0, 1 计算 res2=a*c: 1, 2 计算 d=res1+res2: 2, 3 section 内存分配 res1 : 0, 2 res2 : 1, 2 d : 2, 3
  • 时间 0-1 :计算 res1,为 res1 分配内存。
  • 时间 1-2 :计算 res2,为 res2 分配另一块内存。此时 res1res2 都必须在内存中。
  • 时间 2-3 :计算 d。一旦 d 计算完成,res1res2 就无用了。ggml_gallocr 发现这一点后,可以立即将 res1res2 的内存回收,并分配给 d。这样,峰值内存占用就是 sizeof(res1) + sizeof(res2),而不是 sizeof(res1) + sizeof(res2) + sizeof(d)

ggml-alloc.c中,你可以找到ggml_dyn_tallocr_allocggml_dyn_tallocr_free_tensor这两个函数,它们实现了一个简单的动态内存管理器,维护一个空闲块列表,并支持合并相邻的空闲块,以减少内存碎片。

总结

在本章中,我们认识了 ggml 中强大的内存规划专家------图分配器 (ggml_gallocr)

  • 它像一个精打细算的仓管员 ,通过分析计算图中张量的生命周期来工作。
  • 它的核心能力是智能地复用内存,将已经"死亡"的临时张量的空间回收,给新生的张量使用。
  • 这极大地降低了执行复杂计算图所需的峰值内存,让在有限资源(如消费级显卡)上运行大型模型成为可能。
  • 我们学会了标准的使用流程:ggml_gallocr_new -> ggml_gallocr_reserve -> ggml_gallocr_alloc_graph

我们已经走过了 ggml 的核心组件之旅:从数据(张量)、内存(上下文、缓冲区)到计算计算图、后端),最后到内存的自动规划(图分配器)。现在,所有的准备工作都已就绪。

但是,当一个模型非常大,一部分权重在 CPU 内存中,另一部分在 GPU 显存中时,我们该如何协调这两个"厨师"共同完成一份"菜谱"呢?

相关推荐
TDengine (老段)6 分钟前
TDengine IDMP 高级功能(4. 元素引用)
大数据·数据库·人工智能·物联网·数据分析·时序数据库·tdengine
curdcv_po7 分钟前
😲AI 💪🏻超级 整合时代 已经 到来~
人工智能·trae
*星星之火*13 分钟前
【GPT入门】第47课 大模型量化中 float32/float16/uint8/int4 的区别解析:从位数到应用场景
人工智能·gpt
雨落倾城夏未凉28 分钟前
8.被free回收的内存是立即返还给操作系统吗?为什么?
c++·后端
雨落倾城夏未凉31 分钟前
6.new和malloc的区别
c++·后端
郝学胜-神的一滴35 分钟前
深入理解QFlags:Qt中的位标志管理工具
开发语言·c++·qt·程序人生
aneasystone本尊1 小时前
学习 Coze Studio 的工作流执行逻辑
人工智能
aneasystone本尊1 小时前
再学 Coze Studio 的智能体执行逻辑
人工智能
xuanwuziyou1 小时前
LangChain 多任务应用开发
人工智能·langchain
INS_KF1 小时前
【C++知识杂记2】free和delete区别
c++·笔记·学习