Skip to content
BranchesTags
 
 

Repository files navigation

RWKV-Lab

tests Python 3.10+ License: MIT

A reproducible research workbench for RWKV language models, multimodal captioning, recurrent-depth training, memory, post-training, kernels, and cross-architecture conversion.

RWKV-Lab turns research ideas into isolated, measurable implementations. Most language-model levers are off by default, carry an explicit configuration, and have a CPU correctness test or reference oracle. Hardware-specific paths must pass parity before they are allowed to claim a speedup. Training runs preserve their configuration, data fingerprints, checkpoints, evaluation artifacts, and resume state so that a promising graph is evidence—not just a screenshot.

Important

This is an active research repository, not a packaged pretrained model. Weights, datasets, generated manifests, caches, and run directories are intentionally excluded from Git. Reproducing a model run requires supplying those artifacts and using the same fingerprints recorded by the launcher.

At a glance

Track What is implemented Maturity
Multimodal RWKV Frozen MoonViT and RWKV backbones; trainable visual bridge/resampler; optional SigLIP2 + DINOv2 + SAM fusion; grounding losses; recurrent loops; NextLat; Engram; qualitative eval artifacts Active experiment pipeline
Cross-architecture conversion Exact Gated DeltaNet → RWKV-7 remap plus isolated attention-layer distillation, assembly, and consolidation GDN remap validated; attention conversion experimental
RWKV training research Recurrent depth, latent prediction, lexical and online memory, Muon-family optimizers, mixed-context training, and typed post-training Reference implementations and A/B tooling
Training systems Exact-resume checkpoints, immutable receipts, watchdogs, kernel qualification, and a Go/SQLite/Pixi dashboard Operational research infrastructure

The long-term multimodal design—compress several frozen vision teachers into one RWKV-native vision student—is documented in MULTI_TEACHER_VISION_DISTILLATION.md.

Start here

Goal Entry point
Understand the available research levers What's in the box · TRAINING_LEVERS.md
Train the image captioner Multimodal captioning · vision_train.py
Run a conversion experiment Conversion pipeline · validated conversion result
Monitor training dashboard/README.md
Prepare or deduplicate image data UNLABELED_IMAGE_DEDUP.md · scripts/
Run tests Quick start · scripts/test_parallel.sh
Check project claims and maturity Status · References

Quick start

CPU installation is sufficient for the Python correctness suite. Install the PyTorch build appropriate for your platform first; CUDA training additionally requires a compatible CUDA PyTorch stack and fla/flash-linear-attention. The last command below also requires Go and validates the dashboard separately.

git clone https://github.com/sirus20x6/rwkv-lab.git
cd rwkv-lab

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

# Example CPU-only PyTorch installation. Use the matching CUDA build for training.
python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
python -m pip install -e '.[vision,test]'

pytest -q -m 'not gpu'
go -C dashboard test ./...

The package uses a src/ layout, so installed entry points run as python -m rwkv_lab.<module>. For development without an editable install, prefix commands with PYTHONPATH=src.

Multimodal captioning

The current captioner adapts Kimi K2.6's MoonViT representation to a 2.9B RWKV G1H language model. Both pretrained backbones remain frozen. The baseline trains only a visual prefix projector; research configurations can add a resampler, intermediate-layer injection, multi-teacher residual fusion, recurrent-depth adapters, NextLat, grounding objectives, and Engram memory.

flowchart LR
    I[Image] --> M[MoonViT<br/>frozen]
    I --> S[SigLIP2<br/>optional, frozen]
    I --> D[DINOv2<br/>optional, frozen]
    I --> A[SAM<br/>optional, frozen]
    M --> P[Visual bridge / resampler]
    S --> F[Teacher fusion adapter]
    D --> F
    A --> F
    F --> P
    P --> R[RWKV G1H 2.9B<br/>frozen backbone]
    R --> C[Caption tokens]
    N[NextLat · loops · Engram<br/>optional trainable levers] --> R
Loading

What trains and what stays frozen

Component Baseline Advanced configuration
MoonViT Frozen Frozen
RWKV G1H backbone and LM head Frozen Frozen
Prefix projector Trainable Trainable
Vision resampler and layer/deep-vision adapters Disabled Trainable when enabled
SigLIP2 + DINOv2 + SAM towers Disabled Frozen; features may be cached
Fusion residual, loop adapters, loop-index embedding Disabled Trainable when enabled
NextLat head and Engram memory Disabled Trainable when enabled

Data contract

Training and explicit evaluation data are JSONL files with one image-caption pair per line:

{"image":"images/000001.jpg","text":"Two red kites fly above a calm blue sea.","source":"example"}
  • image may be absolute or relative to the repository root.
  • text is the target caption; prompts and image-prefix positions are masked from caption cross-entropy.
  • An optional per-row prompt overrides the default Describe this image:.
  • Explicit evaluation manifests must be image-disjoint from training data.
  • Exact image-caption duplicates are removed before sampling.

Generated manifests are local artifacts and are not checked into Git. The curation scripts under scripts/ build them from source datasets.

Minimal bridge baseline

Supply the RWKV World vocabulary, the two model checkpoints, and your own train/eval manifests. This command disables the experimental recurrent and latent levers so the result is a clean visual-bridge baseline:

export VOCAB=/path/to/rwkv_vocab_v20230424.txt

python -m rwkv_lab.vision_train \
  --data /path/to/train.jsonl \
  --eval-data /path/to/eval.jsonl \
  --rwkv /path/to/rwkv7-g1h-2.9b-20260710-ctx10240.pth \
  --moonvit /path/to/model-00064-of-000064.safetensors \
  --feature-cache /path/to/moonvit-cache \
  --out runs/vision-bridge-baseline \
  --prefix-tokens 64 --loop-count 1 --no-loop-index \
  --nextlat-weight 0 --resume auto

--resume auto restores last.pt, including optimizer, sampler, and RNG state. The trainer refuses to overwrite an existing run that lacks a recoverable checkpoint; --fresh --resume none is the explicit destructive-new-run path.

Advanced multi-teacher run

prepare_vision_next_levers.sh caches MoonViT taps and aligned SigLIP2/DINOv2/SAM features, verifies capacity, and writes a manifest-bound receipt. run_vision_next_levers.sh refuses to launch if that receipt, its cache counts, fingerprints, or mtimes do not match.

# Set VISION_NEXT_* paths for your machine before running these scripts.
scripts/prepare_vision_next_levers.sh
scripts/run_vision_next_levers.sh

The checked-in advanced launcher uses 128 visual tokens, three MoonViT taps, three RWKV injection sites, full-frame plus quadrant views, a two-layer visual resampler, early-token grounding weight, image-text contrastive loss, two factored recurrent passes, loop-index embeddings, NextLat, and two Engram sites. These are experimental choices, not universal defaults.

Evaluation and observability

Evaluation loss and perplexity are teacher-forced monitoring metrics; they do not update model weights. Qualitative evaluation runs greedy decoding on a stable, source-filtered held-out set and stores the image, prompt, generated caption, reference, termination reason, and token count as run artifacts. Trainboard ingests those artifacts so each evaluation point can be inspected alongside loop, NextLat, Engram, throughput, GPU, and checkpoint telemetry.

go -C dashboard run ./cmd/trainboard \
  -repo "$PWD" -runs "$PWD/runs" -addr 127.0.0.1:9124
# open http://127.0.0.1:9124

Optional Kimi teacher captions

kimi_teacher.py can rank a clean candidate queue and request long-form Kimi K2.6 captions through OpenRouter. Existing captions influence selection only; they are not sent as hints. Network execution is deliberately gated by both an API key and --execute.

python -m rwkv_lab.kimi_teacher select

OPENROUTER_API_KEY=... python -m rwkv_lab.kimi_teacher caption \
  --budget-usd 3.80 --execute

Every response is written atomically and preserves model/provider metadata, usage, cost, chosen-token log probabilities, and returned alternatives. The completion limit is a runaway-spend guard rather than a desired caption length; guard-truncated responses remain available for diagnosis but are excluded from the accepted training manifest.

Image acquisition and deduplication

Reproducibility and artifact boundaries

The repository contains code, tests, launch contracts, and documentation. It does not contain third-party datasets, model weights, generated captions, feature caches, checkpoints, logs, dashboard databases, or review packs. Those paths are ignored because they are large, machine-specific, or may carry local absolute paths and upstream license constraints.

For reproducible experiments, preserve the run's config.json, train.jsonl, status.json, last.pt, best/, dataset/cache receipts, and qualitative eval artifacts together. A checkpoint without its data and cache fingerprints is not enough to reconstruct the experiment.

What's in the box

The table below is the fastest map from a research question to its implementation. Most experimental levers are off by default; paper links, reference oracles, and tests are cataloged in References. A listed implementation is not automatically a claim of production readiness—see Status for the evidence boundary.

