基于GPT2的LLM预训练

训练大语言模型往往被认为是很复杂很神秘的技术。预训练是训练LLM的第一步通常是通过让模型学习海量的文本知识得到一个可以对用户输入文本提示词进行续写的模型。如果你有transformer基础和一张独立显卡对LLM进行预训练并不是一件很难的事情。接下来我首先会简要介绍一下LLM模型的结构然后用一个例子给大家演示如何对GPT-2进行预训练让它能说中文LLM结构介绍基于GPT结构的LLM又称Decoder-only Autoregressive Language Model纯解码器自回归模型去除了原始transformer的Encoder编码器部分并对Decoder解码器的结构进行了增强下图是一个39层的LLM结构输入数据在底部输出在顶部*注图片引用自论文《Qinkai Zheng, Xiao Xia, et, CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X》数据准备数据集我使用的是Skypile-150B语料数据集的一部分。Skypile天工150B数据集是2023年10月由昆仑万维开源的600GB、150B Tokens的超大高质量开源中文文本数据集。Skypile数据下载地址https://modelscope.cn/datasets/modelscope/SkyPile-150B/tree/master/data整个数据集由多个数据文件组成一个数据文件包含的内容参考下图图中展示了一个数据文件中的前2条数据内容这里我们只下载第一个数据文件2020-40_zh_head_0000.jsonl。如果你跑通了这段代码、想提高输出质量可以下载更多的数据集文件进行训练。预训练代码详解接下来我们来看一下预训练代码完整代码链接在文章末尾依赖库加载fromdatasetsimportload_datasetimporttorchfromtransformersimport(GPT2TokenizerFast,GPT2Config,GPT2LMHeadModel,DataCollatorForLanguageModeling,Trainer,TrainingArguments,)这里加载了3个依赖库datasets用来加载数据集torch即深度学习框架pytorch用来构建网络transformers是Huggingface推出的一个transformer模型常用功能库包含transformer模型的加载、推理、训练等功能。模型等相关参数设置接下来是设置一些构造和训练模型相关参数num_layers12emb_size768num_heads16seq_length1024epochs1batch_size4output_dir/path/to/your/saving/directory主要参数如下num_layers指的是Decoder的层数emb_size指的是输入向量会被映射到多少维的空间中num_heads指的是注意力头的数量seq_length指进入模型的语句数据长度这里设置成1024表示如语句长度大于1024个token则只保留前1024个token如语句长度小于1024个token则对缺失的位置补0。epochs1指训练时只遍历一遍数据集。和普通神经网络训练时不一样LLM由于训练语料库大、消耗算力多一般不会遍历很多次数据集batch_size指一次训练时进入LLM的语句数量这里是4句适合3060这种显存小的GPU如果你的GPU显存多可以把这个数值调大如32G显存可以设置为32可以加快训练速度output_dir是输出目录会保存模型训练的中间结果比如优化器状态等如果你选择保存模型程序也会把模型保存在这个目录下一定要指定好可以不用创建最后一级文件夹程序运行时会自动创建。分词tokenizerGPT2TokenizerFast.from_pretrained(gpt2)tokenizer.pad_tokentokenizer.eos_tokendeftokenize_batch(batch):returntokenizer(batch[text],truncationTrue,max_lengthseq_length)这段代码首先初始化了一个GPT2的分词器。分词器的作用是把文本语句按GPT2的词表转化为类似[112, 24, 53, ……]这样的整型数组方便进入模型的embedding层处理即模型结构图中底部的tokenization层的作用。加载数据集dsload_dataset(json,data_files./data/2020-40_zh_head_0000.jsonl,splittrain)dsds.map(tokenize_batch,batchedTrue,remove_columnsds.column_names)ds.set_format(typetorch,columns[input_ids,attention_mask])这里首先使用load_dataset()方法将2020-40_zh_head_0000.jsonl文件里的数据加载到dataset对象中。再调用刚才的tokenize_batch()方法进行map()操作即按GPT2的词表转化为类似[112, 24, 53, ……]这样的整型数组。map()操作对CPU的要求很高可以用来测试CPU并发处理能力。构建模型configGPT2Config(vocab_sizetokenizer.vocab_size,n_embdemb_size,n_layernum_layers,n_headnum_heads,n_positionsseq_length,bos_token_idtokenizer.bos_token_id,eos_token_idtokenizer.eos_token_id,)modelGPT2LMHeadModel(config)这段代码是按照GPT-2-Small模型构建的要求配置了初始化GPT-2模型需要的参数然后再从这些参数初始化出一个GPT2LMHeadModel。其中vocab_size是指GPT-2的词表大小约为5万多构造完成的模型结构如下GPT2LMHeadModel( (transformer): GPT2Model( (wte): Embedding(50257, 768) (wpe): Embedding(1024, 768) (drop): Dropout(p0.1, inplaceFalse) (h): ModuleList( (0-11): 12 x GPT2Block( (ln_1): LayerNorm((768,), eps1e-05, elementwise_affineTrue) (attn): GPT2Attention( (c_attn): Conv1D(nf2304, nx768) (c_proj): Conv1D(nf768, nx768) (attn_dropout): Dropout(p0.1, inplaceFalse) (resid_dropout): Dropout(p0.1, inplaceFalse) ) (ln_2): LayerNorm((768,), eps1e-05, elementwise_affineTrue) (mlp): GPT2MLP( (c_fc): Conv1D(nf3072, nx768) (c_proj): Conv1D(nf768, nx3072) (act): NewGELUActivation() (dropout): Dropout(p0.1, inplaceFalse) ) ) ) (ln_f): LayerNorm((768,), eps1e-05, elementwise_affineTrue) ) (lm_head): Linear(in_features768, out_features50257, biasFalse) )该模型是GPT2系列模型中参数量最小的模型只用了12个Decoder模块总参数量为: 124439808个124M。数据从最开始的Embedding层进入从最后的(lm_head): Linear层流出LayerNorm、Attention、MLP在文章开始的结构图里都有展示可以对照着图看结构。设置训练参数training_argsTrainingArguments(output_diroutput_dir,overwrite_output_dirTrue,num_train_epochsepochs,per_device_train_batch_sizebatch_size,save_steps1000,save_total_limit2,logging_steps200,bf16torch.cuda.is_available(),remove_unused_columnsFalse,)这部分代码利用transformer库的TrainingArguments类一个对模型训练各个参数的封装类设置了模型训练常见的一些参数output_dir保存训练中间产出物等的目录overwrite_output_dir设置为True时每次训练产生的输出都会覆盖前一次的输出save_steps1000每训练1000步存储一次中间状态包括模型权重、优化器状态等为一个checkpoint检查点步数数据集数据总数/batch_sizesave_total_limit2保存在磁盘上的checkpoint总数。因为在大数据集上训练LLM会产生大量的中间结果全部保存很容易导致磁盘被写满、无法继续训练因此transformer库提供一项删除无用checkpoint的功能。这里设置为2表示只保留最近的2个检查点例如训练超过4000步时只会保留第3000步和4000步的检查点这样可以显著降低保存中间结果占用的磁盘大小logging_steps200每训练200步打印一次日志会输出学习率等一些训练相关数据bf16 torch.cuda.is_available()在支持bf16精度的GPUAmpere或更新的核心上开启混合精度训练。我使用的3060就是Ampere核心使用bf16混合精度训练速度要比默认的fp32精度更快保存的模型也更小。模型训练trainerTrainer(modelmodel,argstraining_args,train_datasetds,data_collatordata_collator,)print(-------------------start training-------------------)trainer.train()print(-------------------train fininshed-------------------)模型训练是通过创建一个transformer库Trainer对象进行训练。创建该对象提供这几个参数model指定要训练的模型就是刚才的创建的modelargs指定训练相关的参数就是上一步创建的training_args对象train_dataset指定训练使用的数据集同样是之前创建的数据集对象。最关键的是在创建完之后要调用该对象的.train()方法训练如代码所示看到控制台显示“–train fininshed–”字样说明训练完成。如果使用的是Skypile的数据集文件2020-40_zh_head_0000.jsonl训练时间会比较长3060 12G要6个多小时训练的时候需要耐心等待。我在完整代码中还保留了另一个数据集的调用方法dsload_dataset(cornell-movie-review-data/rotten_tomatoes,splittrain)这是另一个小型文本数据集rotten_tomatoes里面大概只有8千条电影评论文本用cpu 40分钟左右就可以训完可以在上GPU之前用来测试代码能否正常运行注需要你的机器能访问Huggingface或Huggingface的镜像站点。保存模型可选trainer.save_model(output_dir)tokenizer.save_pretrained(output_dir)这两行代码可以在训练结束以后保存模型文件和分词器。之所以写成“可选”是因为不写这两行代码程序也会自动把模型保存在最后一个检查点文件夹中不过因为检查点里的文件还包括优化器状态等临时文件刚接触LLM的人区分这两种文件可能有点困难所以执行这两行代码可以直接把模型相关的文件保存在你之前指定的输出目录下方便接下来推理使用。推理训练好了模型接下来就要把模型用起来推理让它能够根据输入的中文进行续写。推理代码如下首先导入transformers库中的两个模块GPT2Tokenizer和GPT2LMHeadModel用于加载分词器和模型再导入pytorch用于支撑transformers库进行计算transformers库需要基于pytorch这样的深度学习框架fromtransformersimportGPT2Tokenizer,GPT2LMHeadModelimporttorch接下来从刚刚预训练的结果中加载分词器和模型model_pathpath/to/your/output/directorytokenizerGPT2Tokenizer.from_pretrained(model_path)modelGPT2LMHeadModel.from_pretrained(model_path)其中model_path就是你刚刚保存模型的目录。接下来编写提示词再用分词器将其转为token词元intput_text山东警方出警迅速inputstokenizer(intput_text,return_tensorspt)print(inputs)如果你运行了print(inputs)程序会打印出词元化后的数组{input_ids: tensor([[ 161, 109, 109, 10310, 250, 164, 255, 99, 43095, 49035, 118, 164, 255, 99, 32573, 227, 34460, 253, 171, 120, 234]]), attention_mask: tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}其中161等数字就是“山东警方出警迅速”词元化后对应的id。接下来将该向量输入模型调用模型的generate()方法得到词元化后的结果outputsmodel.generate(inputs.input_ids,attention_maskinputs.attention_mask,# max_length100,max_new_tokens200,num_return_sequences2,pad_token_idtokenizer.eos_token_id,do_sampleTrue,# 开启采样top_k50,# top-k采样top_p0.95,# 核采样temperature1.0,# 温度no_repeat_ngram_size2,# 避免重复n-gram)由于模型输出是二维数字张量无法直接阅读接下来用tokenizer的decode()方法将该向量转为文字fori,outputinenumerate(outputs):generated_texttokenizer.decode(output,skip_special_tokensTrue)print(fSequence{i1}:{generated_text}\n)输出因为之前generate时指定num_return_sequences2因此模型会输出2段例句现在GPT2可以说中文了。输出文不达义是因为训练语料和训练轮数都不够多通过增加训练语料和训练轮数可以提高GPT2表达中文的能力。*限于时间和水平有限文中可能会有疏漏错误欢迎各位专家批评指正补充说明完整源代码地址运行本文中的代码需要有python编程基础下载代码后还请给我个star哈Github Gitee如果你想看各个显卡在本训练任务上的对比表现欢迎查看我的这篇文章各款显卡性能对比如果你对搭建分布式训练网络感兴趣欢迎查看我的另一篇文章《搭建模型训练RDMA网络上》