| --- |
| 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). |
