Version: Commit 5488fb78
Feature: x86_64 CPU Only, mmap,no blas. 略过了所有有关KV缓存的部分.以下文字可能因代码更新或本人认知水平而不再适用.
LLaMA.cpp or GGML
GGML是LLaMa.cpp同一作者的作品,作者将LLaMa和Whisper等cpp实现库中用到的运算相关函数提取到了一个库中,命名为GGML(A Tensor Lib for ML), 因此GGML并不是类似MNN等的推理框架,而仅仅是一个操作算子的加速计算库而已,有关模型序列化\子图划分等上层功能仍要由调用方实现,因此这里还是选择阅读LLaMa.cpp的代码,来梳理一个简单的推理库要做什么工作.
Weights File Structrue
LLama.cpp(下文简称Lc)没有像其他ML框架一样借助Proto或者FlatBuf这种序列化框架来实现权重的序列化，而是简单采用二进制顺序读写来自定义序列化，比起框架方案缺少了向前兼容和透明迁移等特性，但是毫无疑问简单了很多。
下面介绍文件结构:
第一个 u32 是Magic Number，用于识别文件格式与版本类型(LLAMA_FILE_VERSION_GGMF_V1 ,LLAMA_FILE_VERSION_GGJT_V1 etc.) ，具体可参考下面的enum.
enum llama_file_version {
LLAMA_FILE_VERSION_GGML,
LLAMA_FILE_VERSION_GGMF_V1, // added version field and scores in vocab
LLAMA_FILE_VERSION_GGJT_V1, // added padding
LLAMA_FILE_VERSION_GGJT_V2, // changed quantization format
LLAMA_FILE_VERSION_GGJT_V3, // changed Q4 and Q8 quantization format
};接下来的7个u32是模型的参数，如layer数目,head数目等，是与模型类型紧密结合的，此外还有一个n_vocab 参数，用于决定接下来的vocab存储的数目。
再接下来LLama.cpp会从文件中读取vocab，每个vocab由字符串长度len(u32) vocab字符串word(string 长度为len) score (float) 组成,并被按顺序从文件中依次读出.
vocab表读取完之后,文件剩下的部分就是Tensor Map了,Lc使用一个while循环一直读到文件末尾来读取出所有的Tensor信息(包括名称和权重).每个Tensor的存储格式为维度数目n_nums(u32) 名称长度 name_len(u32) Tensor type(u32) Tensor 维度数组(长为n_nums的u32数组,记录了每一维的大小 Tensor名称(长为name_len的字符串 ,读出以上metadata之后,如果文件存储格式版本大于LLAMA_FILE_VERSION_GGJT_V1,则metadata之后会有一段padding来对齐到32字节.从下一个32字节开始存储Tensor的weights数据 ,但是在最开始的load过程中不做读取,仅仅计算出应当读取的字节数并记录在llama_load_tensor 结构体中,后期再使用mmap或其他方式读入内存.读取metadata完毕后,会将llama_load_tensor放入一个列表中,并生成一个name_to_tensor_idx的索引便于查找.
至此 模型文件的加载就完成了.
struct llama_load_tensor {
std::string name;
enum ggml_type type = GGML_TYPE_F32;
std::vector<uint32_t> ne; //维度数据
size_t file_off; //计算在文件中的偏移量
size_t size; //计算字节数
struct ggml_tensor * ggml_tensor = NULL; //指向真正的 可执行计算的 Tensor对象
uint8_t * data;
};Lc计算weights的字节数:
对于Float32和FP16型的Tensor而言,使用一个相应的Float类型来存储即可,Lc计算出的字节数即为sizeof(float)但是对于量化类型,每个FP32元素会用较低精度的类型来表示(如4bit int),为了弥补精度下降带来的计算误差,必须额外携带纠正数据,因此可能使用一个结构体来存储一组数据.一个典型结构如下: 以Q4_0量化为例,量化后的数据每32个元素分为一个block,每个数据以4bit(半个int8)来存储,占用
sizeof(fp16)+sizeof(uint8)*16=18bit存储,因此n个元素的Tensor占用的字节数目为sizeof(block_q4_0) *n /32依照此可以算出Tensor的weight在文件中的偏移量.这也是为什么4Bit/8Bit量化可以显著降低内存占用的原因.
#define QK4_0 32
typedef struct {
ggml_fp16_t d; // delta
uint8_t qs[QK4_0 / 2]; // nibbles / quants
} block_q4_0;
static_assert(sizeof(block_q4_0) == sizeof(ggml_fp16_t) + QK4_0 / 2, "wrong q4_0 block size/padding");int qk =32
for (int j = 0; j < qk/2; ++j) {
//同时取两个数字进行存储
const float x0 = x[i*qk + 0 + j]*id;
const float x1 = x[i*qk + qk/2 + j]*id;
const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f));
const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f));
//分别占据一个uint8的前4位和后四位
y[i].qs[j] = xi0;
y[i].qs[j] |= xi1 << 4;
}讨论:Lc中的权重格式是高度定制化的,而且也没有数据校验设计,只能用于相对简单且格式更迭不剧烈的情况下.
Memory Buffer in GGML
Lc与底层GGML各有内存池实现,不过都比较简单,是使用静态数组来作为内存池并且预先分配完成.Lc的内存池为llama_context结构体中的buf_compute buf_scratch* ,类型为llama_buffer.此外,还用llama_model 中的buf(,但仅仅用于存储从文件中读取的权重Tensor结构体信息),llama_context中的work_buffer (算子计算时的中间结果缓冲区,是每个算子进行重用的). Lc创建的内存池都会在不同时期通过以引用形式传递给下层ggml的方式来进行使用,lc本身不进行内存池上对象的分配.GGML的内存池位于ggml_context中的mem_buffer ,此外还有ggml_scratch scratch 来持有上述中llama buf_scratch内存池列表中内存池的引用,便于GGML进行对象的分配.各内存池都随相应的context来进行创建(或引用已有内存池).
buf_compute内存池用于放置推理时生成的对象(ggml_object),内存池的大小在context初始化时就已经决定了,而且没有动态扩容功能.当llama_context加载模型时,会根据模型的类型从预设中取出预定义的buf_compute大小.buf_scratch内存池在LLAMA_USE_SCRATCH宏启用后,会用于Tensor中的实际Data数据的分配存储.Lc会维护一个buf_scratch列表,内部包含两个内存池,大小也是预先静态定义的.在推理过程中,Lc会切换两个内存池来存储Tensor Data,如self-attention放置在第1块内存池,而feed-forward层放置在第2块内存池.这样做的设计我还没有理解.
以上两个内存池的初始化代码如下:
buf_compute.resize(MEM_REQ_EVAL().at(ctx->model.type) + ggml_graph_overhead());
buf_scratch[0].resize(MEM_REQ_SCRATCH0(hparams.n_ctx).at(ctx->model.type));
buf_scratch[1].resize(MEM_REQ_SCRATCH1().at(ctx->model.type));- 在ggml中的两个内存池指针中,
ggml_scratch仅仅为了持有上层框架传入的内存池引用而设计(ggml_set_scratch),ggml不负责其创建和销毁.而mem_buffer是由ggml来进行维护的.在创建新的ggml_context时,可以选择传入mem_buffer的地址来引用预先创建的内存池,或交由ggml来创建指定大小的内存池(仅在创建KV存储部分见到过这种情形).mem_buffer在启用LLAMA_USE_SCRATCH特性后,仅仅用于存储ggml_object链表和ggml_tensor对象,而ggml_tensor并不实际存储Tensor中的data,data可能会被存储于mmap区域.计算一个模型所需的mem_buffer的方法如下(注意 这仅仅是加载模型所需Tensor的空间大小.运算时的mem_buffer会指向lc的buf_compute内存池,计算图和计算Tensor都在其中分配):
void calc_sizes(size_t * ctx_size_p, size_t * mmapped_size_p) const {
*ctx_size_p = *mmapped_size_p = 0;
for (const llama_load_tensor & lt : tensors_map.tensors) {
*ctx_size_p += sizeof(struct ggml_tensor) + GGML_OBJECT_SIZE;
*(use_mmap ? mmapped_size_p : ctx_size_p) += lt.size + 16; //如果开启mmap,实际上ctx中的mem_buffer就不需要存储Tensor中的数据.
}
}我们借助一个Tensor的创建过程来实际描述Lc中内存池的使用:
static struct ggml_tensor * ggml_new_tensor_impl(
struct ggml_context * ctx,
enum ggml_type type,
int n_dims,
const int64_t* ne,
void* data) {
size_t data_size = 0;
if (data == NULL && !ctx->no_alloc) { //no_alloc只有在非CPU后端时才会为true
data_size += GGML_TYPE_SIZE[type]*(ne[0]/GGML_BLCK_SIZE[type]);
for (int i = 1; i < n_dims; i++) {
data_size *= ne[i];
}
}
if (ctx->scratch.data != NULL && data == NULL) {
// allocate tensor data in the scratch buffer
if (ctx->scratch.offs + data_size > ctx->scratch.size) {
GGML_PRINT("%s: not enough space in the scratch memory pool (needed %zu, available %zu)\n",
__func__, ctx->scratch.offs + data_size, ctx->scratch.size);
assert(false);
return NULL;
}
data = (char * const) ctx->scratch.data + ctx->scratch.offs;
ctx->scratch.offs += data_size;
data_size = 0;
}
struct ggml_object * const obj_new = ggml_new_object(ctx, GGML_OBJECT_TENSOR, GGML_TENSOR_SIZE + data_size);
struct ggml_tensor * const result = (struct ggml_tensor *)((char *)ctx->mem_buffer + obj_new->offs);
*result = (struct ggml_tensor) {
/*.type =*/ type,
/*.backend =*/ GGML_BACKEND_CPU,
/*.n_dims =*/ n_dims,
/*.ne =*/ { 1, 1, 1, 1 },
/*.nb =*/ { 0, 0, 0, 0 },
/*.op =*/ GGML_OP_NONE, //GGML_OP_NONE代表constant Tensor 不执行任何计算.
// ... something we do not care!...
/*.data =*/ (data == NULL && !ctx->no_alloc) ? (void *)(result + 1) : data,
/*.name =*/ { 0 }, //注意 这里的名字也是空白的,等待后续使用set_name来命名.
};
// 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++;
return result;
}在 ggml_new_tensor_impl 中,当要求创建一个Tensor且没有为它传入data指针时,证明需要创建一个指定维度的空白Tensor.首先ggml根据维度数据计算出内部的data所需的内存大小.当启用了LLAMA_USE_SCRATCH后,data会被分配到scracth内存池中.如果内存池空间不足,则直接报错退出,没有扩容操作.scracth内存池为很简单的offset模型,获取一定长度的内存创建data后,增加已经使用的offset数值即可.
ggml_new_object 函数用于维护ggml_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 pool
struct 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_ALIGN
size_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_PRINT("%s: not enough space in the context's memory pool (needed %zu, available %zu)\n",
__func__, cur_end + size_needed, ctx->mem_size);
assert(false);
return 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 context
ctx->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;
}
//ggml object
enum ggml_object_type {
GGML_OBJECT_TENSOR,
GGML_OBJECT_GRAPH,
GGML_OBJECT_WORK_BUFFER
};
// ggml object
struct ggml_object {
size_t offs;
size_t size;
struct ggml_object * next;
enum ggml_object_type type;
char padding[4];
};这里是一个标准的链表模式,ctx中记录了object链表的头和尾元素.object元素自身也携带着指向下个元素的指针.但是实际上object本身并不承担任何功能,仅仅为了标记内存池的分配情况.ctx->objects_end 指向了链表的末尾元素,也是当前内存池中最后一个元素.当需要分配新的元素时,通过最后一个元素的首地址(offset)+元素自身大小(size)就可以得到下一个空闲地址在内存池中的位置了.当我们真正要为一个Object(以 ggml_tensor为例子)分配内存时,首先计算出Object所需的内存大小(在这里为 sizeof(ggml_tensor)+ datasize)并进行内存对齐,在加上一个ggml_object 对象的大小,即为需分配的内存大小.从内存池中先创建一个ggml_object ,并写入Object 的offset和size,插入链表并返回.调用者即可利用返回的ggml_object 中的offset内存地址和size来创建上层的Object.
此外,所有的内存池数组都已经Aligned,在32位系统上为4bit对齐,64位为16bit对齐.(但是在Metal中可能有不同的对齐方式,这里不做研究),视需要也可能使用mlock来pin住内存.
1.这里有个小问题: 如果我们delete了某个池内对象,造成的内存碎片怎么处理呢?
答案是:不delete. GGML没有设计delete_object的相关函数(最起码我还没有找到).当GGML不需要一个Tensor的时候,极有可能整张运算图都不需要了,直接把整个context连同Memory Buffer扔掉就好了.
- Object 链表是拿来做什么的呢?我一开始以为内存池的分配会是类似OS中首次适应算法的设计,但是实际上并没有那么复杂.目前看来 链表的唯一作用是给定Context和Tensor名称的时候可以查找对应的Tensor(但是这个功能看上去完全没有需求,包括Lc本身的examples里面也没有调用这种方案)或者按顺序PPrint Obejcts.
Tensor in GGML
在GGML中,从Weights文件中反序列化出的Tensor仅仅具有名字,维度和内容所在的文件偏移地址,并不包含模型结构与计算图信息,因此无法参与运算.正如上文所提到的,这种Tensor的类型为llama_load_tensor ,仅仅用于读取模型时候使用.在model加载完毕后,lc会根据tensor_map 中的大小来分配内存池.分配完成后,则开始根据事先硬编码的Tensor名称来查询并为llama_load_tensor 创建真正的Tensor,包括norm等层的权重向量,model的output等.创建过程类似上文所述ggml_new_tensor_impl .[参见 llama.cpp line_{1164,1224}].Tensor全部创建完成之后,lc会从文件中加载实际Tensor的data. 加载过程如下:
void load_data_for(llama_load_tensor & lt) {
if (use_mmap) {
lt.data = (uint8_t *) mapping->addr + lt.file_off; //根据文件偏移值直接从mmap中读出data起始地址
} else {
llama_file & file = file_loader->file;
file.seek(lt.file_off, SEEK_SET);
file.read_raw(lt.data, lt.size);//若未启用mmap,lt.data为uint8数组,从文件中拷贝进去
}若启用了mmap,则不需要手动进行读入,只需要计算出mmap后的地址赋值即可,此外根据Tensor位于的Backend不同,加载data后也执行不同操作,CPU backend会使用Mlock来Pin住这一块data的内存,防止被换页造成性能损耗,GPU会转移到对应的硬件内存中并释放data内存.
ggml_tensor的结构如下:
struct ggml_tensor {
enum ggml_type type; //数据类型 FP32 FP16 or Quant
enum ggml_backend backend; //位于的计算后端
int n_dims; // N维矩阵
int64_t ne[GGML_MAX_DIMS]; // number of elements of each dim
size_t nb[GGML_MAX_DIMS]; // stride in bytes:
// nb[0] = sizeof(type)
// nb[1] = nb[0] * ne[0] + padding
// nb[i] = nb[i-1] * ne[i-1]
// stride用来快速通过维度坐标取出底层元素
// compute data
enum ggml_op op; // Operation Code,标明这个向量是什么操作的运算结果
// op params - allocated as int32_t for alignment
int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];
bool is_param; //用于标识Trainable的(猜测) Infer用不到
struct ggml_tensor * grad; //梯度Tensor 记录反向传播用的 在Infer用不到
struct ggml_tensor * src[GGML_MAX_SRC]; // Tensor在图中的数据来源,即OP的运算数
// performance 性能benchamrk用的
int perf_runs;
int64_t perf_cycles;
int64_t perf_time_us;
// 底层数组
void * data;
char name[GGML_MAX_NAME];
void * extra; // extra things e.g. for ggml-cuda.cu
// 没找到干嘛的↓
char padding[4];
};在模型的运算(Inference)中,大部分框架遵循的一个基本的模式即为数据流从Input 开始流动,(Input可能是模型的输入Tensor,也可能是权重文件中固化的Weights常量,也可能是配置常量),流经Operation节点对数据进行操作,并输出新的数据,新的数据沿着计算图设定好的模型结构继续流向下一个Operation节点,最终直至流到输出节点进行结果输出.每个Operation通常不止一个Input(一般有一个或多个来自模型输入的变量Tensor,也有多个常量作为Weights来参与计算),但一般情况下可能只有一个Output(::反正Lc只考虑了有一个Output的情况,在不考虑RNN\LSTM这些算子和控制流的情况下,大概也是够用的?::).因此 Lc的Tensor做出了以下设计:- Tensor为Constant,在模型中可能对应权重常量或模型输入,仅作为Operation的数据来源,不执行计算,op属性为
GGML_OP_NONE,src为{NULL}.在计算图中处于叶子结点.其data属性即为携带的底层内容数据. - Tensor为
Intermediate Tensor,即计算节点.为了实现方便,Lc中的Operation与此Operation计算的结果都存储在此Tensor中.此Tensor会执行op属性的算子计算,计算所需的操作数为src属性指向的多个Tensor id,计算得出的结果为data属性所指向的数据. op_params属性为执行计算所需的其他parameter.
通过src属性,每个计算Tensor都可以找到自己的数据源Tensor,同时也会被下一级计算的Tensor引用,直到追溯到Constant,这便组成了一个DAG计算图.Lc在执行阶段便沿着这条计算图来进行计算和数据流动.
Computation Graph
参见Lc中的llama_eval_internal 函数,在约{1488->1715}行,使用CPP代码手动构建了一遍LLaMa的结构模型,包括norm,self-Attention,feed_forward等部分.至此,DAG算子图建立完毕.但是仅仅有算子之间的关联并不足以构建执行模型所需的执行顺序等信息,ggml通过ggml_cgraph 来创建执行模型的计算顺序信息. ggml_cgraph 的结构如下:
// computation graph
struct ggml_cgraph {
int n_nodes;
int n_leafs;
struct ggml_tensor * nodes[GGML_MAX_NODES];
struct ggml_tensor * grads[GGML_MAX_NODES];
struct ggml_tensor * leafs[GGML_MAX_NODES];
void * visited_hash_table[GGML_GRAPH_HASHTABLE_SIZE];
// performance
int perf_runs;
int64_t perf_cycles;
int64_t perf_time_us;
};ggml通过ggml_build_forward_expand 函数来通过Tensor组成的DAG图来生成计算图,阅读下ggml_build_forward_impl 的核心代码就知道,这是个典型的拓扑排序.
static void ggml_build_forward_impl(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, bool expand) {
if (!expand) {
cgraph->n_nodes = 0;
cgraph->n_leafs = 0;
}
const int n0 = cgraph->n_nodes;
UNUSED(n0);
ggml_visit_parents(cgraph, tensor);
const int n_new = cgraph->n_nodes - n0;
GGML_PRINT_DEBUG("%s: visited %d new nodes\n", __func__, n_new);
if (n_new > 0) {
// the last added node should always be starting point
GGML_ASSERT(cgraph->nodes[cgraph->n_nodes - 1] == tensor);
}
}
static void ggml_visit_parents(struct ggml_cgraph * cgraph, struct ggml_tensor * node) {
if (node->grad == NULL) {
// this usually happens when we generate intermediate nodes from constants in the backward pass
// it can also happen during forward pass, if the user performs computations with constants
if (node->op != GGML_OP_NONE) {
//GGML_PRINT_DEBUG("%s: warning: node %p has no grad, but op %d\n", __func__, (void *) node, node->op);
}
}
// check if already visited
if (hash_insert(cgraph->visited_hash_table, node)) {
return;
}
for (int i = 0; i < GGML_MAX_SRC; ++i) {
if (node->src[i]) {
ggml_visit_parents(cgraph, node->src[i]);
}
}
if (node->op == GGML_OP_NONE && node->grad == NULL) {
// reached a leaf node, not part of the gradient graph (e.g. a constant)
GGML_ASSERT(cgraph->n_leafs < GGML_MAX_NODES);
if (strlen(node->name) == 0) {
ggml_format_name(node, "leaf_%d", cgraph->n_leafs);
}
cgraph->leafs[cgraph->n_leafs] = node;
cgraph->n_leafs++;
} else {
GGML_ASSERT(cgraph->n_nodes < GGML_MAX_NODES);
if (strlen(node->name) == 0) {
ggml_format_name(node, "node_%d", cgraph->n_nodes);
}
cgraph->nodes[cgraph->n_nodes] = node;
cgraph->grads[cgraph->n_nodes] = node->grad;
cgraph->n_nodes++;
}
}这个拓扑排序保证了以下几点:
- nodes数组中的元素应当都是计算Tensor.leafs数组中应当都是常量.
- nodes数组中的元素索引从小到大按照拓扑顺序排序的,这也就意味着某个计算Tensor所依赖的所有数据源Tensor一定在nodes数组中排在此Tensor之前.当GGML按照nodes索引从小到大的顺序依次执行计算Tensor时,保证了每个Tensor的数据源一定已经计算完毕了,从而避免了复杂的依赖问题.
成功构建计算图后,GGML通过ggml_graph_compute_helper 函数来编排执行顺序ggml_cplan .这个执行plan更多的是根据算子类型来考虑每个算子能否拆分成多线程并行操作,并计算每个算子计算过程中所需要多少内存,因为算子计算过程中的中间结果会存储在work_buffer 中,因此这些内存需求的最大值会决定work_buffer 的大小.
最后,在ggml_graph_compute 函数中,GGML会创建多个线程,分别执行ggml_graph_compute_thread 计算函数来多线程推理模型,他们之间通过 ggml_compute_state_shared 的引用来共享状态.
struct ggml_compute_state_shared {
const struct ggml_cgraph * cgraph; //计算图
const struct ggml_cplan * cplan; //执行计划,负责提供执行过程中的Context,包含中间变量的内存池,那些算子可以多线程并行,以及cancel_callback.
int64_t perf_node_start_cycles;
int64_t perf_node_start_time_us;
const int n_threads;
// synchronization primitives 用于线程间状态同步
atomic_int n_active; // num active threads
atomic_int node_n; // active graph node
bool (*abort_callback)(void * data); // abort ggml_graph_compute when true
void * abort_callback_data;
};而ggml_graph_compute_thread 的核心思想即为: 所有线程共同计算一个算子(假设这个算子可以多线程并行).
- 在线程执行的最开始,每个线程都会去查询计算图中的算子列表,并尝试从0开始执行.
- 当取出第一个计算算子后,线程会首先执行算子的INIT子任务.大部分算子的
INIT子任务是空白的,但也有些是初始化内存等操作.(注意,此时可能所有线程都会执行第一个算子的INIT子任务,不过这种任务一般不会有副作用) INIT执行完毕后,线程会检查此算子能否并行执行,
- 如果不能,则会自己完成算子的
COMPUTE子任务,即实际上的算子计算任务. - 如果可以并行执行,那么他会尝试将任务分配给cplan规划的指定数目的进程,即设置共享状态.:共享状态(简述为State)中n_active值为应当并行的线程数目, node_n为当前计算的算子ID.并开始执行自身分配到的一部分并行计算工作.
- 如果不能,则会自己完成算子的
- 如果某个线程先完成了自己的一份计算工作,则他会将State中的n_active -1(Atomically),并挂起等待,直到node_n变更,代表下一个可以并行的算子计算工作的开始.
- 当此算子最后一块计算工作也被完成时,线程在-1时发现自己为最后一个完成的线程,会执行算子的
Finalize子任务,即收集结果等,同样也不是所有任务的Finalize都有内容. - 最后一个进程执行
Finalize完毕后,他会检查下一个算子能否并行并执行INIT,若不能并行,则自己完成计算任务后继续检查下一个算子,直到某个算子可以多线程并行,此时跳转到第三步并依次循环.
static thread_ret_t ggml_graph_compute_thread(void * data) {
struct ggml_compute_state * state = (struct ggml_compute_state *) data;
const struct ggml_cgraph * cgraph = state->shared->cgraph;
const struct ggml_cplan * cplan = state->shared->cplan;
const int * n_tasks_arr = cplan->n_tasks;
const int n_threads = state->shared->n_threads;
set_numa_thread_affinity(state->ith, n_threads);
int node_n = -1;
while (true) {
// 判断是否调用callback取消了推理
if (cplan->abort_callback && cplan->abort_callback(cplan->abort_callback_data)) {
state->shared->node_n += 1;
return (thread_ret_t) GGML_EXIT_ABORTED;
}
if (atomic_fetch_sub(&state->shared->n_active, 1) == 1) {
// all other threads are finished and spinning
// do finalize and init here so we don't have synchronize again
struct ggml_compute_params params = {
/*.type =*/ GGML_TASK_FINALIZE,
/*.ith =*/ 0,
/*.nth =*/ 0,
/*.wsize =*/ cplan->work_size,
/*.wdata =*/ cplan->work_data,
};
if (node_n != -1) {
/* FINALIZE */
struct ggml_tensor * node = state->shared->cgraph->nodes[node_n];
if (GGML_OP_HAS_FINALIZE[node->op]) {
params.nth = n_tasks_arr[node_n];
ggml_compute_forward(¶ms, node);
}
ggml_graph_compute_perf_stats_node(node, state->shared);
}
// distribute new work or execute it direct if 1T
while (++node_n < cgraph->n_nodes) {
GGML_PRINT_DEBUG_5("%s: %d/%d\n", __func__, node_n, cgraph->n_nodes);
struct ggml_tensor * node = cgraph->nodes[node_n];
const int n_tasks = n_tasks_arr[node_n];
state->shared->perf_node_start_cycles = ggml_perf_cycles();
state->shared->perf_node_start_time_us = ggml_perf_time_us();
params.nth = n_tasks;
/* INIT */
if (GGML_OP_HAS_INIT[node->op]) {
params.type = GGML_TASK_INIT;
ggml_compute_forward(¶ms, node);
}
if (n_tasks == 1) {
// TODO: maybe push node_n to the atomic but if other threads see n_tasks is 1,
// they do something more efficient than spinning (?)
params.type = GGML_TASK_COMPUTE;
ggml_compute_forward(¶ms, node);
if (GGML_OP_HAS_FINALIZE[node->op]) {
params.type = GGML_TASK_FINALIZE;
ggml_compute_forward(¶ms, node);
}
ggml_graph_compute_perf_stats_node(node, state->shared);
} else {
break;
}
if (cplan->abort_callback && cplan->abort_callback(cplan->abort_callback_data)) {
break;
}
}
atomic_store(&state->shared->n_active, n_threads);
atomic_store(&state->shared->node_n, node_n);
} else {
// wait for other threads to finish
const int last = node_n;
do {
//sched_yield();
node_n = atomic_load(&state->shared->node_n);
} while (node_n == last);
}
// check if we should stop
if (node_n >= cgraph->n_nodes) break;
/* COMPUTE */
struct ggml_tensor * node = cgraph->nodes[node_n];
const int n_tasks = n_tasks_arr[node_n];
struct ggml_compute_params params = {
/*.type =*/ GGML_TASK_COMPUTE,
/*.ith =*/ state->ith,
/*.nth =*/ n_tasks,
/*.wsize =*/ cplan->work_size,
/*.wdata =*/ cplan->work_data,
};
if (state->ith < n_tasks) {
ggml_compute_forward(¶ms, node);
}
}
return GGML_EXIT_SUCCESS;
}当所有算子都执行完毕(node_n数目大于等于graph中算子数目)后,线程退出.GGML取出output Tensor的结果,执行完毕.
以上就是Lc的Inference过程的大致思路,但是还有一些细节有待探究(GPU Offload的转移,KV的使用等).
Lc in 1 Weekend?
为什么LLama.cpp能在短短一个周末的时间写完?我认为除了作者对于MLSys很深的了解和基础外,Lc的很多设计实际上被大幅简化了,从而大幅度减少了设计和实现的工作量:
- Memory Pool为静态且不回收. Lc通过模型的类型直接从预设的几种常数中选取出内存池的大小,内存池为最简单的Offset模型,且没有动态扩容功能.内存池中的Object仅可创建,不可销毁.
- 模型结构为预先定义且不可变.这实际上免去了大部分推理框架所面临的麻烦: 模型结构预定义,则不需要考虑计算图的有效性,不需要编写很多算子,甚至反序列化权重的时候都能心中有数.不需要从第三方模型格式转换,也避免了很多容错和算子转换代码.
- Transformer结构模型自身的算子比较简单,多为线性计算,相比RNN等复杂算子,实现和优化起来的难度大大减小,也没有单算子多个输出的问题. 同样 NLP模型可以完全把NHWC和NCHW的问题抛到一边.
- Lc面向的是Cuda和Metal,两者都有较为完善的生态支持,不存在某些算子无法被计算的问题,但是这点在移动端的加速推理场景中经常遇见,不得不采用分割子图等方式规避,实际上带来了很高的工作量.