RWKV-7-Arithmetic-0.1B 是基于 RWKV-LM 仓库的 RWKV-7 架构预训练的 0.1B(L12-D768)模型,可以实现基础加减法运算和加减法方程求解功能。
该模型能够处理简单的数字加减、混合中英文,甚至是额外混合全半角数字的加减法任务。其中整数部分允许 1-12 位、小数部分允许 0-6 位,支持中英文数字、全半角格式以及大小写字符等多种混合表示形式。
2025 年 7 月 24 日更新模型 ,支持在题目中使用科学计数法,例如
95-7.16219e+5
。
测试结果
2025 年 7 月 24 日新增:添加了 3M 的科学计数法数据进行训练,然后测试了 440 条科学计数法数据,该类数据正确率达到 94.55%。
我们利用 RWKV Runner 提供的接口对数据集进行批量测试。测试采用零样本(0-shot)方式,每条数据独立输入模型进行推理。
测试结果表明,模型整体正确率超过 99.2%,大部分测试文档达到 100% 准确率,即使表现最差的文档正确率也超过 97% 94.5%。
文件名 | 测试数据条数 | 正确条数 | 正确率 (%) |
---|---|---|---|
X_ch_en_mix_test_v2 | 500 | 499 | 99.80% |
X_ch_en_mix_test_v1 | 100 | 100 | 100.00% |
ADD_use-connect_test | 500 | 487 | 97.40% |
X_ch_mix_test | 500 | 499 | 99.80% |
X_ch_mix_test_hard | 500 | 500 | 100.00% |
ADD_en_mix_test | 488 | 487 | 99.80% |
ADD_en_test | 491 | 488 | 99.39% |
ADD_base_test | 500 | 500 | 100.00% |
ADD_en_base_test | 500 | 500 | 100.00% |
ADD_many0_test_v1 | 50 | 50 | 100.00% |
ADD_many0_test_v2 | 500 | 498 | 99.60% |
ADD_test | 499 | 497 | 99.60% |
ADD_test_hard | 463 | 457 | 98.70% |
ADD_random_test | 1000 | 997 | 99.70% |
Test_sci.jsonl | 440 | 416 | 94.55% |
全部文件 | 7131 | 7075 | 99.21% |
详细的测试实现可参考 RWKV-7-Arithmetic 数据集中的 Test.ipynb
测试脚本。
模型训练过程
训练数据集
我们使用一套数据生成代码生成了 15 种类型,23 个文件,共 30000587 行单轮加减法数据,然后使用脚本合并,最终构建成共 2297887 行的 10~20 轮随机多轮对话。
最终训练数据集共 1B(1014434168) token,文件大小约 2.33 GB。
数据集的细节可查看 RWKV-7-Arithmetic 数据集仓库。
训练方法(参数)
训练基于 RWKV-LM 仓库,完整的训练教程可见:RWKV 预训练教程
具体的训练参数设置:
MODEL_TYPE="x070"
N_LAYER="12"
N_EMBD="768"
CTX_LEN="1024"
PROJ_DIR="out/L"$N_LAYER"-D"$N_EMBD"-"$MODEL_TYPE
M_BSZ="16"
LR_INIT="6e-4"
LR_FINAL="1e-5"
GRAD_CP=1
EPOCH_SAVE=1
N_NODE=1
GPU_PER_NODE=1
DS_BUCKET_MB=2
python train.py --load_model "0" --wandb "ArithmeticTest" --proj_dir $PROJ_DIR --my_testing $MODEL_TYPE \
--ctx_len $CTX_LEN --train_stage 3 --epoch_count 999999 --epoch_begin 0 \
--data_file "/home/rwkv/RWKV-LM/RWKV-v5/qa_add_v2.2-30M" --my_exit_tokens 1014434168 --magic_prime 990599 \
--num_nodes $N_NODE --micro_bsz $M_BSZ --n_layer $N_LAYER --n_embd $N_EMBD \
--lr_init $LR_INIT --lr_final $LR_FINAL --warmup_steps 10 --beta1 0.9 --beta2 0.99 --adam_eps 1e-18 --data_type "binidx" --vocab_size 65536 \
--weight_decay 0.001 --epoch_save $EPOCH_SAVE --head_size 64 \
--accelerator gpu --devices $GPU_PER_NODE --precision bf16 --strategy deepspeed_stage_2 --grad_cp $GRAD_CP --enable_progress_bar True --ds_bucket_mb $DS_BUCKET_MB
训练 Loss 图:
如何使用此模型
推荐使用 CUDA + FP32 精度进行推理,可以在 RWKV Runner 或者 Ai00 中推理。
由于这是一个只需要准确率,不需要任何创造性的模型。因此我们推荐使用如下解码参数:
- Temperature=0
- Top_P= 0
- Presence Penalty=0
- Frequency Penalty=0