Area Techniques Modules
Recurrent-depth loops weight-tied loops · hyper-connections · per-iterate readout · PonderNet halt · CART contractive gate · HRM DEQ / Neumann-k gradient · FPRM fixed-point halt · RWKV-Product multi-substep looped_rwkv · rwkv_product
Latent prediction MTP · MuToR · TOP · NextLat · ConceptLM · FSP · L-MTP · Belief-State · JTP · LLM-JEPA · Coconut continuous-thought lookahead_module · llm_jepa · coconut
Memory / retrieval Engram lexical bank · ROSA suffix-automaton (+ golden reference) · Fast-weight Product-Key Memory · L³ large-lookup · WriteSAE state autoencoder engram_lmb · rosa_sam · fwpkm · l3_lookup · write_sae
Scale & online adaptation (P0) u-μP · Titans/MIRAS/ATLAS/Nested memory · offline sleep and guarded adapter consolidation · GSPO/Dr.GRPO/DAPO RLVR · bounded Reasoning Cache u_mup · online_memory · adapter_consolidation · reasoning_cache
Post-training platform typed SFT/preference/feedback/PRM/RLVR · named LoRA/NF4 QLoRA · initial-state and per-token state-offset adapters · routed state bank · paired confirmation · safe exports posttrain_data · adapters · state_tuning · state_bank
Training systems (P1) RegMix/MDE · NVFP4 · Decoupled DiLoCo · stable triangular delta-rule inversion oracle/qualification data_mixture · nvfp4 · diloco · triangular_delta
Representation, decoding, serving & tracing (P2) BLT · byte-aware/SuperBPE experiments · decoder matrix · typed policies · EAGLE-3 · paged state forks · B³D-RWKV · HiLS · recurrent attribution byte_patches · tokenizer_experiments · decoding_eval · decoding_policy · recurrent_serving
Optimizers & dynamics Muon (+ MuonClip) · 12 spectral-Muon levers (Muonᵖ, Aurora, MONA, DDC, RSAV, Hierarchical, Distance-Aware, ARO…) · PC-Layer preconditioning · layerwise-LR · grokking probes spectral_muon · muon_helpers · pc_layer
Cross-arch conversion GDN ⊂ RWKV-7 lossless remap · RADLADS · QRWKV7 balance_state · Taylor-Calibrate · Comba · Attention-to-Mamba convert_gdn_lossless · convert_train · rwkv8_deltanet
From-scratch lab Future-Seed cross-layer state chaining · DeepEmbed per-token FFN gates (output / BlinkDL-exact hidden / +shift / +emb-residual) · Engram-as-lever · semantic context-bucket packing + mixed-context training (reciprocal batch) · grad-accum / EMA / fp8 / 8-bit optimizers rwkv_pretrain · experiment · build_corpus

All Python lives under src/rwkv_lab/ (python -m rwkv_lab.<module>); a from-scratch Go + SQLite + Pixi.js dashboard (dashboard/) drives and monitors runs.

RWKV community research additions and implementation notes

RWKV community research additions

The following off-by-default capabilities came from a July 2026 review of the RWKV community's technical channels. Discord posts are treated as leads; paper and repository links below are the implementation sources of record.

  • Balanced conversion state: --balance-state implements the alternate normalized-key and write/forget scaling path in Recursal's QRWKV7 implementation, motivated by its reported large-model GroupNorm stability behavior. Existing recurrence remains the default.
  • RWKV state adapters: state_tuning.py learns named initial WKV and token-shift states with frozen weights, inspired by rwkv-rlhf and OpenMOSE/RWKV-LM-RLHF.
  • Paged recurrent serving: recurrent_serving.py provides request-isolated CPU paging, prefix-state forks, continuous-batch stack/split, and asynchronous device transfer primitives inspired by AUXStar/RWKV-Server.
  • Stable delta-rule inversion: triangular_delta.py implements direct and unit-lower Neumann oracles plus parity/speed receipts from Fast and Stable Triangular Inversion for Delta-Rule Linear Transformers.
  • Offline consolidation: online_memory.py adds explicitly budgeted, bounded sleep passes following Do Language Models Need Sleep?.
  • Reasoning Cache: reasoning_cache.py implements bounded response→summary→response iteration and auditable short-horizon training pairs following Reasoning Cache. It does not execute generated code.
  • Parallel decoding: diffusion_rwkv.py is an isolated triplet-block bidirectional diffusion head following B³D-RWKV and its official repository.
  • Sparse hybrid layers: hils_attention.py is a causal, CPU-readable compressed-landmark/chunk-fusion oracle following HiLS-Attention and its official implementation. Production adoption waits for a qualified sparse kernel and long-context evidence.
  • Decoder-as-evaluation: decoding_eval.py runs paired, deterministic greedy/top-k/top-p/typical/Mirostat tapes and records task score, entropy, loops, throughput, and recurrent-state divergence. The matrix follows A Thorough Examination of Decoding Methods in the Era of LLMs; the RWKV-specific state-drift criterion came from the community decoding discussion.
  • State-offset tuning: state_tuning.py adds frozen-base FP32 matrix/shift offsets at every recurrent token, with an optional scheduled-offset ablation and an exact slow chunk oracle, following the ACL 2025 State-offset Tuning paper. It is launchable as the scratch-LM state_offset arm or --state-offset 1.
  • Routed state bank: state_bank.py learns soft/hard routing over reusable constant-size state slots, optionally adds a bounded hypernetwork residual, and reports entropy/collapse telemetry. This is a correctness oracle for the community's dynamic-state and document-state retrieval proposals, not a performance claim.
  • Byte-aware and superword experiments: tokenizer_experiments.py provides exact-no-op token-length and position-specific UTF-8 byte embeddings based on BlinkDL's proposal, plus an auditable builder for SuperBPE and Faster Superword Tokenization corpus arms. The actual tokenizer trainer/runtime is implemented in ztok as ztok train --kind superbpe; RWKV-Lab owns tokenizer fingerprints and model A/B evidence.
  • Guarded adapter consolidation: adapter_consolidation.py turns the community's day-state/SDFT proposal into immutable daily snapshots, a hard training-token cap, held-out improvement/regression gates, preserved receipts, and an explicit human-only promotion step.
  • Typed decoding policies: decoding_policy.py narrows the community's self-controlled sampling proposal to allowlisted named modes, bounded nesting/transitions, and optional grammar labels; generated tokens cannot invent parameters or bypass operator policy.
  • Post-training state expansion: state_expansion.py implements StateX's uniform layer selection, actual single-head RWKV TimeMix replacement, block-diagonal runtime-state oracle, paper-recommended mixer reinitialization, and auditable capacity receipt from StateX; the generic kernel can execute the expanded geometry, while production adoption still requires long-context parity, memory, and throughput qualification. Community lead: RWKV Discord.
  • Supervised Memory Training: supervised_memory.py provides the time-parallel one-step (memory, next input) → next memory loss, predictive-future composition, anti-collapse uniformity term, detached offline-label mode, and rollout-drift diagnostic from Pretraining Recurrent Networks without Recurrence. Community lead: RWKV Discord.
  • Routing-Free MoE: routing_free_moe.py replaces the FFN with independently self-activating experts—no router, softmax, or top-k—and exposes the paper's interpolated token/expert balancing objective as the routing_free_moe arm or --routing-free-moe 1. Sources: paper, official code, community lead.
  • Rapid dense-to-sparse transfer: sparse_transfer.py implements retrieval-head calibration, trainable 16-D pre-RoPE indexers, query-adaptive top-p support, exact full-dimensional attention on that support, and teacher-top-logit distillation from RTPurbo. Community lead: RWKV Discord. The module is an exact small oracle; production sparse decoding still needs a qualified block kernel.
  • External kernel-candidate qualification: kernel_candidates.py accepts already-imported candidates from systems such as Mirage, then gates adoption on randomized output/gradient parity, determinism, and median speedup. It deliberately never executes generated source or shell commands (community lead).
  • ROSA backend qualification: rosa_backends.py registers CPU/CUDA/JAX/FPGA callables behind one exact suffix-matching contract and compares each with the independent CPU oracle before adoption. The semantics come from ROSA-Tuning; the hardware interface is inspired by ROSA-FPGA and its community report. Reported energy/throughput metadata is not treated as locally reproduced evidence.
  • Action-conditioned JEPA diagnostics: llm_jepa.py now includes a multi-step action-conditioned latent predictor, variance regularization, and SVD rank/dimension residual curves based on A Generalization Theory for JEPA-Based World Models. The diagnostic reports empirical low-rank approximation trade-offs without claiming planning-regret guarantees outside the paper's assumptions.
  • Key-Value Means: key_value_means.py provides fixed, square-root-growing, and saturating state budgets; least-redundant append; winner-take-all cosine merging; fixed value radii; JIT normalization; and joint state/window readout from KVM and its official code (community lead).
  • Neural Procedural Memory: procedural_memory.py stores inter/intra-trajectory activation contrasts, retrieves them by task, synthesizes a consensus direction, and applies an explicitly bounded residual intervention following NPM (community lead).
  • Mamba-3 recurrence ablations: mamba3_recurrence.py isolates exponential SSM discretization, complex-valued state rotation, and MIMO state lanes from Mamba-3 (community lead).
  • Nonlinear matrix state: m2rnn.py is a small matrix-to-matrix nonlinear recurrence suitable for isolated hybrid-layer experiments, based on M²RNN (community lead).
  • Compositional Muon: compositional_muon.py implements partner-whitened steepest-descent updates for paired factors such as QK and OV, with an operator update receipt, from the Apache-2.0 reference and community lead.
  • Distillation merge stage: distillation_merge.py adds weighted expert-state merging, ensemble-logit alignment, and tolerance-corrected win/tie evaluation from Effective Distillation to Hybrid xLSTM Architectures (community lead).
  • Guarded test-time training: test_time_training.py wraps inference-time next-token adaptation in an explicit parameter/step budget, held-out regression gate, immutable snapshot, and automatic rollback. It follows TTT-E2E and its official code. Baseline inference never mutates weights.
  • ROSA+ fallback: rosa_plus.py supplies an independently written, token-generic interpolated Witten–Bell distribution only when exact ROSA has no match, inspired by bcml-labs/rosa-plus without copying its GPL code (community lead).
  • Tool-use length generalization: tool_length_generalization.py builds inert, allowlisted tool tapes and extrapolation curves following Malach et al. (community lead). Adamaton exclusively owns execution.
  • Energy-based refinement: energy_refinement.py adds a contrastive compatibility-energy head and norm-bounded candidate-latent descent from Energy-Based Transformers (community discussion). Model weights remain fixed during refinement.
  • Compressed Convolutional Attention: cca_attention.py is a causal oracle that projects Q/K/V and performs convolution plus attention entirely at latent width, following CCA (community lead).
  • Compute-aware data filtering: data_filter_audit.py permits raw-versus-filter comparisons only inside matched parameter/token/compute cells, based on A Bitter Lesson for Data Filtering and Apple's quality-filter analysis (community lead).
  • Runtime backend matrix: runtime_backends.py applies the common parity/gradient/determinism/speed gate to Albatross, vLLM RWKV, Tenstorrent WKV7, and ROSA-JAX; merely being installed never adopts a backend.
  • Native RWKV megakernel backend: megakernel.py executes the one-token RWKV-7 DPLR state transition in-place with mutation-safe autotuning over batch/head/state geometry. A separately timed candidate fuses the fp32 state update directly through GroupNorm, the RWKV bonus, and output gating. megakernel_ops.py supplies ln1 + six TimeMix branches, attention add + ln2 + ChannelMix, and the final residual plus output norm. All are compiler-visible triton_op operations and are adopted independently only when whole-layer timing improves. Fixed-budget greedy plans unroll the entire device-side token/EOS loop into one CUDA Graph replay instead of one host replay per token. Exact-shape prompt plans are compiled/captured, and both decode and prompt shapes can be serialized with experimental PyTorch AOT compilation. Following the Apache-2.0 Albatross faster3b implementation, block 0 normalization is folded once into an immutable inference embedding table; the prepared buffer is part of the plan hash. Its GPU-specific GEMV layouts, worker-role executor, sparse FFN, and reduced-precision recurrent-state variants inspired separately qualified candidates in megakernel_linear.py; they remain on cuBLAS unless local parity, determinism, and speed gates pass. A serving-only preparation can replace the raw embedding with its folded table and release the duplicate, while refusing DeepEmbed modes that still need it. sm120_kernels.py adds an RTX PRO 6000-specific lane: its persistent RWKV state scheduler caps the resident grid at the card's 188 SMs and consumes multiple disjoint state tiles per CTA; its CC12.0 plan uses 2-warp, 8-column tiles against Blackwell's 48-warp/SM occupancy envelope; and its CUDA access-policy controller can reserve and target the device-reported persisting-L2 window on the packed R/K/V or folded-embedding buffer. The window is reapplied to private CUDA Graph capture streams and retained only when end-to-end timing improves. These settings follow NVIDIA's Blackwell tuning guide and L2 persistence guidance. An optional pip install -e '.[sm120]' enables NVIDIA's CUTLASS Operator API projection adapter. It rejects portable SM80 results: the receipt calls a projection native only when CUTLASS metadata explicitly advertises SM120. As of Operator API 0.1.0, BF16 discovery on this machine returns portable SM80 kernels, so projection execution correctly remains on cuBLAS while the adapter is ready for a native registry release. Plans are cached by model/device/state geometry and can be requested with rwkv_lab.generate --engine megakernel; auto adoption requires token parity, warm throughput, and profiler-measured launch reduction from rwkv_lab.production_kernels. The design follows HazyResearch Megakernels and TileRT, but uses a native RWKV implementation because their published runtimes target different model layouts. execution_plan.py remains the generic fail-closed gate for other externally supplied plans (community lead).

