优化边缘设备上的大型语言模型（LLM）--tinychat

文章目录

一、项目启动
- 1.背景：针对不同操作系统架构的4bit权重重排
- 2.初始环境配置
- - 下载LLaMA2-7B-chat模型
- 3.项目启动
- - 项目结构说明
  - 评估不同优化技术
  - 可能遇到的bug以及措施
  - - 1.macOS上部署
二、各种优化技术实现
- 1.前置条件
- 2.优化----循环展开
- 3.优化----多线程
- 4.优化----单指令多数据
三、各个优化技术实验结果
- 1.评估基准
- - Profiler的使用
  - 指标的计算
- 2.效果对比
- - 评估实验数据
  - 使用同一个问题“What is Newton's First Law?”测试效果
四、思考体会

一、项目启动

git clone --recursive https://github.com/mit-han-lab/tinychat-tutorial

1.背景：针对不同操作系统架构的4bit权重重排

为了减少与权重重排相关的运行时开销，TinyChatEngine在模型转换过程中离线执行此过程。本项目为ARM和x86 CPU量身定制的QM_ARM和QM_x86权重布局，分别支持128位SIMD和256位SIMD操作。

QM_ARM布局：
对于QM_ARM，考虑一个128位的权重向量[w0, w1, … , w30, w31]，其中每个wi都是一个4位量化权重。TinyChatEngine通过交错下半部分和上半部分的权重，将这些权重重新排列为[w0, w16, w1, w17, …, w15, w31]。这种新排列方式便于使用128位AND和移位操作同时解码下半部分和上半部分，如后续图所示。这将消除运行时重排开销并提高性能。
QM_x86布局：
类似地，对于QM_x86，TinyChatEngine将一个256位的权重向量[w0, w1, … , w62, w63]重新排列为[w0, w32, w1, w33, …, w31, w63]。这种布局确保了与256位AVX2 SIMD指令的最佳兼容性，以实现高效的去量化。

2.初始环境配置

请按照以下说明在你的电脑上进行设置。

MacOS： 使用以下命令通过Homebrew安装boost和llvm。
```
brew install boost
brew install llvm
```
确保你的电脑上有Python。如果没有，建议使用Anaconda安装。
基于Linux的操作系统（如Ubuntu）： 确保你的电脑上有Python。如果没有，建议使用Anaconda安装。
Windows： 确保以下命令在你的电脑上可用：
- g++
- make
- unzip
- git
- Python
如果没有，你可以从官方网站下载Python，并使用MSYS2安装其他依赖项。请按照本教程进行安装，然后使用以下命令安装所需的依赖项：
```
pacman -S --needed base-devel mingw-w64-x86_64-toolchain make unzip git
```

下载LLaMA2-7B-chat模型

从我们的模型库下载LLaMA2-7B-chat模型。模型将被存储在INT4/models/LLaMA_7B_2_chat/目录中。这个下载过程需要一段时间……
- 对于x86设备（Intel/AMD）：
```
cd <path-to-repo>/transformer
python download_model.py --model LLaMA_7B_2_chat --QM QM_x86
```
- 对于ARM设备（M1/M2）：
```
cd <path-to-repo>/transformer
python download_model.py --model LLaMA_7B_2_chat --QM QM_ARM
```
- 对于离线下载的用户：
  离线下载模型，在transformer目录下解压。

3.项目启动

TinyChatEngine是一个专为在边缘设备上高效部署量化的大型语言模型（LLMs）而设计的强大神经网络库。在kernel/starter_code中它采用以下优化技术加速推理并最小化内存占用：

循环展开： 一种循环转换技术，尝试通过牺牲二进制文件大小来优化程序的执行速度，这是一种空间与时间的权衡方法。（参考讲座TinyEngine：循环优化）
多线程： 多线程是一种允许单个进程同时管理多个执行线程的编程概念。每个线程共享进程的资源，但可以独立执行，从而提高性能、改善用户体验和更好地利用资源。（参考讲座TinyEngine：多线程）
SIMD（单指令多数据）指令： 一种通过单条指令处理多个数据的计算方法。（参考讲座TinyEngine：SIMD编程）

项目结构说明

文件结构： 所有起始代码都位于kernels/starter_code目录中。以下是详细结构：

kernels/
│     ...
└───starter_code/
    ├── reference.cc
    ├── loop_unrolling.cc
    ├── multithreading.cc
    ├── simd_programming.cc
    ├── multithreading_loop_unrolling.cc
    └── all_techniques.cc
transformer/
└─── evaluate.sh

