这次打算用 Hugging Face 的 API 来写一份预训练大(小)模型的代码,也就是用 Trainer 来做预训练。由于只是想练习一下,因此打算选一个极小模型 + 小数据集。为了贴近主流,于是打算预训练一个 LLaMA 3——不过是超迷你版本,大小仅不到 20M。
想起来曾经看到过的微软的工作 TinyStories,探索的是语言模型在多小的情况下还能流利地讲故事,工作非常直白、有趣,刚好也契合我的练习想法,于是这次来复现一下。
代码放在这里了:GitHub - Mxoder/TinyStories: 从头预训练一只超迷你 LLaMA 3——复现 TinyStories。
1. 前期准备
让我们先来想一想大概需要做什么。
首先是模型架构的选择。原工作用的是 GPT Neo 架构(可以看他们的 config),这个算是很老的模型了,最初是 EleutherAI 用来复现追踪 GPT-3 的工作的,现在用的也比较少了。我打算选用 LLaMA 架构,也算是符合研究主流、便于推广。LLaMA 3 主要多了个 GQA,也是现在模型的主流,我这里也用一下。
其次是数据的选择。既然是复现,就直接贯彻拿来主义,用原工作开源的数据集(主要是从头生成要花不少 api 费用)。原工作第一版的时候用的是 GPT-3.5 生成的数据,后面社区有人更新了第二版,是用 GPT-4 生成的,比原数据更好,就用它了。
最后是训练。其实我手上就两张 3060 12G 和 4060 Ti 16G,训这个确实是绰绰有余,但我还是不想在桌前吵我自己,于是继续用 Colab。现在 Colab 可以直接看到剩余使用时长了,虽然已经被砍到只有 3h 左右的用卡时间,但至少心里有个底,况且 3h 训我们这个也完全够了。
我们这次用到的 Hugging Face 的库如下:
理论上比较新的版本都没问题,但如果你很久没更新了,最好用 pip install -U 来升级一下。
我这里的用到的库版本如下,供参考:
另外,接下来的步骤讲解主要是以 jupyter notebook 的形式展开的,并不是 .py 文件的形式,也就是说前面执行的变量会在中间储存下来。
2. 原工作简介
虽然是练习,但既然打着复现工作的名头,还是来简要回顾一下原工作究竟做了什么吧。
原工作探索的问题是语言模型(LM)在文本连贯性上的表现。像早期的一些语言模型如 GPT-2,即使在一些 Common Craw 这样的语料库上大量预训练后,也很难生成长的、连贯的文本。比如前几年有一种 AI 玩具类型是做文本续写,例如彩云小梦,可以写写作文、小说什么的,如果大家玩过就知道效果其实一言难尽,和今天的大模型完全没法比,其实这就是 GPT-2 level 的续写能力。
作者就在想,会不会是因为训练的语料库太多、太宽泛,需要学习各种语法元素、词汇、知识、推理等等,才导致小语言模型(SLM)没法有一个很好的表现。作者决定专注于一个任务——短篇故事续写,来探索一下 LM 的性能边界。
作者用 GPT-4 和 GPT-3.5 构建了一个英文短篇小说数据集 TinyStories,将内容限制在三四岁儿童也能轻松理解的程度,并且使用不同的关键词来让故事的主题足够丰富。此外,他们还加入了额外的关键词,来控制故事有更曲折的走向、不同的结局等等。
作者用的模型基座架构是 GPT Neo,词表大小约为 50k,并且他们尝试了不同的模型参数,调整了隐藏层维度(hidden_size)、隐藏层数(num_hidden_layers)等,来探索不同参数对于模型性能的影响。
作者的评估方式是经典的 GPT-4 监督打分模式,就是让不同的 SLM 根据提示生成故事,然后 GPT-4 从设定好的不同维度来打分,主要有 Creativity、Grammar、Consistency 三项,分别代表创造性、语法正确性、上下文一致性。此外,作者额外加入了一套 TinyStories-Instruct 数据集,来训练一批指令微调的 SLM,并测试他们的指令跟随能力,也就是第四项 Instruct。
作者主要和 GPT-Neo 以及 GPT-2 的小中大杯进行了对比。
3. 模型初始化
让我们正式开始复现!
3.1 决定模型的参数
首先是定义我们自己的模型。由于 LLaMA 3 的架构早就集成于 transformers 库中,因此我们可以直接用 AutoConfig 初始化一个模型配置,传入参数 model_type='llama' 即可。
架构确定了,那么现在来探讨一下模型具体参数,比如隐藏层大小、隐藏层数等等。我们先来看看 TinyStories 原工作的实验结果:
可以看到,隐藏层维度从 64 增长到 256 时的收益是比较大的,往后收益就逐渐放缓了。而层数的影响并不如隐藏层维度那么大,大而浅的网络也能有不错的表现(例如 hidden_size=1024, num_hidden_layers=1 的模型)。综合考虑,我这里选择 hidden_size=256 和 num_hidden_layers=4。
其他参数方面,我们遵循现在主流的研究表现,将 FFN 的维度从传统的 4 倍隐藏层维度设为 8/3 倍(按 128 向上取整)。头的数目我们设为 16,并应用 GQA 机制。GQA 的实现在 transformers 中非常简单,只需要配置 num_key_value_heads 即可。num_key_value_heads 取值和 num_attention_heads 相同时即为 MHA 机制,取值为 1 时即为 MQA 机制。
综上,我们的配置如下:
3.2 分词器 Tokenizer
我这里选用 LLaMA 2 的分词器,因为二代的词表比较小(32k),LLaMA 3 的词表太大了(128k),在 SLM 中会占用太多的参数比重,并且这只是个专有任务数据训练,没必要用太大的词表。
另外注意这里 padding_side='left',如果不是的话需要设置 tokenizer.padding_side='left',即批量填充的时候从左边开始填充,这对于 decoder-only 的模型做生成任务是必要的,因为我们本质上做的是 next token prediction,如果 pad 挡在了生成序列的右边,会影响到模型生成。
吃果冻不吐果冻皮
专注于AI工程化(LLM、MLOps、LLMOps、RAG、Agent)落地。
150篇原创内容
公众号
3.3 模型实例化
接下来就是实例化模型,这里就不用从预训练模型加载 from_pretrained() 了,而是从配置加载 from_config():
可以看到,k_proj 和 v_proj 的 out_features 从 256 变为了 128,这即是 GQA 机制。
此时,模型已经初始化了,让我们来打印一下看看参数:
得到结果如下:
可以看到,我们的模型只有不到 20M!非常非常小,并且其中 Embedding 占了大头。
尽管模型还没有训练,但我们仍然可以测试一下推理:
嗯,的确是胡言乱语呢,不过可以正常推理,说明模型没问题!
但现在模型是随机初始化的,为了让模型更好地收敛,我们最好给模型一个更好的初始化方法,我这里选用 Kaiming 初始化,比较适用于 ReLU 类的激活,当然也可以选用高斯初始化、Xavier 初始化等等。
现在,我们的模型真正初始化完成了!如果你愿意,可以先将这个初始化好的模型保存到本地,用 save_pretrained() 即可。
4. 数据集
让我们继续!
4.1 加载数据集
我们接下来需要从 Hugging Face 加载数据集,我这里是建立在网络畅通的基础上的,如果你没有用 Colab 或者网络无法直连 Hugging Face,那么也可以先下载到本地某个文件夹中,load_dataset 也可以直接读取本地文件夹。我们要用的数据集路径如下:noanabeshima/TinyStoriesV2 · Datasets at Hugging Face
我们来看看数据长什么样子:
这里需要注意,datasets 加载后的数据是 Dict[str, List[str]] 的形式的,并非 List[Dict[str, str]]。
4.2 数据预处理
接下来,我们要将数据预处理一下,也就是用 tokenizer 进行 tokenize。让我们来写一个处理函数:
我们来解析一下其中的一些点:
tokenizer 的 encode
根据前面的示例,我们知道这里 examples['text'] 其实是一个 List[str],当一个 List 传入 tokenizer() 时,tokenizer 会自动进行 batch encode,得到的是 {'input_ids': List[int], 'attention_mask': List[int]}(当然,如果设置了 return_tensors='pt' 就会得到 Tensor)。
add_special_tokens=False 则是让 tokenizer 不要加上特殊 token,在 LLaMA 中就是不会在句首加上 bos_token <s>。
填充还是截断?
在这里,我采用直接截断的方式,最大截取当前输入序列的后 (max_token - 1) 位,再加上一个 eos_token_id,组成总长度不超过 max_token 的序列。attention_mask 的长度保持一致,全为 1。
这里利用到了 list 的切片特性,input_ids[-max_token+1:] 可以获取 min(max-token, len(input_ids)) - 1 的序列。
当然,也可以采取将超出长度部分再按照 max_token 来分块,重新组装。
应用在所有数据上
接下来,我们用 map() 函数,来将 process_func() 应用到 ds_train 和 ds_val 中的每个样本:
数据预处理成功!
4.3 数据批处理——DataCollator
4.3.1 两行代码
我们在训练的时候往往不会一条一条训练,而是成批次地训练,那么我们就需要对数据做批处理。因此我们需要用到 transformers 中的一个工具系列——DataCollator[1]。
既然是预训练,那么就是让模型在语料上做无监督学习,也就是我们熟知的 next token prediction,即根据前面的所有输入来预测下一个 token,然后把新的 token 拼接在已有输入上作为下一输入,如此往复,直到触发停止设定(例如触发 max_new_tokens)。
所以我们的训练目标——或者说是 label——显而易见,就是把输入偏移一位当作预测目标,我们计算的就是输出和这个目标之间的 loss:
所以我们只需要把 input_ids 复制一份、再偏移一位,就可以作为 labels 了。
……再等等,让我们看看这条问题:Shifting ids to the right when training GPT-2 on text generation? - Beginners - Hugging Face Forums[2]:
这里 sgugger 提到,在 Hugging Face 的实现里,training 时已经实现好了偏移一位的逻辑,不需要我们再手动实现了。我们也可以在 transformers 的源码里看到这一点,例如 LLaMA[3] 的实现。
所以,我们只需要将 input_ids 直接复制一份作为 labels 即可。
那么怎么做呢?我们可以用 DataCollatorForLanguageModeling,并设置 mlm=False:
两行代码,非常简单!
不过我们可以稍微多讲一点,这个 data collator 是如何发挥作用的?为什么选的是它而不是在微调中更常见的 DataCollatorForSeq2Seq?
吃果冻不吐果冻皮
专注于AI工程化(LLM、MLOps、LLMOps、RAG、Agent)落地。
150篇原创内容
公众号
4.3.2 More things……
实际上,就像 mlm 这个参数所显示的一样,DataCollatorForLanguageModeling 一开始其实是设计给 Bert 的 MLM 任务的。MLM 任务就是 Masked Language Modeling,掩码语言建模,就是在一整个序列中挑选一部分 token 用 [mask] 给盖住,让模型去根据上下文预测被盖住的是什么 token。
可以看到,这样建立的 label 就是原来的 input 的 copy,它是将 input 中随机 mask 一部分,label 不变。
这几乎就是我们想要的——只是不需要 mask。所以,我们设置 mlm=False 后,就可以直接得到 input_ids 的 copy 了。
让我们继续看看 DataCollatorForLanguageModeling 怎么作用的:
可以看到:
- • 由于长度不一,所以 data collator 做了 padding,padding 的方向就是我们 3.2 中提到的 tokenizer.padding_size
- • labels 确实是 input_ids 的原位复制,区别在于 input_ids 里用 pad_token_id 来填充,labels 里对应的是 -100、表示不计算 loss
那么微调里常用的 DataCollatorForSeq2Seq 又是如何作用的呢?我们仍然用刚刚的数据例子:
可以发现,DataCollatorForSeq2Seq 和 DataCollatorForLanguageModeling 一样,做了批处理和 padding,但是没有标签 labels。原因是:DataCollatorForSeq2Seq 设计之初用于的任务和它的名字一样,是序列到序列(seq2seq)任务,放到文本任务上,就是要有两个 seq:输入 text 和 输出 label,比如下面的例子:
可以看到:
- • DataCollatorForSeq2Seq 需要指定当前输入的文本和后面需要生成的文本,即 text 和 label,如果像 DataCollatorForLanguageModeling 那样处理会得不到 labels 字段
- • 因此,DataCollatorForSeq2Seq 适合有监督微调(SFT),输入是 text,输出是 label,非常合理
5. 超迷你 LLaMA,启动!
5.1 配置训练参数
我们需要用到 transformers 的 TrainingArguments 来配置训练参数,具体参数说明可以看这里。
如果你之前用了 wandb,现在想禁用掉,可以设置环境变量:
5.2 配置 Trainer
同样地,具体参数说明可以看这里。
5.3 训练与保存
配置好 Trainer 后,通过下列代码即可启动训练:
接下来只需要等待训练完成。我用一个半小时训练了 2 epochs,loss 达到了 1.6 左右。
训练完成后,如果用的是 jupyter notebook,那么此时 model 已经是训练好的状态了。我们可以再次推理试试看:
得到如下结果:
可以看到:
- • 20M 模型确实能够流畅续写故事了
- • 20M 模型写出的故事的语法、流畅度都不错,但是一致性欠佳,特别是故事主题、人名的前后连贯性不高
总之,这个超迷你 LLaMA 3 确实训练完成了!我们可以将它保存到本地:
也可以推送到 Hugging Face:
6. 结尾
这次尝试用 Trainer 来做一个模型的预训练,以往都是用 Trainer 来做微调,这次也算是学习了一下吧。TinyStories 这个工作之前就有关注过,但一直没顾上来复现一下,这次也算是简单复现了个小模型出来,和原工作的丰富度确实是比不了,但也算完成一个 todo。
后面这个小模型可以继续做 SFT,也就是做指令微调,可以和原工作一样,给定故事背景、关键词、开头让小模型续写,也可以迁移到别的任务。不过由于我们的预训练任务只针对了讲短篇故事这一类任务,加上参数又特别少,如果直接迁移其它指令任务估计表现不会很好。
本篇文章也是希望尽可能写得详细一点,欢迎大家交流。