Highlight result — GDN ⊂ RWKV-7 (lossless conversion)

The conversion track's anchor result. Qwen3.5's linear-attention layers are gated DeltaNet (GDN). We proved — algebraically and end-to-end — that GDN's gated-delta recurrence is an exact special case of the RWKV-7 wkv7 kernel at matched head dimensions:

Given GDN kernel inputs (q, k, v, g, β), with q/k L2-normalized:
    r        = normalize(q)
    gk       = g                       # GDN's scalar log-decay, broadcast over the key dim
    k_write  = β · normalize(k)
    a        = −normalize(k)           # delta-rule removal key
    b        = normalize(k) · exp(g)·β # in-context learning rate
    out      = wkv7(r, gk, k_write, v, a, b) · (1/√head_dim)

Feeding a GDN layer's own activations through this map reproduces its output at cosine 0.999995. Patching all 24 GDN layers of the full 9B model changes perplexity by +0.013% (8.4898 → 8.4908) — with zero training. See convert_gdn_lossless.py.

That collapses the conversion problem to just the 8 full-attention layers, which are not a linear-attention subset and need distillation (RADLADS-style block-alignment + logit-KD). The attn_L3_poc.py proof-of-concept and the convert_train.py per-layer trainer target exactly those.

Why this matters: an earlier version of this project built the RWKV core at head-size 64 against GDN's 32×128, a self-imposed 2:1 state compression that forced a whole distillation-and-codec pipeline. The matched-dimension remap makes 24 of 32 layers free. The "RWKV-7 decay floor" that dogged early runs turned out to be a parametrization artifact, not a kernel limit.

Historical per-layer conversion map with 22 of 32 layers accepted
Historical dashboard snapshot at 22/32 accepted layers. The 8 dashed cells are the full-attention layers still being distilled; the green gated-delta-net layers were subsequently validated as a 24-layer lossless remap.

Validated lever results — from-scratch lab

The first A/B-validated wins, measured with the lab's multi-seed harness on small from-scratch models (d256 · 4 layers) — synthetic diagnostics plus a 388M-token real corpus (Open-PerfectBlend, chat/math/code, flattened + ztok-tokenized; ~25M tokens per run, so nothing repeats):

Lever Benchmark Result
Future-Seed --seed-chain (layer L's state scan starts from layer L−1's final state) Open-PerfectBlend, 3k steps −9.2% val ppl (86.2 vs 94.9) at identical params/FLOPs; on synthetic recall, length-gen acc@2× 0.82 → 0.98 with seed variance collapsing to zero
Engram LMB --engram (token suffix-automaton recall → learned table + copy head) induction:64, 4 seeds 0.029 → 0.442 (baseline is at chance 1/32; Δ+0.414 SIGNIFICANT); recall:16 perfect incl. 2× length
DeepEmbed --deepembed --de-mode hidden --de-shift (BlinkDL-exact FFN-hidden gate + separate gate token-shift) both corpora −2.4% / −2.1% ppl, replicated across corpora; the variant ordering (output gate loses, hidden ≈ parity, +shift wins) independently reproduces BlinkDL's report
Semantic context-bucket packing (whole docs, best-fit-decreasing into 512…32k) Open-PerfectBlend 0.1% padding (worst bucket 0.37%); mixed-context training holds tokens/step and activation VRAM constant via reciprocal batch (B = budget/T)

LM ppl rows are single-seed so far; the synthetic results are 4-seed with significance calls. All levers remain off-by-default and compose (seedchain + engram + de_shift launch together from the dashboard).

Screenshots

