从零预训练139M中文大模型:混合精度、数据采样、LoRA微调与全量SFT实战

发布时间:2026/6/11 10:24:53
从零预训练139M中文大模型:混合精度、数据采样、LoRA微调与全量SFT实战 项目开源说明本项目的代码数据权重均已开源可以在单张5090显卡上跑通模型预训练全量SFT和LoRA微调的全流程。同时项目还提供便捷的WebUI可以轻松体验本项目的Base模型和微调模型。重新预训练书接上文。在预训练了一个305M的模型之后我尝试使用全量SFT和LoRA微调两种方式进行微调但是效果都很差甚至没法输出一段逻辑连贯的话。与之相对比的tiny-llm-zh里面的92M模型能够流畅地完成问答和对话。于是我开始停下来反思其中的原因。由于两者在模型架构上几乎相同所以问题只可能是数据和训练量。在数据方面tiny-llm-zh用的是中文百科数据百度百科和Wiki CN数据质量会干净很多也天然的适合做常识类的QA任务而我们的Open Web Text数据集相对而言更杂乱数据量比百科要少数据分布也更为稀疏更不利于模型学习到一些常识类的知识。在训练量方面tiny-llm-zh的92M模型训练了9B个token这直接是模型参数量的100倍了相比之下我们的模型训练量还不到参数量的20倍这里的差距是很明显的。所以我决定使用tiny-llm-zh提供的中文数据来重新训练一个更小的模型。也就是遵循Chinchilla经验法则的指导在算力资源有限的情况下用更小的模型参数来换取更大的训练量。最终模型的参数量从305M缩减到了139M这里主要是缩减了d_model的大小从原来的1024调整到了768并且tokenizer也换成了tiny-llm-zh同款的chatglm3 tokenizer。此外还有一个比较大的改动就是模型的LM Head和Embedding层共享权重了。权重共享所谓权重共享就是指模型最开始的Embedding层和最后的LM Head层使用同一个矩阵作为参数。如下图所示Embedding层的参数是一个vocab_size * d_model的矩阵在输入阶段一个token会被编码成一个vocab_size大小的one-hot向量然后与vocab_size * d_model大小的Embedding矩阵相乘得到d_model大小的向量。然后这个向量会经过一系列Transformer层最终结果还是一个d_model维向量。最后这个d_model向量和之前那个Embedding矩阵的转置相乘便能回到vocab_size维向量。LM Head和Embedding层共享权重已经是业界一个比较常见的做法了。GPT系列LLAMA系列Qwen系列的部分模型都有使用权重共享。在transformers库中可以配置config.tie_word_embeddings True来启用权重共享。至于为什么能进行权重共享从直觉上来看Embedding和LM Head本质上都是在做查字典的操作Embedding负责把某一个具体的token映射到模型空间中而LM Head则负责把模型空间中的词映射回token空间。既然是字典那英译中和中译英使用同样的字典自然是很合理的。当然在深度学习的世界里只有直觉是不够的。《Using the Output Embedding to Improve Language Models》这篇文章通过实验验证了权重共享能起到一种正则化的效果能够提升模型的泛化性。从训练成本的角度来看权重共享本身也能减少参数数量从而减少训练成本。这个收益在小模型上更为明显我们模型的vocab_size是64798d_model是768那Embedding矩阵的大小就是49M占了模型参数量的35%如果能把这部分参数量节省下来就可以在同等参数量的情况下提升模型的深度和d_model数量从而提升模型的知识积累能力。混合精度训练接下来的几个小节会展示几个重新预训练时发现的新的小技巧以及遇到的一些坑。首先是使用混合精度进行训练。在之前的训练中我们一直都是在使用FP32精度来进行训练即一个浮点数占用4个字节。但是对于某些精度不敏感的操作例如LinearReLU在这些算子上面使用FP32会在增加显存占用的同时拖慢计算速度。❝这里有个参考数据在NVIDIA B200上FP32的最大吞吐量是80TFLOPS在FP16或者BF16上的最大吞吐量是2500TFLOPS这个差距是很大的这时我们就可以借助PyTorch提供的autocast来进行混合精度训练。同样的在PyTorch里接入混合精度和DDPtorch.compile一样简单不需要对模型做任何修改只需要使用一行代码套上一个Context即可示例代码如下所示with torch.autocast(device_typecuda, dtypetorch.bfloat16): logits model(input_ids, None) ce_loss cross_entropy_loss(logits, labels)在autocast上下文里PyTorch会根据算子类型自动调整算子的数据类型对于精度不敏感的类型会使用BF16对于精度敏感的类型例如包含求和规约的LayerNorm等则会使用FP32。在使用了混合精度之后训练速度提升是很明显的训练速度几乎是原来的2倍。与此同时显存占用也从原来的26GB降到了18GB这使得我们可以进一步减少gradient checkpoint的数量来提升训练速度。序列数据的采样方式这是这次预训练前期遇到的一个坑这个问题可以描述为数据集里是一系列连续的token但是训练时需要的是长度为1024的连续的token序列应该怎么采样这1024个token呢一个很直观的思路就是每次采样的时候随机选择一个起点然后从这个起点出发采样1024个token。这种方法很直观但是会有token利用率的问题即多次采样到的片段可能有重叠Overlap的部分这会导致实际训练到的不同的token数量小于理论值。下图是一个模拟实验的数据图这里假定总token数量为3.6B橙色曲线表示随机起点采样训练到的不同的token占数据集总量的比值即真实训练token覆盖率而蓝色图则是理论最大值即采样到的总token数量与总token数量的比值。❝注这里每一步采样16组长度为1024的token序列和预训练实验里的batch size对齐图中数据是蒙特卡洛出来的可以观察到在20w步的时候理论上模型应该已经见过88%的训练集数据了但是在随机起点采样的情况下模型才只见过58%的数据有相当一部分的算力被浪费在了训练已经见过的token上面。从第一性原理的角度来看模型没法回答出它没见过的问题因此这种情况下模型的表现必然会更差。说完反面案例接下来我们来看看正面案例。现在主流的方法是先分块再进行不放回采样具体两者的区别如下图所示。由于每次选择Block都是不放回采样所以这种方法的训练token覆盖率就严格等于理论最大值即不会有算力浪费在训练已经见过的token上面。这种方法的确也会存在一些问题例如一段连续的关键知识可能会被切分到两个block里。这个问题完全可以靠暴力scale来解决一方面现在大模型的预训练上下文长度在提升另一方面同一个知识在海量训练数据中肯定会出现多次这就导致模型训练时一次都见不到这个完整知识的概率是微乎其微的。DataLoader的shuffle的代价这是上面那个小节的一个补充也是踩到的一个坑。在309M那个时期我的Dataset的实现方式是把getitem里传来的index作为序列第一个token在数据集中的下标然后以此往后采样1024个token并使用了DataLoader提供的shuffleTrue来实现随机采样。当时遇到的一个很大的问题就是训练启动很慢而且很占内存。当时以为这个慢是因为torch.compile占内存是因为加载数据所以也没太在意。直到后面尝试把随机起点采样放到Dataset侧并取消了shuffleTrue时我才发现原来训练启动能这么快并且内存占用也比之前低了一个数量级。出现这个现象的原因在于shuffleTrue的实现方式DataLoader为了实现不放回采样会构造一个和Dataset相同大小的下标数组然后把这个数组shuffle。也就是说这里其实显式地构造了一个和数据集大小相同的下标映射。我们可以算一下构造这个映射的代价有多大由于数据集大小远大于序列长度1024所以这里直接认为数据集大小就是3.6B如果以Python的list来存储每个元素会占用28字节的int对象和8字节的指针那么总内存消耗能达到惊人的129GB。这里给我带来的启发是在使用DataLoader自带的shuffle之前要考虑数据集的大小。对于很大量的数据能在Dataset里直接做随机采样自然是最好的。如果不能那最好先对数据做分片在分片内进行shuffle以避免构造整个数据集的shuffle映射。在项目实际实现中我还是使用了Dataset分块DataLoader shuffle这是因为在分块之后Dataset的大小缩小为了以前的1/1024下标映射数组的大小也同步的下降了3个数量级这种代价已经完全是可接受的了。Tensorboard可视化这是一个很有用的数据记录工具。它可以在每个step中记录某个标量值例如当前的lrvalid loss等然后根据这些记录日志绘图。更重要的是Tensorboard的日志记录是不断追加的也就是说即使是中途从之前的某个checkpoint开始重新训练step发生了回退Tensorboard终端也会智能的处理这种情况并且之前记录的数据也不会丢掉。Tensorboard的一个常见用法如下所示from torch.utils.tensorboard import SummaryWriterwriter SummaryWriter(log_dir./runs/exp1)global_step 0for epoch in range(num_epochs): for batch in dataloader: loss train_step(batch) writer.add_scalar(train/loss, loss.item(), global_step) global_step 1writer.close()如果是在autodl上训练可以把日志目录设置为/root/tf-logs然后打开AutoPanel的TensorBoard就能直接查看相关数据了。如果是运行项目自带的webui这个webui还会把本机的TensorBoard端口转发到本机的6006端口这使得我们可以直接使用autodl提供的自定义服务链接来访问。过拟合实验排查明显错误我们经常会遇到这种情况在训练了一段时间之后模型的表现仍然不好此时可能会纠结要不要继续训练下去。这时候可以考虑先用少量重复数据来做一个过拟合实验以确认模型效果不佳不是代码实现错误导致的当然在这之前最好先用AI检查一轮看看有没有很明显的实现错误。例如可以收集10条“法国首都是巴黎”这种数据然后反复训练到过拟合确认loss能下降到0.1以下的水平然后用数据集里存在的Prompt做自回归推理。经过这一步检查之后基本就可以排除大部分代码实现错误了。虽说这样检查之后也不能保证继续训练下去效果会变好但是至少能够避免事后发现Label没有移位之类的低级错误而血压飙升。观察参数更新率在遇到train loss震荡但是valid loss还在缓慢下降时可以考虑观察一下参数的实际更新率以此来判断当前的学习率是否合适有没有出现死掉的参数。参数更新率Optimizer Update Ratio的计算方法是其中norm是取的二阶范数。本项目自带了打印参数更新率的功能把json里的print_optimizer_update_ratio设置为true即可。一般而言在预训练中参数更新率在1e-4左右是个比较健康的范围如果太小了就需要考虑调大一些学习率。不同的情况可能会有所变化具体的可以考虑和GPT老师交流一下这里只提供一个思路结果展示最终在训练了9.3B个token之后模型的valid loss来到了2.68。训练时的lr曲线和loss曲线如下图所示并且模型在few shot的情况下能够表现出一定的知识储备和泛化能力。以下是一些实际的例子日期推理英文翻译地理常识分类常识不过当前依然有很多case效果表现不好一个典型的例子就是模型几乎没有算数能力并且现在base模型还不能进行QA问答如果Prompt是类似于“介绍一下地球”的指令模型会答非所问。小结预训练的主要目的是让模型具有语言能力和世界知识从这个角度来看我们的139M中文模型已经具备这种能力了。在预训练过程中我最深的一个感悟就是数据才是决定模型效果的第一因数据的重要性远远超过炼丹本身。即使是大模型也依旧遵循Rubbish in, Rubbish out的法则。之前的305M模型的训练数据是OWT模型能够接一些日常对话但是没法对一些知识类问答做出响应这次139M模型的训练数据是中文百科所以这次模型就有了回答一些百科知识的能力。如果按照这个思路推广开来一个足够体量的模型加上足够强大的算力基建和全网大量的数据就能够训练出一个表现出“涌现”能力的大模型。LoRA微调预训练后的Base模型只是有一定的语言能力和知识还不能直接和人类对话。这其中的本质原因在于预训练阶段模型只见过大量连续的文本几乎没有见过日常对话的示例。例如预训练数据里面可能有“北京是中国的首都”但是没有“问中国的首都是哪 答北京”。所以我们需要通过微调来修正模型在面对各种指令时的答案分布。这里我首先尝试了使用LoRA技术来进行微调尝试之后发现效果比较差于是最后使用了全量微调。开源的代码仓库里有两次微调的结果权重感兴趣的读者可以自行体验一下两者的差别。这一章节主要介绍一下LoRA微调的原理实现以及效果。原理在介绍LoRA之前首先需要回顾一下传统的全量微调是如何做到的。传统的全量微调流程和预训练几乎相同微调的数据集是由一组一组的问题-答案对组成如下所示{question: 为以下文本生成一个摘要在明天的会议上我们将讨论公司的财务状况并提出几个建议来改进状况。\n在此会议上我们会讨论公司目前的财务状况以及如何改进它。,answer: 会议的议程将会涉及公司的财务状况和如何改进它。}微调的时候首先把数据组织成|user| question |assistant| answer |eos|的形式其中形如|xxx|的是一个special token。然后把数据送入模型得到预测结果把label的问题部分设置为Ignore Label计算交叉熵反向传播更新参数over。关于Ignore Label在上一篇文章里有过详细介绍本质是为了让模型预测答案而不是复读问题可以参考下面这张图理解一下这里不再赘述通过上面的描述我们可以看出全量微调所需要的显存等资源和预训练几乎是相同的。如果我们想要微调Qwen 32B这种体量的模型即使是在使用了ZeRO-3的情况下也需要4-8张A100-80G。于是LoRA诞生了。LoRA的全称是Low-Rank Adaptation其基本思想是在原模型的参数基础上再外挂一个参数量更小的参数然后在微调时冻结模型原本的参数只训练这个参数量更小的外挂参数。一个具体的例子如下图所示对于需要微调的Linear层LoRA会额外新增两个参数矩阵A和B其中B的形状是(in_features, r)A的形状是(r, out_features)两个矩阵相乘正好能得到和参数矩阵形状相同的(in_features, out_features)所以可以把BA认为是一个外挂的参数参与线性层的计算。在现实中一般in_features和out_features的数量级都是几千r一般取8或者16此时LoRA引入的增量参数A和B的大小是远小于原始参数矩阵的参数大小的。这就是LoRA微调节省显存的秘密所在。实现在实现层面需要做两件事。首先是把模型里需要微调的线性层换成上述的LoRA线性层。对于Transformer模型一般的处理方式是对Multi Head Attention处的QKV投影层进行微调如下图所示。然后就是需要把不需要参与训练的原始模型参数进行冻结这一步在PyTorch里实现起来很简单只需要把不需要更新的参数的requires_grad设置为False即可用下面这两行代码就能完成for name, param in self.named_parameters(): param.requires_grad False是否冻结Embedding我们在微调的时候会引入预训练数据中没有的一些special token例如|user|在这种情况下如果依然冻结Embedding显然是不合理的因为模型在预训练时没有见过这些special token其对应的Embedding层自然也没有学会关于这些token的语义。但是另一方面我们又不希望在微调时改动其他token对应的Embedding层参数也就是说我们其实希望只训练这些新增的token对应的Embedding。这里HuggingFace的transformers库提出了一种只训练部分Embedding参数的解决方案把Embedding层参数的requires_grad设置为True但是更新参数之前把除了新增token以外的行的梯度置零。这个方法非常简单粗暴在得知这个之前我一直以为transformers库的实现是像LoRA那样新增了一个小的Embedding参数在训练时把这个小参数的requires_grad设置为True最后把两个参数合并。不过在本项目的代码实现中Embedding层依然是全量微调的。感兴趣的读者可以自行尝试一下如果在LoRA微调时只微调那几个special token对应的Embedding效果会不会更好。效果展示LoRA微调之后其在SFT验证集上的loss为2.83训练过程的valid loss曲线如下图所示虽然loss很好看但是模型的实际表现很糟糕。对于一些在全量微调下能很好回答的问题LoRA微调之后无法正常回答例如下图中的“介绍一下玫瑰花”模型直接复读了玫瑰花然后就输出eos了。但是很奇怪的是如果在问题后面加一个问号那模型的表现又不同了还有一个勉强算是输出了一些有意义的内容的例子小结LoRA微调的效果并不是很好可能主要原因是Base模型本身能力不够。一般来说我们都是对Qwen 32B这类比较强参数量也比较大的模型进行LoRA微调的这里LoRA微调也只是起到一个演示作用了。全参数量微调在LoRA微调失败以后我又尝试了全参数量微调最终valid loss降到了2.63训练时的valid loss和lr曲线图如下所示可以发现valid loss在1w步左右的时候出现了一个peak这是在微调过程中需要非常小心的一个点微调时特别容易过拟合所以需要用更小的学习率来训练一般是训练时的1/20到1/30。结合右图的lr曲线可以看得更明白我在1w步以前还是用的cosine lr scheduler在1w步的时候观察到了loss反弹于是手动把学习率固定在了5e-7。最终的自回归效果如下可以发现模型已经能够用逻辑连贯的流畅的语言回答问题并且也学会了用言简意赅的短语来回答问题但是依然还存在一些复读的问题小结目前全量微调之后的模型效果已经足够惊艳了“训练出一个能说人话的模型”是我这个项目的最终目标现在它无疑是已经做到了。学AI大模型的正确顺序千万不要搞错了2026年AI风口已来各行各业的AI渗透肉眼可见超多公司要么转型做AI相关产品要么高薪挖AI技术人才机遇直接摆在眼前有往AI方向发展或者本身有后端编程基础的朋友直接冲AI大模型应用开发转岗超合适就算暂时不打算转岗了解大模型、RAG、Prompt、Agent这些热门概念能上手做简单项目也绝对是求职加分王给大家整理了超全最新的AI大模型应用开发学习清单和资料手把手帮你快速入门学习路线:✅大模型基础认知—大模型核心原理、发展历程、主流模型GPT、文心一言等特点解析✅核心技术模块—RAG检索增强生成、Prompt工程实战、Agent智能体开发逻辑✅开发基础能力—Python进阶、API接口调用、大模型开发框架LangChain等实操✅应用场景开发—智能问答系统、企业知识库、AIGC内容生成工具、行业定制化大模型应用✅项目落地流程—需求拆解、技术选型、模型调优、测试上线、运维迭代✅面试求职冲刺—岗位JD解析、简历AI项目包装、高频面试题汇总、模拟面经以上6大模块看似清晰好上手实则每个部分都有扎实的核心内容需要吃透我把大模型的学习全流程已经整理好了抓住AI时代风口轻松解锁职业新可能希望大家都能把握机遇实现薪资/职业跃迁这份完整版的大模型 AI 学习资料已经上传CSDN朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】