🚀 让我们来从头训练一个属于自己的Mini-RWKV-7吧~ 小小的鹅也能飞得很高喔~

🤗 Pure RWKV V7 HuggingFace Space~~~😜 Have a Try?

🤗 RWKV V7 with Deep Eembedding HuggingFace Space~~~😜 Have a Try?

🤗 我自己练的模型权重现已开源~

本模型是基于 RWKV-V7 架构 训练的一个 34M 参数量 的语言模型Mini-RWKV-V7-LM-34M。它在保持轻量的同时，具备良好的语言理解和生成能力，非常适合资源极其有限的设备部署和快速迭代开发。

💫 现已经支持训练

• RWKV V7架构(Pure RWKV V7)
• RWKV V7a DE架构(RWKV V7 with Deep Eembedding)
• RWKV V7b DEA架构(RWKV V7 with Deep Eembedding Attention)

只要在demo-training-prepare.sh | demo-training-run.sh | demo-training-run-sft.shbash脚本里面的

MODEL_TYPE="x070"中修改，可以调整模型架构

# Note: "x070" is Pure RWKV-V7
#       "x070#DE" is RWKV-V7 with Deep Eembedding
#       "x070#DE#DEA" is RWKV-V7 with Deep Eembedding Attention

• Vocab 和MiniMind的保持一致

INFO:pytorch_lightning.callbacks.model_summary:
  | Name   | Type       | Params | Mode  | FLOPs
------------------------------------------------------
0 | emb    | Embedding  | 3.3 M  | train | 0    
1 | blocks | ModuleList | 27.6 M | train | 0    
2 | ln_out | LayerNorm  | 1.0 K  | train | 0    
3 | head   | Linear     | 3.3 M  | train | 0    
------------------------------------------------------
34.2 M    Trainable params
0         Non-trainable params
34.2 M    Total params

• Vocab 和MiniMind的保持一致

INFO:pytorch_lightning.callbacks.model_summary:
  | Name   | Type       | Params | Mode  | FLOPs
------------------------------------------------------
0 | emb    | Embedding  | 3.3 M  | train | 0    
1 | blocks | ModuleList | 84.9 M | train | 0    
2 | ln_out | LayerNorm  | 1.0 K  | train | 0    
3 | head   | Linear     | 3.3 M  | train | 0    
------------------------------------------------------
91.5 M    Trainable params
0         Non-trainable params
91.5 M    Total params

• Vocab 和MiniMind的保持一致

INFO:pytorch_lightning.callbacks.model_summary:
  | Name   | Type       | Params | Mode  | FLOPs
------------------------------------------------------
0 | emb    | Embedding  | 3.3 M  | train | 0    
1 | blocks | ModuleList | 128 M  | train | 0    
2 | ln_out | LayerNorm  | 1.0 K  | train | 0    
3 | head   | Linear     | 3.3 M  | train | 0    
------------------------------------------------------
135 M     Trainable params
0         Non-trainable params
135 M     Total params

• 🪿 架构：RWKV-V7
• 📚 数据源：minimind_dataset 特别感谢MiniMind的作者 @jingyaogong开源了训练数据集 🤗
• 📈 学习率：动态调整
• 🖥️ 硬件：可以使用4060laptop等显卡进行训练，甚至Radeon 780M 核显也可以在轻薄本上进行训练 😜
• 👀我是在AMD Instinct MI300X 上快速复现的(十分感谢AMD公司的对我个人以及RWKV的云算力赞助)😊
• 📦 模型大小：68.4MB 参数量 34.2M Params
• 📊 预损失曲线：预训练收敛稳定 loss = 2.12左右波动(因为预训练数据量比较少)
• 📊 SFT训练损失曲线 SFT训练最终loss=0.5左右波动

• 📝 预训练（Pre-training）
• 📚 监督微调训练（Supervised fine-tuning [SFT]）

pip install -r requirements.txt