Experiments are driven and monitored through trainboard, a from-scratch Go + SQLite + Datastar + Pixi.js dashboard (dashboard/) that ingests every run's train.jsonl and paints run state live — loss/PPL curves, per-layer conversion maps, and ablation sweeps across the levers above.

Single-layer conversion loss curve
A single GDN→RWKV layer conversion (block-relative distillation): train loss 1.09 → 0.22, next-token top-1 88%, converging to the frozen-teacher reference line (green).

Run detail view with KPIs
Per-run view: KPI tiles (step, loss, eval PPL, top-1), the conversion map, and per-layer status.

Run leaderboard
Run leaderboard — the conversion is an experiment sweep: hundreds of isolated per-layer runs, ablations (looped vs. control, neg-eigval, schedule-free), and optimizer studies, all sortable by best PPL / top-1 / recency.

Experiments card: config-driven A/B builder + registry results
Experiments card — the current config-driven builder exposes task/init/budget/model sizing, optimizer, precision, compilation, and batch controls before the lever matrix and accumulated registry evidence below. Campaigns retain every seed and rung, paired confidence intervals, learning curves, measured throughput/memory/energy, Pareto status, lineage, and fresh-seed confirmation.

Verifiable-reward training campaign builder, evidence, and recursive lineage
Verifiable-reward training — configure cold-start and recurrent GSPO/Dr.GRPO/DAPO campaigns, rollout and evaluation budgets, promotion gates, and distributed rollout devices; inspect per-algorithm held-out evidence and the recursive checkpoint lineage. Promotion remains evidence-gated and never overwrites the parent checkpoint.

Post-training campaigns and data panel
Post-training — validate/version typed JSONL and role-aware masks; launch equal-token paired-seed campaigns with fresh confirmation; inspect confidence intervals, promotion receipts, and adapter-recursive lineage; compare checkpoints and save explicit training-only preferences. Paths stay repository-confined, and Trainboard itself cannot promote a checkpoint.

Production kernel qualification panel with megakernel tuning and regression gates
Production qualification — choose fast compilation or full megakernel plan autotuning; launch parity-before-speed kernel and serving checks; inspect adopted backends, CUDA launch reduction, cold compile cost, and regression gates; and retain checkpoint-bound machine-readable receipts under runs/. Backend adoption remains fail-closed; the dashboard cannot promote or publish a model.

Research capability inventory with readiness levels, entry points, and sources
Research capability inventory — one auditable index of the community-inspired decoding, recurrent-state, tokenizer, memory, optimizer, attention, data-filter, and runtime experiments in this repository, including each capability's readiness, executable entry point, and source. Experimental and oracle paths remain off by default.

Conclusive experiment campaigns

python -m rwkv_lab.experiment now treats a sweep as a reproducible campaign rather than a table of final averages:

python -m rwkv_lab.experiment --task recall:16 \
  --configs baseline,loop3,engram --factorial \
  --seeds 4 --steps 3000 --confirm-seeds 8
  • Baseline and candidate seeds consume identical deterministic batch/evaluation tapes.
  • Arms pass through configurable successive-halving rungs; promoted model, optimizer, RNG, data-tape position, and learning curve resume from persistent rung checkpoints instead of restarting.
  • The registry stores every trial, curve, profile, RNG hash, resolved config, environment, package set, dataset identity, and compressed dirty Git patch.
  • Comparisons use paired bootstrap intervals, sign-flip permutation tests, effect sizes, Holm correction, next-seed power guidance, and pre-registered O'Brien–Fleming alpha spending across interim rung looks.
  • Evaluation covers a length grid, corruption stress, NLL, calibration, loop engagement, and actual time/memory/energy—not only approximate FLOPs.
  • Exploratory winners are rerun on unused seeds in a child confirmation campaign; only those results can receive the dashboard's CONFIRMED badge.

The same normalized registry now owns synthetic, language-model, and conversion experiments. LM launches record each seed, JSONL curve, final checkpoint, and paired −validation-loss comparison. A declarative conversion campaign uses each layer/seed combination as a paired unit:

name: loop-gate conversion A/B
conversion:
  model_dir: Qwen3.5-9B-Base
  data: /path/to/qwen-token-cache
  layers: [0, 1, 2]
  args: {w_lmce: 1.0, w_block: 20.0, w_smt: 0.0, w_dmt: 0.0}
seeds: 2
train: {steps: 4000, seq_len: 1024, optimizer: schedulefree, eval_every: 100}
configs:
  baseline: {loop_count: 1}
  loop4: {loop_count: 4, loop_gate: factored}

Run it with python -m rwkv_lab.config run experiments/my_conversion.yaml; it produces the same campaign, trial, comparison, artifact, and reproducibility records shown by trainboard.

Verifiable-reward post-training

rlvr_train.py closes the model-side RLVR loop: grouped rollouts, deterministic rewards, GSPO/Dr.GRPO/DAPO policy updates, a fixed or rollout reference policy, held-out evaluation, resumable optimizer state, and lineage-bearing checkpoints. Sparse-reward cold starts can use trusted-answer SFT and staged curricula before a reward-diversity preflight, following DeepSeek-R1. The objectives follow Dr.GRPO, DAPO, GSPO, and the programmatic-verifier boundary from Absolute Zero.

python -m rwkv_lab.rlvr_train \
  --ckpt runs/lm/ckpt.pt --out runs/rlvr-gspo \
  --algorithm gspo --steps 100 --prompts-per-step 2 --group-size 8 \
  --curriculum-stages 1,2 --sft-steps 16 --preflight-prompts 8

--rollout-engine auto batches each rollout group and uses RWKV-7 constant-size recurrent state when the checkpoint is causally compatible. Future-Seed, looped, online-memory, and Engram checkpoints automatically use batched full-prefix recomputation so their semantics are not silently changed. Response scoring combines nearby lengths into token-budgeted padded forwards, uses fused selected-token cross entropy instead of a full-vocabulary log-softmax buffer, and computes group-relative statistics with device-side scatter reductions. The immutable scoring layout is built once per rollout set and reused for old-policy, frozen-reference, and every current-policy epoch. External verification overlaps that CPU preparation. Optional --rollout-devices cuda:0,cuda:1,... creates fresh inference replicas from the current policy each step, deterministically shards prompt groups, and reports wall-clock tokens/s.

For an equal-budget comparison, rlvr_campaign.py runs isolated GSPO, Dr.GRPO, and DAPO arms over paired seeds, then atomically aggregates their held-out reward, variance, update count, and promotion decisions into campaign.json:

python -m rwkv_lab.rlvr_campaign \
  --ckpt runs/lm/ckpt.pt --out runs/rlvr-comparison \
  --algorithms gspo,dr_grpo,dapo --seeds 0,1,2 \
  --steps 100 --prompts-per-step 2 --group-size 8

Trainboard's verifiable-reward training panel launches the same versioned campaign contract and renders the per-algorithm held-out mean/standard deviation, delta from the frozen parent, applied RL/SFT updates, preflight passes, rollout budget, and promotion count. It also discovers bounded recursive-loop lineage. It does not execute generated code: external verifier tasks retain the Adamaton sandbox boundary described below.

Without --tasks, the trainer creates disjoint deterministic arithmetic train/eval curricula. An Adamaton task producer can instead write rlvr_arithmetic.example.jsonl:

{"id":"t1","split":"train","prompt":"Compute 17 * 9.","verifier":{"kind":"numeric","expected":153},"metadata":{}}

Local verifiers support normalized exact answers, bounded arithmetic expressions, and final numeric answers. Generated code is never executed by RWKV-Lab. A task with verifier.kind="external" is sent to --verifier-command in a versioned batch JSON request; Adamaton owns that command's sandbox, private tests, timeout policy, and verifier independence. The final result.json, manifest.json, train.jsonl, and rlvr.pt form the machine-readable return contract. The parent checkpoint is never overwritten. A separate --heldout-tasks file can keep evaluation prompts outside the proposal process; rlvr_heldout.example.jsonl only demonstrates the format and is not a secret benchmark. Exact ID/prompt leakage fails before training. Promotion requires an informative update, point improvement, a paired bootstrap lower bound, no task-family regression beyond policy, and token/time budget compliance. Any failed gate names the untouched parent as the rollback target.

Structured adapter and preference post-training

posttrain_data.py defines one versioned JSONL contract for pretraining text, multi-turn SFT, preference pairs, binary feedback, step-labeled PRM, and existing RLVR tasks. Chat rendering retains system/user/assistant/tool roles, emits an explicit token-level loss mask, and canonicalizes structured assistant tool calls, and records both source and template hashes. Unlike the legacy corpus flattener, SFT trains only the assistant spans. The contract follows the typed data/template direction in LLaMA-Factory, without importing its Transformer model registry.

