（注：本专栏只讨论相关技术细节，不涉及任何其他目的，如侵删） 

提醒：本专栏全文较长，部分是选做内容，可自行取舍。

摘要

2023年3月，清华大学NLP团队开源了对话模型ChatGLM-6B，该模型具有对话流畅、部署门槛较低等优点。同年6月，该模型的升级版本ChatGLM2-6B发布，二代模型具有更强大的性能和推理能力。本项目利用大模型的低秩适配（Low-Rank Adaptation，LoRA）技术对INT4量化ChatGLM-6B系列模型进行微调，在包含~7k段对话的微信聊天对话数据集上进行训练，以实现一个微信聊天机器人。

本专栏讨论了微调相关的实现细节，包括数据集预处理、量化模型加载、并行计算（受限于Kaggle CPU的显存，实际上并未实现，但正文中有相关的讨论）、LoRA实现和训练等，同时包含部分ChatGLM的源码解析。需要指出的是，由于微信聊天对话数据的处理比较粗糙，不少对话上下文不具有明显的逻辑性，因此模型最后的推理性能并不理想。想要提高模型的推理回答质量，需要进行更严格的数据清洗工作。

本专栏是对我上一期视频的补充：

相应的Kaggle Notebook链接为https://www.kaggle.com/code/littleweakweak/finetune-lora-chatglm2，可以与本专栏结合起来看。

目录

• 前言（废话）

• 数据集预处理

• 模型量化与加载

• 并行计算

• Monkey Patch（选做）

• LoRA

• 模型训练（Kaggle + Colab）

• 模型推理

• 链接汇总

前言（废话）

首先感谢ChatGLM团队提供的开源模型和官方微调教程：

https://huggingface.co/THUDM/chatglm-6b

https://huggingface.co/THUDM/chatglm2-6b

（PS：前两天ChatGLM官方账号私聊我，问我能不能出一期“智谱清言”的使用体验，我打算下一期视频就出。作为一个AI初学者 & B站小透明新人，能被官方翻牌子还是很受宠若惊的。）

训练一个微信聊天机器人的想法实际上去年11月就有了，当时为了帮女朋友写计算语言学的大作业，提取了几千段微信对话一通分析，想不到现在就已经变成前女友了，真是感慨万千（小丑见大丑）。

今年2月份，我尝试利用中文GPT2训练了一个微信聊天机器人，最后的效果还挺出乎我意料的。后来ChatGLM-6B和ChatGLM2-6B相继发布，我很好奇更大的模型在同样的数据集上能表现出怎样的性能。然而Kaggle GPU的显存过小，全参数微调不现实，而LoRA可以实现较低的微调显存消耗和较好的模型推理性能，因此我基于LoRA对模型进行微调。

这个项目做了差不多有一个月，最后的实现难度大大超出了我原来的设想。在Kaggle那两块破GPU上做微调真就是“戴着镣铐跳舞”，好在最后也算是做完了。在整个过程中，我也是从零开始学习LoRA、模型量化、并行计算（Accelerate库）等相关知识，时间有限，能力不足，项目还是有不少缺陷。有任何问题或者意见都欢迎评论区留言或者私聊我，大家一起学习、一起进步。

（以下将ChatGLM-6B和ChatGLM2-6B分别简称为GLM1和GLM2）

数据集预处理

微信聊天记录的提取可以看我之前的专栏：

这里我们只需要把微信文本提取到txt文件（data.txt）中，格式如下（每段对话之间由\n分隔）：

然后利用下面的代码划分数据集，指定训练集和验证集的存储路径（train.txt / val.txt）：

然后把train.txt和val.txt上传到Kaggle的Dataset中，具体操作见视频。

下一步需要构建训练和测试用的Dataset和定义collate_fn，这里我主要参考了下面两个链接：

https://github.com/yuanzhoulvpi2017/zero_nlp

https://github.com/liucongg/ChatGLM-Finetuning

同时，根据GLM2源码中的chat函数以及tokenizer中的build_prompt函数：

模型推理生成回答时，需要在问答前后加上“[Round 0]\n问”或“\n答”。因此我们需要转换一下输入数据的格式，加上相应的prompt，即convert_txt函数：

根据一些实验测试，最后我选用的prompt方式是，在所有对话前后加了“Round”和“问” 或“答”，同时在最开始加了一段对话：“\n问：你现在是一个微信聊天机器人，接下来你要根据我的提问，作出相应的回答。\n答：好的，请开始你的提问。”

convert_txt函数会对输入文本进行问答划分，最后将处理后的数据集保存在一个字典中。