参考实现： 我们提供了一个使用标准for循环的基本参考实现（reference.cc）。这将作为性能比较的基准。
优化技术： 你将填充代码并完成其他实现，包括：
- 循环展开（loop_unrolling.cc）
- 多线程（multithreading.cc）
- SIMD编程（simd_programming.cc）
- 循环展开与多线程结合（multithreading_loop_unrolling.cc）
- 所有技术的组合（all_techniques.cc）

评估不同优化技术

评估脚本： 提供的bash脚本evaluate.sh将用于测试实现的正确性并测量学生实现的加速效果。脚本可以测试单个优化技术或所有技术。
使用evaluate.sh脚本测试你的实现。脚本将编译并运行你的代码，然后提供关于正确性和性能（GOPs）的反馈。
测试所有实现： ./evaluate.sh
你可以指定要测试的优化技术或测试所有技术。支持的参数： reference（基准）、 loop_unrolling、 multithreading、simd_programming、multithreading_loop_unrolling、 all_techniques
例如：要测试特定的实现，例如循环展开技术，可以使用以下命令：

./evaluate.sh reference

这还会生成一个名为chat的可执行文件，你可以运行该文件以使用特定实现部署本地聊天机器人。要使用本地聊天机器人，可以运行：

./chat

在这里插入图片描述

可能遇到的bug以及措施

1.macOS上部署

1️⃣opt_smooth_expoter.py文件Cannot find reference ‘opt’ in ‘init.py’ 在这里插入图片描述

不要使用pip安装，仍会报错。
去https://github.com/mit-han-lab/smoothquant将子文件夹smoothquant复制到transformer下面

2️⃣make 时候./include/OPTTokenizer.h:22:10: fatal error: ‘nlohmann/json.hpp’ file not found 22 | #include <nlohmann/json.hpp>

1.brew install nlohmann-json
2.在transformer/makefile文件第28行添加INCLUDE_DIRS += -I/opt/homebrew/include

3️⃣./evaluate.sh loop_unrolling ./evaluate.sh: line 36: bear: command not found Compilation failed!

1.brew install bear
2.将evaluate.sh脚本中 
 bear make chat test_linear -j IMP="$arg"替换为
 bear -- make chat test_linear -j IMP="$arg"

二、各种优化技术实现

1.前置条件