{"id":"s1","kind":"sft","split":"train","messages":[{"role":"user","content":"Compute 17*9."},{"role":"assistant","content":"153"}],"metadata":{}}
{"id":"p1","kind":"preference","split":"train","messages":[{"role":"user","content":"Explain the result."}],"chosen":"17 groups of 9 equals 153.","rejected":"It is 162.","metadata":{}}
{"id":"k1","kind":"feedback","split":"train","messages":[{"role":"user","content":"Be concise."}],"response":"Okay.","label":true,"metadata":{}}
{"id":"r1","kind":"prm","split":"eval","messages":[{"role":"user","content":"Show the proof."}],"steps":[{"text":"First valid step.","label":true},{"text":"Unsupported leap.","label":false}],"adversarial_steps":[{"text":"Reordered invalid step.","label":false}],"metadata":{"family":"proof"}}

Validate and preview the exact rendered variants before using them:

python -m rwkv_lab.posttrain_data datasets/posttrain.jsonl --json \
  > datasets/posttrain.manifest.json

Multiple validated sources can be merged into an immutable content-addressed version; duplicate content and train/eval/test overlap are rejected:

python -m rwkv_lab.posttrain_data data/sft.jsonl data/preferences.jsonl \
  --version-root datasets/versions --json

adapters.py implements frozen-base, zero-output-init LoRA for explicitly selected RWKV time/channel-mix linears. It supports multiple named adapters, activation/deactivation, deterministic merge/unmerge, base-model fingerprints, and safetensors artifacts. quantization.py adds an opt-in native QLoRA correctness backend: packed NF4 weights, per-block scales, a frozen recurrent base, adapter-gradient and zero-init checks, dense-merge parity, and measured stored bytes. Its NF4 table matches the primary bitsandbytes implementation. The portable path dequantizes for F.linear and is deliberately correctness-scale only; the optional TorchAO NF4Tensor path uses its accelerator/compile dispatch and double-quantized scales. --quant-backend auto adopts TorchAO only after representative-layer output, input-gradient, storage, and median-throughput qualification and fails closed when no accelerated backend passes. The dequantizing oracle must be requested explicitly with --quant-backend portable, preventing an apparently production run from silently taking the measured ~8×-slower reference path.

The executable trainer supports SFT, four preference paths, outcome rewards, and full process-reward training. Tokenization may be streamed into a safe content-addressed cache; its receipt reports token length percentiles and truncation counts. --packing reset best-fit packs examples, resets RWKV's matrix state, TimeMix/ChannelMix/DeepEmbed token shifts, and RoPE position at every boundary, and uses FLA cumulative sequence lengths on the accelerated kernel. Before the first optimizer step, the exact objective must pass unpacked-vs-packed loss and trainable-gradient parity. Unsupported stateful levers fail closed; audit remains available to inspect utilization without executing a pack. Immutable tokenized examples are tensorized once and transferred as one padded batch. DPO/KTO frozen-base log-probabilities are cached once per example rather than recomputed every optimizer step; PRM calibration is evaluated at reporting cadence rather than introducing per-bin GPU synchronization into every update. --log-every N also batches scalar loss/non-finite telemetry; gradients are sanitized on-device each step so a bad update cannot poison adapter state before the next fail-fast telemetry boundary.

python -m rwkv_lab.posttrain_train \
  --checkpoint runs/lm/ckpt.pt --data datasets/posttrain.jsonl \
  --eval-data datasets/posttrain-heldout.jsonl --objective dpo \
  --output runs/posttrain-dpo --rank 16 --alpha 32 --steps 500 \
  --token-cache caches/posttrain --packing reset \
  --base-quantization nf4 --quant-backend auto

The loss implementations map directly to DPO, KTO, ORPO, and SimPO. preference.py also contains pairwise outcome-reward heads/losses following InstructGPT and step-position process-reward heads/masked losses following Let's Verify Step by Step. The prm objective trains the head, reports held-out accuracy/Brier/ECE, evaluates explicitly supplied adversarial steps, and versions the reward head with base/data/weight hashes. Learned rewards do not bypass hidden-split confidence, calibration, adversarial, or family-regression promotion gates.

Post-training campaigns put SFT/DPO/KTO/ORPO/SimPO/RM/PRM behind the same evidence contract:

python -m rwkv_lab.posttrain_campaign \
  --checkpoint runs/lm/ckpt.pt --data datasets/posttrain-train.jsonl \
  --eval-data datasets/posttrain-heldout.jsonl --output runs/posttrain-campaign \
  --objectives sft,dpo,kto,orpo,simpo --seeds 0,1,2 \
  --confirm-seeds 100,101,102 --token-budget 100000 \
  --devices cuda:0,cuda:1 --arm-timeout 14400 --retries 2

Every arm receives the same scored-input-token budget and frozen-parent held-out examples. Per-example loss deltas are paired, bootstrapped, and checked for minimum gain, confidence lower bound, content/ID leakage, and task-family regression. Unused confirmation seeds are a distinct registry phase. Only a successful confirmation writes an eligible rwkv-lab.posttrain-promotion.v1 receipt; each trained adapter remains preserved regardless of the decision. Every arm also writes an atomic command-hashed state/attempt history. Restarting the same campaign skips verified completed arms, retries failed or timed-out attempts with the original seed, refuses configuration drift, schedules at most one arm per listed device slot, and completes exploration before launching fresh-seed confirmation.

adapter_recursive.py applies that receipt to recursive improvement. Following the proposer/solver lineage explored by Absolute Zero and Self-Rewarding Language Models, Adamaton may propose bounded train-only records and a small allowlisted configuration. Each round trains an isolated adapter; rejections stay on disk, and only a confirmed adapter is densely materialized into a new immutable parent with checkpoint, adapter, parent, and receipt hashes. Held-out data and promotion authority are never sent to the proposal command.

posttrain_kernels.py benchmarks compiled LoRA branches, outcome/process reward heads, and batched recurrent preference scoring. A candidate is adopted only when output and gradient parity pass and median timing improves; activation offload uses PyTorch's save_on_cpu.

Production kernel qualification

production_kernels.py applies the same parity-before-speed policy to training and serving backends. It records the exact CUDA/device capability, backend availability, output/state/gradient errors, median timings, and adoption decision in one JSON receipt:

python -m rwkv_lab.production_kernels --device cuda --output runs/kernel-qualification.json
python -m rwkv_lab.production_kernels --device cuda \
  --baseline runs/kernel-qualification.baseline.json \
  --max-throughput-regression 0.05 --max-memory-regression 0.10 \
  --max-kernel-regression 0.10 \
  --output runs/kernel-qualification.json
python -m rwkv_lab.production_kernels --device cuda --checkpoint runs/lm/ckpt.pt \
  --prompt-ids 1,123,456 --max-new 64 \
  --megakernel-artifact runs/qualification/rwkv-megakernel.pt2 \
  --output runs/serving-qualification.json