然后构建Dataset，这里我用的是Hugging Face的datasets库，可以方便地构建Dataset对象。具体用法见https://huggingface.co/docs/datasets/index。

我们需要定义preprocess函数来处理数据集，其功能是对数据集中的文本做tokenize，同时保存每段对话中提问的起始位置，以便后续进行mask操作：

这里我用的是GLM2的tokenizer，tokenizer会在每段文本前加上[gMASK]和sop两个标记：

其中64790和64792分别对应[gMASK]和sop：

（注意：GLM1和GLM2用的tokenizer的词表大小不一样，如果后续训练中出现tokenizer报错数组越界或是下面这种ValueError，可能是tokenizer用错了）

下一步是定义collate_fn，目的是把所有提问的部分都mask掉（-100表示不对该位置的元素计算loss）：

在官方推荐的LoRA-GLM1的微调代码中（也就是上面出现的zero_nlp），collate_fn还返回了处理后的attention_mask：

具体来说，模型在生成提问部分对应的输出时，可以看到整个提问部分；在生成回答部分对应的输出时，才和原始的GPT模型一致，只看到当前已经生成的内容。

而在GLM2的微调代码中，collate_fn没有返回预先定义的attention_mask，而是由模型在计算Attention时自动创建一个上三角形的mask（与原始GPT一致）：

按照我个人的理解，预定义atttention_mask似乎有些多此一举了，因为返回的label已经把提问部分给mask掉了，即loss不会传给提问部分的元素，所以模型如何预测提问部分并不重要。

Anyway，最终的collate_fn只返回了input和label，这里可以验证一下collate_fn的结果是否正确（主要看label中mask的位置）：

然后设置Batch size，定义dataloader：

模型量化与加载

（这一部分由于我也是初学者，下面的理解可能会有不少错误，欢迎大家指正）

接下来是模型量化与加载，这是最麻烦的一步。Kaggle的CPU只有~13GB的内存，用AutoModel.from_pretrained加载模型时需要先加载到CPU上，而GLM1的模型大小为~14GB，GLM2的模型大小为~13GB，因此GLM1不能直接加载，GLM2有时可以直接加载（取决于之前占用的内存大小）。此外，如果像我之前部署GLM1推理一样，使用load_checkpoint_and_dispatch函数，之后可能会报错tensor的device不一样。考虑到GPU也只有~15GB的显存，最终我使用的是INT4量化的GLM2进行微调。

加载量化模型主要有两种方式：

1. 基于Accelerate和bitsandbytes库，在AutoModel.from_pretrained中指定load_in_4bit=True，即

模型量化的原理可以参考https://huggingface.co/blog/hf-bitsandbytes-integration，我觉得写的很清楚。这种方式似乎需要先加载一遍完整的模型到CPU中，然后在移动到目标device的过程中完成量化，因此只能用这种方式加载相对小一些的GLM2，加载完成后大约占用显存5GB，不过很多时候会加载失败。

量化完成后，模型中的nn.Linear层会被替换为bitsandbytes中的Linear4bit层，为了数值稳定性，输出层没有进行量化：

这种加载方式的好处是Peft库可以支持基于bitsandbytes的量化模型，后续在LoRA微调时，代码比较简洁。Peft库官方提供的INT8微调demo都用的是这种方式。

2. 直接加载官方的量化模型，即

ChatGLM官方提供了模型量化方式，见https://huggingface.co/THUDM/chatglm2-6b/blob/main/quantization.py，其中定义了量化线性层类QuantizedLinear：

然而，Peft库不知道这个QuantizedLinear是什么东西，所以如果要使用Peft库，需要把QuantizedLinear换成Peft支持的Linear4bit层或LinearNF4层，之后在LoRA部分会详细说明。

加载时最好不要指定device_map='auto'，后续可能会报错。

（PS：我还尝试过使用low_cpu_mem_usage和tensor_parallel来加载，都失败了，感兴趣的朋友可以参考这两个链接试一下：

https://huggingface.co/docs/transformers/main_classes/model

https://www.kaggle.com/code/blacksamorez/tensor-parallel-int4-llm

）

并行计算

Kaggle免费提供的GPU比Colab的更好：1块P100或者2块T4。然而P100不支持量化模型（见https://github.com/THUDM/ChatGLM-6B/issues/1322），因此我们很自然地想用2块T4做并行计算。

Hugging Face提供了用于并行计算的Accelerate库，下面是一些介绍：

https://huggingface.co/docs/accelerate/index

https://blog.csdn.net/cxx654/article/details/131817042

