Skip to content

feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM)#1521

Merged
jlarson4 merged 10 commits into
TransformerLensOrg:devfrom
mukund1985:feat/rwkv7-adapter
Jul 20, 2026
Merged

feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM)#1521
jlarson4 merged 10 commits into
TransformerLensOrg:devfrom
mukund1985:feat/rwkv7-adapter

Conversation

@mukund1985

Copy link
Copy Markdown
Contributor

Closes #1462.

Adds a TransformerBridge adapter for RWKV7ForCausalLM (RWKV-7 "Goose") from fla-org/flash-linear-attention, loaded via trust_remote_code=True.

Architecture overview

RWKV-7 is an attention-free recurrent model:

model.embeddings → model.layers[i] (RWKV7Block) → model.norm → lm_head

Each RWKV7Block contains:

  • attn_norm — standard biased LayerNorm (NOT RMSNorm)
  • attnRWKV7Attention: generalized delta-rule time-mixing with matrix-valued state; 4-tuple return (y, v_first, state, ...)
  • ffn_norm — fused add-and-norm when fuse_norm=True (default)
  • ffnRWKV7FeedForward: token-shift + sqrelu channel-mixing (not a gated MLP)

v_first (cross-block state) is managed internally by RWKV7Model.forward — no adapter handling required.

Adapter decisions

SSMBlockBridge for model.layers — delegates the whole recurrent block; exposes hook_in/hook_out on the residual stream. BlockBridge is not usable: its hook aliases hardcode a standard pre-norm flow that RWKV-7 does not follow.

SSM2MixerBridge for both attn and ffn sublayers — tuple-aware passthrough that correctly preserves attn's 4-tuple return (carrying v_first) and ffn's 2-tuple return without re-implementing either.

NormalizationBridge (biased LayerNorm) for attn_norm and ln_finaluses_rms_norm=False, normalization_type="LN".

ffn_norm as plain GeneralizedComponent — when fuse_norm=True, the block calls ffn_norm(hidden_states, residual, True), a fused 3-arg signature that NormalizationBridge's single-input contract cannot express.

head_dim excluded from _HF_PASSTHROUGH_ATTRS — read-only property on TransformerBridgeConfig (alias of d_head); a passthrough setattr would crash boot.

applicable_phases = [] — correctness covered by integration tests.

prepare_loading mirrors Raven's transformers v5 guards (_tied_weights_keys list→dict, _init_weights re-randomization guard), applied defensively via try/except.

Files changed

File Change
transformer_lens/model_bridge/supported_architectures/rwkv7.py New adapter
tests/unit/model_bridge/supported_architectures/test_rwkv7_adapter.py 31 unit tests (synthetic config)
tests/integration/model_bridge/test_rwkv7_adapter.py Integration tests (gated: slow + skipif(CI) + importorskip("fla"))
transformer_lens/model_bridge/supported_architectures/__init__.py Register RWKV7ArchitectureAdapter
transformer_lens/factories/architecture_adapter_factory.py Register RWKV7ForCausalLM
transformer_lens/tools/model_registry/__init__.py HF_SUPPORTED_ARCHITECTURES + CANONICAL_AUTHORS_BY_ARCH
transformer_lens/tools/model_registry/generate_report.py Architecture description
transformer_lens/model_bridge/sources/_bridge_builder.py RWKV-7 passthrough attrs

Testing

Unit tests (31) pass on synthetic config — no HF Hub access:

uv run pytest tests/unit/model_bridge/supported_architectures/test_rwkv7_adapter.py -v

Integration tests require flash-linear-attention (~191 MB checkpoint download):

pip install flash-linear-attention
uv run pytest tests/integration/model_bridge/test_rwkv7_adapter.py -v -m slow

@mukund1985

Copy link
Copy Markdown
Contributor Author

@jlarson4 ready for review

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
@mukund1985
mukund1985 force-pushed the feat/rwkv7-adapter branch from 7835baf to ea4e0c2 Compare July 17, 2026 06:43
@jlarson4

Copy link
Copy Markdown
Collaborator

Same note about reviewing the failing CI checks here @mukund1985. Thanks for taking both of these adapters on!

@jlarson4

Copy link
Copy Markdown
Collaborator

@mukund1985 Very similar comment here as on #1520. We cannot use the SSM-based components for non-SSM recurrent architectures. I recommend adopting the same OpaqueBlockBridge in place of the SSMBlockBridge as recommended in #1520. For your usage of SSM2MixerBridge, we should be able to replace it with a GeneralizedComponent without running into any issues. Let me know if you run into any issues with those adjustments.

mukund1985 and others added 8 commits July 20, 2026 15:52
…ge base class

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
…7 attn/ffn sublayers

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
@mukund1985

Copy link
Copy Markdown
Contributor Author