输入矩阵A，B；相乘输出矩阵C;
A 矩阵的维度是 m x k，B 矩阵的维度是 k x n，因此 C 矩阵的维度为 m x n。
reference.cc使用三个嵌套的循环遍历矩阵：外层的 row 和 col 循环遍历矩阵 C 的每一个元素，内层的 ch 循环则用于计算。

 for (int row = 0; row < m; row++) {
        for (int col = 0; col < n; col++) {
            float acc = 0;
            // Compute each block
            for (int ch = 0; ch < k;) {

2.优化----循环展开

循环展开带来的运算数变化
假设没有使用循环展开技术，那么原始代码在处理每一列时会逐列进行（例如 col += 1）。使用循环展开后，一次处理 4 列，意味着：

  for (int row = 0; row < m; row++) {
        for (int col = 0; col < n; col += 4) {
            float acc0 = 0;
            float acc1 = 0;
            float acc2 = 0;
            float acc3 = 0;
            // Compute each block
            for (int ch = 0; ch < k;) {}

每次循环展开后：处理的计算量从 1 列增加到了 4 列。由于每次处理了更多列，因此循环的总次数减少了约 4 倍。
循环展开带来了运算数的变化：每次循环处理的工作量增加（如处理 4 列而不是 1 列），减少了总的循环次数。
延迟减少：循环展开减少了循环控制的开销，增加了指令级并行性，改善了流水线效率，从而减少了延迟。

3.优化----多线程

核心思想是将矩阵乘法的任务按列划分给多个线程并行处理。
每个线程独立计算某一范围列内的所有结果，大大加快了矩阵乘法的计算速度，尤其在多核处理器上能显著提高性能。
列的划分使用了简单的均匀分配策略，且通过 pthread_create 创建线程，主线程通过 pthread_join 等待所有线程完成任务。

  for (int row = 0; row < m; row++) {
        for (int col = mat_args->start; col < mat_args->end; col++) {
            float acc = 0;
            // Compute each block

4.优化----单指令多数据

ARM 和 x86 的默认指令集（ARM A32/Thumb T32 和 x86/x86-64）都不是SIMD指令集。默认指令集用于处理单个数据操作，而SIMD指令集（如ARM的NEON和x86的SSE/AVX）则是这些架构的扩展指令集，用于在一条指令中并行处理多个数据元素。
SIMD的核心思想是并行处理多个数据，这在ARM和x86分支中都有体现，在不同的分支他使用不同的指令集合：

#ifdef QM_ARM
#include <arm_neon.h>
#endif
#ifdef QM_x86
#include <immintrin.h>
#endif

并行数据加载：通过指令一次性加载128位或256位数据（例如 _mm256_loadu_si256、vld1q_u8），同时处理多个数据元素。
并行运算： ARM中的 vdotq_s32，x86中的 _mm256_maddubs_epi16
都是典型的SIMD点积运算，这些指令可以在单条指令中对多个数据元素同时进行乘法和加法操作。
并行累加： ARM中的 vaddvq_f32 和 x86中的 _mm256_fmadd_ps 能够在一次指令中将多个累加结果同时计算出来

三、各个优化技术实验结果

1.评估基准

Profiler的使用

在该项目中，性能分析与剖析是通过 Profiler 类实现的。考虑到 Profiler 类的结构，尤其是在计算 FLOPs（浮点运算次数）时，我们需要关注代码中如何调用 Profiler 相关的宏或函数。

test_linear.cc文件中STATS_FLOPS 和 STATS_END：
```
STATS_FLOPS(int4_op.profile_name, flops);
int4_op.forward(hidden_states, outputQ_fast);
STATS_END(int4_op.profile_name);
```
这些宏的作用通常是用来标记某一段代码的开始和结束，并将性能数据记录到 Profiler 中。具体来说，它们执行以下操作：
- STATS_FLOPS(name, flops)：
  - 记录开始时间。
  - 将 flops 作为该操作的浮点操作数数量进行记录。
  - name 用于标识这段代码块（函数或操作）的名称。
- STATS_END(name)：
  - 记录结束时间。
  - 计算这段代码块的执行时间。
  - 将计算结果存储在 Profiler 中，与 name 相关联。

指标的计算

基于 Profiler 类的实现，以下是主要指标的计算方法：

总时间 (Total time):
- 通过 stop 方法计算并累积从 start 到 stop 的时间差，以微秒为单位记录。
平均时间 (Average time):
- 总时间除以这段代码块的调用次数。
调用次数 (Count):
- 每调用一次 STATS_FLOPS 和 STATS_END 对应的代码块，调用次数增加 1。
Giga FLOPs (GOPs):
- 利用传递给 STATS_FLOPS 的 flops 值，结合时间差，计算每秒执行的 FLOPs 数量。
- 如果总 FLOPs 为 F, 时间为 T 微秒，那么 GOPs 计算为 (F / T) / 1,000,000。

2.效果对比

评估实验数据

在这里插入图片描述

使用同一个问题“What is Newton’s First Law?”测试效果

其实大体分布类似上图，就是给个问题使得延迟更加具象化。

1.reference(基准)

Section, Total time(ms), Average time(ms), Count, GOPs
Inference latency, 996283.937500, 13647.724609, 73, N/A

在这里插入图片描述
2.loop_unrolling

Section, Total time(ms), Average time(ms), Count, GOPs
Inference latency, 583191.125000, 7988.918945, 73, N/A

在这里插入图片描述

3.multithreading

Section, Total time(ms), Average time(ms), Count, GOPs
Inference latency, 272030.093750, 3726.438965, 73, N/A

在这里插入图片描述

4.simd_programming

Section, Total time(ms), Average time(ms), Count, GOPs
Inference latency, 480822.906250, 6678.095215, 72, N/A

在这里插入图片描述

5.multithreading_loop_unrolling

Section, Total time(ms), Average time(ms), Count, GOPs
Inference latency, 197656.734375, 2707.625977, 73, N/A

在这里插入图片描述

6.all_techniques

Section, Total time(ms), Average time(ms), Count, GOPs
Inference latency, 107534.695312, 1493.536987, 72, N/A

在这里插入图片描述

四、思考体会

目前完成starter_code 的几个技术发现上述技术只是从矩阵乘法方面优化的：

还有一个着重要关注的点就是量化用到了哪些方法，能否进一步优化？
内核编译是不是相当于一个小小的操作系统，假如使用微控制器操作系统开源有源码的话，是不是可以从更底层优化？
以及没有用到pytorchAI编译框架的话，用的是啥，是否可以相互借鉴？
x86效果是咋样的，和arm一样吗？
还有其他针对硬件结构的代码没有细看，后续可以继续揣摩。