具体来说，我们需要修改训练代码，去掉所有的.to(device)或者.cuda()操作，而是让accelerate来分配device。然后把训练代码写在一个train函数里，调用notebook_launcher函数进行分布式训练：

可参考下面的链接：

https://huggingface.co/learn/nlp-course/chapter3/4#supercharge-your-training-loop-with-accelerate

https://huggingface.co/docs/accelerate/basic_tutorials/notebook

然而，notebook_launcher似乎是要在CPU上加载两次完整的模型，而加载一次INT4模型最多需要~9GB的内存消耗，因此CPU根本吃不消。之后我又换了launch py脚本的方式（用Accelerate或DeepSpeed），也出现了同样的问题，所以并行计算的尝试宣告失败。之后我可能会尝试学一下DDP，看看能不能做。

这里有两个小坑点，用Accelerate的xd可以注意一下，以免出现nan梯度。（本来是在一篇CSDN上看到的，后面找不到原文链接了）：

1. 量化模型加载时，需要选择参数类型为float32：

2. 混合精度训练时，不要使用accelerate.autocast，而是用accelerate.accumulate，训练for loop大概长下面这样：

调用Accelerate库的小demo：https://www.kaggle.com/code/littleweakweak/test-accelerate-chatglm

Monkey Patch（选做）

（这一部分是选做，有兴趣的xd可以自行尝试，不感兴趣的可以直接跳到LoRA部分）

Monkey Patch（猴子补丁）是一种在程序运行时动态修改对象属性的方法，像这样：

可参考下面两个链接：

https://zhuanlan.zhihu.com/p/71181926/

https://stackoverflow.com/questions/5626193/what-is-monkey-patching

由于模型文件是直接从Hugging Face加载，想要修改内部代码并不容易，此时就可以运用Monkey Patch来魔改内部实现。我这里使用Monkey Patch主要用于简单并行，即把模型的不同层放在两个GPU上，以节约显存。我在尝试的时候为了简便，只把embedding和output_layer（均为65024 * 4096）放在了cuda 0，其它层放在cuda 1。

GLM2的model是一个ChatGLMForConditionalGeneration对象，而model的transformer属性是一个ChatGLMModel对象：

在GLM2的源码中，ChatGLMForConditionalGeneration的forward函数（简称为外层forward）会进一步调用transformer（ChatGLMModel）的forward函数（简称为内层forward），其中embedding层在内层forward中被调用，output_layer在外层forward中被调用。

ChatGLMForConditionalGeneration的forward函数（外层forward）：

ChatGLMModel的forward函数（内层forward）：

所以，为了实现简易版“模型并行”，我们需要利用Monkey Patch重写内外层的forward函数，即对embedding和output_layer的相关tensor做device转换。

内层forward改写：

外层forward改写：

Monkey Patch和模型加载的代码：

（这里提醒一下，想要移动nn.Module对象到特定device，需要直接对nn.Module对象调用.cuda()或.to()，移动它的parameters是没用的。参考

https://discuss.pytorch.org/t/moving-a-module-to-a-device/142175

https://github.com/pytorch/pytorch/issues/7460

）

经过测试，移动embedding和output_layer可以节省~1GB的显存（INT4量化操作会消耗cuda 0的~1GB显存）：

综上，移动模型的部分层确实可以显著减小显存占用，提高GPU的利用效率。只是后面训练的时候我懒得加了，大家可以按照这个模式自行魔改内部函数，以实现更高效的模型并行（比如前一半layer在cuda 0，后一半layer在cuda 1）。上面的代码汇总在https://www.kaggle.com/littleweakweak/test-multi-device-glm。

LoRA

低秩适配（Low-Rank Adaption，LoRA）是一种最近常用的大模型微调技术，其原理非常简单巧妙。它认为权重矩阵的更新本质上是低秩的，所以可以用两个低秩矩阵的乘积来近似。按我个人的理解，LoRA假设权重矩阵中只有极少元素的更新是独立有效的，而大部分元素的更新是冗余的，那么就可以用低秩矩阵去近似其中的“有效更新”。原理可见：

https://baijiahao.baidu.com/s?id=1771735489327080446&wfr=spider&for=pc

运用LoRA微调时，需要冻结原始权重的梯度，然后对目标线性层添加低秩适配器（矩阵A、B），在训练过程中只更新低秩适配器的权重，从而大大降低了训练参数量。

实现LoRA可以通过Peft库，也可以自己手写（LoRA from scratch）。

1. 通过Peft库

Peft库实现了多种大模型的adapter微调方式，同时支持对bitsandbytes量化模型的微调（量化层目前还不支持梯度更新，但LoRA等方法只需要更新adapter层的参数），可参考：

