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设备代码进行加快的，那就毫不犹豫地去完成它吧！算法的加快并不是局限在怎么运用别人的结构和现有的功用上的。