python -m rwkv_lab.generate --ckpt runs/lm/ckpt.pt --engine auto \
  --megakernel-receipt runs/serving-qualification.json \
  --megakernel-artifact runs/qualification/rwkv-megakernel.pt2 \
  --megakernel-serving-prepare \
  --prompt "Explain RWKV."

  • NF4: TorchAO remains opt-in and is adopted only after representative-layer output, input-gradient, storage, and throughput checks against the packed NF4 oracle.

  • NVFP4: --nvfp4-backend transformer_engine uses NVIDIA Transformer Engine's native Blackwell E2M1 tensor-core recipe, including hierarchical scales and optional RHT, and fails closed unless it beats the fake-quant oracle on parity and speed. This implements the production path described by NVFP4 pretraining and TetraJet-v2 through the documented Transformer Engine NVFP4 API.

  • Online memory: fixed-size Titans/MIRAS/ATLAS/Nested scans can be qualified under torch.compile. Recurrent inference carries the memory matrix, momentum, and ATLAS key/value window across chunks and must match a full-prefix scan. --online-memory-kernel auto installs the compiled live-parameter path only after parity and throughput pass; the reference module remains the stateful oracle.

  • Triangular delta rule: direct float64/float32 triangular solves remain the correctness oracle; iterative or hardware implementations inspired by Sobczyk et al. are adopted only after inverse/residual parity and measured speed pass.

  • Hard ROSA: the ROSA-Tuning online suffix automaton runs entirely on the current CUDA stream, reuses a stream/shape-isolated workspace, and retains the Numba CPU implementation as its exact oracle. Qualification includes throughput and projected long-context workspace fraction; allocation fails before launch when it would consume over half of currently free device memory.

  • Megakernel serving: --engine megakernel explicitly compiles a native RWKV plan; production qualification separately records cold compilation, warm token throughput, exact tokens, and a four-stage eager → fused-state/epilogue → compiled-fullgraph → CUDA-Graph latency/launch ablation. --engine auto --megakernel-receipt ... adopts only an approved receipt whose SHA-256 matches the checkpoint bytes and whose GPU compute capability, PyTorch, and Triton versions match the current runtime, preventing stale evidence from authorizing another model or machine. Inductor and Triton reuse their on-disk compiler caches; the process retains CUDA Graphs by device, batch, dtype, and recurrent-state geometry. The implementation cites and follows the execution-plan direction of HazyResearch Megakernels and TileRT, plus Albatross's folded block-0 normalization, without loading their model-specific binaries. --megakernel-artifact loads an adjacent hash-checked .pt2.json manifest and AOT plan whose checkpoint, plan geometry, compute capability, PyTorch, and Triton identities must all match. Exact prompt shapes use adjacent *.prefill-bN-tN.pt2 artifacts when available. Qualification receipts include per-path GPU time, allocation traffic, and the hottest CUDA kernels—not just launch counts—plus separately gated row-one GEMV, distinct-input packed R/K/V, and exact squared-ReLU/value candidates. --megakernel-serving-prepare is deliberately destructive and requires an adopted receipt; use it only for an immutable serving process.

  • SM120 evidence: on compute capability 12.0, the megakernel receipt also includes the physical-SM persistent-scheduler plan, CC12 occupancy geometry, native-versus-portable CUTLASS metadata, and the measured persisting-L2 window. State parity and median latency gate the scheduler, full-model parity and latency gate the L2 window, and projection requires an actual native-SM120 operator—being on a Blackwell GPU alone adopts nothing.

    On the local RTX PRO 6000 Blackwell qualification model (batch 1, d256, four layers, 1,024-token synthetic vocabulary), exact greedy decode measured 6,710 → 5,379 → 2,116 → 130 µs across those four paths and 153 → 11,556 end-to-end tokens/s with one-replay 32-token generation. Profiler events fell from 313 to 101. The combined state/epilogue passed at 1.02×; portable layer-boundary fusion (0.93×) and generic row/RKV/FFN candidates (0.34–0.50×) correctly remained disabled. This is a reproducible systems smoke test, not a quality or full-scale-model throughput claim; every real checkpoint must produce its own bound receipt.

  • Serving fallback: without an adopted megakernel receipt, generate --engine auto selects RWKV-7 constant-state decoding only for compatible checkpoints and reports tokens/s; otherwise it records the reason for exact full-prefix fallback. The EAGLE-3 verifier validates a whole draft in one target call, reports acceptance/target-call throughput, and cannot be adopted unless its output tokens exactly match ordinary target-greedy decoding. Its direct draft positions share one packed vocabulary GEMM (legacy per-position head checkpoints migrate on load).

Trainboard's post-training panel validates repository-confined JSONL, previews the rendered text and trainable spans, reports duplicates/split leakage, materializes validated content-addressed versions/merges, and performs base-versus-candidate generation with identical prompt, seed, temperature, and token budget. An explicit operator choice can append a training-only preference to datasets/trainboard_preferences.jsonl; it never modifies held-out evaluation data. It also launches the allowlisted campaign command and reads campaign phase, paired interval, promotion, and adapter-loop lineage receipts from runs/. It cannot run an Adamaton proposal command, merge an adapter, overwrite a parent, or publish a model.

FSDP2 checkpoints and safe exports

From-scratch pretraining has an opt-in PyTorch FSDP2 path with bottom-up RWKV block sharding, collective gradient clipping, per-block activation checkpointing, optional CPU offload, and Distributed Checkpoint model/optimizer/per-rank RNG state. Resuming at the same world size is exact; model and optimizer state can be resharded by DCP when the world size changes.

torchrun --standalone --nproc-per-node=4 -m rwkv_lab.rwkv_pretrain \
  --distributed fsdp2 --activation-checkpointing \
  --data models/corpus.bin --out runs/fsdp2 --steps 1000 --save runs/fsdp2.dcp

torchrun --standalone --nproc-per-node=4 -m rwkv_lab.rwkv_pretrain \
  --distributed fsdp2 --activation-checkpointing \
  --data models/corpus.bin --out runs/fsdp2-resume --steps 2000 \
  --resume runs/fsdp2.dcp --save runs/fsdp2-next.dcp

export_bundle.py converts a trusted self-describing single-process checkpoint to safetensors plus architecture, chat-template, tokenizer, adapter, dataset, lineage, and promotion receipts. Every file is hashed and re-opened for verification. Export does not publish or execute anything; external hub publication remains a separate manual operator action.

python -m rwkv_lab.export_bundle --checkpoint runs/lm/ckpt.pt \
  --adapter runs/posttrain-dpo/adapter --dataset-manifest datasets/posttrain.manifest.json \
  --promotion-receipt runs/posttrain-dpo/promotion.json --output exports/candidate

Bounded recursive improvement

recursive_improve.py implements the first closed, reversible iteration inspired by Absolute Zero and Self-Rewarding Language Models, while keeping rewards independent:

python -m rwkv_lab.recursive_improve \
  --ckpt runs/lm/ckpt.pt --out runs/recursive-rlvr \
  --heldout-tasks /secure/eval_tasks.jsonl \
  --proposal-command '/path/to/propose-rwkv-tasks' \
  --verifier-command '/path/to/verify-rwkv-batch' \
  --rounds 3 --max-total-rollout-tokens 1000000

The versioned proposal request reveals the parent hash, accepted/rejected history, and immutable operator caps—but never the held-out file or private verifier results. Adamaton may return training tasks and a small allowlisted, cap-checked configuration. Each candidate runs in an isolated directory; only rlvr_train's independent multi-gate decision can advance the parent pointer. The controller stops at its round, rollout-token, wall-clock, proposal-size, or consecutive-rejection limits and atomically persists loop.json for audited resume.

Conversion track — target model

The conversion levers are developed against Qwen3.5-9B-Base — a 32-layer, hidden-size-4096 hybrid (the loop / latent-prediction / memory / optimizer levers are model-agnostic and drop onto any RWKV-7/8 core):

Layers Mechanism Geometry
Linear 24 (all except every 4th) Gated DeltaNet (GDN) 32 value heads × 128, 16 key heads × 128
Full attention 8 (indices 3, 7, 11, 15, 19, 23, 27, 31) Gated GQA + RoPE + per-head q/k-norm 16 query heads × 256, 4 KV heads (GQA rep 4)

A second track targets Qwen3.6-35B-A3B (a Mixture-of-Experts model) for the MLA / Engram experiments — the origin of this repo's old moe-mla name.

What you need locally

This repo does not contain model weights, datasets, generated manifests, feature/token caches, run logs, or checkpoints. Supply the artifacts required by the track you are running and keep their upstream licenses with them:

Input Used by Notes
RWKV G1H checkpoint + World vocabulary multimodal captioning Pass the checkpoint with --rwkv; set VOCAB to rwkv_vocab_v20230424.txt. The checked-in experiment defaults target the 2.9B G1H model.
Kimi K2.6 MoonViT shard multimodal captioning and vision caches Pass model-00064-of-000064.safetensors with --moonvit; the other Kimi language-model shards are not required for the frozen vision tower.
Image-caption JSONL multimodal training/evaluation Use the documented schema, with disjoint explicit train/eval manifests. Images remain external to Git.
Optional SigLIP2, DINOv2, and SAM weights multi-teacher vision experiments Frozen towers used by the advanced cache/fusion path; set the corresponding VISION_NEXT_* paths before preparation.
Qwen3.5-9B-Base weights conversion, baseline eval, target extraction Pass with --model-dir, or put the HF snapshot at Qwen3.5-9B-Base.
Tokenized eval/train stream eval_baseline.py, build_memory_targets.py, convert_train.py --data may be a cache directory or a flat tokens.bin accepted by build_memory_targets.load_token_stream.
CUDA Torch + fla RWKV-7 kernel path Install these for your CUDA stack; requirements.txt only covers the regular Python deps.
Go toolchain dashboard/ Needed only for trainboard.

For a small data-format smoke test, build_qwen35_data.py --max-docs 1000 --out_root /tmp/qwen35-cache writes the same flat-cache format without pulling the full corpus.

Repository layout

Python source lives under src/rwkv_lab/; entry points run as python -m rwkv_lab.<module>. Conversion components are generally drop-in linear_attn or attention-module swaps, while the vision, pretraining, post-training, and serving tracks provide their own explicit trainers and contracts. Model weights, datasets, checkpoints, generated manifests, caches, and local paper archives are git-ignored—the repository publishes source and reproducibility machinery, not third-party artifacts.

Conversion core

