README.md

---
license: apache-2.0
base_model:
- BlinkDL/rwkv7-g1
- allenai/Olmo-3-1025-7B
---
# RWKV-7 "g1d" 0.1B — HuggingFace port (OLMo tokenizer)
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
# If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.
A `trust_remote_code=True` HuggingFace wrapper around the BlinkDL
**RWKV-7 "Goose" g1d 0.1B** checkpoint
(`rwkv7-g1d-0.1b-20260129-ctx8192.pth`), re-headed for the **OLMo tokenizer**.

- The 12 transformer-style RWKV-7 blocks keep their **pretrained** weights.
- The **embedding** and **lm-head** are **re-initialized** (from RWKV's 65536-token
  vocab to OLMo's 100278-token vocab), so they are *untrained* and need fine-tuning.
- Time-mixing runs on a **fused CUDA kernel** (forward **and** backward) when a GPU
  + CUDA toolchain are available, and transparently **falls back to pure PyTorch**
  otherwise.
- `generate()` runs in **RNN mode with a recurrent state cache** — after the prompt
  is absorbed into the state, each new token is a single-token forward instead of a
  full-sequence recompute. Stateful inference runs on a second, **forward-only
  "wkv7s" CUDA kernel** (state in/out), so prefill and decode are kernel-speed too.

> Because emb/head are freshly initialized, `generate()` produces gibberish until
> you fine-tune on OLMo-tokenized data. The model *body* is pretrained; only the
> vocabulary projection is new.

---

## Files

| File | Purpose |
|------|---------|
| `config.json` | Serialized `RWKV7Config` (dims + `auto_map` to the remote-code classes). |
| `configuration_rwkv7.py` | `RWKV7Config` — all architecture hyperparameters. |
| `modeling_rwkv7.py` | Model code: CUDA-kernel dispatch + PyTorch fallback, recurrent state cache for generation, `RWKV7Model`, `RWKV7ForCausalLM`. |
| `cuda/wkv7_cuda.cu`, `cuda/wkv7_op.cpp` | The fused bf16 "wind_backstepping" RWKV-7 kernel (forward + backward), copied from `RWKV-v7/train_temp/cuda/`. Used on the stateless (training) path. |
| `cuda/wkv7s.cu`, `cuda/wkv7s_op.cpp` | The stateful, forward-only "wkv7s" inference kernel (fp32 state in/out, arbitrary `T`, no chunk padding), adapted from RWKV-LM's inference kernel — patched here for **bf16** (upstream typedef was fp16) and **B > 1** (batch-aware state indexing). Used on the stateful (generation) path. |
| `model.safetensors` | Converted weights (bf16, ~244M params). |
| `tokenizer.json`, `tokenizer_config.json` | OLMo tokenizer (vocab 100278, GPT2-style BPE). |
| `generation_config.json` | Default generation settings (eos/pad ids). |
| `convert.py` | Reproduces `model.safetensors` from the original `.pth`. |
| `verify.py` | End-to-end smoke test (load / forward / parity / backward / generate). |

## Architecture (from `config.json`)

| field | value | meaning |
|-------|-------|---------|
| `num_hidden_layers` | 12 | RWKV-7 blocks |
| `hidden_size` | 768 | embedding dim `C` |
| `head_size` | 64 | → 12 heads (`H = C / head_size`) |
| `intermediate_size` | 3072 | channel-mix (FFN) hidden |
| `decay_lora` / `aaa_lora` / `mv_lora` / `gate_lora` | 64 / 64 / 32 / 128 | LoRA ranks for `w` / `a` / `v` / `g` |
| `vocab_size` | 100278 | OLMo tokenizer size (re-initialized emb/head) |
| `chunk_len` | 16 | CUDA kernel chunk length; sequence is padded to a multiple of this |
| `use_cuda_kernel` | true | prefer the fused kernel when possible |

---

## Usage

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

PATH = "/workspace/rwkv7-g1d-olmo"

tok = AutoTokenizer.from_pretrained(PATH, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    PATH, trust_remote_code=True, dtype=torch.bfloat16
).cuda().eval()

ids = tok("The Eiffel tower is in the city of", return_tensors="pt").input_ids.cuda()
with torch.no_grad():
    logits = model(ids).logits          # (1, T, 100278)

# training / backward
model.train
#🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨labels[labels == tokenizer.pad_token_id] = -100🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
out = model(ids, labels=ids)            # shifted cross-entropy loss
out.loss.backward()                     # gradients flow through the CUDA kernel
```

### Forcing the PyTorch fallback

Set the flag on the config (useful on CPU, non-bf16, or to debug the kernel):

```python
model.config.use_cuda_kernel = False    # every RWKV-7 op now runs in pure PyTorch
```

The kernel is also skipped automatically when: no CUDA device, input dtype is not
bfloat16, or the kernel fails to compile — in all cases the fallback is used and a
one-line notice is printed.

### Stateful generation (RNN mode)

`generate()` uses the state cache automatically — nothing to configure:

```python
out_ids = model.generate(ids, max_new_tokens=64, do_sample=False)
```

You can also drive the state manually (streaming, chunked prefill, state
save/reuse, "infinite context" style loops):

```python
with torch.no_grad():
    out = model(ids, use_cache=True)          # prefill: prompt -> state
    next_tok = out.logits[:, -1].argmax(-1, keepdim=True)
    out = model(next_tok, state=out.state)    # passing `state` implies use_cache
```

`outputs.state` is a list of per-layer tuples `(att_x_prev, wkv_state, ffn_x_prev)`:

| tensor | shape | dtype | role |
|--------|-------|-------|------|
| `att_x_prev` | `(B, C)` | input dtype | last post-`ln1` token, time-shift for the time-mix |
| `wkv_state` | `(B, H, N, N)` | fp32 | the WKV matrix state |
| `ffn_x_prev` | `(B, C)` | input dtype | last post-`ln2` token, time-shift for the channel-mix |

States are ordinary tensors — you can `clone()` them to snapshot/branch a
context, or persist them to disk and resume later.

> **Which implementation runs when?** Stateful calls under `torch.no_grad()`
> (i.e. normal generation) use the forward-only **wkv7s** kernel — one kernel
> launch per forward, any `T`, no chunk padding — so both long-prompt prefill
> and per-token decode are fast. wkv7s has **no backward**, so a stateful call
> with autograd enabled (e.g. TBPTT-style training that carries state) silently
> drops to the sequential PyTorch path instead. The kernel mutates its state
> buffer in place; the wrapper always clones the incoming state first, so a
> `state` you're holding on to is never corrupted (snapshot/branching stays safe).

Zero token-shift state + zero WKV state is exactly equivalent to the zero-padded
non-cached forward, so prefill→decode with the cache reproduces the full-recompute
logits.

---

## How the CUDA kernels + fallback work (`modeling_rwkv7.py`)

The time-mixing recurrence is dispatched by
`run_rwkv7(r, w, k, v, a, b, config, initial_state=None, output_final_state=False)`,
which returns `(y, final_state_or_None)`. Dispatch is three-way:

| call | implementation |
|------|----------------|
| stateless (training / plain forward), CUDA + bf16 | `wind_backstepping` (fused fwd + bwd) |
| stateful (cache), CUDA + bf16, **autograd off** | `wkv7s` (fwd-only, state in/out) |
| everything else (CPU, non-bf16, grad-enabled stateful, compile failure) | pure PyTorch |

1. **Training kernel** (`_rwkv7_cuda`, "wind_backstepping") — used when
   `config.use_cuda_kernel` is set, the tensors are CUDA + bf16, the call is
   **stateless** (`initial_state is None` and no final state requested), and
   `_try_load_cuda_kernel()` succeeds. On first call it
   JIT-compiles `cuda/wkv7_op.cpp` + `cuda/wkv7_cuda.cu` via
   `torch.utils.cpp_extension.load` (flags `-D_C_=head_size`, `-D_CHUNK_LEN_=chunk_len`)
   and registers the `wind_backstepping` op. Compilation happens once per process and
   is cached; any failure is caught and flips the model to the fallback.
   - `_WindBackstepping` is a `torch.autograd.Function`:
     - **forward** → `wind_backstepping.forward` (produces `y` plus the saved state
       `s` and `sa` needed for backprop).
     - **backward** → `wind_backstepping.backward` (returns gradients for all six
       inputs `w, q, k, v, z, b`), so training works end-to-end on the kernel.
   - The sequence length is padded to a multiple of `chunk_len` (16) and sliced back.

2. **Inference kernel** (`_rwkv7_cuda_stateful`, "wkv7s") — used for **stateful**
   calls (an incoming state and/or `output_final_state=True`) when the tensors are
   CUDA + bf16 **and autograd is disabled** (`torch.no_grad()` — the kernel has no
   backward). JIT-compiled from `cuda/wkv7s_op.cpp` + `cuda/wkv7s.cu`
   (flag `-D_N_=head_size`), registered as the `wkv7s` op.
   - The time loop lives inside the kernel: a single launch processes any `T`
     (no `chunk_len` padding), fully parallel over `B·H` blocks × `N` threads.
   - The fp32 state `(B, H, N, N)` is read at the start and written back at the
     end (in-place); the Python wrapper clones the incoming state so the caller's
     tensor is never mutated.
   - Patched relative to the upstream RWKV-LM inference kernel: `bf16` typedef
     was actually `at::Half` (fp16) upstream → switched to `at::BFloat16`, and
     the state indexing was `B = 1`-only → made batch-aware.

3. **Fallback path** (`_rwkv7_pytorch`) — a plain sequential-over-time
   implementation of the same recurrence
   `state = state*exp(-exp(w)) + state·aᵀ·b + vᵀ·k`, `y = state·r`, in fp32.
   It is fully differentiable through ordinary autograd (no custom backward needed),
   accepts an `initial_state`, and can return the final state — it covers every
   case the kernels can't (CPU, non-bf16, grad-enabled stateful calls).

All three paths take the **raw (pre-exp) decay `w`** and apply `exp(-exp(w))`
internally, and implement the identical per-head recurrence
(`s[i,j] = s[i,j]·w_j + v_i·k_j + (Σⱼa_j·s[i,j])·b_j`), so they are numerically
interchangeable. `verify.py` confirms kernel-vs-fallback
parity (identical bf16 logits and top-1 prediction).

### Model structure

`RWKV7ForCausalLM` → `.rwkv` (`RWKV7Model`) + `.head` (lm-head). `RWKV7Model` holds
`emb`, `blocks[0..11]` (`RWKV7Block` = `ln1` + `att` time-mix, `ln2` + `ffn`
channel-mix; block 0 also has `ln0`), and `ln_out`. State-dict keys mirror the
original RWKV layout under the `rwkv.` prefix (e.g. `rwkv.blocks.0.att.receptance.weight`).

Generation runs in **RNN mode**: the model carries a per-layer recurrent state
(see the state table above) through the `state=` kwarg / `outputs.state` field —
the same convention as `transformers`' Rwkv and Mamba models, so
`GenerationMixin` propagates it between steps automatically (`state` is in
`ALL_CACHE_NAMES`; a `_update_model_kwargs_for_generation` override covers older
versions). `prepare_inputs_for_generation` feeds only the last token once a state
exists. Passing `state=` implies `use_cache=True`; gradient checkpointing forces
it off during training. Stateful forwards run on the wkv7s kernel under
`no_grad`, PyTorch otherwise.

---

## Reproducing the conversion (`convert.py`)

```bash
# 1) download the original checkpoint (already done in /workspace)
wget https://huggingface.co/BlinkDL/rwkv7-g1/resolve/main/rwkv7-g1d-0.1b-20260129-ctx8192.pth \
     -O /workspace/rwkv7-g1d-0.1b-20260129-ctx8192.pth

# 2) convert -> /workspace/rwkv7-g1d-olmo
python convert.py
```

What `convert.py` does:

1. Loads the OLMo tokenizer to read the target vocab size (100278).
2. Loads the `.pth` and infers all dims from the tensor shapes → builds `RWKV7Config`.
3. Remaps the original RWKV keys to the HF module layout (`rwkv.` prefix), **dropping**
   `emb.weight` / `head.weight`, and loads them with `strict=False`
   (asserts there are **no** unexpected or unmatched keys besides emb/head).
4. **Re-initializes** the embedding (`uniform(±1e-4)`) and head
   (`orthogonal`, gain `0.5·√(vocab/hidden)`) for the new vocabulary — RWKV's own
   init scheme.
5. Saves weights (`safetensors`), config, tokenizer, and copies the remote-code files.

## Verifying (`verify.py`)

```bash
python verify.py
```

Checks: load via `AutoModelForCausalLM(trust_remote_code=True)`, CUDA-kernel forward,
kernel-vs-fallback parity, cached-vs-uncached parity (prefill + stateful decode vs
full recompute), a backward pass (gradients on emb / attention / decay-LoRA), and a
short greedy `generate()` through the state cache.

> Note: 399/402 parameters receive gradients — the 3 without are
> `blocks.0.att.v0/v1/v2` (the value-residual params are unused in layer 0 by design).

---

## Requirements

- PyTorch with CUDA (bf16-capable GPU; tested on RTX 3060 / CUDA 13.0) for the kernel;
  CPU/other works via the fallback.
- `transformers >= 5`, `safetensors`.
- A working CUDA toolchain (`nvcc`) for first-call kernel JIT compilation; if absent,
  the model still runs on the PyTorch fallback.

## Credits

RWKV-7 architecture and the original checkpoint/kernels by **BlinkDL** —
<https://github.com/BlinkDL/RWKV-LM>. Training kernel copied from
`RWKV-v7/train_temp/cuda/`; the stateful `wkv7s` inference kernel adapted from
RWKV-LM's inference code (patched here for bf16 + batched state).