SII-Jiaquan/StateDiffRWKV-2.9B-trajectory4096-pretrained

StateDiffRWKV 2.9B Trajectory4096 Pretrained

This repository contains the best currently recorded trajectory4096 pretrained / pre-SFT checkpoint for StateDiffRWKV. It is uploaded with pretrained in the name because it is not a response-SFT model.

Checkpoint

• File: model.pt
• Source path: outputs_relay/s3a-4096-s2-align-kl-4h200-dai-20260701-highutil-v2/step_00005000/model.pt
• Checkpoint step: 5,000
• Backbone: RWKV/RWKV7-Goose-World3-2.9B-HF
• Context: 4096 tokens
• Trajectory layout: 64 x 64 (trajectory_chunk_size=64, trajectory_horizon=64)
• Training data: preprocessed_data/fineweb_4096_packed_full

What this is

This is the best scored 4096-token trajectory checkpoint available at release time:

• S1 mode: rwkv
• S2 denoiser: birwkv
• objective: DDPM S2 with decoded-logit KL alignment (s2_align_loss_weight=0.01)
• trajectory state blend: 0.5
• SFT: no

This is a pretrained/pre-SFT checkpoint from the older 4096 route. The newer 4096 independent S1 + birwkv S2 + condboundary recipe had been submitted but was still queued at release time, so it is not included here.

Recorded eval result

Protocol: CoLA-style generate-then-match MMLU proxy, 200 samples, cfg_scale=1.

Checkpoint	MMLU proxy	Average
step5k	43.50	43.50

Evidence path in the source repo:

baseline/Cola-DLM/eval_output/tasks_s3a_dai_v2_step5000_mmlu200_cfg1/accuracy_summary.csv

Important: this is a repo CoLA-style generative MMLU proxy, not official Hendrycks MMLU 5-shot logprob scoring.

Recommended sampling command

Use these parameters as the safe default for this pretrained 4096 checkpoint. Keep cfg_scale=1 and trajectory_state_blend=0.5; higher CFG was not the recorded best setting for this model. The stronger repetition penalty is intentional for long 4096-token samples.

CUDA_VISIBLE_DEVICES=0 python scripts/eval/sample_prefix_suffix_trajectory_cfg.py \
  --ckpt_dir <local_dir_with_model_pt> \
  --prompt "The history of artificial intelligence" \
  --output /tmp/statediffrwkv_trajectory4096_pretrained_sample.json \
  --steps 100 --cfg_scale 1 --max_new_tokens 4096 \
  --trajectory_s1_mode rwkv --trajectory_state_blend 0.5 \
  --temperature 0.5 --top_k 10 --top_p 0.75 --repetition_penalty 1.5 \
  --seed 42

For quick smoke tests, reduce only --max_new_tokens (for example to 512); keep the other parameters fixed.

Suggested eval command

python scripts/eval/run_cola_dlm_tasks_prefix_suffix_trajectory_cfg.py \
  --ckpt_dir <local_dir_with_model_pt> \
  --task_data_dir baseline/Cola-DLM/eval_output/tasks_default \
  --output_dir /tmp/diffrwkv_trajectory4096_pretrained_mmlu200_cfg1 \
  --tasks mmlu \
  --max_samples 200 --steps 100 --cfg_scale 1 \
  --max_new_tokens 32 --temperature 0.0 \
  --trajectory_s1_mode rwkv --trajectory_state_blend 0.5

See config.yaml and release_metadata.json for full training and release metadata.