File Role
rwkv8_deltanet.py RWKV-7/8 time-mix + channel-mix modules (port of BlinkDL's RWKV_Tmix_x070), using fla's Triton wkv7 kernel with a Python reference fallback. The swap target.
convert_gdn_lossless.py The lossless GDN→RWKV-7 kernel remap (proven above). Weight-preserving, zero training.
convert_train.py Single-layer conversion trainer: block-MSE + logit-KL + SMT/DMT state distillation, device-side non-finite gradient guards with cadence-based host telemetry, spectral-optimizer levers. See TRAINING_LEVERS.md.
smt_dmt.py Supervised (one-step) + Dynamical (closed-loop rollout) Memory Training + the bilinear state codec.
distill_objectives.py Alignment-invariant / relational distillation losses (CKA, relative-L2).
attn_L3_poc.py Full-attention→RWKV proof-of-concept (RADLADS init + freeze-most), self-contained.
layer_swap.py, svd_init.py Hot-swap a decoder layer's mixer; SVD-based weight transfer.
build_memory_targets.py Extract frozen-teacher GDN state/block targets for the SMT/DMT caches.
assemble_looped.py, drive_isolation.py, distill_consolidate.py Assemble independently-converted layers into one looped model; drive the per-layer sweep; joint consolidation pass.
load_converted.py, eval_baseline.py Load a converted stack; evaluate the untouched base for reference PPL.

From-scratch pretraining lab

File Role
rwkv_pretrain.py From-scratch RWKV-7 LM trainer where every lever attaches natively: seed-chain, Engram, DeepEmbed (all variants), loops, aux heads; mixed context-length training (--ctx-buckets, reciprocal batch); grad-accum, fp32 EMA shadow weights, fp8, 8-bit optimizers; GPU-resident window sampling; fused LM-head CE.
experiment.py / experiment_analysis.py N-seed lever A/Bs (named LEVERS) with preflight gates, length-gen / noise eval grids, and paired bootstrap CIs + Holm correction.
synthetic_tasks.py GPU-native copy / associative-recall / induction diagnostics (accuracy, not ±0.1-nat ppl noise).
build_corpus.py ztok corpora from local files or streamed HF datasets (chat records flattened to role-tagged text); doc offsets; semantic context-bucket packing (whole docs, best-fit-decreasing, ~0% padding).
config.py Declarative YAML campaigns; the corpora registry (local / blend / blend-mix) behind the dashboard's LM tasks.
registry.py / fused_ce.py Campaign/trial SQLite registry; fused LM-head cross-entropy (flash CE, pad-masking).
rlvr.py / rlvr_train.py / rlvr_evaluation.py / rlvr_campaign.py / recursive_improve.py GSPO/Dr.GRPO/DAPO objectives; cold-start/preflight training; recurrent batched rollouts; leakage, confidence, family-regression, and budget gates; equal-budget campaigns; and bounded Adamaton proposal lineage.
posttrain_data.py / posttrain_train.py Typed tool-aware post-training JSONL, streamed content-addressed token caches, qualified reset-mask recurrent multipacking, and executable native-RWKV SFT/DPO/KTO/ORPO/SimPO/ORM/PRM training.
adapters.py / quantization.py / preference.py Named RWKV-aware LoRA lifecycle, portable packed-NF4 QLoRA qualification, and reference-tested preference/outcome/process-reward losses.
posttrain_campaign.py / adapter_recursive.py / posttrain_kernels.py / production_kernels.py Equal-token paired/confirmation campaigns and receipts; adapter-first immutable recursive parents; parity-before-speed training and serving kernel qualification.
distributed.py / export_bundle.py FSDP2 + DCP exact-resume state and verified safetensors export packages with lineage/promotion receipts.

Looped recurrence (recurrent depth)

File Role
looped_rwkv.py Weight-tied N-loop refinement wrapper (pre-norm + zero-init residual gate ⇒ identity at init). Factored head/channel gates, spectral-radius cap, hyper-connection lanes.
loop_probe.py Loop-iterate diagnostics + depth-usefulness sweep.
looped_rwkv_rosa_engram_v3.py The integrated looped + ROSA + Engram core.

Memory: Engram + ROSA

File Role
engram_lmb.py, engram_lmb_build.py Lexical Memory Bank — a suffix-automaton-recalled embedding memory (offline builder + runtime module).
engram_integration.py, build_engram_patch.py, gpu_engram_prefill.py Wiring, patch builder, GPU prefill of the memory table.
rosa.py, rosa_sam.py, rosa_soft_layer.py ROSA suffix-matching retrieval (v1 drop-in, device-native online suffix-automaton kernel with CPU oracle, soft-retrieval layer).
verify_engram.py, load_mla_engram.py Verification + combined MLA+Engram loader.

MLA (Multi-head Latent Attention)

File Role
mla_module.py DeepSeek-V2/V3-style MLA attention module (hot-swappable).
train_mla.py, train_mla_engram.py GQA→MLA finetune trainers (frozen backbone, MLA-only params).

Prediction objectives (research modules, training-only, default-off)

File Role
mtp_module.py, parallel_heads_module.py Multi-token-prediction auxiliary heads (Gloeckle-style).
mutor_module.py MuToR register-token auxiliary MTP.
lookahead_module.py, fsp_module.py Latent-lookahead + future-summary auxiliary objectives.
pc_layer.py PC-Layer polynomial weight preconditioning.

Optimizers & diagnostics

File Role
muon_helpers.py, spectral_muon.py MuonClip helpers; one configurable Muon-family optimizer collecting the 2026 spectral-optimizer literature (Muonᵖ, DDC, distance-aware, hierarchical, …).
llr.py Heavy-tail layerwise learning rate.
grokking_metrics.py, grok_autopilot.py Memorization-vs-grokking diagnostics + reactive recovery.

Infra

File Role
dashboard/ trainboard — Go + SQLite + Datastar + Pixi.js real-time training dashboard (see its README).
live_controls.py Trainer-side consumer of the dashboard's live-tuning panel.
safe_torch.py Safer torch-serialization load wrappers.
build_qwen35_data.py Build Qwen3.5-tokenized DCLM + FineWeb-Edu caches.
tests/ CPU/GPU invariant + feature tests (loops, lookahead, Engram, ROSA, the SOTA levers). Run scripts/test_parallel.sh for process-parallel CPU tests with bounded native thread pools plus serialized CUDA tests (PYTEST_WORKERS and PYTEST_NATIVE_THREADS tune the split). Set RWKV_GPU_STRESS=1 to append the idle-GPU compile-core and DMT graph programs.
scripts/ Overnight sweep / A-B drivers (gate_ab.sh, gdn_sweep.sh, rel_sweep.sh, supervisor_night.sh).

Pipeline

# 0. install Python deps, then install CUDA-specific torch + fla separately
pip install -r requirements.txt
pip install -e '.[test]'  # package + pytest/xdist parallel test runner
#                          use a venv (e.g. python -m venv --system-site-packages .venv) for the
#                          CUDA torch/fla stack; tests also run without this via tests/conftest.py

MODEL_DIR=/path/to/Qwen3.5-9B-Base
DATA=/path/to/qwen3.5-token-cache-or-tokens.bin

# 1. baseline eval on the same windows used by conversion runs
python -m rwkv_lab.eval_baseline --model-dir "$MODEL_DIR" --data "$DATA" --out runs/_baseline.json

# 2. GDN layers — lossless, no training
python -c "from transformers import AutoModelForCausalLM; \
           from rwkv_lab.convert_gdn_lossless import install_lossless_wkv7; \
           m = AutoModelForCausalLM.from_pretrained('$MODEL_DIR'); \
           print(install_lossless_wkv7(m), 'GDN layers converted')"

# 3. attention layers — per-layer distillation against the frozen original
python -m rwkv_lab.build_memory_targets --model-dir "$MODEL_DIR" --data "$DATA" --layer 3 --out mem_targets/L3
python -m rwkv_lab.convert_train --model-dir "$MODEL_DIR" --data "$DATA" \
  --layer 3 --codec-cache mem_targets/L3 --out runs/iso_L3 --steps 10000

# 4. assemble accepted isolated layers, then consolidate
python -m rwkv_lab.assemble_looped runs/iso_L*/best/ckpt.pt --out Qwen3.5-9B-RWKV/rwkv_layers_looped.pt
python -m rwkv_lab.distill_consolidate --model-dir "$MODEL_DIR" --data "$DATA" \
  --rwkv-ckpt Qwen3.5-9B-RWKV/rwkv_layers_looped.pt \
  --kl-weight 1.0 --out Qwen3.5-9B-RWKV/rwkv_layers_distilled.pt

# 5. watch it (separate terminal)
go -C dashboard run ./cmd/trainboard   # http://127.0.0.1:9124

Some research scripts retain machine-local defaults as convenience examples; pass explicit paths or the documented environment variables on another host. Conversion levers described below default off, while specialized experiment launchers explicitly opt into named bundles. See TRAINING_LEVERS.md.

Levers & flags

python -m rwkv_lab.convert_train exposes a large set of optional research levers; its documented baseline leaves them disabled so each can be evaluated in a controlled A/B. Many are live-tunable from trainboard without restarting a run. The complete manual—defaults, sources, compatibility constraints, and when to use each lever—is TRAINING_LEVERS.md. Headline levers:

Recurrent-depth loops — wrap the RWKV layer in LoopedRWKV

Flag Turns on Paper
--loop-count N N weight-tied refinement passes Iso-Depth Scaling Laws
--loop-hyper K K hyper-connection lanes at the loop boundary Hyper-Connections
--loop-iter-readout supervise every loop iterate toward the teacher Readout Blind Spot
--loop-adaptive-halt PonderNet per-token adaptive depth PonderNet
--loop-cart-anchor contractive LTI gate (bounds the deep loop) CART
--loop-deq (--loop-deq-window k) DEQ 1-step / Neumann-k gradient (O(1) memory) HRM · FPRM
--loop-fp-halt fixed-point-residual halting FPRM

Optimizer — --optimizer spectral_muon (12 levers, all off; the flagships)

Flag Turns on Paper
--sm-spectral-power p Muonᵖ fractional-power orthogonalization 2606.13867
--sm-mona MONA momentum-Nesterov 2605.26842
--sm-rsav SpecMuon gradient-energy adaptation 2602.16167
--sm-tile-size T Hierarchical / tiled Newton–Schulz 2606.27216
--sm-da-muon Distance-Aware adaptive radius 2605.18999
--sm-aro ARO-Sinkhorn (replaces orthogonalization) 2602.09006
--sm-aro-compile 1 parity-tested compiled ARO tensor subgraph; foreach Adam fallback stays exact-resumable 2602.09006
--sm-ddc-strength Dead-Direction Conditioner 2606.29176

Distillation & grokking

Flag Turns on Paper
--block-loss rel per-token relative-L2 block match (OpenMOSE) —
--nuc-weight nuclear-norm generalization penalty 2606.04405
--grokfast Grokfast slow-gradient amplification 2405.20233
--logit-kl (attn PoC) top-k logit self-distillation RADLADS

From-scratch pretraining — python -m rwkv_lab.rwkv_pretrain (all off by default; also launchable as dashboard lever checkboxes)

Flag Turns on Source
--seed-chain Future-Seed: layer L's wkv scan starts from layer L−1's final state yanghu819/future-seed
--engram (+--engram-sites/-drow/-rows/-boundary-id) Engram lexical memory bank as a native lever (token-SAM recall, gated injection, copy head) DeepSeek Engram
--deepembed · --de-mode hidden · --de-shift · --de-emb-res DeepEmbed per-token FFN gates — v1 output gate, BlinkDL-exact bilinear hidden gate, separate gate token-shift, emb-residual fold BlinkDL RWKV-LM (rwkv_v7a)
--ctx-buckets meta.json mixed context-length training over packed 512…32k buckets, reciprocal batch (B = budget/T), pad-masked loss, per-bucket val —
--grad-accum N / --ema d / --fp8 / --optimizer adamw8bit large-batch simulation · fp32 EMA shadow weights · torchao fp8 GEMMs · bitsandbytes 8-bit moments —
--u-mup-base-width W (+ --u-mup-base-depth L) u-μP initialization and AdamW LR groups for width/depth scale transfer u-μP
--online-memory 1 (+ --online-memory-mode, --online-memory-kernel) in-forward associative memory; Titans, MIRAS, ATLAS, and learned nested-controller modes; parity-gated compiled scan Titans · MIRAS · ATLAS · Nested Learning
--nvfp4 (+ --nvfp4-rht, --nvfp4-backend) E2M1 fake-quant oracle or parity-gated native Transformer Engine NVFP4 on Blackwell NVFP4 pretraining · TetraJet-v2

Prediction & memory objectives are aux heads / standalone modules, not convert_train flags: the lookahead heads (L-MTP, Belief-State, JTP, TOP, NextLat, …) are wired via lookahead_module.lookahead_from_args; LLM-JEPA, Coconut, L³, FwPKM, and WriteSAE are standalone modules for a paired-data / SFT stage — see References.

Status

Area Evidence boundary
Multimodal MoonViT → RWKV training Implemented with exact resume, explicit train/eval separation, qualitative artifacts, cache receipts, and CPU-tested contracts; caption quality and large-scale convergence remain experimental
Multi-teacher vision compressor/student Architecture and cache contracts designed; the first teacher shard is an experiment input, not evidence that the proposed deployable student has been trained
GDN → RWKV-7 conversion (24 layers) Validated exact remap: cosine 0.999995 and +0.013% full-model perplexity change
Full-attention → RWKV conversion (8 layers) Active distillation work; the RADLADS proof of concept and two-stage logit-KL path do not yet constitute an end-to-end converted release
From-scratch lever validation Lab-scale wins include seed-chain −9.2% PPL, Engram 15× induction accuracy across four seeds, and DeepEmbed de_shift −2% replicated across corpora (results)
Dashboard and run infrastructure Go tests, ingestion/schema tests, eval-artifact tests, exact-resume contracts, and fail-closed launcher checks are implemented; operational behavior still depends on the local CUDA/filesystem stack
Large-model composition Open problem: isolated or small-model wins must still survive joint training and substantially larger compute budgets

The rule throughout the repository is to label an implementation, a local measurement, and a production-ready result as three different things. Claims above are intentionally bounded by the evidence that is checked into source or described by a reproducible receipt.

References

Papers we build on

Only papers with a concrete implementation or adopted design decision in this repo are listed (each maps to the module named, arXiv id linked). The wider reading pile is intentionally not committed.

Architecture conversion

Looped / recurrent depth

From-scratch pretraining levers

  • Future-Seed — cross-layer recurrent-state chaining (s₀ of layer L = s_T of layer L−1); validated here for length generalization and −9.2% LM ppl → rwkv_pretrain.py (--seed-chain)
  • DeepEmbed (BlinkDL RWKV-LM, rwkv_v7a) — per-layer per-token FFN gates; our A/B independently reproduces BlinkDL's variant ordering (separate gate token-shift is the win) → rwkv_pretrain.py (--deepembed, --de-mode hidden, --de-shift, --de-emb-res)
  • Fewer Truncations Improve Language Modeling — the best-fit packing idea behind our semantic context-bucket packer (whole docs, standard context sizes, ~0% padding) → build_corpus.py (pack_context_buckets)

P0 — scale transfer, online learning, and verifiable rewards

P1 — data and distributed/numerical training systems

P2 — bytes, speculative serving, and interpretability

Memory (Engram / ROSA)

Latent attention & prediction objectives

Optimizers & training dynamics

Code & upstream references

Project Use here
BlinkDL/RWKV-LM The RWKV-7/8 time-mix design; run_rwkv7_qwen35.py validated our GDN→RWKV mapping.
fla-org/flash-linear-attention Triton chunk_rwkv7 / wkv7 kernels used throughout.
RADLADS (Recursal) · recursal/QRWKV The attention→RWKV distillation protocol (block-align → logit-KL → CE) and RAD-RWKV7 init.
OpenMOSE Normalized-MSE + logit-KL conversion recipe guidance.
DeepSeek-V2/V3 MLA formulation (mla_module.py).
Qwen3.5 / Qwen · HF Transformers Base model + modeling code.
Muon (Keller Jordan) · schedule-free Optimizer bases.
Open-PerfectBlend (mlabonne, Apache 2.0) The lab's real LM corpus — 788k chat/math/code conversations → 388M ztok tokens (blend / blend-mix).
bitsandbytes · torchao · Transformer Engine 8-bit optimizer states · FP8/NF4 · native Blackwell NVFP4. Hardware backends are adopted only through qualification receipts.
Datastar · Pixi.js trainboard front-end (hypermedia SSE + WebGL charts).

Acknowledgments

Special thanks to BlinkDL (creator of RWKV) and OpenMOSE for their direct guidance on this work — BlinkDL for the RWKV-7 architecture and sharing the run_rwkv7_qwen35.py reference that confirmed our GDN→RWKV kernel mapping, and OpenMOSE for the normalized-MSE + logit-KL conversion recipe and RADLADS pointers that shaped the distillation pipeline. This project would not have gotten off the ground without their generosity.

License

MIT — original code in this repository. The referenced papers, base model weights (Qwen3.5/3.6), and upstream projects retain their own licenses; this repo contains no model weights or copyrighted PDFs.

RWKV-Lab is independent research and is not affiliated with the RWKV project, Recursal, or Alibaba/Qwen.

About

Cross-architecture LLM conversion: distilling pretrained Transformer/GDN models into RWKV linear-attention cores, with looped recurrence and memory (Engram/ROSA).

Resources

Readme
License
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages