2023年的深度学习入门指南(23) - ChatGLM2
在《在你的电脑上运行大模型》这一节，我们曾经介绍过ChatGLM模型，它是当时最好的中文大模型之一。现在，它又更新到了第二代，即ChatGLM2。
当时，我们的技术储备还不足，只能让它运行起来，还不敢讲解它的原理和代码。
现在，经过LLaMA 2和百川的代码的狂轰滥炸，大家已经适应了看代码的节奏了。现在，是时候来看看ChatGLM2的原理和代码了。
运行ChatGLM2
首先我们还是将ChatGLM2的代码运行起来。在大于13GB显存的机器上，ChatGLM2都可以顺利运行起来。比如我是在一个15G的T4上运行的。
第一步还是将安装相关的库：
pip install protobuf
pip install transformers==4.30.2
pip install cpm_kernels
pip install torch>=2.0
pip install gradio
pip install mdtex2html
pip install sentencepiece
pip install accelerate
pip install sse-starlette
pip install streamlit>=1.24.0
第二步就可以用Transformers的标准接口来调用ChatGLM2了：
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
model = model.eval()
response, history = model.chat(tokenizer, "生成scala语言的快速排序", history=[])
print(response)
输出如下：
def quickSort(arr: Int[]): Int[] = {
val pivot = arr(arr.length / 2)
val left = 0
val right = arr.length - 1
while (left <= right) {
while (arr(left) < pivot) {
left = left + 1
}
arr(left) = pivot
while (arr(right) > pivot) {
right = right - 1
}
arr(right) = pivot
left = left + 1
right = right - 1
}
return arr
}
如果在更小显存的显卡上运行，我们可以使用4位量化后的结果：
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).cuda()
model = model.eval()
response, history = model.chat(tokenizer, "生成Kotlin语言编写的快速排序", history=[])
print(response)
这是我在3060上运行的结果：
fun quickSort(arr: IntArray): IntArray {
val left = 0
val right = arr.size - 1
val quicksortFactor = arr.size / 2
while (left < right) {
quicksortFactor--.let {
let x = left
let y = right
let temp = arr[x]
if (temp < arr[y]) {
x++
} else {
y--
}
if (x == y) {
break
}
quicksortFactor++.let {
arr[x] = arr[y]
arr[y] = temp
}
}
}
return arr
}
量化的CUDA代码解析
之前讲过不少多头注意力的代码实现了，后面也还会讲。在本节中我们讲一个之前没有讲到的内容，量化所用的CUDA代码。LLaMA 2部分没讲是因为它还没有量化部分，而百川是CUDA核代码暂时还没开源。所以我们就先借着GLM的代码来讲一下。
我们先看一下CUDA核部分的Makefile:
NVCC=nvcc
OPTIONS=-gencode arch=compute_61,code=sm_61 \
-gencode arch=compute_62,code=sm_62 \
-gencode arch=compute_70,code=sm_70 \
-gencode arch=compute_72,code=sm_72 \
-gencode arch=compute_75,code=sm_75 \
-gencode arch=compute_80,code=sm_80 \
-gencode arch=compute_86,code=sm_86
TARGETS=$(patsubst %.cu, %.fatbin, $(wildcard *.cu))
all: $(TARGETS)
%.fatbin: %.cu
$(NVCC) -fatbin $^ $(OPTIONS) -o $@
.PHONY : clean, copy
clean:
rm $(TARGETS)
copy:
cp $(TARGETS) ../kernels/
我们可以看到，这里的代码是支持多个CUDA架构的，包括了6.1、6.2、7.0、7.2、7.5、8.0、8.6。这里的架构是指GPU的架构，比如RTX 3090的架构是8.6，RTX 3060的架构是8.0。
- 6.1和6.2对应的是Pascal架构，比如P100, GTX 1060
- 7.0是Volta架构，比如V100
- 7.5是Turing架构，比如RTX 2080, T4
- 8.0和8.6是Ampere架构，比如A100, RTX 3090
别看已经支持这么多架构了，但是更早的Maxwell和Kepler等更老的架构已经随风而去了。
要支持这么多架构，就需要引入一个新的知识点 - fatbin.
.fatbin 文件是 CUDA 二进制格式（CUDA Fat Binary Format）的文件。这是 NVIDIA 的 CUDA 平台使用的一种特殊的二进制文件格式。fatbin文件包含了针对多种 GPU 架构和计算能力的代码，可以在多种不同类型的处理器上运行。
在 CUDA 编程中，GPU 代码（通常称为 kernel）经常以类似于内联汇编的方式存储在主机代码中。然而，这种方法在实际应用中存在一些困难，主要是因为不同的 GPU 架构和设备可能需要不同的 GPU 代码版本。CUDA Fat Binary 解决了这个问题，它包含了多个版本的 GPU 代码，每个版本都针对一个特定的 GPU 架构进行优化。
当 CUDA 程序运行时，CUDA 运行时系统会检查正在运行的设备，并从 fat binary 文件中选择最适合该设备的 GPU 代码版本。这样就可以用一个.fatbin文件,使同一个CUDA程序可以在不同算力的GPU上运行。不需要为不同GPU单独编译。
下面我们就来看8位量化的实现，这个代码完全就是一个如何写最简单CUDA代码的例子：
template<typename T>
__device__ void
int8WeightExtractionDevice(const int8_t* weight,
const T* scale_list,
T* output,
const int n,
const int k)
{
for(int i = blockIdx.x * k + threadIdx.x; i < blockIdx.x * k + k; i += blockDim.x){
output[i] = T(weight[i]) * scale_list[blockIdx.x];
}
}
在GPU线程中：
- 计算当前线程读取的weight索引:blockIdx.x代表块id,_k代表每个块处理k个值,threadIdx.x代表线程id
- 读取weight数组当前索引处的int8值
- 对其缩放:用scale_list中对应块id的缩放因子乘以weight值
- 结果输出到output数组对应索引处
最后，通过blockDim.x线程并行完成weight数组到output数组的整体复制和缩放计算。
如果大家忘了CUDA那一节的内容，我们来复习一下blockIdx，threadIdx和blockDim的概念：
- blockIdx: CUDA将线程组织成块(block),每个块有一个id,称为blockIdx。可以有多个块,通过blockIdx区分不同块
- threadIdx: 每个块里面有多个线程,通过threadIdx区分同一块中的不同线程。线程id从0开始计数
- blockDim: 指明每个块中含有的线程数目
在启动核函数时指定,例如调用核函数时执行<<<32, 128>>>表示有32个块,每个块中有128个线程。
在函数体内部，有一个for循环，请注意，这个循环不是像在单核CPU上那样串行的，而是在CUDA的每个线程上执行的！
循环变量i的初始值是 blockIdx.x * k + threadIdx.x，这是一个常用的模式，用于将数据的不同部分分配给不同的CUDA线程。每轮循环中，i增加 blockDim.x，这表示每个线程处理的数据间隔是一个block的大小。
在for循环中，函数将权重乘以对应的缩放因子，并将结果存储在输出数组中。这里，权重的类型被转换为T，然后乘以对应的scale_list元素。注意scale_list[blockIdx.x]的使用，这表示对于同一个block内的所有线程，它们使用的是同一个缩放因子。
CUDA的核被封装在host的函数里：
extern "C" __global__ void int8WeightExtractionHalf(const int8_t* weight,
const half* scale_list,
half* output,
const int n,
const int k){
int8WeightExtractionDevice<half>(weight, scale_list, output, n, k);
}
extern "C" __global__ void int8WeightExtractionFloat(const int8_t* weight,
const float* scale_list,
float* output,
const int n,
const int k){
int8WeightExtractionDevice<float>(weight, scale_list, output, n, k);
}
我们来看一下在Python中如何调用这个函数的：
def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
if source_bit_width == 8:
func = kernels.int8WeightExtractionHalf
elif source_bit_width == 4:
func = kernels.int4WeightExtractionHalf
else:
assert False, "Unsupported bit-width"
with torch.cuda.device(weight.device):
n, m = weight.size(0), weight.size(1)
out = torch.empty(n, m * (8 // source_bit_width), dtype=torch.half, device="cuda")
stream = torch.cuda.current_stream()
gridDim = (n, 1, 1)
blockDim = (min(round_up(m, 32), 1024), 1, 1)
func(
gridDim,
blockDim,
0,
stream,
[
ctypes.c_void_p(weight.data_ptr()),
ctypes.c_void_p(scale_list.data_ptr()),
ctypes.c_void_p(out.data_ptr()),
ctypes.c_int32(n),
ctypes.c_int32(m),
],
)
return out
我们可以看到，如果是针对8位量化的时候，就会调用 kernels.int8WeightExtractionHalf 函数，这个函数对应的就是我们上面写的那个函数。
下面我们讲解一下是如何划分并行度的。
gridDim这个变量表示 CUDA 内核函数的网格大小。在这段代码中，它被设置为 (n, 1, 1)，其中 n 是权重张量的第一维大小。这意味着网格中有 n 个块，每个块负责处理权重张量的一行。
blockDim是一个三元组，用于指定每个线程块的维度。在这个代码中，blockDim 被设置为 (min(round_up(m, 32), 1024), 1, 1)。这表示每个线程块中的线程数量为 min(round_up(m, 32), 1024)，这个数量是 m（weight 张量的第二维的大小）向上取到最近的32的倍数，但最大不超过1024。这是因为CUDA架构的限制，每个线程块的线程数量不能超过1024。
下面我们再讲一下4位压缩的：
__device__ void
int4WeightCompressionDevice(const int8_t* input,
int8_t* output,
const int n,
const int k)
{
for(int i = blockIdx.x * k + threadIdx.x; i < blockIdx.x * k + k; i += blockDim.x){
output[i] = (input[i * 2] << 4) | (input[i * 2 + 1] & 0b00001111);
}
}
int4WeightCompressionDevice里面，对于每个线程，它会计算出它应该处理的元素的索引 i。然后，它会将输入数组中索引为 i * 2 和 i * 2 + 1 的两个元素压缩成一个元素。压缩方法是将第一个元素左移 4 位，然后与第二个元素进行按位或运算。最后，将结果存储在输出数组中索引为 i 的位置。
虽然可以高度并行，但是其实GPU上的代码写起来跟CPU上也并没有太大的不同，不需要学新的语句。
同理，我们看下将4位压缩的权重转换为8位的：
template<typename T>
__device__ void
int4WeightExtractionDevice(const int8_t* weight,
const T* scale_list,
T* output,
const int n,
const int k)
{
for(int i = blockIdx.x * k + threadIdx.x; i < blockIdx.x * k + k; i += blockDim.x){
int8_t original = weight[i];
int8_t high = original >> 4;
int8_t low = original << 4; low = low >> 4;
output[i * 2] = T(high) * scale_list[blockIdx.x];
output[i * 2 + 1] = T(low) * scale_list[blockIdx.x];
}
}
有了上面的知识，这里不需要额外讲解了吧？
量化层的实现
最后，我们来看看量化如何在神经网络中使用：
解释下面的代码：
import torch
from kernels import extract_weight_to_half
class W8A16Linear(torch.autograd.Function):
@staticmethod
def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width):
ctx.inp_shape = inp.size()
ctx.weight_shape = quant_w.size()
ctx.weight_bit_width = weight_bit_width
out_features = quant_w.size(0)
inp = inp.contiguous().view(-1, inp.size(-1))
weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width)
output = inp.mm(weight.t())
ctx.save_for_backward(inp, quant_w, scale_w)
return output.view(*(ctx.inp_shape[:-1] + (out_features,)))
@staticmethod
def backward(ctx, grad_output: torch.Tensor):
inp, quant_w, scale_w = ctx.saved_tensors
weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width)
grad_output = grad_output.contiguous().view(-1, weight.size(0))
grad_input = grad_output.mm(weight)
grad_weight = grad_output.t().mm(inp)
return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None
forward 方法接受四个参数：inp 是一个输入张量，quant_w 是一个量化后的权重张量，scale_w 是一个权重缩放张量，weight_bit_width 是权重的位宽。这个方法首先保存输入张量和权重张量的形状，然后将输入张量转换为连续的并调整形状。接着，它使用我们刚才讲过的 extract_weight_to_half 函数从量化后的权重和权重缩放中提取出半精度的权重。然后，它使用矩阵乘法计算输出，并将结果调整为正确的形状。最后，它将输入、量化后的权重和权重缩放保存起来，以便在反向传播时使用。
backward 方法接受一个参数：grad_output 是一个梯度输出张量。这个方法首先从上下文中获取保存的输入、量化后的权重和权重缩放，并从中提取出半精度的权重。然后，它将梯度输出转换为连续的并调整形状。接着，它使用矩阵乘法计算输入梯度和权重梯度。最后，它返回调整形状后的输入梯度和权重梯度。
这段代码实现了一个自定义的线性层，它使用半精度的权重进行计算，并支持 PyTorch 的自动求导机制。
这还没有完，为了实现更大规模并行化，量化层还可以进一步的封装：
import torch
from torch.nn.parameter import Parameter
from SwissArmyTransformer.mpu import copy_to_model_parallel_region
from SwissArmyTransformer.mpu import gather_from_model_parallel_region
from SwissArmyTransformer.mpu import reduce_from_model_parallel_region
from SwissArmyTransformer.mpu import scatter_to_model_parallel_region
from SwissArmyTransformer.mpu import ColumnParallelLinear, RowParallelLinear
from .functional import W8A16Linear
from kernels import compress_int4_weight
class QuantizedColumnParallelLinear(ColumnParallelLinear):
def __init__(self, weight_bit_width: int, weight=None, *args, **kwargs):
super(QuantizedColumnParallelLinear, self).__init__(*args, **kwargs)
self.weight_bit_width = weight_bit_width
shape = self.weight.shape
del self.weight
if weight is None:
self.weight = torch.empty(
shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=kwargs["device"]
)
self.weight_scale = torch.empty(shape[0], dtype=kwargs["params_dtype"], device=kwargs["device"])
else:
self.weight_scale = (weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).half()
self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8)
if weight_bit_width == 4:
self.weight = compress_int4_weight(self.weight)
self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False)
self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False)
def forward(self, input_):
# Set up backprop all-reduce.
input_parallel = copy_to_model_parallel_region(input_)
# Matrix multiply.
output_parallel = W8A16Linear.apply(input_parallel, self.weight, self.weight_scale, self.weight_bit_width)
if self.bias is not None:
output_parallel = output_parallel + self.bias
if self.gather_output:
# All-gather across the partitions.
output = gather_from_model_parallel_region(output_parallel)
else:
output = output_parallel
return output
这段代码定义了一个名为 QuantizedColumnParallelLinear 的类，它继承自 ColumnParallelLinear 类。这个类实现了一个量化的列并行线性层。
__init__ 方法接受若干个参数，其中 weight_bit_width 是权重的位宽，weight 是一个可选的权重张量。这个方法首先调用父类的构造函数，然后保存权重的位宽。接着，它获取权重的形状并删除权重属性。如果没有提供权重，则创建一个空的权重张量和一个空的权重缩放张量。否则，根据提供的权重计算权重缩放，并将权重量化为整数。如果位宽为 4，则使用 compress_int4_weight 函数对权重进行压缩。最后，将权重和权重缩放转换为 PyTorch 参数并保存。
forward 方法接受一个参数：input_ 是一个输入张量。这个方法首先使用 copy_to_model_parallel_region 函数将输入复制到模型并行区域。然后，使用 W8A16Linear.apply 函数计算输出。如果有偏置，则将偏置加到输出上。如果需要收集输出，则使用 gather_from_model_parallel_region 函数收集输出。否则，直接返回输出。
这段代码实现了一个量化的列并行线性层，它可以在多个 GPU 上并行计算。
小结
本节我们借着讲ChatGLM2功能的机会，顺便把从CUDA一直到多GPU并行时要用到的量化方法完整地介绍了一遍。
如果你有哪些功能可以用CUDA设备代码进行加速的，那就毫不犹豫地去实现它吧！算法的加速并不是局限在如何使用别人的框架和现有的功能上的。
