Liu Yao's attic
0%

ABQ-LLM源码分析（1) - llama_exapmle.cc

发表于 更新于 分类于

前言

本系列将分析ByteDance项目ABQ-LLM的源码，并将重点放在llama模型端到端性能测试上，ABQ-LLM的模型端到端性能测试是基于NVIDIA的FasterTransformer项目搭建的。
ABQ-LLM的github仓库：https://github.com/bytedance/ABQ-LLM
ABQ-LLM论文链接：https://arxiv.org/abs/2408.08554
FasterTransformer的github仓库：https://github.com/NVIDIA/FasterTransformer

在构建FasterTransformer时，由于需克隆第三方的cutlass和googletestgithub仓库用于性能对比，而github的访问不太稳定，导致构建缓慢或失败，可选的解决方案包括挂代理和使用github镜像，这里推荐使用github镜像，自己用下来效果比较好的是githubfast.com。在开始构建前在CMakeLists.txt中把相关git仓库的链接修改即可。

GIT_REPOSITORY https://github.com/google/googletest.git
改为
GIT_REPOSITORY https://githubfast.com/google/googletest.git

GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
改为
GIT_REPOSITORY https://githubfast.com/NVIDIA/cutlass.git

ABQ-LLM为FasterTransformer添加了量化版llama和llama2模型的支持，模型的端到端性能测试文件在fastertransformer/example/cpp/llama和fastertransformer/example/cpp/llamaV2中，具体的main函数分别在llama_example.cc和llamaV2_example.cc中。

llama_example.cc - main()

直接看main函数。

mpi::initialize(&argc, &argv);

mpi是用于多GPU运算的库，如果不开启多GPU的支持，可以在FasterTransformer项目的build.sh中调整-DBUILD_MULTI_GPU=OFF，此时mpi::initialize是一个空函数，不会执行任何操作。

接下来的代码读取.ini文件中的配置项，并根据其中data_type的值决定用float还是用half实例化并执行llama_example函数。这里感觉有点奇怪，明明执行的是低精度INT量化的模型，为何还会使用浮点类型呢？留待后续解决。

llama_example.cc - llama_example()

这是一个模板函数，带有一个可变类型T，回忆一下，这个类型是由配置文件里的data_type决定的。

超参数

进到函数里首先是读了一大堆参数，分别看下这些参数是干嘛的。
tensor_para_size：张量并行度
pipeline_para_size：流水线并行度
sparse：是否采用稀疏化
release_dense比较抽象，追踪进去发现代码段

1
2
3
4
if (release_old && dense_weights[i] != nullptr){
    T* release_point = const_cast<T*>(dense_weights[i]);
    deviceFree(release_point);
}

这就能理解了，当release_dense为true时，压缩后的稀疏权重会替代原始的密集权重，并且原始的密集矩阵会被释放，从而减少内存占用；release_dense为false时，原始的密集权重不会被释放，仍然保留在内存中。
在实际实现中，llama_config.ini中未指定这两项，即这两项会取默认值，不开启稀疏。
int8_mode指定了量化的形式，取值及其代表的意义可查ABQ-LLM的README.md。

模型结构参数

head_num：表示注意力头的数量。在多头自注意力机制（Multi-head Attention）中，head_num 控制了注意力头的数量，每个注意力头会学习不同的表示。
size_per_head：每个注意力头的维度。它决定了每个注意力头的输出向量的大小。和head_num一起决定了总的隐藏层大小。
vocab_size：模型的词汇表大小。
decoder_layers：解码器层的数量。
rotary_embedding_dim：旋转嵌入的维度大小。
layernorm_eps：Layer Normalization的epsilon值。epsilon是一个非常小的常数，用来避免在进行除法时出现除零错误。
start_id：模型输入的开始标记ID。
end_id：模型输入的结束标记ID。
hidden_units：模型的隐藏单元数。它由head_num * size_per_head计算得来，表示每个层的维度大小。
inter_size：前馈神经网络的中间层大小。

生成策略参数

