Add RWKV-7 (Goose) - #47780
Conversation
|
Update on the ONNX known issue — narrowed, still open, and the evidence now points Two no-op views in the model are gone (f3c38d5): the chunked recurrence padded to The export test still fails, and that is the interesting part. After the The failing node is I do not think that is something I can fix from the model side, and I would rather |
|
Correction to my previous comment: this was fixable from the model side, and it is fixed. I was right that the failing Both are replaced in 6869cd7. Export tests on this branch: Worth flagging for anyone reading the diff, because I got it wrong first. The matrix is nilpotent, so the Neumann series terminates and Newton doubling reaches the inverse in No maintainer input needed on this any more. |
13aa294 to
9be4ac3
Compare
|
Rebased onto current main and pushed a follow-up commit; summary for reviewers:
Re-verified on the official 1.5B after the change: the fp32 and bf16 cross-runtime statistics reported in #47787 reproduce bit-for-bit (the changes are numerics-neutral), the slow integration suite passes, and compiled decode throughput is unchanged (~16× eager on an RTX 3090; |
Adds the RWKV-7 attention-free recurrent language model as Rwkv7Model / Rwkv7ForCausalLM: portable PyTorch with no third-party runtime dependency, a chunk-parallel prefill form of the recurrence, sequential single-token decode, and an Rwkv7Cache built on LinearAttentionLayer with a separately configurable WKV state dtype. Parameter names follow the upstream RWKV reference implementation, so converting a native .pth checkpoint is a prefix rename; the converter also reads the flash-linear-attention layout. Tests include the common mixins, an integration test matched token-for-token against BlinkDL's own runtime, and an independent numpy re-derivation of the forward pass that shares no code with the model.
Generated by utils/add_dates.py; fixes the repository-consistency check.
The examples_torch shard failed on a self-hosted runner infrastructure
error ("Executing the custom container implementation failed"), not on
anything in this PR - it does not touch examples/.
The chunked recurrence pads to a multiple of the chunk size and computes in a wider dtype, then unconditionally slices back to seq_len and casts back. When the input needed neither, both are no-ops that still return views, and a view becomes aten.alias under export. The exported graph now has zero alias nodes.
Seven ONNX export subtests failed under RUN_SLOW, and both causes are the same shape. `aten.cumprod` and `aten.linalg_solve_triangular` each decompose into a graph carrying a scalar tensor constant; `aten.lift_fresh` on that constant decomposes to `aten.alias`, functionalization rewrites the alias to the in-place `aten.detach_`, and aot_autograd rejects a graph that mutates its input. Neither op is reached by the recurrent path, so decode-only exports were green and only the prefill ones failed. `cumprod(exp(x))` is `exp(cumsum(x))`, and the decay is only ever wanted exponentiated, so the sum is taken first. That is also the better arithmetic: a cumulative product of factors below one underflows the longer it runs, and `c_prev` becomes a subtraction rather than a division by a possibly-tiny value. The solve is replaced by a batched inverse of the same unit lower triangular matrix, computed by block forward substitution. The tempting alternative is a series -- the matrix is nilpotent, so `I - x + x^2 - ...` terminates and Newton doubling reaches it in ceil(log2(span)) steps, exactly. Exactly in exact arithmetic: the intermediate powers are not bounded by the answer. On a real chunk the matrix has entries at most 0.977 and its inverse has entries at most 1.0, while the 32nd power reaches 1.3e11, so float32 cancels away every digit it has and the result comes back with entries of 1e4 where the answer is 1. Block forward substitution never forms a power above `block`; against float64 `solve_triangular` on that chunk it lands at 1.3e-7 for block 4, 6.6e-7 for block 8 and 1.9e-5 for 16. Worth recording that random triangular matrices cannot detect any of this: their own inverses are as large as the intermediate powers, so a series looks accurate to 1e-6 on exactly the input a test reaches for first. The test added here is end-to-end and seeded, at a length of several chunks, because the failure needs chunks to compound -- a fifth of random initialisations at T=1024, one in forty at T=256, none at T=128 or below.
The note said a third slower on CPU end to end and left the GPU unmeasured, which is the number anyone reading it would want. Against the per-chunk solve it replaces, loop included, on an RTX 5090 at 1.5B shapes: 2.5x at T=1024, where 16 chunks do not fill the card, falling to 1.2x at T=4096 and flat in batch and head count.
The Checkpoints section opened with "No checkpoint on the Hub loads into this
implementation as-is". That was true when it was written and is not any more:
the RWKV organisation's 20260805 release is a plain safetensors directory whose
config.json maps onto Rwkv7Config key for key, and it loads with no conversion
and no trust_remote_code.
Verified on an A800 against the 1.5B release rather than assumed:
- config.json: every key is a declared parameter of Rwkv7Config except
model_type, which the base class handles.
- 798 tensors, exact name and shape match against the state dict.
- The release is byte-identical to BlinkDL's source .pth -- 798/798 tensors,
same bf16 storage dtype, matched by content hash so no naming convention
has to be agreed first. The conversion is lossless.
- fp32 greedy output matches BlinkDL's own numpy forward on 24/24 lambada
prompts, and is invariant to batching (0/24 differing at batch 4 and 8) and
to left padding (0/8).
- bf16 differs on 2 of those 24 continuations. Both are near-ties: the first
divergence is at generated token 9 with a top1-top2 margin of 0.023 logits,
which several recurrent steps of 8-bit mantissa are enough to tip. The fp32
arm being exact is what identifies this as precision rather than an error.
The usage example moves to the official 1.5B, which also removes the
trust_remote_code line: these repos ship the World vocabulary as a tokenizer.json,
so AutoTokenizer loads it as an ordinary fast tokenizer.
The slow integration test is deliberately not touched. It converts BlinkDL's raw
0.1B .pth and compares against values from BlinkDL's own runtime, so it does not
certify this implementation with this implementation; and at 768 wide with
head_dim 64 the head count and head width differ, which the 1.5B and larger
releases do not give (7.2B and 13.3B are 64 and 64).
…e, adopt standard names - Drop the private RWKV7_WKV_FUNCTIONS registry and its config knob: it duplicated the kernel-substitution mechanism huggingface#47630 just unified across linear-attention models. The torch reference is called directly; an optimised kernel arrives via hub-kernels when one exists, not a per-model registry. The doc section demonstrating the registry goes with it. - Gate generate's input truncation on is_first_iteration instead of the state's existence: a pre-allocated state (the compile contract) or a warm state (a resumed chat turn) silently dropped every prompt token but the last one. Same gate Mamba uses; bare callers keep the old behaviour. - Rename cu_seq_lens to the ecosystem's cu_seq_lens_q, so padding-free collators reach the packed path instead of being swallowed by kwargs. - Refuse Rwkv7Cache.crop() loudly (the state is not invertible), fix keys_to_ignore_at_inference to name the output this model actually has, add the missing explicit models/__init__ re-export, and drop the attentions=None output field and ignored output_attentions parameter. - Tests: generate-with-state equalities, end-to-end left-padded batch generate, compiled decode step over allocate_state, the chunked-prefill mask-slicing contract, a multi-chunk numpy-oracle case; pin the numpy reference to its external sources and the rwkv package version used for the integration expectations; let the shared gradient-checkpointing trainings run via test_all_params_have_gradient. - Docs: compile guidance now reflects measurement (reduce-overhead with an eager prefill; max-autotune measured slower on this kernel chain).
utils/add_dates.py derives the expected date from the first commit of the model's files on main; for a model that is not merged yet there is no such commit, so it falls back to today and the check fails once the branch is more than a day old. Re-stamped with the value the script itself produces.
9be4ac3 to
86af266
Compare
… needs The time-mix and the channel-mix both have a `key` and a `value`, so a `target_modules` list carried over from an attention model attaches to twice as many projections as its author expects, silently. Name both sets and show the qualified form that selects the time-mix alone. Also record that a chunked-loss trainer may look for the recurrence under `past_key_values`: TRL's `SFTTrainer` defaults to `loss_type="chunked_nll"`, whose head reads that field, and this model returns `state` — the same shape of mismatch `MambaOutput` has with `cache_params`. `loss_type="nll"` skips it.
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, rwkv7 |
CI recapDashboard: View test results in Grafana |
2 participants
Adds RWKV-7 ("Goose") as
Rwkv7Model/Rwkv7ForCausalLM.Follows up on #46984, where the integration was declined because the published checkpoints did not follow Transformers conventions. That objection was correct at the time; it has since been resolved upstream — see below.
Checkpoints (updated 2026-08-07)
On 2026-08-05 the RWKV org published official checkpoints in the standard layout:
RWKV/RWKV7-1.5B-20260805,2.9B,7.2Band13.3B—safetensorsonly, noauto_map, no remote code, a standardconfig.json, and the World vocabulary shipped as a plaintokenizer.jsonthatAutoTokenizerloads as an ordinary fast tokenizer.This implementation loads them unchanged:
config.jsonis a declaredRwkv7Configparameter, exceptmodel_type, which the base class handles;BlinkDL/rwkv7-g1.pthit was packed from — 798/798 tensors matched by content hash, a bijection (measured in Add RWKV-7 (Goose) #47787).The docs and the
configuration_rwkv7.pydoc checkpoint point at the official repos. (Historical note, since it was the reason #46984 was closed: the olderRWKV/*-PTHreleases ship nosafetensors, and the olderRWKV/*-HFones requiretrust_remote_code. The 20260805 line is what resolved this.)Implementation
Portable PyTorch, no third-party runtime dependency:
Rwkv7Cacheis built onLinearAttentionLayer. The WKV state dtype is configurable independently of the model dtype, since the recurrent state is more precision-sensitive than the weights..pthis a prefix rename. The converter also reads the flash-linear-attention layout.Correctness
Three independent checks, in addition to the common mixins:
tests/models/rwkv7/test_modeling_rwkv7.py— an integration test matched token-for-token against BlinkDL's own runtime (therwkvpackage at cpu fp32), through this PR's converter, on a checkpoint whose head count (12) and head width (64) differ, so a quantity indexed by the wrong one cannot pass by coincidence.tests/models/rwkv7/test_modeling_rwkv7_numpy_reference.py— a numpy re-derivation of the forward pass that shares no code with the model, so an error in the PyTorch path cannot be masked by the same error in the reference.rwkvpackage's PyTorch path, 24 lambada prompts, same forced tokens, per-position logits and per-layer recurrent states, prefill (T>1) and decode (T=1) legs both. fp32: worst relative deviation 8.2e-6 across everything compared. bf16 vs bf16, same protocol: median deviation 6e-3 (bf16's own precision), with signed bias 20–180× smaller than the deviation itself — symmetric rounding, with no directional offset at the state level, which is where the products the dtype question concerns are formed.Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. These often are low-quality, or fix extremely minor issues that occur rarely or never in practice.
As a result, we're instituting a rule that first-time contributors should not use code agents to submit PRs or issues.
We'd also ask autonomous "OpenClaw"-like agents not to open any PRs or issues.
Issues/PRs from first-time contributors that violate this rule will probably just be closed without review, and we
might block you, especially if you open more than one or appear to be deliberately ignoring this. We especially do not
want new contributors to jump in on random issues to contribute an agent-written fix. This creates lots of noise
for reviewers and other users and will almost certainly get you blocked.
For more information, please read
CONTRIBUTING.md.Before submitting
Who can review?
@Cyrilvallez — you raised the conventions objection on #46984; the checkpoint section above is meant to answer it directly.
@ArthurZucker @vasqu