个人从零预训练 1B LLM 实践与总结

训练框架

训练代码这块主要借鉴了 TinyLlama。有群友问为何不用 Megatron，最根本原因是 Steel LLM 是个 1B 的模型，比较小，没有必要用 Megatron。首先，Megatron 支持 pipeline 并行、tensor 并行，对训练大尺寸模型友好，但是在 A100 上，1B 的模型使用原生 PyTorch 的 FSDP 甚至是 DDP 就能稳定训练。其次，Megatron 提供的算子融合、dataloader 等能力，其他项目也有，而且简洁，更好集成和修改。

基于 TinyLlama 训练代码进行了如下几点改进：

• 兼容 HuggingFace 格式的模型：TinyLlama 使用的模型结构是定义在项目里的，如果使用 TinyLlama 项目训练其他模型不是很方便。将其修改为使用 HuggingFace 的 Transformers 库的模式定义模型，开发者可以方便的更换自己想训练的模型结构。
• 数据训练进度恢复：预训练时间长，难免会出现中断的情况。想要继续训练的话，不仅要恢复模型和优化器的状态，还需要恢复数据训练的进度。TinyLlama main 分支提供的恢复数据训练进度的方法比较简单粗暴，保存模型 checkpoint 时候也记录下迭代轮数，加载 checkpoint 之后直接跳过之前迭代过的轮数的数据，这种方式的数据恢复要求预训练过程数据不能有变动。进一步改进为保存所有数据的文件名、每条数据的索引，以实现数据进度的精准恢复。
• 支持训练过程中追加数据：预训练时间比较长，不免会有追加新数据的需求。基于第 2 点'数据训练进度恢复'改动，实现了新追加的数据索引会和老数据中未训练的数据索引重新 shuffle 的功能，防止加入新数据后，新老数据分布差异过大，影响后续的模型训练效果。具体原理如下图所示：

• 防止数据重复：为了防止数据块被意外的重复追加到训练过程中，Steel-LLM 还实现了使用 hash 值检测数据内容是否重复的功能，具体使用的是 MD5 哈希算法。对一个 1.6G 的文件进行哈希大约要 9s，效率太低。因此选择只用一块数据的头部和尾部数据计算哈希值。

对于预训练来说，进行模型的训练效率优化是必不可少的。卡少的情况下，即使预训练个小模型也得跑好几十天，那么训练效率优化个 10%，就能省下好几天的电钱呢。个人或者小机构训练模型很难去优化底层的分布式训练机制，那么性价比最高的优化方式就是对模型的部分算子用算子融合。模型在训练时候，除了 gpu 的计算时间外，从显存把数据搬运到缓存也占用了很多时间。举个例子，算一个最简单的标准化 (x-mean)/var，需要涉及到如下步骤：1.读 x+ 读 mean 2.写 x-mean 3.读 x-mean，读 var 4.写 (x-mean)/var，看起来十分的啰嗦，如果想让 gpu 端到端的计算结果（只读一次写一次），就需要做算子融合了，用 cuda 编程或者 triton 编程都行。在 LLM 中，RMSNorm、RoPE、self-attention、交叉熵损失函数都是频繁读取数据的大户，因此能把这些操作融合掉能提高不少的训练时间。对 self-attention 进行算子融合的实现就是大名鼎鼎的 flash attention。在我的实验中，对各个算子的融合进行了消融，使用算子融合整体上能提升 50% 的训练效率，以及节省 12% 的显存，MFU（指模型一次前反向计算消耗的算力与机器能够提供的算力的比值）也能从 43% 提高到 63%，大大提高了 GPU 的利用率。

分布式训练方式上，沿用了 FSDP（Fully-sharded data-parallel，完全分片数据并行）方法。它的原理和 deepspeed 的 ZeRO stage3 一致，FSDP 是 pytorch 的官方实现，将优化器参数、模型参数和梯度分布存储到不同的 GPU 上，节省分布式训练时候的显存，模型越大，显存节省的越多。

模型结构

公司训练模型大部分都是沿用传统的 LLM 结构（self attention、RMSnorm、RoPE 位置编码等），一来是 LLM 的效果主要是靠数据，二来是 team leader 也没必要去冒风险搞一些不一定有确定收益的创新。笔者训练 Steel LLM 没有任何的 KPI 压力，也愿意尝试进行一些结构上的修改，比较遗憾的是，没有算力再训练一遍传统结构的 LLM，去消融做的结构改动是否是好的。Steel LLM 的 tokenizer 直接使用的是 qwen 的，并没有重新训练。

最开始应群友建议，尝试了下训练 recurrent gemma 结构的模型，只不过它的 pytorch 实现训练效率太差了，和训 transformer 结构相比慢了几十倍，遂放弃。gemma 是 google 训的模型，他们在 TPU 上训练时候做了不少的工程优化，训练效率才是可接受的。对于主流的 LLM 结构来说，self attention 和 FFN 是两大核心模块。目前 self attention 这块的实现基本都用 flash attention，修改 self attention 的实现比较困难。之前测试，如果不使用 flash attention 会掉 25% 左右的训练效率。因此，魔改结构的重点放到 FFN 上边。Steel LLM 在 FFN 上的修改有两部分，soft MOE 和 SENet。

Soft MOE

很多大型号的 LLM 使用 hard MOE（每个 token 选择 top k 个 FFN），目的是减少训练和推理时的计算量，但是 hard MOE 并不省显存，在训练和推理时仍然需要将完整的模型加载到显存中。考虑到只有 1 台机器训练 Steel LLM，显存并不富裕，模型也比较小，想充分训练每一份参数，并且 hard MOE 并不好训练，需要考虑高效的 token 路由实现、训练稳定性以及负载均衡等问题。hard MOE 存在种种问题，但笔者仍然想尝试一下 MOE 结构，因此使用了 soft 形式的 MOE，这在搜广推多任务学习上被广泛使用。