• 如果你使用的是AMD Instinct MI300X等别的AMD显卡，请不要安装requirements.txt中的torch，请安装对应最新版本的torch
• 比如说pip3 install torch --index-url https://download.pytorch.org/whl/rocm6.3
• 具体安装指令可以参考Pytorch官网下载链接

python3 ./API_DEMO_CHAT.py

modelscope download --dataset gongjy/minimind_dataset sft_512.jsonl sft_1024.jsonl sft_2048.jsonl pretrain_hq.jsonl sft_mini_512.jsonl --local_dir ./minimind_dataset

python3 ./data/make_data.py pretrain_hq.jsonl 1 512

• 这里的1是只重复一遍数据集, 512是上下文长度

sh ./demo-training-prepare.sh

• 执行后会生成一个./out/L8-D512-x070文件夹，里面有初始化的模型权重文件,名为rwkv-init.pth

sh ./demo-training-run.sh

• 修改第28行的M_BSZ="16"可以调整batch size
• 修改第52行的--wandb ""可以启动wandb记录训练过程
• 训练完成后会生成一个最终权重./out/L8-D512-x070文件夹里面有训练好的模型权重文件,名为rwkv-final.pth

• 训练之前删除./out/L8-D512-x070文件夹里面的rwkv-init.pth初始权重文件
• 然后将./out/L8-D512-x070文件夹里面的rwkv-final.pth权重文件改名为rwkv-init.pth来作为初始权重文件进行Post-training继续训练
• 按照数据集Token数对照表的数据集信息修改对应的配置参数

CTX_LEN="512"
MY_EXIT_TOKENS="1843253579"
MAGIC_PRIME="3600053"
DATA_FILE=[your data file path]

sh ./demo-training-run-sft.sh

• 训练之前删除./out/L8-D512-x070文件夹里面的rwkv-init.pth初始权重文件
• 然后将./out/L8-D512-x070文件夹里面的rwkv-final.pth权重文件改名为rwkv-init.pth来作为初始权重文件进行Post-training继续训练
• 按照数据集Token数对照表的数据集信息修改对应的配置参数

CTX_LEN="1024"
MY_EXIT_TOKENS="1430197426"
MAGIC_PRIME="1396673"
DATA_FILE=[you data file path]

sh ./demo-training-run-sft.sh

• 训练之前删除./out/L8-D512-x070文件夹里面的rwkv-init.pth初始权重文件
• 然后将./out/L8-D512-x070文件夹里面的rwkv-final.pth权重文件改名为rwkv-init.pth来作为初始权重文件进行Post-training继续训练
• 按照数据集Token数对照表的数据集信息修改对应的配置参数

CTX_LEN="2048"
MY_EXIT_TOKENS="2398644915"
MAGIC_PRIME="1171199"
DATA_FILE=[you data file path]

sh ./demo-training-run-sft.sh

• LR_INIT=6e-4 LR_FINAL=2e-5

• LR_INIT=2e-5 LR_FINAL=1e-6

Total tokens in sft_512.jsonl:

• --my_exit_tokens 1843253579 --magic_prime 3600053 --ctx_len 512

Total tokens in sft_1024.jsonl:

• --my_exit_tokens 1430197426 --magic_prime 1396673 --ctx_len 1024

Total tokens in sft_2048.jsonl:

• --my_exit_tokens 2398644915 --magic_prime 1171199 --ctx_len 2048

• 🖥️ 感谢AMD公司的对我个人以及RWKV的云算力赞助
• 🙌 感谢 RWKV 社区提供的开源代码和训练框架！
• 🚀 感谢 MiniMind提供灵感！
• 如发现 bug 或有任何建议，欢迎提交 issue 或 PR 🛠️

• MiniMind:一个轻量级 LLM 教学项目 📚
• RWKV_v7_G1_Translate:一个基于RWKV-V7-G1全量微调的超轻量端侧中英翻译模型 💫
• RWKV-LM:RWKV 官方架构仓库 🧠

🎉 感谢小伙伴们使用 Mini_RWKV_7！如果你喜欢这个项目，欢迎推给大家一起来玩！🌟