planckgpt

PlanckGPT是一个教育项目，演示了如何在消费级硬件上从零训练GPT风格的语言模型。该模型约有5.4亿参数，在来自Fineweb-edu数据集的约20亿个token上进行预训练，并可选使用Smol-smoltalk数据进行聊天微调阶段。 该项目旨在实现最佳的质量与训练时间比，软限制为50小时的训练时间。在作者的笔记本电脑（Intel Core i7 14700HX，24GB RAM，RTX 5070 Mobile 8GB）上，预训练约需28小时20分钟，聊天微调约需4小时10分钟。 架构使用14层Transformer解码器，7个查询头，896维嵌入，旋转位置嵌入，交替值嵌入，平方ReLU激活，以及RMSNorm。分词器是使用GPT-2编码的Tiktoken（50,257词汇量）。训练采用NorMuon优化器用于Transformer权重，8位Adam用于嵌入，BF16混合精度，梯度检查点，以及torch.compile。 该项目提供预训练、评估、聊天微调和推理脚本。预训练性能接近GPT-2-medium，在Fineweb-edu上的平均验证损失约为2.97。聊天微调后的性能在Smol-smoltalk上达到约1.13的平均验证损失。推理使用top-k和top-p采样，并带有KV缓存以加速生成。 该项目受modded-nanogpt和nanochat启发，采用Apache 2.0许可证。