Mini-LLM 项目旨在基于有限的算力，复现主流开源模型架构，从零手撕一个 100-200M 参数量的迷你 LLM，并完成预训练和微调，使其具备基本的对话能力。 本次从零复现迷你版 DeepSeekV4架构，总参数量大约160M，激活参数量大约80M。不得不说，虽然论文读起来CSA和HCA的原理还是容易理解的，但是上手实现就会发现中间有很多细节。代码实现参考了官方inference源码和transformers近期更新的DeepSeekV4代码。这期间也发现了transformers官方代码实现上的一些问题并提了issue。相比官方代码，本项目的手撕代码大概有以下几个特点： 1.官方inference源码和transformers源码的Compressor都不支持padding，但是如果我们要从零训练，尤其是SFT，就需要让一个batch的数据pad到相同长度。因此，本项目专门为Compressor加入了padding逻辑。 2.为了兼容transformers，Attention部分主要参考了transformers实现，但transformers实现的HCA、CSA，包括它们的CacheLayer是分开的，本项目将它们统一起来。MoE部分主要参考官方inference源码实现，也跟我们之前手撕的Mini DeepSeekV3差别不大，并且均实现了无辅助损失负载均衡、序列级辅助损失和MTP。 3.Indexer由于topk的逻辑，端到端的训练导致梯度无法回传。DeepSeek在预训练时，应该是分为了几个阶段，先让indexer进行warm up与dense attention对齐，然后再继续加入到训练中（V3.2的DSA对这个过程介绍更详细些）。本项目为了复用固有的pretrain代码从而简化实现，为topk添加了score bias和随机选择槽位，让indexer得以收到梯度。 4.前3层使用hash routing的MoE，这里本项目为了简单，没有从语料中统计token id分布，而是直接将词表中的token id按均衡的分布分配给各个专家。 P.S. 由于最近算力有点紧张，这次用了4卡A100（40GB）训练了2B token，用了13小时左右，大概训出个差不多的效果就收尾吧。另外，自己实现的Mini DeepSeekV4确实推理速度感觉会慢一些，因为中间涉及到的Compressor、Indexer都没有优化过，infra确实很重要啊。 项目地址：https://github.com/WKQ9411/Mini-LLM https://huggingface.co/WKQ9411 如果对你有帮助，欢迎给个star🌟 #大模型 #大模型算法 #deepseek #人工智能