题解：洛谷 P1059 [NOIP 2006 普及组] 明明的随机数
2026/10/8 7:45:35
- 人工智能
- 大模型
- 预训练
- 模型评测
- 模型量化
【免费下载链接】RWKV-LM
RWKV (pronounced RwaKuv) is an RNN with great LLM performance, which can also be directly trained like a GPT transformer (parallelizable). We are at RWKV-7 "Goose". So it's combining the best of RNN and transformer - great performance, linear time, constant space (no kv-cache), fast training, infinite ctx_len, and free sentence embedding.
项目地址：https://gitcode.com/gh_mirrors/rw/RWKV-LM
本文以 RWKV-LM 仓库中 RWKV-v7/train_temp/README.md 为核心，系统讲解 RWKV-7（代号 "Goose"，模型类型
x070）的完整训练流程：如何用一张 7GB 显存的 GPU 在 MiniPile（约 15 亿 token）上从零训练、如何解读期望的 loss 曲线、如何升级高吞吐 CUDA kernel、以及如何理解 RWKV-7 为每个参数精心设计的初始化与优化器分组。读完本文，你将能够复现 MiniPile 训练、把配置扩展到 Pile 20B 等更大语料，并能读懂train_temp中数百行训练代码进而改造适配自己的任务。
一、环境要求与准备工作
官方给出的参考环境如下（版本均为"参考值"，代码内部对 torch 与 CUDA 无硬性上限）：
- Python 3.10+
- PyTorch 2.5+
- CUDA 12.5+
- 最新版 DeepSpeed
- pytorch-lightning 固定为 1.9.5（切勿升级，训练代码依赖该版本的 API 行为）
安装命令（README 原文）:
# 可以直接使用最新 torch + 最新 cuda（不限于 cu121 构建） pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu121 pip install pytorch-lightning==1.9.5 deepspeed wandb ninja --upgrade其中ninja用于加速 CUDA kernel 的即时编译（src/model.py 通过torch.utils.cpp_extension.load加载多个自定义算子），wandb用于训练日志可视化。
一个重要的排障提示：两个 demo 脚本头部都注明如果启动卡住，请清理TORCH_EXTENSIONS_DIR中的锁文件（见 demo-training-prepare.sh）。这是因为多个 kernel 并发编译时产生的锁文件可能残留，导致后续load()永久等待。
二、快速开始：MiniPile 单卡训练（两步流程）
整个训练分为两个阶段，对应两个 shell 脚本。默认配置只需1 张 7GB 显存的 GPU（显存不足时可进一步调小micro_bsz），非常适合先在本地验证全流程。
第一步：下载数据
脚本 demo-training-prepare.sh 中给出的是 HuggingFace 上的 MiniPile tokenized 数据集（约 15 亿 token，rwkv_vocab_v20230424词表，约 3GB）。下载后把minipile.bin与minipile.idx放入RWKV-v7/train_temp/data/目录：
cd RWKV-v7/train_temp/ mkdir -p data # wget --continue -O data/minipile.idx https://huggingface.co/datasets/BlinkDL/minipile-tokenized/resolve/main/rwkv_vocab_v20230424/minipile.idx # wget --continue -O data/minipile.bin https://huggingface.co/datasets/BlinkDL/minipile-tokenized/resolve/main/rwkv_vocab_v20230424/minipile.bin第二步：生成初始权重（stage 1）
sh ./demo-training-prepare.sh该脚本在CPU上运行train.py的 stage 1，生成初始权重rwkv-init.pth并保存到out/……/输出目录中。脚本顶部可直接修改模型规格：
MODEL_TYPE="x070" # x070 => rwkv-7.0 N_LAYER="12" N_EMBD="768" CTX_LEN="512" # !!! 修改 ctx_len 时必须同时修改 magic_prime !!! PROJ_DIR="out/L"$N_LAYER"-D"$N_EMBD"-"$MODEL_TYPE第三步：正式训练（stage 3）
sh ./demo-training-run.sh该脚本会加载rwkv-init.pth（默认通过--load_model "0"让 train.py 自动从PROJ_DIR中找出最新 checkpoint 续训）并开始 GPU 训练。建议先登录 wandb，以便实时查看 loss / lr / wd / 吞吐量曲线。训练前脚本会先删除输出目录中的旧 checkpoint（rwkv-*0.pth、rwkv-71.pth、rwkv-final.pth），避免自动续训时误加载过期权重——如果你不希望这种清理行为，可以注释掉 demo-training-run.sh 中对应的rm行。
训练完成后，out/....../train_log.txt中的 loss 应与参考值一致（误差必须在 ±0.01 以内，否则说明配置或环境有问题）：
0 4.875856 131.0863 0.00059975 2025-04-24 02:23:42.481256 0 1 4.028621 56.1834 0.00059899 2025-04-24 02:28:16.674463 1 2 3.801625 44.7739 0.00059773 2025-04-24 02:32:51.059568 2 3 3.663070 38.9808 0.00059597 2025-04-24 02:37:25.409892 3 4 3.578974 35.8368 0.00059371 2025-04-24 02:41:59.711315 4 5 3.510906 33.4786 0.00059096 2025-04-24 02:46:33.990839 5 6 3.462345 31.8917 0.00058771 2025-04-24 02:51:08.378331 6 7 3.412196 30.3318 0.00058399 2025-04-24 02:55:42.927474 7 8 3.376724 29.2747 0.00057978 2025-04-24 03:00:17.504665 8 9 3.336911 28.1321 0.00057511 2025-04-24 03:04:52.006063 9 10 3.313411 27.4787 0.00056999 2025-04-24 03:09:27.563336 10 11 3.295895 27.0016 0.00056441 2025-04-24 03:14:01.786079 11每行格式为：miniepoch号 loss 困惑度(exp(loss)) 当前lr 时间戳 current_epoch。该日志由 src/trainer.py 中的train_callback在每轮结束时写入，同时以REAL it/s与Kt/s形式记录吞吐。
三、核心训练参数详解（train.py 全参数表）
两个 demo 脚本最终都调用同一个入口 RWKV-v7/train_temp/train.py，其中 L16-L60 定义了全部命令行参数。下表按功能分组整理，供改造配置时查阅：
| 参数 | 默认值 | 说明 |
|---|---|---|
--load_model | "" | 权重文件完整路径（含.pth）。demo 中用"0"表示自动加载PROJ_DIR下最新 checkpoint |
--wandb | "" | wandb 项目名；为空则关闭 wandb（prepare 阶段即用""） |
--proj_dir | "out" | 输出目录，保存rwkv-init.pth、rwkv-{epoch}.pth、rwkv-final.pth与train_log.txt |
--random_seed | -1 | ≥0 时设置全局随机种子（会显著影响多 GPU 采样，README 警告慎用） |
--data_file | "" | 数据文件前缀（binidx 格式的.bin+.idx） |
--data_type | "utf-8" | 数据格式，本项目固定为"binidx" |
--vocab_size | 0 | MiniPile 用 65536；Pile 20B 用 50304；0 表示自动（字符级 / .txt 数据） |
--ctx_len | 1024 | 上下文长度。修改它必须重新计算magic_prime |
--epoch_steps/--epoch_count | 1000/500 | 一个 miniepoch 的步数 / 总 miniepoch 数；stage 2+ 由代码根据magic_prime自动推导（见下文） |
--epoch_begin | 0 | 续训时起始 miniepoch 号 |
--epoch_save | 5 | 每 N 个 miniepoch 保存一次rwkv-*.pth |
--micro_bsz | 12 | 每 GPU 的微批大小；MiniPile demo 用 16（≈7G VRAM），显存不足可调小 |
--n_layer/--n_embd | 6/512 | 层数与嵌入维度；demo 用 L12-D768 |
--dim_att/--dim_ffn | 0 | 0 时自动取n_embd与n_embd*3.5（32 对齐），见 train.py |
--lr_init | 6e-4 | 初始学习率。注释给出的经验值：L12-D768 用 6e-4，L24-D1024 用 4e-4，L24-D2048 用 3e-4 |
--lr_final | 1e-5 | 最终学习率（cosine 退火终点） |
--warmup_steps | -1 | 预热步数；-1表示续训时自动设为 10（见 train.py） |
--beta1/--beta2 | 0.9/0.99 | Adam 动量系数 |
--adam_eps | 1e-18 | Adam epsilon（MiniPile 训练阶段用1e-18） |
--grad_cp | 0 | 梯度检查点：1 更省显存但更慢；0 更快更耗显存 |
--weight_decay | 0 | 仅作用于带wdecay标记的大张量（详见第四节） |
--grad_clip | 1.0 | 梯度裁剪阈值；遇到问题样本可降至 0.7 / 0.5 / 0.3 / 0.2 |
--train_stage | 0 | 1=生成初始权重；≥2 自动寻找最新 checkpoint 并续训 |
--head_size | 64 | RWKV 注意力头维度，默认 64 |
--head_chunk | 0 | LM head 分块计算：0 最快但最耗显存；65536 省 ~70% 显存；4096 省 ~80%（更慢） |
--magic_prime | 0 | 数据遍历序列的素数周期（详见下文） |
--kernel | "" | 自定义 kernel 选择；"@rwkv3"表示启用 v3 系 fused kernel（H100 上通常更快） |
--my_exit_tokens | 0 | 训练总 token 数；达到后保存rwkv-final.pth并退出 |
--my_testing | 'x070' | 模型变体标记，本项目固定为x070（RWKV-7.0） |
--ds_bucket_mb | 200 | DeepSpeed allgather/reduce bucket 大小（MB）；脚本注释提示消费级 GPU 可设 2、A100/H100 设 200，但新版本 DeepSpeed 中该功能有 bug，已被注释禁用 |
--accelerator/--devices | 见脚本 | prepare 阶段用cpu，训练阶段用gpu |
--precision | 见脚本 | 本项目统一用bf16（fp32 极慢，fp16 可能溢出） |
--strategy | 见脚本 | deepspeed_stage_2；stage 3 与 torch.jit 不兼容（train.py 会自动关闭RWKV_JIT_ON） |
关于 magic_prime
magic_prime是小于datalen/ctxlen - 1的最大 3n+2 型素数，用于控制数据分片与每个 GPU 的采样游走，保证多卡、多轮遍历时数据交错且不重复。MiniPile 的计算过程在脚本注释中给出：
datalen = 1498226207（token 数），ctxlen = 512 1498226207/512 - 1 = 2926222.06 小于它的最大 3n+2 素数 = 2926181修改CTX_LEN时必须用素数工具重新计算该值，否则数据采样会出错。在 train.py 中，epoch_count = magic_prime // 40320、epoch_steps = 40320 // real_bsz，即每个 miniepoch 固定消费40320 × ctx_len个 token（脚本注释与 demo-training-run.sh 中的EPOCH_SAVE说明可相互印证）。
四、升级 kernel：让训练更快
README 给出了一组比仓库当前代码更新的 kernel，用于把已有训练代码升级为更快的版本：
rwkv7_clampw_v3 (可同时尝试 rwkv7_clampw_v3_for_h100 与 rwkv7_clampw_v3_for_h100_alt) rwkv7_cmix_bf16_v5 rwkv7_tmix_mix6_bf16_v5 rwkv7_tmix_kk_pre_bf16_v5 rwkv7_tmix_lnx_rkvres_xg_bf16_v1 rwkv7_tmix_a_gate_bf16 rwkv7_tmix_vres_gate_bf16_v1 rwkv7_l2wrap_ce_bf16_v2这些 kernel 的源码已包含在仓库中，分布在 RWKV-v7/train_temp/cuda/ 目录（.cu+.cpp成对出现）。从 src/model.py 的加载逻辑可以看到当前实现如何选择：
- 默认编译
rwkv7_clampw.cu； - 当命令行传入
--kernel @rwkv3时，改用rwkv7_clampw_v3_for_h100.cu（源码注释说明其对 H100 以及部分消费级 GPU 的某些Bsz × Headcount组合，前向+反向 kernel 可快约 20%；rwkv7_clampw_v3_for_h100_alt有时更快）。
同时注意 src/model.py 中的断言：HEAD_SIZE必须为 64，且T % CHUNK_LEN == 0（CHUNK_LEN=16），训练时ctx_len需是 16 的倍数。
五、RWKV-7 的权重设计：初始化表与优化器分组
README 强调了一个关键设计哲学：RWKV-7 是"整模型"而非"单层"——它包含为每个参数精心设置的初始化（init）、权重衰减（wd）与学习率（lr），因此天然可扩展且训练稳定（无 loss spike）。代价是：不存在开箱即用的简单 "RWKV-7 layer"，因为一个 PyTorch 层无法自行保证使用了正确的 init 与超参。所以你要把 RWKV-7 用于其他任务时，请研究train_temp代码（仅几百行）并按其思路改造。
5.1 参考权重表（1.5B 规模：L24-D2048，vocab 65536）
README 给出完整权重表，其中wdecay注释表示该张量是"大张量"，只有它们才应用权重衰减——这是性能关键，若对全部参数做 wd，性能会明显变差：
| name | shape | comment | initialization |
|---|---|---|---|
| emb.weight | [65536, 2048] | wdecay | see code |
| blocks.0.ln0.weight | [2048] | for layer 0 | 1 |
| blocks.0.ln0.bias | [2048] | for layer 0 | 0 |
| blocks.*.ln1.weight | [2048] | 1 | |
| blocks.*.ln1.bias | [2048] | 0 | |
| blocks.*.att.x_r | [1, 1, 2048] | see code | |
| blocks.*.att.x_w | [1, 1, 2048] | see code | |
| blocks.*.att.x_k | [1, 1, 2048] | see code | |
| blocks.*.att.x_v | [1, 1, 2048] | see code | |
| blocks.*.att.x_a | [1, 1, 2048] | see code | |
| blocks.*.att.x_g | [1, 1, 2048] | see code | |
| blocks.*.att.w0 | [1, 1, 2048] | lr 2x | see code |
| blocks.*.att.w1 | [2048, 96] | 0 | |
| blocks.*.att.w2 | [96, 2048] | see code | |
| blocks.*.att.a0 | [1, 1, 2048] | 0 | |
| blocks.*.att.a1 | [2048, 96] | 0 | |
| blocks.*.att.a2 | [96, 2048] | see code | |
| blocks.*.att.v0 | [1, 1, 2048] | for layer 1+ | 1 |
| blocks.*.att.v1 | [2048, 64] | for layer 1+ | 0 |
| blocks.*.att.v2 | [64, 2048] | for layer 1+ | see code |
| blocks.*.att.g1 | [2048, 256] | 0 | |
| blocks.*.att.g2 | [256, 2048] | see code | |
| blocks.*.att.k_k | [1, 1, 2048] | 1 | |
| blocks.*.att.k_a | [1, 1, 2048] | 1 | |
| blocks.*.att.r_k | [32, 64] | 0 | |
| blocks.*.att.receptance.weight | [2048, 2048] | wdecay | see code |
| blocks.*.att.key.weight | [2048, 2048] | wdecay | see code |
| blocks.*.att.value.weight | [2048, 2048] | wdecay | see code |
| blocks.*.att.output.weight | [2048, 2048] | wdecay | 0 |
| blocks.*.att.ln_x.weight | [2048] | see code | |
| blocks.*.att.ln_x.bias | [2048] | 0 | |
| blocks.*.ln2.weight | [2048] | 1 | |
| blocks.*.ln2.bias | [2048] | 0 | |
| blocks.*.ffn.x_k | [1, 1, 2048] | see code | |
| blocks.*.ffn.key.weight | [8192, 2048] | wdecay | see code |
| blocks.*.ffn.value.weight | [2048, 8192] | wdecay | 0 |
| ln_out.weight | [2048] | 1 | |
| ln_out.bias | [2048] | 0 | |
| head.weight | [65536, 2048] | wdecay | see code |
其中att.w1/w2、a1/a2、v1/v2、g1/g2是 LoRA 形式的低秩门控系数；简化 demo 中使用的对应秩为D_DECAY_LORA、D_AAA_LORA、D_MV_LORA、D_GATE_LORA，且注释指出LM 训练时应放大为max(32, round(2.5*C^0.5))级别（见 rwkv7_train_simplified.py）。
5.2 源码级初始化策略（generate_init_weight）
表内 "see code" 的具体实现在 src/model.py 的generate_init_weight中，要点如下：
emb.weight：均匀初始化，scale =±1e-4；head.weight：当vocab_size > n_embd时用正交初始化，gain =0.5 * sqrt(vocab_size / n_embd)；att.output、ffn.value等投影：零初始化（保证深层网络起点是恒等映射）；att.key：正交初始化，gain = 0.1；ln_x.weight：按层比例缩放，layer_scale = (1+layer_id)/n_layer，最终ln_x.weight = layer_scale^0.7——这是 RWKV 特有的逐层深度缩放；- 各
x_r / x_w / x_k / x_v / x_a / x_g（time-mix 系数）与w0 / a0 / v0等在RWKV_Tmix_x070.__init__中构造：通过1 - ddd^pow、zigzag（与 head_size 相关的锯齿波）与 linear 组合给出先验（rwkv7_train_simplified.py 中有最清晰的参考实现）。
5.3 优化器分组：wdecay 与 lr 2x
configure_optimizers（src/model.py）将参数分成三组：
- lr_1x：其余全部参数，
weight_decay=0，my_lr_scale=1.0； - lr_2x：
att.w0（时间衰减基座），学习率放大 2 倍（my_lr_scale=2.0）且不做 wd； - lr_decay：仅当
weight_decay > 0且参数名含.weight且squeeze后 ≥2 维的大张量（即权重表中标wdecay的项）加入，wd 取命令行--weight_decay值。
优化器根据deepspeed_offload自动选择DeepSpeedCPUAdam（offload）或FusedAdam。这与 README 的警告完全一致："Make sure you only apply wd to large tensors (with 'wdecay' in comment) here，否则性能会差很多"。
5.4 学习率调度
train.py 设置max_epochs=-1（无限续训），实际调度由 src/trainer.py 的train_callback.on_train_batch_start完成：
- 按
real_tokens计算进度progress，执行cosine 退火（lr_mult公式），从lr_init衰减到lr_final； warmup_steps内做 1%→100% 的线性预热；- 每个 miniepoch 结束时按
lr * my_lr_scale写回各参数组，因此att.w0实际始终以 2 倍 lr 训练； - 达到
my_exit_tokens时保存rwkv-final.pth并exit(0)。
六、更大规模：Pile 20B 训练配置
仓库还提供了面向完整 Pile 20B 语料（约 3321 亿 token）的两个脚本，可作为扩展到生产规模的标准模板：
- demo-training-prepare-v7-pile.sh：需要
pile_20B_tokenizer_text_document.bin/.idx，模型为 L12-D768、ctx_len=4096、vocab 50304，magic_prime=81082817（对应 4096 上下文重新计算过）； - demo-training-run-v7-pile.sh：8×80GB GPU（
micro_bsz=30）、lr_init=8e-4、lr_final=3e-5、weight_decay=0.1、beta2=0.99、adam_eps=1e-18，每 50 个 miniepoch 保存一次。
两套配置中 lr 与my_exit_tokens的差异体现了脚本注释总结的经验法则：
- 小数据 → 用较小的 bsz 与稍小的 lr；大数据 → 用较大的 bsz 与稍大的 lr；
- 更大的模型 → 用更小的 lr；微调（finetuning）→ 用很小的 lr，例如
1e-5。
七、简化训练 Demo：理解 RWKV-7 的"解剖"课
仓库在 rwkv7_train_simplified.py 提供了一份简化版 RWKV-7 训练 demo（README 注明：速度慢、结果与正式版不同，仅用于学习）。它把整个体系浓缩为约 300 行：
- 任务：数字反转（把随机整数序列倒序输出），词表仅 13 个 token（0-9、
,、#）； - 模型：2 层
RWKV_Tmix_x070（fp32）+ 2 个 SquaredReLU FFN，V=12, C=32, B=256, T=129, steps=10000，HEAD_SIZE=16（注释强调LM 必须用 64）； - kernel：编译
cuda/wkv7_cuda_fp32.cu+wkv7_op_fp32.cpp（fp32 版 wind backstepping），要求T % CHUNK_LEN == 0； - 前向结构：第一层注意力保存
v_first，后续层通过v0/v1/v2门控加入 value residual（v = v + (v_first - v) * sigmoid(...)）；a是"in-context learning rate"，kk做 L2 归一化后进入RUN_CUDA_RWKV7g(r, w, k, v, -kk, kk*a)，最后叠加r_k项并乘以output * g； - 优化器：
AdamW中只有.weight且非ln的参数组weight_decay=0.1，其余为 0——与正式版的分组逻辑一脉相承； - 训练：
CosineAnnealingLR(T_max=steps, eta_min=1e-6)，loss 用F.cross_entropy，每步clip_grad_norm_(1.0)，并把loss/lr/Mtokens/kt/s记录到 wandb； - 验证：训练结束随机取 5 个样本，把
,之后的回文区作为判定区域计算逐 token 准确率acc，输出in/gold/pred/diff四行诊断信息。
这份 demo 是把第四节"整模型初始化"思想落地的最小可读实现，适合想彻底搞懂 RWKV-7 各组件（time-shift 混合、decay LoRA、value residual、in-context learning rate）的读者。
八、总结：把 RWKV-7 迁移到自己的任务
- 复现基线：按第二节两步脚本在 1×7GB GPU 上即可完整跑通 MiniPile 训练，以
train_log.txt的 ±0.01 误差校验环境正确性； - 扩展规模：换用 demo-training-run-v7-pile.sh 的参数模板并重新计算
magic_prime（改ctx_len必改）与my_exit_tokens； - 性能优化：优先选用 README 列的 8 个升级 kernel，并尝试
--kernel @rwkv3； - 改造迁移：不要试图"取出"一个独立的 RWKV-7 层，而应整体复制
train_temp的训练骨架（train.py+src/），按需替换src/dataset.py的数据格式，并保持generate_init_weight与configure_optimizers的逐参数初始化/分组设计不动——这是 RWKV-7 训练稳定、无 spike 的根本保证。
赞
- 人工智能
- 大模型
- 预训练
- 模型评测
- 模型量化
【免费下载链接】RWKV-LM
RWKV (pronounced RwaKuv) is an RNN with great LLM performance, which can also be directly trained like a GPT transformer (parallelizable). We are at RWKV-7 "Goose". So it's combining the best of RNN and transformer - great performance, linear time, constant space (no kv-cache), fast training, infinite ctx_len, and free sentence embedding.
项目地址：https://gitcode.com/gh_mirrors/rw/RWKV-LM
相关推荐
RWKV-7 "Goose" 实战指南：并行化 RNN 架构解析、MiniPile 训练配置与推理部署
RWKV 7 "Goose" 实战指南：并行化 RNN 架构解析、MiniPile 训练配置与推理部署 RWKV 是一个兼具 RNN 与 Transformer
人工智能大模型预训练模型评测模型量化探秘 RWKV-LM: 一款高效、灵活的预训练语言模型
探秘 RWKV LM: 一款高效、灵活的预训练语言模型 在自然语言处理领域，预训练模型已经成为基础工具，它们能够理解和生成高质量文本，而正是这样一款强大的工具。
人工智能大模型预训练模型评测模型量化10个RWKV-LM常见问题快速解决方案：从安装到训练的完整指南
10个RWKV LM常见问题快速解决方案：从安装到训练的完整指南 RWKV LM是一个结合RNN和Transformer优势的革命性语言模型，它拥有线性时间计算
人工智能大模型预训练模型评测模型量化
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考