soft MOE 的原理如下图所示，input 就是一个向量，通过不同的专家网络计算出结果，然后再通过 input 计算出来的权重进行加权求和。在 LLM 中的 FFN 层使用 soft MOE 原理类似，只不过有多个 input（即多个 token），分别通过各个专家并进行加权求和。

尝试了 3 版 soft MOE，第一版训练效率太差还费显存，第二版 loss 无法收敛到正常水平，最终尝试的第三版实现在效率和效果上才满足要求。这三版实现都在 github 仓库中。

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 = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = intermediate_size
        self.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))

训练过程

预训练输入的最大序列长度是 2048，将收集的数据训练了两个 epoch，大概 1.1T 个 token，需要训练 1M 个 step。如果使用 SXM 版本的 A100 80G*8，大概需要 60 天，H800 需要 30 天左右。训练时的 batchsize 为 8，梯度累计步数为 8。优化器使用的 AdamW，最大学习率时 3e-4，学习率衰减策略是 CosineAnnealingWarmUp，这些都没啥特别的。训练的 wandb 链接如下：

https://api.wandb.ai/links/steel-llm-lab/vqf297nr

因为是单机 8 卡，训练还是比较稳定的，没有出现过啥 error。唯一一次故障重启是机器网断了，wandb 把训练进度给卡住了。

24 年 6 月初，在 Steel LLM 训练了 200k step 的时候，朋友搞到的 8 卡 A100 被收回了。供应方说是临时拿走做下推理测试，但是之后再也没给回来。眼看 Steel LLM 就要成为烂尾项目，但非常的幸运，项目搁置一周多之后，某合作机构借了一个多月的 1 台 H800 用，项目才得以完成。

微调和评估

24 年 8 月中旬，Steel LLM 预训练完成，打算开始做下后训练。目前先只考虑做 sft，强化学习又是个大坑，工作有点卷，sft 也是到了 10 月底才做完的，强化学习如果也要做一下估计得等到明年才能发布了。

微调阶段数据量不大，为了方便直接用 llama factory 训练了，选择如下 4 份数据：

• BAAI/Infinity-Instruct：今年 8 月发布的比较新的中英文 sft 数据，BAAI 又是开源社区比较靠谱的机构，果断用起来。虽然有'做 sft 有少量高质量数据就行'的说法，但是觉的小模型做 sft 还是数据多点好一些，毕竟预训练阶段学到内容相对较少，sft 阶段补一补。Infinity-Instruct 有 700w 条数据，量大管饱。
• wanjuan 中文选择题部分：wanjuan 中文选择题是预训练阶段就有的数据，因为微调之后的模型要在 ceval 等数据集上测试，因此 sft 阶段回炉重造一下，提高一下做题能力。
• ruozhiba：前段时间大火的数据集，从百度贴吧'弱智吧'扒的问题，用 GPT4 生成答案。
• 自我认知数据：希望让训练完的模型知道自己是 Steel LLM，而不是回答自己是智能助手。模板来自于项目。

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 上并没有拿到更好的结果。

Steel LLM 模型和其他部分模型在 ceval 和 cmmlu 上的对比如下所示，其他模型的结果出自 ceval、MiniCPM、MAP-Neo、CT-LLM 的论文或技术报告。

经验总结与展望

随着开源数据的不断增多、算力价格不断降低，个人或小机构正经预训练一个小型号的 LLM 并不是遥不可及，希望 Steel-LLM 系列博客能对您在资源受限的情况下训练 LLM 产生一些启发或帮助。

头一次做耗时这么长的个人项目，精力和资源有限，确实存在一些没做到位的地方，比如：训练 tokenizer、数据配比探究、全局数据清洗、模型英文能力较弱等。后续依算力情况，打算基于 Steel LLM 做一些微调方面的探索（样本筛选等）、强化学习或者是 VLM。

关于显存优化的补充：在实际训练中，显存管理是重中之重。除了使用 FSDP 之外，混合精度训练（AMP）也是基础配置。在 PyTorch 中开启 amp 模式可以自动混合 FP16 和 BF16，通常配合 torch.cuda.amp 使用。对于 1B 模型，BF16 通常能提供比 FP16 更好的数值稳定性，尤其是在训练后期 loss 震荡时。此外，激活检查点（Activation Checkpointing）技术可以在不增加太多计算开销的前提下，显著减少显存占用，允许更大的 batch size 或更深的模型。

关于数据质量的思考：数据质量永远优于数据量。虽然本项目使用了 T 级数据，但在清洗过程中，去除低质、重复、有害内容是关键。简单的规则过滤可能不够，结合 perplexity 筛选或人工抽检能进一步提升数据纯度。特别是在加入 SFT 数据时，确保指令遵循的多样性至关重要，避免模型过拟合某种特定格式。

关于硬件选择的建议：对于个人训练者，A100 80G 是性价比极高的选择，单卡显存大，适合 FSDP 切分。如果是消费级显卡，如 4090，由于显存限制，可能需要更多的卡进行数据并行，或者使用量化技术（如 INT8/INT4）进行训练，但这会增加工程复杂度。H800 虽然性能更强，但受限于供应和成本，并非所有人都能轻易获得。

未来方向：下一步计划探索 LoRA 等高效微调技术在 1B 模型上的表现，看是否能以更小的数据量达到接近全量微调的效果。同时，尝试引入多模态能力，构建 VLM，扩展模型的应用边界。