这些全是控制模型生成结果的特性的一些参数，之前没有接触过，用GPT快速学习一下，基本都能有个大概的理解。
beam_width：束搜索的宽度。在生成任务中，束搜索（Beam Search）是一种用于生成多个候选序列并选择最佳的策略，beam_width 决定了每次生成时考虑的候选序列数量。较大的 beam_width 可以提高生成质量，但也会增加计算开销。
top_k：Top-K 采样的 K 值。在生成文本时，Top-K 采样会从概率分布中选择前 K 个最可能的单词进行采样，其他的单词则被丢弃。较大的 K 会导致更多的多样性，但也可能降低生成质量。
top_p：Top-P 采样的概率阈值。Top-P 采样（也叫 nucleus 采样）会选择累计概率大于 p 的单词集合进行采样，控制生成结果的多样性。top_p 控制了保留多少最可能的单词，较小的 top_p 会限制输出的选择范围。
temperature：温度系数。在采样过程中，温度用于控制概率分布的平滑度。较低的温度会使得输出更加确定性（模型更倾向于选择最高概率的词），较高的温度则增加了生成的多样性（模型选择更多低概率的词）。
repetition_penalty：重复惩罚系数。用于惩罚模型生成重复的单词或短语。repetition_penalty 设置了惩罚的力度，较高的值能减少重复输出的可能性。
presence_penalty：出现惩罚系数。用于惩罚模型生成已经在上下文中出现过的单词。这有助于避免重复使用同一个单词或短语。
len_penalty：长度惩罚系数。用于控制生成文本的长度，防止生成过短或过长的输出。较高的len_penalty值可以使模型生成较长的序列。
beam_search_diversity_rate：束搜索多样性率。这个参数用于控制束搜索中每个候选序列之间的多样性。如果设置较高的值，会鼓励生成更加多样化的输出。
min_length：生成文本的最小长度。这个参数保证了生成的文本至少包含min_length个标记。常用于限制模型生成太短的文本。
request_batch_size：每次请求的批处理大小。这里的值为 1，即每次只处理一个请求，通常在生成任务中用于设置模型每次处理的输入批次大小。
total_output_len：请求的总输出长度。表示希望模型生成的最终输出序列的总长度。
request_output_len：请求的输出长度。指定模型在一次请求中生成的文本长度，可能用于控制生成的文本长度或者用于分页生成。

参数合法性的检查

head_num % tensor_para_size == 0：注意力头必须能被均匀地分到多个张量并行的设备上。
decoder_layers % pipeline_para_size == 0：不同的解码器必须能被均匀地分到多个流水线并行的设备上。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
int rank       = mpi::getCommWorldRank();
int world_size = mpi::getCommWorldSize();
if (rank == 0) {
    printf("Total ranks: %d.\n", world_size);
}
int device, device_count;
check_cuda_error(cudaGetDeviceCount(&device_count));
check_cuda_error(cudaSetDevice(rank % device_count));
check_cuda_error(cudaGetDevice(&device));

struct cudaDeviceProp prop;
check_cuda_error(cudaGetDeviceProperties(&prop, device));
printf("Device %s\n", prop.name);

printf("P%d is running with GPU #%d.\n", rank, device);

同样是多GPU相关的操作。值得注意的是，项目中为了保持外层代码干净，并没有将是否开启多GPU的编译指令写在外层代码中，而是定义了一个叫mpi的命名空间，无论是否开启了多GPU的支持，mpi中的函数都会执行，区别在于在未开启多GPU的情况下mpi中的函数几乎什么都不做，只返回一些默认值（如rank = 0, world_size = 1），而开启的情况下则执行多GPU的相关逻辑。

cudaGetDeviceCount(&device_count)：获取系统中可用的 GPU 设备数量。
cudaSetDevice(rank % device_count)：设置当前设备为根据rank计算得到的GPU设备。设备内存的分配和核函数的启动都会在当前设置的设备上；流与事件的创建也是发生在当前所设置的设备上。如果没有显式的调用设置设备的函数，则默认的当前设备是0。
cudaGetDevice(&device)：获取当前设置的设备ID。