Thanks for the review. Both points addressed in the latest commits:

  • Replaced SSMBlockBridge with OpaqueBlockBridge for the model.layers block list (same refactor as feat: add RavenForCausalLM (Huginn) depth-recurrent adapter #1520).
  • Replaced SSM2MixerBridge with GeneralizedComponent for the attn and ffn sublayers. Unit tests updated accordingly.
    Let me know if there is anything else to adjust.

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
@mukund1985

Copy link
Copy Markdown
Contributor Author

@jlarson4 all CI checks are now green — ready for your re-review when you get a chance.

@jlarson4

Copy link
Copy Markdown
Collaborator

Looks great, thanks @mukund1985! Merging now

@jlarson4
jlarson4 merged commit d96c36a into TransformerLensOrg:dev Jul 20, 2026
25 checks passed
mukund1985 added a commit to mukund1985/TransformerLens that referenced this pull request Jul 20, 2026
… entry

Brings in the RWKV-7 additions from TransformerLensOrg#1521 that landed in dev:
- _bridge_builder.py: added RWKV-7 passthrough attrs (num_heads,
  value_dim, decay/gate/a/v_low_rank_dim, norm_first, norm_bias,
  fuse_norm, attn_mode, hidden_act)
- generate_report.py: added RWKV7ForCausalLM description entry

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
jlarson4 pushed a commit that referenced this pull request Jul 21, 2026
* Add RavenForCausalLM (Huginn) TransformerBridge adapter

Adds a TransformerBridge architecture adapter for RavenForCausalLM
(tomg-group-umd/huginn-0125), a depth-recurrent decoder with a
prelude / weight-tied recurrent core / coda structure, resolving #1469.

- raven.py delegates the recurrence to the remote-code HF forward and maps
  the three physical block lists (prelude / core_block / coda) via
  SSMBlockBridge, with combined-QKV native attention and a gated MLP. Sets
  applicable_phases=[] (a random initial latent state and post-residual
  sandwich norms diverge from the verify_models phases) and
  supports_fold_ln=False (ln_f is reused mid-network). prepare_loading
  patches Huginn's remote code for transformers v5 (tied-weights-keys dict
  form + a weight re-init guard).
- Registers in the adapter factory, the model registry (canonical author and
  description) and supported_models.json; surfaces the recurrence-shape
  config via both _HF_PASSTHROUGH_ATTRS lists.
- Adds synthetic-config unit tests and CI-gated integration tests.

* fix: pass num_steps as int to iterate_forward (0-d tensor has no len)

* fix: remove unused torch import and use setattr for dynamic cfg attrs

* style: apply black formatting to raven.py

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* refactor: introduce OpaqueBlockBridge; use it for Raven instead of SSMBlockBridge

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: apply black formatting to opaque_block, ssm_block, __init__

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* fix: strip dead hook_q/k/v aliases from Raven AttentionBridge

Raven uses a combined Wqkv projection (no separate q/k/v submodules),
so AttentionBridge's default hook_q/hook_k/hook_v aliases (which target
q.hook_out / k.hook_out / v.hook_out) are unresolvable. The upstream
test_every_hook_alias_resolves_to_hookpoint audit catches these as 9
dead aliases across prelude / core_block / coda.

Strip them by setting an instance-level hook_aliases that omits those
three keys, leaving only the aliases that have real HookPoint targets.

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* chore: merge upstream/dev — add RWKV-7 passthrough attrs and registry entry

Brings in the RWKV-7 additions from #1521 that landed in dev:
- _bridge_builder.py: added RWKV-7 passthrough attrs (num_heads,
  value_dim, decay/gate/a/v_low_rank_dim, norm_first, norm_bias,
  fuse_norm, attn_mode, hidden_act)
- generate_report.py: added RWKV7ForCausalLM description entry

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

---------

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
jlarson4 added a commit that referenced this pull request Jul 24, 2026
* Add dedicated GraniteMoe adapter tests (#1302) (#1524)

granite_moe had substantive test coverage, but embedded inside
test_granite_adapter.py rather than in its own file — inconsistent
with every other architecture listed in #1302 (including
granite_moe_hybrid), which each get a dedicated test_<name>_adapter.py.
Moves that coverage into test_granite_moe_adapter.py and adds the one
piece that was missing: an explicit config-flag test class (GraniteMoe
doesn't override config setup, so it inherits dense Granite's flags,
but nothing asserted that directly before).

test_granite_adapter.py goes back to covering only
GraniteArchitectureAdapter, matching its own docstring.

