始智AI wisemodel.cn开源社区
随着开源数据的日益丰富以及算力价格的持续下降，对于个人或小型机构而言，预训练一个小型的 LLM 已逐渐成为可能。开源中文预训练语言模型 Steel - LLM 就是一个典型案例，其模型参数量与数据量并非十分庞大，基本处于参数量为 B 级别、数据量为 T 级别的规模。
本文将着重阐述项目实施过程中所遭遇的问题、复盘后的思考以及相关技术细节，期望能为在资源有限条件下开展 LLM 训练的开发者们提供一定的启发与助力。该模型已上线始智AI-wisemodel开源社区，欢迎大家前去体验。
模型和代码地址：
https://wisemodel.cn/models/zhanshijinwat/Steel-LLM
https://wisemodel.cn/codes/zhanshijinwat/Steel-LLM
01
数据收集与处理
<|im_start|>systemYou are a helpful assistant.<|im_end|><|im_start|>user{问题}<|im_end|><|im_start|>assistant{回答}<|im_end|>
在预训练阶段加入SFT数据的方法在minicpm的训练过程中也使用了，但是只在预训练末期的退火阶段加入了SFT数据，这种方法预训练出来的模型指令跟随能力应该会比较强。如果预训练全程都杂糅SFT数据，这些少量的SFT数据会淹没在海量的原始文本中，起的作用会比较小。
Steel-LLM的预训练没有退火阶段，退火阶段的的核心是“高质量”数据，但是高质量的定义我觉得还是比较模糊的。不同类型的数据配比合理可以理解为是一种“高质量”，但是开源数据往往并没有给出数据的类型。经过各种规则筛选的数据也可以被称为高质量数据，但是开源数据基本都是经过数据清洗pipeline处理的。
02
训练框架
兼容HuggingFace格式的模型：TinyLlama使用的模型结构是定义在项目里边的，如果使用TinyLlama项目训练其他模型不是很方便。笔者将其修改为使用HuggingFace的Transformers库的模式定义模型，开发者可以方便的更换自己想训练的模型结构。 数据训练进度恢复：预训练时间长，难免会出现中断的情况，想要继续训练的话，不仅要恢复模型和优化器的状态，还需要恢复数据训练的进度。TinyLlama main分支提供的恢复数据训练进度的方法比较简单粗暴，保存模型checkpoint时候也记录下迭代轮数，加载checkpoint之后直接跳过之前迭代过的轮数的数据，这种方式的的数据恢复要求预训练过程数据不能有变动。笔者对这块进行了进一步改进，保存所有数据的文件名、每条数据的索引，以实现数据进度的精准恢复。 支持训练过程中追加数据：预训练时间比较长，不免会有追加新数据的需求。基于第2点“数据训练进度恢复”改动，实现了新追加的数据索引会和老数据中未训练的数据索引（图中红色数字表示）重新shuffle的功能，防止加入新数据后，新老数据分布差异过大，影响后续的模型训练效果。具体原理如下图所示：
为了防止数据块被意外的重复追加到训练过程中，Steel-LLM还实现了使用hash值检测数据内容是否重复的功能，具体使用的是MD5哈希算法。对一个1.6g的文件进行哈希大约要9s，效率太低。因此选择只用一块数据的头部和尾部数据计算哈希值。
对self-attention进行算子融合的实现就是大名鼎鼎的flash attention。在我的往期文章中，笔者对各个算子的融合进行了消融，使用算子融合整体上能提升50%的训练效率，以及节省12%的显存，MFU（指模型一次前反向计算消耗的算力与机器能够提供的算力的比值）也能从43%提高到63%，大大提高了GPU的利用率。
03
模型结构
（2）SENet
SENet（Squeeze-and-Excitation Networks）来自于计算机视觉领域ImageNet 2017竞赛的冠军方案，目的是通过学习的方式来获取到每个图像特征通道的重要程度，通过重要程度加强或抑制特征通道，经过简单变换之后可变为如下图所示的形式：
我们再来看看Qwen模型的FFN层实现，大致可分为两层，它的第一层就是SENet的思想，用了gate_proj和up_proj，FFN的第二层也换成了SENet。
class Qwen2MoeMLP(nn.Module):def __init__(self, config, intermediate_size=None):super().__init__()self.config = configself.hidden_size = config.hidden_sizeself.intermediate_size = intermediate_sizeself.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)self.act_fn = ACT2FN[config.hidden_act]def forward(self, x):return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
04
训练过程
预训练输入的最大序列长度是2048，将收集的数据训练了两个epoch，大概1.1T个token，需要训练1M个step。如果使用SXM版本的A100 80G*8，大概需要60天，H800*8需要30天左右。训练时的batchsize 为8，梯度累计步数为8。优化器使用的AdamW，最大学习率时3e-4，学习率衰减策略是CosineAnnealingWarmUp。训练的wandb链接如下：
https://api.wandb.ai/links/steel-llm-lab/vqf297nr因为是单机8卡，训练还是比较稳定的，没有出现过啥error。唯一一次故障重启是机器网断了，wandb把训练进度给卡住了。
05
微调和评估
微调阶段数据量不大，为了方便直接用llama factory训练了，选择如下4份数据：
BAAI/Infinity-Instruct：今年8月发布的比较新的中英文sft数据，BAAI又是开源社区比较靠谱的机构，果断用起来。尽管存在有“做sft有少量高质量数据就行”的说法，但就小模型而言，预训练阶段学到的内容相对匮乏，所以在 sft 阶段，数据多点予以补充完善。Infinity-Instruct 涵盖 700 万条数据，规模颇为可观，足以应对相关需求。
wanjuan中文选择题部分：wanjuan中文选择题是预训练阶段就有的数据，因为微调之后的模型要在ceval等数据集上测试，因此sft阶段回炉重造一下，提高一下做题能力。
ruozhiba：前段时间大火的数据集，从百度贴吧“弱智吧”扒的问题，用GPT4生成答案。
自我认知数据：希望让训练完的模型知道自己是Steel LLM，而不是回答自己是智能助手。模板来自于EmoLLM项目。
Steel LLM的预训练数据里80%以上都是中文的，所以只测了ceval和cmmlu这两个中文benchmark。第一版实验用了700w条全量的Infinity-Instruct数据，ceval能有33%的准确率。
后来发现Infinity-Instruct里90%数据都是英文的，和预训练数据分布严重不符（预训练数据里只有20%英文）。之后从Infinity-Instruct里抽出了70w的中文数据，并糅合其他3个数据集，最终在ceval取得了38%准确率，cmmlu取得了33%准确率。
同时，还做了刷榜测试，直接将cmmlu数据也放到sft数据里。在cmmlu上的正确率从33%提高了36%，在ceval上的准确率几乎没变化。这说明让模型去死记硬背答案也没那么容易。
cmmlu在sft训练时候的标签只有一个选项字母，如果标签中有答案的解释，应该会效果更好一些。除此之外，还尝试了一下让模型以COT的方式进行答案生成，即先输出解释再输出答案，发现在benchmark上并没有拿到更好的结果。
本次项目属于个人项目，在精力和资源有限的情况下，还存在一些没做到位的地方，比如训练tokenizer、数据配比探究、全局数据清洗、模型英文能力较弱等。后续依算力情况，还将基于Steel LLM做一些微调方面的探索（样本筛选等）、强化学习或者是VLM，欢迎持续关注。
关于wisemodel更多