1
2
3
4
5
6
7
8
if (tensor_para_size * pipeline_para_size != world_size) {
    if (world_size % pipeline_para_size) {
        printf("[ERROR] tensor_para_size * pipeline_para_size should equal to world_size \n");
        exit(-1);
    }
    tensor_para_size = world_size / pipeline_para_size;
    printf("[INFO] Setting tensor_para_size to %d \n", tensor_para_size);
}

接下来检查tensor_para_size * pipeline_para_size == world_size是否满足。如果不满足，那么检查pipeline_para_size能否被world_size整除，如果可以则据此计算pipeline_para_size，否则报错退出。

1
2
3
4
5
6
7
8
const int layers_per_group = decoder_layers / pipeline_para_size;
if (layers_per_group * pipeline_para_size != (int)decoder_layers) {
    printf("[ERROR] layers_per_group (%d) * pipeline_para_size (%d) should equal to decoder_layers (%ld) \n",
           layers_per_group,
           pipeline_para_size,
           decoder_layers);
    exit(-1);
}

这里是检查解码器能否均分到几个流水线上，但是其实之前就检查过了。

1
2
3
4
5
6
// assume gpu_num = k * n,
// tensor parallelism group size is n
// pipeline parallelism group size is k
NcclParam tensor_para;
NcclParam pipeline_para;
ftNcclInitialize(tensor_para, pipeline_para, tensor_para_size, pipeline_para_size);

ftNcclInitialize()中对nccl进行了一些初始化，由于不打算进行多GPU的实验，此处暂不做研究。

读取停止词和坏词

接下来代码读取坏词和停止词的列表，用到了函数read_word_list()：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
int read_word_list(const std::string& filename, std::vector<int>& file_data)
{
    std::ifstream word_list_file(filename, std::ios::in);

    std::string line_buf;
    int         line_count   = 0;
    size_t      id_counts[2] = {0, 0};
    while (std::getline(word_list_file, line_buf)) {

        std::stringstream line_stream(line_buf);
        std::string       vals;
        while (std::getline(line_stream, vals, ',')) {
            file_data.push_back(std::stoi(vals));
            id_counts[line_count]++;
        }
        line_count++;

        if (line_count > 1) {
            break;
        }
    }
    assert(id_counts[0] == id_counts[1]);

    return 0;
}

代码通过std::getline逐行读取文件，每一行存储在line_buf中，然后将该行的内容通过std::stringstream按逗号分隔成多个值。对于每个分隔出的值，使用std::stoi将其转换为整数，并存储到file_data向量中，然后通过id_counts[line_count]++记录当前行的列数，line_count增加，直到读取第二行数据后跳出循环。也就是说，代码只管csv文件的前两行，并且要求两行的列数是相等的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
std::vector<int> bad_words;
read_word_list("../examples/cpp/llamaV2/bad_words.csv", bad_words);

int* d_bad_words = nullptr;
deviceMalloc(&d_bad_words, bad_words.size(), false);
cudaH2Dcpy(d_bad_words, bad_words.data(), bad_words.size());

// Handle stop_words dictionary
std::vector<int> stop_words;
read_word_list("../examples/cpp/llamaV2/stop_words.csv", stop_words);

const size_t stop_words_len = stop_words.size() / 2;
// Tile with same dict for each element
std::vector<int> tiled_stop_words;
for (int i = 0; i < request_batch_size; i++) {
    tiled_stop_words.insert(tiled_stop_words.end(), stop_words.begin(), stop_words.end());
}

int* d_stop_words = nullptr;
deviceMalloc(&d_stop_words, tiled_stop_words.size(), false);
cudaH2Dcpy(d_stop_words, tiled_stop_words.data(), tiled_stop_words.size());

这样，停止词的数量就为stop_words_len = stop_words.size() / 2。停止词会被拷贝request_batch_size次，至于为什么，等待后续研究吧。

读取开始词