* Updating documentation for future deprecation (#1523)

* feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM) (#1521)

* feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* refactor: use OpaqueBlockBridge for RWKV-7; introduce OpaqueBlockBridge base class

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* added hook_out_is_single_residual_stream default to OpaqueBlockBridge

* style: apply black formatting to opaque_block, ssm_block, __init__

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* refactor: replace SSM2MixerBridge with GeneralizedComponent for RWKV-7 attn/ffn sublayers

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix black formatting in generate_report.py

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

---------

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Co-authored-by: jlarson4 <jonahalarson@comcast.net>

* feat: add RavenForCausalLM (Huginn) depth-recurrent adapter (#1520)

* Add RavenForCausalLM (Huginn) TransformerBridge adapter

Adds a TransformerBridge architecture adapter for RavenForCausalLM
(tomg-group-umd/huginn-0125), a depth-recurrent decoder with a
prelude / weight-tied recurrent core / coda structure, resolving #1469.

- raven.py delegates the recurrence to the remote-code HF forward and maps
  the three physical block lists (prelude / core_block / coda) via
  SSMBlockBridge, with combined-QKV native attention and a gated MLP. Sets
  applicable_phases=[] (a random initial latent state and post-residual
  sandwich norms diverge from the verify_models phases) and
  supports_fold_ln=False (ln_f is reused mid-network). prepare_loading
  patches Huginn's remote code for transformers v5 (tied-weights-keys dict
  form + a weight re-init guard).
- Registers in the adapter factory, the model registry (canonical author and
  description) and supported_models.json; surfaces the recurrence-shape
  config via both _HF_PASSTHROUGH_ATTRS lists.
- Adds synthetic-config unit tests and CI-gated integration tests.

* fix: pass num_steps as int to iterate_forward (0-d tensor has no len)

* fix: remove unused torch import and use setattr for dynamic cfg attrs

* style: apply black formatting to raven.py

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* refactor: introduce OpaqueBlockBridge; use it for Raven instead of SSMBlockBridge

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: apply black formatting to opaque_block, ssm_block, __init__

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* fix: strip dead hook_q/k/v aliases from Raven AttentionBridge

Raven uses a combined Wqkv projection (no separate q/k/v submodules),
so AttentionBridge's default hook_q/hook_k/hook_v aliases (which target
q.hook_out / k.hook_out / v.hook_out) are unresolvable. The upstream
test_every_hook_alias_resolves_to_hookpoint audit catches these as 9
dead aliases across prelude / core_block / coda.

Strip them by setting an instance-level hook_aliases that omits those
three keys, leaving only the aliases that have real HookPoint targets.

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* chore: merge upstream/dev — add RWKV-7 passthrough attrs and registry entry

Brings in the RWKV-7 additions from #1521 that landed in dev:
- _bridge_builder.py: added RWKV-7 passthrough attrs (num_heads,
  value_dim, decay/gate/a/v_low_rank_dim, norm_first, norm_bias,
  fuse_norm, attn_mode, hidden_act)
- generate_report.py: added RWKV7ForCausalLM description entry

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