https://huggingface.co/docs/peft/conceptual_guides/lora

https://github.com/huggingface/peft

注意安装的Peft库的版本应该为0.4.0，否则在import的时候会报错。

前面提到Peft库不支持对QuantizedLinear层添加adapter，因此需要把QuantizedLinear层转换为LinearNF4层：

需要注意的是，LinearNF4的实际量化发生在.cuda()或.to()调用，模型权重会从float类型转为int8类型。在量化完成后，不能再对量化层进行任何.cuda()或.to()操作，即使不改变device。因为每一次.cuda()或.to()操作似乎都会执行一次量化，而多次量化后在矩阵相乘时会报错：

下一步利用Peft库配置LoRA，这里我们选择adapter的秩 r = 8，比例系数 α = 32，对每个Attention模块的QKV变换层（LinearNF4）和输出层（Linear）添加adapter：

其中prepare_model_for_kbit_training函数的作用是开启输入梯度（有利于微调adapter）和梯度检查点（节省显存，但是训练耗时更长），具体见https://blog.csdn.net/BIT_666/article/details/131675165；target_modules指定了需要添加adapter的module，注意在GLM中最后一层输出层名为lm_head，GLM2中则为output_layer。最后调用get_peft_model完成adapter的添加：

输出层的LoRA模块：

2. LoRA from scratch

我自己手写的LoRA参考了BERT-LoRA-TensorRT项目，项目链接：

https://medium.com/@alexmriggio/lora-low-rank-adaptation-from-scratch-code-and-theory-f31509106650

https://github.com/alexriggio/BERT-LoRA-TensorRT

具体地，我们需要将所有要添加adapter的线性层替换为线性LoRA模块，共包含3个线性层（pretrained线性层，矩阵A、B），然后将原线性层的权重复制到pretrained线性层中。由于原始QuantizedLinear是对原始权重进行INT4量化，因此量化权重的输入维度应该是原权重输入维度的一半（猜测：一个byte表示两个“INT8”，这里weight_bit_width=4）：

这里我们直接复制INT4权重，所以权重维度不需要除2，为了方便后续构造和复制权重，我稍微修改了一下QuantizedLinear的定义：（其实不改貌似也行，可以直接把weight_bit_width设为8）。

接下来调用qlinear_conversion函数，把模型原有的QuantizedLinear层换成自定义的QuantizedLinear；开启梯度检查点和输入梯度（相当于peft中的prepare_model_for_kbit_training）。最后可以自行决定是否要把最后一层参数变成float32：

接下来定义针对量化线性层的QLinearLoRA模块：

类似地，我们也要定义针对普通量化层的LinearLoRA模块。然后定义添加LoRA模块的函数create_qlora：

调用add_lora_layers对模型添加LoRA模块：

最后，冻结除LoRA的adapter以外所有层的梯度，至此模型部分的修改已经全部完成。

模型训练（Kaggle + Colab）

经过了前面繁琐的模型修改，接下来是模型训练，这一部分的逻辑基本上没什么大的改动，所以代码部分我就不细讲了。前面我们通过两种方式实现了LoRA，基于Peft库的和LoRA from scratch。在实际训练的时候，我发现基于Peft库的训练微调要远远慢于LoRA from scratch。基于Peft库的微调耗时大概是3h / epoch（不开梯度检查点~2h）：

而LoRA from scratch的耗时大概是50min ~ 1h / epoch（参考原视频，耗时似乎与网速有关，最快可以30min / epoch，平均~1h），训练耗时受batch size影响不大。这里我还不太清楚训练耗时差异的原因，有知道的大佬可以在评论区说一下。我猜测是Kaggle的T4 GPU太辣鸡了，可能对持Peft库内部的一些优化兼容性较差。总之，由于基于Peft库的微调实在太久了，我最后只用了LoRA from scratch来微调。

另外，由于使用AutoModel.from_pretrained加载的模型参数均为float16类型，而后续添加的LoRA矩阵A、B均为float32，因此直接训练会报错：

所以我们需要利用torch.autocast开启混合精度训练，让torch帮我们自动做类型转换（autocast只用wrap forward部分）：

关于自动混合精度（Automatic Mixed Precision，AMP）的介绍可进一步参考：

https://pytorch.org/docs/stable/amp.html

https://www.cs.toronto.edu/ecosystem/documents/AMP-Tutorial.pdf

然而，如果在加载模型的时候指定参数类型为float32，则在量化层的backward过程会报类型错误：