根据大模型的工作机制，开始词就是我们给到模型的输入。下面是读取开始词的代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
// Read ids of request from file.
size_t           max_input_len = -1;
std::vector<int> v_start_lengths;
std::vector<int> v_start_ids;
read_start_ids(request_batch_size,
                &v_start_lengths,
                &v_start_ids,
                max_input_len,
                end_id,
                1,
                "../examples/cpp/llamaV2/start_ids.csv");

int* d_input_ids;
int* d_input_lengths;
if (max_input_len == 0) {
    // unconditional case, no input ids, so do nothing.
    d_input_ids     = nullptr;
    d_input_lengths = nullptr;
}
else {
    // conditional case.
    deviceMalloc(&d_input_ids, request_batch_size * max_input_len, false);
    deviceMalloc(&d_input_lengths, request_batch_size, false);
    cudaH2Dcpy(d_input_ids, v_start_ids.data(), request_batch_size * max_input_len);
    cudaH2Dcpy(d_input_lengths, v_start_lengths.data(), request_batch_size);
}
std::vector<int> start_ids(request_batch_size, start_id);
std::vector<int> end_ids(request_batch_size, end_id);

其中主要的函数是read_start_ids()，又是一个很长的函数。

1
2
3
4
5
6
7
int read_start_ids(size_t            batch_size,
                   std::vector<int>* v_start_lengths,
                   std::vector<int>* v_start_ids,
                   size_t&           max_input_len,
                   const int         end_id,
                   const int         beam_width,
                   std::string       file_name);

下面逐步分析。

1
2
std::vector<std::vector<int>> tmp_start_ids;
std::vector<int>              tmp_start_lengths;

tmp_start_ids用于存储读到的所有start_id，因此是一个二维向量；tmp_start_lengths用于存储每行的长度。

开始读取后，代码按行读取文件，把单词存入tmp_start_ids的同时计数，把每行的长度存入tmp_start_lengths。

如果batch_size == 0，那么batch_size就设为start_id的行数；

1
2
3
4
max_input_len = tmp_start_lengths.data()[0];
for (uint i = 1; i < (uint)tmp_start_lengths.size(); i++) {
    max_input_len = max_input_len > tmp_start_lengths.data()[i] ? max_input_len : tmp_start_lengths.data()[i];
}

遍历所有行，将最长的一行的长度存在变量max_input_len中。

1
2
3
4
5
6
7
8
while ((int)tmp_start_lengths.size() < batch_size) {
    std::vector<int> padding_ids;
    for (int i = 0; i < max_input_len; i++) {
        padding_ids.push_back(end_id);
    }
    tmp_start_ids.push_back(padding_ids);
    tmp_start_lengths.push_back(max_input_len);
}

如果当前的输入行数比batch_size小，那么进行填充。填充的方法为，填充到行数等于batch_size，每一行都是max_input_len个结束词。

1
2
3
4
5
6
// Add padding
for (int i = 0; i < (int)tmp_start_ids.size(); i++) {
    for (int j = (int)tmp_start_ids[i].size(); j < max_input_len; j++) {
        tmp_start_ids[i].push_back(end_id);
    }
}

接下来还要对每行使用结束词填充，保证每行的长度一致。

1
2
3
4
5
6
7
8
for (int i = 0; i < (int)tmp_start_ids.size(); i++) {
    for (int b = 0; b < beam_width; b++) {
        for (int j = 0; j < (int)tmp_start_ids[i].size(); j++) {
            v_start_ids->push_back(tmp_start_ids[i][j]);
        }
        v_start_lengths->push_back(tmp_start_lengths[i]);
    }
}

最后，如果开启了束搜索策略，还要对每行数据进行beam_width扩展，复制每行的ID和长度beam_width次，并将其存入v_start_ids和v_start_lengths中。
这样，最终喂给模型的输入就生成了，接下来就是用cudaH2Dcpy()把输入拷贝到设备内存中。
这里start_ids和end_ids是干嘛的，暂未搞懂[todo]。

prompt_learning

代码指定默认prompt_learning_type为0，即不开启prompt learning。

多task