---------

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* Fix Tracr demo output_label: compiled labels are auto-numbered (#1525)

The Tracr demo passes output_label="reverse", but Tracr builds labels as
f"{name}_{unique_id}", so the compiled expression's residual labels are
reverse_1:* and the demo raises ValueError at the state-dict cell. Pass
reverse.label instead of hardcoding the name, and correct the same wrong
label shape in the unit-test fixture (real Tracr never emits "reverse:1")
and in the infer_tracr_output_label docstring example.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

* Add TransformerBridge adapter for lightweight decoder-only pretraining models (#1519)

* Add TransformerLens pretrain bridge adapter

* Add TransformerLens pretrain bridge adapter

* Setup registry for hookedtransformer deprecation

* bug fixes and rebase to 4.x

* Hook management improvements

* implemented `stop_at_layer` on TransformerBridge

* input_to_embed and pos_slice added to bridge

* prevent corruption, check_hooks_to_add placed in all locations that need it

* Adjusted head detector to be HookedTransformer agnostic, added key value injection to bridge

* Add specific migration differences section

* Attempting to let pytest run in parallel

* Attempting to fix the OOM on the multi-device Full Coverage run

---------

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Co-authored-by: Delaida Muminovic <delaida.muminovic.22@size.ba>
Co-authored-by: Mukund Pandey <mukund.pandey@gmail.com>
Co-authored-by: Joseph Quevedo <41499530+dewstend@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cacapice <kombrellaro@gmail.com>
jlarson4 added a commit that referenced this pull request Jul 24, 2026
* Add dedicated GraniteMoe adapter tests (#1302) (#1524)

granite_moe had substantive test coverage, but embedded inside
test_granite_adapter.py rather than in its own file — inconsistent
with every other architecture listed in #1302 (including
granite_moe_hybrid), which each get a dedicated test_<name>_adapter.py.
Moves that coverage into test_granite_moe_adapter.py and adds the one
piece that was missing: an explicit config-flag test class (GraniteMoe
doesn't override config setup, so it inherits dense Granite's flags,
but nothing asserted that directly before).

test_granite_adapter.py goes back to covering only
GraniteArchitectureAdapter, matching its own docstring.

* Updating documentation for future deprecation (#1523)

* feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM) (#1521)

* feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* refactor: use OpaqueBlockBridge for RWKV-7; introduce OpaqueBlockBridge base class

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* added hook_out_is_single_residual_stream default to OpaqueBlockBridge

* style: apply black formatting to opaque_block, ssm_block, __init__

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* refactor: replace SSM2MixerBridge with GeneralizedComponent for RWKV-7 attn/ffn sublayers

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix black formatting in generate_report.py

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

---------

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Co-authored-by: jlarson4 <jonahalarson@comcast.net>

* feat: add RavenForCausalLM (Huginn) depth-recurrent adapter (#1520)

* Add RavenForCausalLM (Huginn) TransformerBridge adapter

Adds a TransformerBridge architecture adapter for RavenForCausalLM
(tomg-group-umd/huginn-0125), a depth-recurrent decoder with a
prelude / weight-tied recurrent core / coda structure, resolving #1469.

- raven.py delegates the recurrence to the remote-code HF forward and maps
  the three physical block lists (prelude / core_block / coda) via
  SSMBlockBridge, with combined-QKV native attention and a gated MLP. Sets
  applicable_phases=[] (a random initial latent state and post-residual
  sandwich norms diverge from the verify_models phases) and
  supports_fold_ln=False (ln_f is reused mid-network). prepare_loading
  patches Huginn's remote code for transformers v5 (tied-weights-keys dict
  form + a weight re-init guard).
- Registers in the adapter factory, the model registry (canonical author and
  description) and supported_models.json; surfaces the recurrence-shape
  config via both _HF_PASSTHROUGH_ATTRS lists.
- Adds synthetic-config unit tests and CI-gated integration tests.

* fix: pass num_steps as int to iterate_forward (0-d tensor has no len)

* fix: remove unused torch import and use setattr for dynamic cfg attrs

* style: apply black formatting to raven.py

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* refactor: introduce OpaqueBlockBridge; use it for Raven instead of SSMBlockBridge

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: apply black formatting to opaque_block, ssm_block, __init__

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* style: fix import order in generalized_components/__init__.py (isort)

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* fix: strip dead hook_q/k/v aliases from Raven AttentionBridge

Raven uses a combined Wqkv projection (no separate q/k/v submodules),
so AttentionBridge's default hook_q/hook_k/hook_v aliases (which target
q.hook_out / k.hook_out / v.hook_out) are unresolvable. The upstream
test_every_hook_alias_resolves_to_hookpoint audit catches these as 9
dead aliases across prelude / core_block / coda.

Strip them by setting an instance-level hook_aliases that omits those
three keys, leaving only the aliases that have real HookPoint targets.

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* chore: merge upstream/dev — add RWKV-7 passthrough attrs and registry entry

Brings in the RWKV-7 additions from #1521 that landed in dev:
- _bridge_builder.py: added RWKV-7 passthrough attrs (num_heads,
  value_dim, decay/gate/a/v_low_rank_dim, norm_first, norm_bias,
  fuse_norm, attn_mode, hidden_act)
- generate_report.py: added RWKV7ForCausalLM description entry

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

---------

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>

* Fix Tracr demo output_label: compiled labels are auto-numbered (#1525)

The Tracr demo passes output_label="reverse", but Tracr builds labels as
f"{name}_{unique_id}", so the compiled expression's residual labels are
reverse_1:* and the demo raises ValueError at the state-dict cell. Pass
reverse.label instead of hardcoding the name, and correct the same wrong
label shape in the unit-test fixture (real Tracr never emits "reverse:1")
and in the infer_tracr_output_label docstring example.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

* Add TransformerBridge adapter for lightweight decoder-only pretraining models (#1519)

* Add TransformerLens pretrain bridge adapter

* Add TransformerLens pretrain bridge adapter

* Setup registry for hookedtransformer deprecation

* bug fixes and rebase to 4.x

* Hook management improvements

* implemented `stop_at_layer` on TransformerBridge

* input_to_embed and pos_slice added to bridge

* prevent corruption, check_hooks_to_add placed in all locations that need it

* Adjusted head detector to be HookedTransformer agnostic, added key value injection to bridge

* Add specific migration differences section

* Attempting to let pytest run in parallel

* Attempting to fix the OOM on the multi-device Full Coverage run

* Big comment sweep

* Comment fix

---------

Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com>
Co-authored-by: Delaida Muminovic <delaida.muminovic.22@size.ba>
Co-authored-by: Mukund Pandey <mukund.pandey@gmail.com>
Co-authored-by: Joseph Quevedo <41499530+dewstend@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cacapice <kombrellaro@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants