StateDiffRWKV-2.9B-512-pretrained

DiffRWKV State-Hijacking RELAY 2.9B 512 Trajectory Champion (from-scratch joint co-adapt + condboundary, pretrained/no-SFT, 57.67)

This release contains a DiffRWKV State-Hijacking RELAY checkpoint plus the minimal inference code needed to load it.

Files

• model.pt: trainable RELAY checkpoint only. It does not include the frozen RWKV backbone weights.
• configs/rwkv_relay_2.9B_state_hijack_dit_vae32_traj32x16.yaml: training/inference config used by the checkpoint.
• models/state_hijacking_dit.py: RELAY model implementation.
• scripts/eval/relay_utils.py: checkpoint loader.
• scripts/eval/sample_prefix_suffix_trajectory_cfg.py: inference entrypoint.
• manifest.json: size, SHA256, source path, and step metadata.

Download (all assets)

# 1. This repo: planner checkpoint (model.pt) + inference code
hf download humanlong/StateDiffRWKV-2.9B-512-pretrained --local-dir ./champion

# 2. Frozen RWKV-7 2.9B backbone (official RWKV repo, downloaded separately)
hf download RWKV/RWKV7-Goose-World3-2.9B-HF --local-dir ./models/RWKV7-Goose-World3-2.9B-HF

# 3. Preprocessed 512-token training data (OpenWebText, RWKV tokenizer)
hf download humanlong/laces-owt-rwkv-tokens --repo-type dataset --local-dir ./data
tar xzf ./data/owt_rwkv_tokens_train.tar.gz -C preprocessed_data/owt_rwkv_tokens/

Base Model

RWKV/RWKV7-Goose-World3-2.9B-HF

Important — backbone path: the shipped config contains an absolute rwkv_local_path from the training machine. After downloading the backbone, either place it at that path or edit rwkv_local_path in configs/rwkv_relay_2.9B_state_hijack_dit_vae32_traj32x16.yaml (or patch cfg.model.rwkv_local_path in scripts/eval/relay_utils.py after loading) before constructing the model.

Training data

The 512-token champion recipe trains on humanlong/laces-owt-rwkv-tokens (304,052 OpenWebText samples, max length 512, RWKV7-Goose-World3 tokenizer). Extract it so that data.token_dir points at preprocessed_data/owt_rwkv_tokens/train.

The 4096-token variant uses a separate FineWeb packed dataset (~45 GB) that is not yet published here; the 512 recipe above is the champion line.

Install

pip install -r requirements.txt
pip install flash-linear-attention fla-core

Use the same FLA build as the training environment when reporting numbers.

Inference

CUDA_VISIBLE_DEVICES=0 python scripts/eval/sample_prefix_suffix_trajectory_cfg.py --ckpt_dir . --prompt 'The history of artificial intelligence' --output sample.json --steps 100 --cfg_scale 3 --max_new_tokens 512 --trajectory_s1_mode independent --trajectory_state_blend 0.7 --temperature 0.5 --top_k 5 --top_p 0.7 --repetition_penalty 1.2