1
2
3
4
5
6
7
const int num_tasks = reader.GetInteger(model_name, "num_tasks", 0);
for (int task_name_id = 0; task_name_id < num_tasks; task_name_id++) {
    std::string config_task_name = model_name + "_task_" + std::to_string(task_name_id);
    std::string task_name        = reader.Get(config_task_name, "task_name");
    const int   prompt_length    = reader.GetInteger(config_task_name, "prompt_length", 0);
    prefix_prompt_table_pair.insert({task_name, {task_name_id, prompt_length}});
}

代码从配置文件中读取task的数目，并构造每个task的名称，读取每个task具体的prompt_length，并{task_name, (task_name_id, prompt_length)}的KV对存放到prefix_prompt_table_pair中。
然后代码定义了std::vector<int> prefix_prompt_task_ids(request_batch_size, 0);，用于存储每个request对应的task_id，各个request会被平均地分给几个task。
最后计算总输出长度，等于输入长度和请求输出长度的和。

stream与handle的创建

接下来创建cudaStream_t stream cublasHandle_t cublas_handle和cublasLtHandle_t cublaslt_handle，并利用一系列初始化函数进行初始化，并把cublas_handle绑定在新创建的cuda流stream上。

cublas在运行时需要加载GEMM的配置，注意cublasAlgoMap的构造函数只进行对应变量的赋值，不会进行实际的文件读取等操作。

Allocator内存管理器

接下来定义了一个fastertransformer里的自有类Allocator，并用AllocatorType::CUDA进行实例化。内联函数getDevice()通过调用cuda的cudaGetDevice获取当前进程的设备id。

成员变量

1
2
3
4
5
6
class Allocator<AllocatorType::CUDA>: public IAllocator {
    const int                          device_id_;
    cudaStream_t                       stream_ = 0;  // initialize as default stream
    std::unordered_map<void*, size_t>* pointer_mapping_;
    ...
};

上面是Allocator类的成员变量。

构造函数

Allocator类的构造函数接收一个参数device_id，并用之初始化成员变量device_id_。
然后构造一个无序map：

