Skip to content

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

Open
mukund1985 wants to merge 3 commits into
TransformerLensOrg:devfrom
mukund1985:feat/raven-adapter
Open

feat: add RavenForCausalLM (Huginn) depth-recurrent adapter#1520
mukund1985 wants to merge 3 commits into
TransformerLensOrg:devfrom
mukund1985:feat/raven-adapter

Conversation

@mukund1985

Copy link
Copy Markdown
Contributor

Summary

Adds a TransformerBridge adapter for RavenForCausalLM (tomg-group-umd/huginn-0125), a depth-recurrent transformer that performs latent reasoning by iterating a weight-tied core block a variable number of times at inference.

Closes #1469.

Architecture

Huginn has a three-phase structure under model.transformer:

  • prelude — 2 SandwichBlocks, run once
  • core_block — 4 SandwichBlocks, weight-tied and iterated num_steps times at runtime (default mean_recurrence=32). Each step re-injects the prelude output via a learned adapter linear
  • coda — 2 SandwichBlocks, run once

SandwichBlock uses a post-residual norm pattern (x = norm_2(attn(norm_1(x)) + x)). BlockBridge enforces a pre-norm hook flow that does not exist here, so SSMBlockBridge is used for all three block lists to delegate the full forward unchanged.

Key implementation notes

transformers v5 compatibility. The remote code targets transformers 4.44 where _tied_weights_keys is a list. Under v5, tie_weights() expects a dict. Fixed in prepare_loading() by rewriting it to {"lm_head.weight": "transformer.wte.weight"}, plus the standard _init_weights re-randomization guard.

Non-deterministic forward. iterate_forward seeds the recurrence with torch.randn_like. Integration tests pin the RNG identically before bridge and HF calls.

Hook exposure. core_block.{i}.hook_in / hook_out fire once per recurrence step, enabling cross-step logit lens, fixed-point analysis, and cross-step activation patching.

Registration

All four required sites updated per AGENTS.md: supported_architectures/__init__.py, factories/architecture_adapter_factory.py (key: RavenForCausalLM), tools/model_registry/__init__.py, tools/model_registry/generate_report.py, tools/model_registry/data/supported_models.json.

Tests

  • Unit (test_raven_adapter.py): 58 tests, synthetic config, no weight loading
  • Integration (test_raven_adapter.py): 16/16 pass locally against the full fp32 checkpoint, including forward logit parity vs HF eager, hook shapes across all three phases, and recurrent-core hook firing count. CI-gated on network/memory budget.
  • TestRegistrySyncedWithFactory (4 invariants): all pass

@mukund1985

mukund1985 commented Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

@jlarson4 ready for review. Rebased against dev — conflicts resolved.

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 TransformerLensOrg#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.
@mukund1985
mukund1985 force-pushed the feat/raven-adapter branch from 5653cbc to 9336e87 Compare July 17, 2026 05:41
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.

1 participant