原因可能是W8A16Linear是自定义的torch.autograd.Function子类，其中的forward和backward都是自定义的，torch.autocast对其不起作用。

其中，forward函数调用了extract_weight_to_half函数，其作用是把INT4量化的权重（int8）恢复为float16类型，然后保存输入张量inp供backward过程使用。而此时模型的非量化层（如embedding）的参数类型是float32，所以经过前面一系列模型层的inp的类型也是float32，也就出现了类型冲突。

关于自定义torch.autograd.Function子类可以参考：

https://pytorch.org/docs/stable/autograd.html

https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html

https://zhuanlan.zhihu.com/p/574119247

因此，我们需要稍微修改W8A16Linear的定义，将inp的类型转换为float16（其实也可以把weight转换为float32，但是这么做训练会非常耗时）：

我自己只尝试了float16加载，感兴趣的xd可以试一下float32加载，看看训练之后的效果。

确保训练循环无误后，就可以用Save version在后台训练了。我自己训练了30个epoch，最终的训练loss大概收敛在2.6左右（很早就收敛了，大概是数据集的缘故）。

Colab的代码基本和Kaggle一样，链接如下：

https://colab.research.google.com/drive/1jHvQXG_SH4aY1fZo2H7degxmszBs2rfN?usp=sharing

注意修改相关路径：

修改完成后，可以直接选择“全部运行”。

模型推理

训练完成后，可以进行模型推理。一般的LoRA微调在推理时为了提高速度，在模型加载后会先将adapter的权重merge到原权重上：

但在量化模型中，W0是量化后的权重，需要先恢复为float16才可以merge，而恢复后的权重会占用大量显存，因此我们这里不预先进行权重的merge，只把模型转为eval模式。

推理的Kaggle Notebook为：https://www.kaggle.com/code/littleweakweak/infer-chatglm-lora。

首先需要上传checkpoint文件，然后打开推理Notebook，指定checkpoint路径，设置推理参数（top_p和temperature），具体可以参考原视频里的操作。

关于top_p和temperature可以参考https://zhuanlan.zhihu.com/p/613428710。在GLM2的chat函数中，top_p和temperature的默认值均为0.8，以使模型的回答具有足够的多样性。然而微调之后，模型的回答在一定程度上会被局限在训练语料的分布中。根据我的测试，应该选取一个较低的top_p值和较高的temperature值，这样模型的回答才“相对”合理一些。最终我选取的参数为top_p=0.15，temperature=20.0。（如果temperature过低，模型根本不能生成正常的回答）

而在对话过程中，我将prompt中的“\n问：你现在是一个微信聊天机器人，接下来你要根据我的提问，作出相应的回答。\n答：好的，请开始你的提问。”作为history传入模型，同时限制模型每次生成回答的长度为20，以免模型会生成过长的回答。（似乎模型在微调后忘记了如何终止回答）

视频中也呈现了推理的效果，由于训练语料的清洗比较粗糙，所以很多时候模型生成的回答根本没法看，我只是选取了几段稍微合理的回答。想要提高模型的推理回答质量，需要进行更严格的数据清洗工作。

链接汇总

1. ChatGLM项目链接：

https://huggingface.co/THUDM/chatglm-6b

https://huggingface.co/THUDM/chatglm2-6b

2. 模型训练Notebook：

https://www.kaggle.com/code/littleweakweak/finetune-lora-chatglm2

3. 调用Accelerate库的demo：

https://www.kaggle.com/code/littleweakweak/test-accelerate-chatglm

4. 手动模型并行demo：

https://www.kaggle.com/littleweakweak/test-multi-device-glm

5. 模型推理Notebook：

https://www.kaggle.com/code/littleweakweak/infer-chatglm-lora

结语

这个项目断断续续做了有小一个月，最后的工作量远远超出我当初的设想。一开始的模型并行就弄了差不多两个星期，最终宣告失败；后面的LoRA实现、模型训练和推理又弄了差不多两个星期，好在最后的效果我个人还能接受，算是努力没有白费吧。

这篇专栏也写了差不多一个星期，中间白天还要上课，也有各种各样的事情，所以有些地方可能观感上不是很连贯，望大家包涵。以后我打算延续这种“视频展示效果、专栏解释细节”的模式，一个原因是我觉得视频时间有限，能传达的信息量也有限；另一个原因是写专栏可以比较系统地整理我的思路，同时便于大家理解其中的细节，我自己以后也会时不时翻看这篇专栏。大家如果觉得有更好的形式，或是有任何的疑问、意见或建议，都欢迎在评论区留言或是私信我。我也是初学者，大家一起进步。