Skip to content

Commit 1712265

Browse files
jlarson4jlarson4
authored
Merge pull request #1366 from TransformerLensOrg/dev
Release v3.4.0
2 parents 535f2d8 + 2db79df commit 1712265

117 files changed

Lines changed: 26471 additions & 6292 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.claude/commands/add-model-support.md‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
---
2+
description: Guided workflow for adding a new architecture adapter to TransformerBridge.
3+
argument-hint: <hf_repo>
4+
---
5+
6+
Adding TransformerBridge support for HF model `$ARGUMENTS`. If empty, ask the user for the HF repo path first.
7+
8+
Each step names the doc to read **when you reach that step** — don't load all up front.
9+
10+
1. **Check registry state and decide whether to verify.**
11+
12+
State:
13+
- Architecture supported? Check `SUPPORTED_ARCHITECTURES` in [`architecture_adapter_factory.py`](../../transformer_lens/factories/architecture_adapter_factory.py).
14+
- Model in registry? Check [`supported_models.json`](../../transformer_lens/tools/model_registry/data/supported_models.json); note `status` (0=unverified, 1=verified, 2=skipped, 3=failed).
15+
16+
Branch:
17+
18+
- **Supported AND `status==1`** → already verified. Ask the user the symptom (bug-report path, not add-support). Stop.
19+
- **Supported, `status != 1`** → proceed to **Confirm before verification**. If `status==3`, read existing `note` for the prior failure mode.
20+
- **Supported, not in registry** → add an entry per [§Adding the HF repo to the registry](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#adding-the-hf-repo-to-the-registry) with `status: 0` and null scores, then proceed.
21+
- **Not supported** → skip to step 2.
22+
23+
### Confirm before verification
24+
25+
Always ask the user first, even for small models:
26+
27+
1. Dry-run to project cost:
28+
```
29+
set -a; source .env; set +a
30+
uv run python -m transformer_lens.tools.model_registry.verify_models --model "$ARGUMENTS" --dry-run
31+
```
32+
2. Show: model ID, architecture class, estimated parameters, projected memory (GB), HF_TOKEN needed?, runtime (30 s–2 min sub-1B, 2–15 min 1B–7B, 15+ min 7B+/multimodal), what verification does (Phases 1–4; updates `supported_models.json` on success).
33+
3. Ask: "Run verification on this machine? (Y/N)"
34+
35+
**Confirm** → `/verify-model $ARGUMENTS`. On pass, done. On fail, see [debugging_numerical_divergence.md](../../docs/source/content/debugging_numerical_divergence.md) (per-sibling adapter bug).
36+
37+
**Reject** → `gh issue create --template verify-model.md` (fill from dry-run output). No `gh`? <https://github.com/TransformerLensOrg/TransformerLens/issues/new?template=verify-model.md>. Stop.
38+
39+
2. **Analyze the HF model.** Read `config.json` and source — identify embedding, attention, MLP, normalization, output-head layouts. Read [§Config-attr propagation](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#config-attr-propagation) and decide which non-standard attrs (`final_logit_softcapping`, `sliding_window`, etc.) need surfacing on `self.cfg`.
40+
41+
3. **Pick a starting adapter.** See [§Starter-adapter table](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#starter-adapter-table). Copy into [`supported_architectures/`](../../transformer_lens/model_bridge/supported_architectures/) as `<arch>.py`. **Tokenizer-policy flags are per-model** — see [§Tokenizer policy](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#tokenizer-policy).
42+
43+
4. **Fill `self.component_mapping`.** Bridge-native hook names. Reference: [§Minimal contract](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#minimal-contract), [§Common gotchas](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#common-gotchas).
44+
45+
5. **Register in all four sites** per [§Registration steps](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#registration-steps). Then run the invariant test: `uv run pytest tests/unit/tools/test_model_registry.py -k TestRegistrySyncedWithFactory`.
46+
47+
6. **Add the HF repo entry** to [`data/supported_models.json`](../../transformer_lens/tools/model_registry/data/supported_models.json) per [§Adding the HF repo to the registry](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#adding-the-hf-repo-to-the-registry). Ask the user about adding canonical sibling variants from `CANONICAL_AUTHORS_BY_ARCH[<HFArchClass>]`.
48+
49+
7. **Verify** end-to-end: `/verify-model $ARGUMENTS`. Read both `status` AND per-phase scores. `STATUS_VERIFIED` means hard gates passed (see [§Phase-score thresholds](../../transformer_lens/tools/model_registry/AGENTS.md#phase-score-thresholds)) — but P4's 50% bar is intentionally lenient. P4 well below 100% on a small parity-test model + `status==1` → suspect missing `preprocess_weights` fold or wrong `default_prepend_bos`; investigate before step 8.
50+
51+
8. **Write tests** per [§Required tests](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#required-tests) (unit + integration). Copy the closest sibling.
52+
53+
9. **`/task-complete`** — comment cleanup, `/format`, standard test tiers, loop until clean.

‎.claude/commands/build-docs.md‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
---
2+
description: Source .env then build the Sphinx docs.
3+
---
4+
5+
Build the documentation locally:
6+
7+
```
8+
set -a; source .env; set +a
9+
uv run build-docs
10+
```
11+
12+
Sourcing `.env` is required so `HF_TOKEN` is available — some doctests and notebook embeddings load gated models. Output goes to [docs/build/](../../docs/build/).
13+
14+
For an interactive live-reloading preview instead, run `uv run docs-hot-reload`.
15+
16+
Docs follow Google docstring style with reST extensions; see [docs/source/content/contributing.md](../../docs/source/content/contributing.md) for the style guide.

‎.claude/commands/format.md‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
---
2+
description: Type-check then format the working tree.
3+
---
4+
5+
Run mypy first, then format. Mypy fixes (`isinstance`, `cast`, signatures) can introduce format drift — running format after means a single pass.
6+
7+
```
8+
uv run mypy .
9+
make format
10+
```
11+
12+
`uv run mypy .` uses the config in [pyproject.toml](../../pyproject.toml). `make format` runs `pycln --all` (unused imports), `isort`, and `black` (line length 100).
13+
14+
If mypy reports errors, fix the underlying typing issue — never add `# type: ignore`. Prefer `isinstance` / `typing.cast` ([AGENTS.md §10](../../AGENTS.md#10-hard-rules)).

‎.claude/commands/task-complete.md‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
---
2+
description: End-of-task gate. Clean up new comments, format, type-check, and run the standard test tiers (unit + docstring + acceptance + integration) — fixing issues along the way.
3+
---
4+
5+
Run the end-of-task gate. Do not declare the task complete until every step below passes cleanly.
6+
7+
### 1. Clean up new comments
8+
9+
Review every comment and docstring **added or modified during this task** against the rules in [AGENTS.md §10](../../AGENTS.md#10-hard-rules):
10+
11+
- Comments should be terse one-liners; docstrings are one-line where possible.
12+
- Inline comments explain WHY, not WHAT — delete any that just restate the code.
13+
- Multi-paragraph explanations belong in PR descriptions or design docs, not source.
14+
- Remove any references to plan files, audit IDs, finding IDs, or "see plan section X" — those rot as the codebase evolves and belong only in the PR description.
15+
16+
Use `git diff` against the merge-base to scope the review to genuinely new comments — do NOT rewrite unrelated comments elsewhere in the file.
17+
18+
### 2. Type-check, then format
19+
20+
Run mypy **before** format. Mypy fixes (`isinstance`, `typing.cast`, signature changes) can introduce format drift — running format after mypy means a single format pass.
21+
22+
```
23+
uv run mypy .
24+
make format
25+
```
26+
27+
If mypy reports new errors, fix the underlying typing issue. Do not add `# type: ignore`.
28+
29+
### 3. Run the standard test tiers
30+
31+
```
32+
set -a; source .env; set +a
33+
make test-pr
34+
```
35+
36+
`make test-pr` runs unit + docstring + acceptance + integration — the tiers that gate PR review for almost every change. Notebook and benchmark suites are intentionally skipped (slow, gated models, CI runs them separately). If your change specifically touched a notebook or a benchmark, also run that file directly (`pytest --nbval-sanitize-with demos/doc_sanitize.cfg demos/<notebook>.ipynb` or `make benchmark-test`).
37+
38+
Investigate every failure. Do not dismiss any failure as "pre-existing" or "unrelated" — fix the underlying issue, even if it predates this task (see [AGENTS.md §10](../../AGENTS.md#10-hard-rules)). Do not add platform skips or `xfail` markers to dodge a failing test.
39+
40+
### 4. Re-loop on failure
41+
42+
If any step surfaces issues, fix them and restart from step 1 — fixes can reintroduce comment, format, type, or test drift.
43+
44+
### 5. Report
45+
46+
Report the actual final command output, not a summary. Reviewers re-run tests; agent self-reports are not evidence ([AGENTS.md §10](../../AGENTS.md#10-hard-rules)).

‎.claude/commands/test-all.md‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
---
2+
description: Run the full test suite (unit + integration + acceptance + benchmark + docstring + notebook). Slow.
3+
---
4+
5+
Run every test tier in TransformerLens via the top-level `make test` target:
6+
7+
```
8+
make test
9+
```
10+
11+
This is slow — it runs unit, integration, acceptance, benchmark, docstring, and notebook tests sequentially. It hits HuggingFace Hub and loads multiple models. Before running, confirm:
12+
13+
1. `.env` is sourced so `HF_TOKEN` is set (`set -a; source .env; set +a`).
14+
2. No other heavy GPU/MPS jobs are running on this machine — model verification cannot run concurrently (see [AGENTS.md §10](../../AGENTS.md#10-hard-rules)).
15+
16+
Report the actual command output. Investigate any failures rather than dismissing them.

‎.claude/commands/test-unit.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
---
2+
description: Run the unit test suite.
3+
---
4+
5+
Run the TransformerLens unit tests:
6+
7+
```
8+
make unit-test
9+
```
10+
11+
If any test fails, investigate the failure rather than dismissing it as "pre-existing" or unrelated — see [AGENTS.md §10](../../AGENTS.md#10-hard-rules). Report the actual command output, not a summary.

‎.claude/commands/typecheck.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
---
2+
description: Run mypy across the project.
3+
---
4+
5+
Run the type checker:
6+
7+
```
8+
uv run mypy .
9+
```
10+
11+
Config lives in `[tool.mypy]` of [pyproject.toml](../../pyproject.toml). If mypy reports errors, fix the underlying typing issue — do not add `# type: ignore`. Prefer `isinstance` assertions or `typing.cast` for narrowing.

‎.claude/commands/verify-model.md‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
---
2+
description: Run verify_models.py against a single model (non-parallel). Always dry-run first.
3+
argument-hint: <model_name_or_hf_repo>
4+
---
5+
6+
Verify model `$ARGUMENTS`. If empty, ask for an HF repo path (e.g. `gpt2`, `meta-llama/Llama-2-7b-hf`) or registry alias.
7+
8+
## Always dry-run first
9+
10+
Verification loads the full model and runs Phases 1–4 — 30 s to 30 min, needs memory to hold the model. **Never invoke the real run blindly.**
11+
12+
```
13+
set -a; source .env; set +a
14+
uv run python -m transformer_lens.tools.model_registry.verify_models --model "$ARGUMENTS" --dry-run
15+
```
16+
17+
Capture: estimated parameter count, projected memory (GB), HF_TOKEN requirement, architecture class.
18+
19+
| Model | Action |
20+
|---|---|
21+
| Cached small (`gpt2`, `attn-only-*`, `tiny-stories-1M`, `distilgpt2`, …) | Proceed; report dry-run in your response so user can intervene |
22+
| ≥1B params, gated, or anything else | Present dry-run, ask before running |
23+
24+
## Run the verification
25+
26+
```
27+
set -a; source .env; set +a
28+
uv run python -m transformer_lens.tools.model_registry.verify_models --model "$ARGUMENTS"
29+
```
30+
31+
## Optional flags
32+
33+
Full reference: [tools/model_registry/AGENTS.md §Flag reference](../../transformer_lens/tools/model_registry/AGENTS.md#flag-reference).
34+
35+
- `--device cpu|cuda|mps` — override device selection
36+
- `--dtype float32|bfloat16` — override dtype
37+
- `--max-memory <gb>` — skip if param estimate exceeds; e.g. `16` on a 24 GB GPU leaves headroom for activations
38+
- `--phases 1 2 3` — restrict (P4 is slowest; restrict when debugging P1 forward parity)
39+
- `--dry-run` — see above; always first
40+
- `--no-hf-reference` / `--no-ht-reference` — skip HF / HT comparison (faster, lower confidence)
41+
- `--reverify` — re-test `status==1`
42+
- `--retry-failed` — re-test `status==3` (read existing `note` first)
43+
44+
Batch flags (`--architectures`, `--per-arch`, `--limit`, `--resume`) don't apply to `--model <repo>` — use [§Canonical invocations](../../transformer_lens/tools/model_registry/AGENTS.md#canonical-invocations).
45+
46+
## Interpreting the output
47+
48+
Hard thresholds (`_MIN_PHASE_SCORES` in `verify_models.py`):
49+
50+
| Phase | Min score | Required tests | Below = |
51+
|---|---|---|---|
52+
| 1 | 100% | — | `STATUS_FAILED` |
53+
| 2 | 75% | `logits_equivalence`, `loss_equivalence` | `STATUS_FAILED` |
54+
| 3 | 75% | `logits_equivalence`, `loss_equivalence` | `STATUS_FAILED` |
55+
| 4 | 50% | — | **Non-gating** — adds `"low text quality"` to `note`; never fails. |
56+
| 7 | 75% | `multimodal_forward` | `STATUS_FAILED`. NULL = fail. |
57+
| 8 | 75% | `audio_forward` | `STATUS_FAILED`. NULL = fail. |
58+
59+
`STATUS_VERIFIED` means hard gates passed. `note` carries quality flags or failure details.
60+
61+
**Adapter-author caveat:** P4's 50% bar is intentionally lenient (coherence, not correctness). P4 well below 100% on a small parity-test model can indicate a real bug the system doesn't gate on — most often a missing [`preprocess_weights` fold](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#when-to-override-preprocess_weights) or wrong [`default_prepend_bos`](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#tokenizer-policy). Investigate even on VERIFIED.
62+
63+
Full reference: [§Phase-score thresholds](../../transformer_lens/tools/model_registry/AGENTS.md#phase-score-thresholds).
64+
65+
## Hard rules
66+
67+
**Use `verify_models`, never `main_benchmark`** — only `verify_models` writes `data/supported_models.json` ([tools/model_registry/AGENTS.md](../../transformer_lens/tools/model_registry/AGENTS.md)).
68+
69+
One model at a time — concurrent loads OOM. Report actual per-phase scores; investigate failures per [AGENTS.md §10](../../AGENTS.md#10-hard-rules).

‎.claude/settings.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"permissions": {
3+
"allow": [],
4+
"deny": []
5+
}
6+
}

‎.cursor/rules/transformerlens.mdc‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
---
2+
description: TransformerLens project conventions for Cursor agents.
3+
alwaysApply: true
4+
---
5+
6+
Read `AGENTS.md` at the repo root before doing any work. It is the single source of truth for project conventions, quickstart commands, repo layout, hook-naming rules, the HookedTransformer ↔ TransformerBridge mirroring rule, PR conventions, and hard rules.
7+
8+
Sub-folder `AGENTS.md` files apply when you're working in those directories — read them too:
9+
10+
- `tests/AGENTS.md` — tier placement, conftest hierarchy, MPS rules
11+
- `transformer_lens/model_bridge/supported_architectures/AGENTS.md` — adapter contract, starter-adapter table, 4-place registration
12+
- `transformer_lens/tools/model_registry/AGENTS.md` — `verify_models` workflow, the `main_benchmark` trap
13+
14+
Quick reminders that override common defaults:
15+
16+
- Use `uv`, not `pip` or `poetry`. Install with `uv sync`; run commands with `uv run …` or `make` targets.
17+
- This repo has two parallel systems (`HookedTransformer` legacy and `TransformerBridge` v3). Changes to HookedTransformer that have equivalents in TransformerBridge must be mirrored to TransformerBridge.
18+
- Base PRs against `dev`, not `main`. Never name a branch `main` or `dev`.
19+
- No pre-commit hook is installed. Run `make format` and `uv run mypy .` manually before push.
20+
- Source `.env` (e.g. `set -a; source .env; set +a`) before any HuggingFace-Hub-hitting command.
21+
- Never add `# type: ignore`, never dismiss failing tests as "pre-existing", never add platform skips to dodge CI, never claim drift is "fp noise" without empirical evidence.

0 commit comments

Comments
 (0)