Add RWKV-7 (Goose) - #47780
Conversation
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.
|
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 |
CI recapDashboard: View test results in Grafana |
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.
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, rwkv7 |
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, and the checkpoint situation is addressed below.
On the checkpoint conventions
To be precise about where the problem actually sits, the checkpoints under
RWKV/come in two kinds:RWKV7-Goose-*-PTH— raw.pth, nosafetensors.RWKV7-Goose-*-HF— these do shipmodel.safetensors, but each also shipsmodeling_rwkv7.pyand anauto_map, so loading one requirestrust_remote_code.So neither kind would load through a library implementation if one existed, which is a fair reason to have closed the earlier request.
There is now a set converted directly from the canonical
BlinkDL/*.pthreleases that follows the standard layout —safetensorsonly, no pickle, no remote code, standardconfig.jsonwitharchitectures/model_type— spanning 0.1B to 7.2B.Hakureirm/rwkv7-168m-pile-hfis the smallest; all 399 of its tensors are verified bit-identical against the source.pthrather than spot-checked.That one is a Pile model, so its tokenizer is the ordinary GPT-NeoX-20B fast tokenizer rather than the World vocabulary — the same reason the existing
rwkvmodel in this repo documentsRWKV/rwkv-4-169m-pile.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
Two 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.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.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