1
pointer_mapping_ = new std::unordered_map<void*, size_t>();`

并获取设备数量：

1
2
int device_count = 1;
check_cuda_error(cudaGetDeviceCount(&device_count));

然后创建内存池。cuda内存池允许在设备上分配和释放内存时复用内存块，从而减少分配和释放内存的开销。cuda设备可以有多个内存池，但每个设备都有一个“默认内存池”，当你没有指定具体的内存池时，默认会使用这个内存池进行内存分配。cudaDeviceGetDefaultMemPool函数能够获取默认的内存池。

1
2
cudaMemPool_t mempool;
check_cuda_error(cudaDeviceGetDefaultMemPool(&mempool, device_id));

接下来的部分又比较抽象了，因为原本的代码是支持多GPU的，虽然ABQ-LLM不包含多GPU的支持，但是多GPU的代码仍然遗留了下来。下面的部分主要是确保不同设备能够相互访问内存池中的内存。并且还设置了内存池的释放阈值，避免内存池在释放时出现不必要的内存回收。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
cudaMemAccessDesc desc = {}; // 创建内存访问描述符对象，初始化为空
int peer_access_available = 0; // 标志变量，表示设备之间的对等访问权限

for (int i = 0; i < device_count; i++) {
    if (i == device_id) {
        continue; // 如果当前是本地设备，则跳过
    }
    
    // 检查当前设备（i）是否可以访问 device_id 所代表的设备的内存
    check_cuda_error(cudaDeviceCanAccessPeer(&peer_access_available, device_id, i));
    
    if (!peer_access_available) {
        // 如果不能访问，则输出警告信息并跳过
        FT_LOG_WARNING("Device " + std::to_string(device_id) + " peer access Device " + std::to_string(i)
                       + " is not available.");
        continue;
    }
    
    // 设置内存访问描述符，指定允许的设备访问权限
    desc.location.type = cudaMemLocationTypeDevice;  // 设置为设备类型
    desc.location.id = i;  // 设备的 ID
    desc.flags = cudaMemAccessFlagsProtReadWrite; // 设置访问权限为读写
    
    // 将内存访问权限设置为描述符指定的内容
    check_cuda_error(cudaMemPoolSetAccess(mempool, &desc, 1));
}

// 设置内存池的释放阈值，避免内存池在不需要时被过早释放
uint64_t setVal = UINT64_MAX;
check_cuda_error(cudaMemPoolSetAttribute(mempool, cudaMemPoolAttrReleaseThreshold, &setVal));

析构函数

析构函数比较简单，就是把pointer_mapping_里面的东西和它本身free掉。

那么这玩意儿分配和释放内存的逻辑是怎样的呢？

malloc()

首先如果请求的内存大小为零直接返回空指针。如果不是，则初始化一个指针ptr，并定义变量o_device用于存储当前设备的ID。
getSetDevice(device_id_, &o_device)的行为则是把当前设备的ID保存到o_device中，并且切换到ID为device_id的设备。这意味着，调用用哪个设备的ID初始化的Allocate就会用切换到哪个设备分配内存，然后再切换回原来的设备。

注意这里并不是size多大就分配多少个字节的内存，而是对size向上取到一个32的倍数，这样能够实现内存对齐提升访存效率。

在is_host == true的情况下会用cudaHostAlloc()在主机上分配内存。用这个函数可以得到“主机锁页内存”。所谓“主机锁页内存”，就是被操作系统标记为不会交换到虚拟内存的固定区域，且GPU可以通过DMA（直接内存访问）直接从锁页内存中读取或写入数据，避免使用cudaMemcpy显式地拷贝数据。

在is_host == false的情况下会在设备上分配内存。如果没有开启内存池技术，那么使用cudaMalloc进行同步的内存分配，否则在自己的流中进行异步分配（要记得Allocator类有一个_stream成员变量，默认初始化为CUDA的默认流）。

另外还提供了is_zero的设置，如果为真则会把新分配的内存置零。

最后把分配取得的内存的地址与大小记录在类的成员变量哈希表中。

free()

与malloc()相同，free()也会切换到分配内存时的设备进行释放，释放完以后再恢复原设备上下文。
在释放时同样是根据是否启用CUDA内存池决定是同步释放还是异步释放，根据是否是主机内存决定调用哪个函数。
在释放时还会检查传入的地址是否在哈希表中有记录，如果没有的话会产生警告日志。

cublas准备

接下来代码进行和cublas相关的准备。

1
2
3
std::mutex*     cublas_wrapper_mutex = new std::mutex();
cublasMMWrapper cublas_wrapper =
    cublasMMWrapper(cublas_handle, cublaslt_handle, stream, cublas_algo_map, cublas_wrapper_mutex, &allocator);

各参数具有如下意义：
cublas_handle：提供 cuBLAS 库的上下文句柄，管理基础矩阵运算（如 FP32/FP16 GEMM），确保计算与 CUDA 流同步。
cublaslt_handle：提供 cuBLASLt 的上下文句柄，支持低精度计算（如 INT8）和灵活算法选择，优化 Tensor Core 利用率。
stream：绑定 CUDA 流，确保 cuBLAS 操作与自定义内核（如注意力计算）顺序执行，避免隐式同步。
cublas_algo_map：存储预配置的 cuBLASLt 算法映射表，针对不同矩阵尺寸（如 M=1 或 M=4096）自动选择最优计算路径。
cublas_wrapper_mutex：提供互斥锁，保证多线程/多 GPU 环境下对 cuBLAS 句柄的线程安全访问，防止状态冲突。
allocator：绑定内存分配器，支持对齐分配和内存池复用，减少显存碎片化，优化大规模模型的内存管理。

然后根据不同的数据类型T为cublas_wrapper设置不同的计算精度配置。

构造gpt_weights

接下来代码对gpt_weights进行构造，并调用loadModel方法加载参数，可以看到这边参数就是普通的二进制序列化的，并不稀奇。