在谈及大语言模型的训练时,大家普遍实践的是模型的微调(SFT)阶段。以一个规模为70亿参数(7b)的语言模型为例,采用QLoRA策略,仅需单张RTX 3090显卡即能实现有效的微调工作。即便是个人研究者或小型团队,也能在有限资源下对大规模预训练模型进行微调。
然而,相比于微调阶段,预训练阶段通常是专业团队依托强大计算资源开展的一项任务。预训练一个大语言模型要求构建并运行复杂的分布式计算系统,涉及成百上千块高性能GPU的协同工作,以应对海量文本数据的处理需求以及模型训练过程中天文数字级别的计算量。
尽管如此,对于学术探究或技术学习目的的个体而言,了解并体验大语言模型的预训练流程并非遥不可及。有不少仓库GitHub仓库提供了可供借鉴的实践路径。
本文以仓库“baby-llama2-chinese”(https://github.com/DLLXW/baby-llama2-chinese)为例,通过实战来了解大语言模型预训练的大致步骤。
高质量的数据对模型性能至关重要。该项目选用五个数据集:中文版Wikipedia、中文版BaiduBaiKe、C4、中文悟道以及shibing624医疗数据。前面四个数据集对于大家来说应该比较熟悉。其中,C4数据集基于Common Crawl项目,经过额外的清理、过滤和处理。
下载好原始数据集后,按道理来说需要经过复杂的预处理,比如处理缺失值、去重、处理异常值、文本规范化等操作。该项目为了简单起见,略过了这些操作。
预训练的数据格式和SFT的数据格式不同。对于SFT的数据,一条数据一般是一个字典,这个字典包括prompt和output两个字段,分别对应着提问和回答。而对于预训练的数据,一般是将所有的文本都合并成一个连续的段落,即将整个语料库组织为一条单一的数据。
以中文Wikipedia数据为例,该数据集由254,547个词条组成。在该项目的预处理方案中,所有词条经过分词,然后连接成一个统一的文本串。在连接时,在相邻词条间插入特定分隔符“[EOS]”。
数据预处理的命令是:
这个文件里的不同函数可以分别处理五个数据集,目的是将离散的、一条一条的数据合并成一个连续的段落。在将数据连接之前,需要先进行分词操作。分词后的数据会被以np.uint16的格式进行存储,uint16是指无符号16位整数,它可以表示的范围从0到(2^16)-1,即最大到65535。
本项目所使用的分词器是来源于chatglm2-6b模型。分词器可以通过以下两种方式进行加载:
或者
tokenizer文件夹下包含三个文件:tokenization_chatglm.py,tokenizer.model和tokenizer_config.json。这些文件是直接从chatglm2-6b模型的文件夹中提取出的。分词器的词汇量大小为64794,这个数量能容纳在uint16的表示范围内。
该项目提供了一个较为基础的预训练实现,没有使用高级的封装(如 Transformers 库中的 Trainer),而是手动定义了 Dataset 和 DataLoader。
在 Dataset 中,所有数据被一次性全部加载进来。我使用的语料是百度百科数据,整体数据规模较小。数据是以np.memmap方式存储的,np.memmap可以允许我们不必将整个数组读入内存就能进行处理,这在处理大数据时特别有用。 整个数据一共有136319488个token,被reshape成266249*512的矩阵,这里的512表示模型的max_length,即最大上下文长度。接下来,每次将 512 个 token 喂给模型。在这里的滑动窗口中没有重叠,即在一个 epoch 中,每个 token 只训练一次。
换言之,在预训练阶段,我们将所有类型的语料(无论是百科数据、小说数据,还是网页数据等)连接在一起,形成一条连续的语料。然后以 max_length 为界限,例如本例中的 max_length 为 512,那么第一条数据包含了第 0 至 511 个 token,第二条数据包含了第 512 至 1023 个 token,依此类推,然后将这些数据送入模型。
在__getitem__方法中,定义了每条数据的X和Y:
这里为什么要从unit16转成int64?这是因为 PyTorch 的交叉熵函数(F.cross_entropy)在计算过程中,需要接收的目标(target)参数是 64 位(int64)或者 32 位(int32)整型。
预训练的命令是:
根据GPU卡的数量来调整nproc_per_node参数的值。
该项目作为入门学习LLM的演示示例,使用到的模型规模不大:hidden_dim=512,layers=8,总参数量是58479616,约等于0.06B。
在训练过程中,每次将shape=32*511的input(即代码中的X)输入到模型中,输出logits,形状为32*511*64793。这里的64793是词汇表大小。接着,将输出的logits与shape=32*511的label(即代码中的Y)做交叉熵计算,这样就得到了loss,用来指导梯度下降。
我使用 A100的显卡,在15分钟就完成了一轮的训练。
以“感冒的症状是”为例,来看看模型是否能生成我想要的结果。
对于完全没有经过任何训练的模型,参数是随机初始化的,产生的输出如同随机生成:
经过了一轮百科数据的训练,产生的输出如下:感染可能会造成病因。21日,北京大禹的传染病药物检查与临床试验开始,严重急性,慢性肝、心核、心痛、肺膜炎和肝素水平皆没有异常作用。但以54型冠状病毒剧疫苗为辅助,使大禹能够安全有效的治疗,因而长期服用疫苗,如服用或治疗多种相关药物。该治疗可能影响中国境内癌症,但相关患者的癌症。
虽然也是胡说八道,但至少能看出回答属于医学话题。