Skip to content

Commit e4d9784

Browse files
bryce13950bryce13950
authored
Merge pull request #1071 from TransformerLensOrg/dev-3.x-folding
Transformer bridge layer norm folding
2 parents 92c5b47 + 5babde2 commit e4d9784

File tree

200 files changed

+36437
-15782
lines changed

Some content is hidden

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

200 files changed

+36437
-15782
lines changed

‎.github/workflows/checks.yml‎

Lines changed: 132 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ on:
3636
required: false
3737
secrets:
3838
HF_TOKEN:
39-
required: false
39+
required: false
4040

4141
permissions:
4242
actions: write
@@ -47,37 +47,52 @@ jobs:
4747
name: Compatibility Checks
4848
runs-on: ubuntu-latest
4949
strategy:
50+
fail-fast: false
5051
matrix:
5152
python-version:
5253
- "3.10"
5354
- "3.11"
5455
- "3.12"
5556
steps:
5657
- uses: actions/checkout@v3
57-
- name: Install Poetry
58-
uses: snok/install-poetry@v1
59-
- name: Set up Python
60-
uses: actions/setup-python@v4
58+
- name: Install uv
59+
uses: astral-sh/setup-uv@v6
6160
with:
6261
python-version: ${{ matrix.python-version }}
63-
cache: "poetry"
62+
activate-environment: true
63+
enable-cache: true
6464
- name: Cache Models used with Tests
6565
uses: actions/cache@v3
6666
with:
6767
path: |
6868
~/.cache/huggingface/hub/models--gpt2
69+
~/.cache/huggingface/hub/models--distilgpt2
70+
~/.cache/huggingface/hub/models--EleutherAI--pythia-70m
71+
~/.cache/huggingface/hub/models--EleutherAI--gpt-neo-125M
72+
~/.cache/huggingface/hub/models--google--gemma-2-2b-it
73+
~/.cache/huggingface/hub/models--bigscience--bloom-560m
74+
~/.cache/huggingface/hub/models--Qwen--Qwen2-0.5B
75+
~/.cache/huggingface/hub/models--bert-base-cased
6976
~/.cache/huggingface/hub/models--NeelNanda--Attn_Only*
7077
~/.cache/huggingface/hub/models--roneneldan--TinyStories-1M*
71-
key: ${{ runner.os }}-huggingface-models
78+
~/.cache/huggingface/hub/models--NeelNanda--SoLU*
79+
~/.cache/huggingface/hub/models--ArthurConmy--redwood_attn_2l
80+
~/.cache/huggingface/hub/models--stas--tiny-random-llama-2
81+
~/.cache/huggingface/hub/models--microsoft--DialoGPT-medium
82+
key: ${{ runner.os }}-huggingface-models-v4
83+
restore-keys: |
84+
${{ runner.os }}-huggingface-models-v2
85+
${{ runner.os }}-huggingface-models-v1
86+
${{ runner.os }}-huggingface-models
7287
- name: Install dependencies
7388
run: |
74-
poetry check --lock
75-
poetry install --with dev
89+
uv lock --check
90+
uv sync
7691
- name: Authenticate HuggingFace CLI
7792
if: env.HF_TOKEN != ''
7893
run: |
79-
pip install huggingface_hub
80-
huggingface-cli login --token "$HF_TOKEN"
94+
uv pip install huggingface_hub
95+
uv run huggingface-cli login --token "$HF_TOKEN"
8196
env:
8297
HF_TOKEN: ${{ secrets.HF_TOKEN }}
8398
- name: Unit Test
@@ -87,51 +102,117 @@ jobs:
87102
- name: Acceptance Test
88103
run: make acceptance-test
89104
- name: Build check
90-
run: poetry build
105+
run: uv build
106+
107+
format-check:
108+
name: Format Check
109+
runs-on: ubuntu-latest
110+
timeout-minutes: 10
111+
steps:
112+
- uses: actions/checkout@v3
113+
- name: Install uv
114+
uses: astral-sh/setup-uv@v6
115+
with:
116+
python-version: "3.12"
117+
activate-environment: true
118+
enable-cache: true
119+
- name: Install dependencies
120+
run: |
121+
uv lock --check
122+
uv sync
123+
- name: Check format
124+
run: make check-format
125+
126+
type-check:
127+
name: Type Check
128+
runs-on: ubuntu-latest
129+
timeout-minutes: 10
130+
steps:
131+
- uses: actions/checkout@v3
132+
- name: Install uv
133+
uses: astral-sh/setup-uv@v6
134+
with:
135+
python-version: "3.12"
136+
activate-environment: true
137+
enable-cache: true
138+
- name: Install dependencies
139+
run: |
140+
uv lock --check
141+
uv sync
142+
- name: Type check
143+
run: uv run mypy .
91144

92-
code-checks:
93-
name: Code Checks
145+
docstring-test:
146+
name: Docstring Test
94147
runs-on: ubuntu-latest
148+
timeout-minutes: 15
95149
steps:
96150
- uses: actions/checkout@v3
97-
- name: Install Poetry
98-
uses: snok/install-poetry@v1
99-
- name: Set up Python
100-
uses: actions/setup-python@v4
151+
- name: Install uv
152+
uses: astral-sh/setup-uv@v6
101153
with:
102154
python-version: "3.12"
103-
cache: "poetry"
155+
activate-environment: true
156+
enable-cache: true
157+
- name: Install dependencies
158+
run: |
159+
uv lock --check
160+
uv sync
161+
- name: Docstring test
162+
run: make docstring-test
163+
164+
coverage-test:
165+
name: Full Code Coverage Test
166+
runs-on: ubuntu-latest
167+
timeout-minutes: 60
168+
steps:
169+
- uses: actions/checkout@v3
170+
- name: Install uv
171+
uses: astral-sh/setup-uv@v6
172+
with:
173+
python-version: "3.12"
174+
activate-environment: true
175+
enable-cache: true
104176
- name: Cache Models used with Tests
105177
uses: actions/cache@v3
106178
with:
107179
path: |
108180
~/.cache/huggingface/hub/models--gpt2
181+
~/.cache/huggingface/hub/models--distilgpt2
182+
~/.cache/huggingface/hub/models--EleutherAI--pythia-70m
183+
~/.cache/huggingface/hub/models--EleutherAI--gpt-neo-125M
184+
~/.cache/huggingface/hub/models--google--gemma-2-2b-it
185+
~/.cache/huggingface/hub/models--bigscience--bloom-560m
186+
~/.cache/huggingface/hub/models--Qwen--Qwen2-0.5B
187+
~/.cache/huggingface/hub/models--bert-base-cased
109188
~/.cache/huggingface/hub/models--NeelNanda--Attn_Only*
110189
~/.cache/huggingface/hub/models--roneneldan--TinyStories-1M*
111-
key: ${{ runner.os }}-huggingface-models
190+
~/.cache/huggingface/hub/models--NeelNanda--SoLU*
191+
~/.cache/huggingface/hub/models--ArthurConmy--redwood_attn_2l
192+
~/.cache/huggingface/hub/models--stas--tiny-random-llama-2
193+
~/.cache/huggingface/hub/models--microsoft--DialoGPT-medium
194+
key: ${{ runner.os }}-huggingface-models-v4
195+
restore-keys: |
196+
${{ runner.os }}-huggingface-models-v2
197+
${{ runner.os }}-huggingface-models-v1
198+
${{ runner.os }}-huggingface-models
112199
- name: Install dependencies
113200
run: |
114-
poetry check --lock
115-
poetry install --with dev
116-
- name: Check format
117-
run: make check-format
118-
- name: Docstring test
119-
run: make docstring-test
120-
- name: Type check
121-
run: poetry run mypy .
201+
uv lock --check
202+
uv sync
122203
- name: Authenticate HuggingFace CLI
123204
if: env.HF_TOKEN != ''
124205
run: |
125-
pip install huggingface_hub
126-
huggingface-cli login --token "$HF_TOKEN"
206+
uv pip install huggingface_hub
207+
uv run huggingface-cli login --token "$HF_TOKEN"
127208
env:
128209
HF_TOKEN: ${{ secrets.HF_TOKEN }}
129210
- name: Test Suite with Coverage Report
130211
run: make coverage-report-test
131212
env:
132213
HF_TOKEN: ${{ secrets.HF_TOKEN }}
133214
- name: Build check
134-
run: poetry build
215+
run: uv build
135216
- name: Upload Coverage Report Artifact
136217
uses: actions/upload-artifact@v4
137218
with:
@@ -142,12 +223,12 @@ jobs:
142223
name: Notebook Checks
143224
runs-on: ubuntu-latest
144225
strategy:
226+
fail-fast: false
145227
matrix:
146228
notebook:
147229
# - "Activation_Patching_in_TL_Demo"
148230
# - "Attribution_Patching_Demo"
149231
- "ARENA_Content"
150-
- "Colab_Compatibility"
151232
- "BERT"
152233
- "Exploratory_Analysis_Demo"
153234
# - "Grokking_Demo"
@@ -162,30 +243,29 @@ jobs:
162243
# - "T5"
163244
steps:
164245
- uses: actions/checkout@v3
165-
- name: Install Poetry
166-
uses: snok/install-poetry@v1
167-
- name: Set up Python
168-
uses: actions/setup-python@v4
246+
- name: Install uv
247+
uses: astral-sh/setup-uv@v6
169248
with:
170249
python-version: "3.11"
171-
cache: "poetry"
250+
activate-environment: true
251+
enable-cache: true
172252
- name: Re-use HuggingFace models cache
173253
uses: actions/cache/restore@v3
174254
with:
175255
path: ~/.cache/huggingface/hub
176256
key: ${{ runner.os }}-huggingface-models
177257
- name: Install dependencies
178258
run: |
179-
poetry check --lock
180-
poetry install --with dev,jupyter
259+
uv lock --check
260+
uv sync
181261
- name: Install pandoc
182262
uses: awalsh128/cache-apt-pkgs-action@latest
183263
with:
184264
packages: pandoc
185265
version: 1.0
186266
- name: Check Notebook Output Consistency
187267
# Note: currently only checks notebooks we have specifically setup for this
188-
run: poetry run pytest --nbval-sanitize-with demos/doc_sanitize.cfg demos/${{ matrix.notebook }}.ipynb
268+
run: pytest --nbval-sanitize-with demos/doc_sanitize.cfg demos/${{ matrix.notebook }}.ipynb
189269

190270

191271
build-docs:
@@ -194,23 +274,24 @@ jobs:
194274
name: 'Build Docs'
195275
runs-on: ubuntu-latest
196276
if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev') || contains(github.head_ref, 'docs')
197-
needs: code-checks
277+
needs: coverage-test
198278
steps:
199279
- uses: actions/checkout@v4
200-
- name: Install Poetry
201-
uses: snok/install-poetry@v1
202-
- name: Set up Python
203-
uses: actions/setup-python@v4
280+
- name: Install uv
281+
uses: astral-sh/setup-uv@v6
204282
with:
205283
python-version: "3.11"
206-
cache: "poetry"
284+
activate-environment: true
285+
enable-cache: true
207286
- name: Install pandoc
208287
uses: awalsh128/cache-apt-pkgs-action@latest
209288
with:
210289
packages: pandoc
211290
version: 1.0
212291
- name: Install dependencies
213-
run: poetry install --with docs
292+
run: |
293+
uv lock --check
294+
uv sync
214295
- name: Download Test Coverage Artifact
215296
uses: actions/download-artifact@v4
216297
with:
@@ -219,12 +300,12 @@ jobs:
219300
- name: Authenticate HuggingFace CLI
220301
if: env.HF_TOKEN != ''
221302
run: |
222-
pip install huggingface_hub
223-
huggingface-cli login --token "$HF_TOKEN"
303+
uv pip install huggingface_hub
304+
uv run huggingface-cli login --token "$HF_TOKEN"
224305
env:
225306
HF_TOKEN: ${{ secrets.HF_TOKEN }}
226307
- name: Build Docs
227-
run: poetry run build-docs
308+
run: make build-docs
228309
env:
229310
HF_TOKEN: ${{ secrets.HF_TOKEN }}
230311
- name: Upload Docs Artifact
@@ -251,4 +332,4 @@ jobs:
251332
with:
252333
folder: docs/build
253334
clean-exclude: |
254-
*.*.*/
335+
*.*.*/

‎.github/workflows/release.yml‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -43,22 +43,25 @@ jobs:
4343
runs-on: ubuntu-latest
4444
steps:
4545
- uses: actions/checkout@v3
46-
- name: Install Poetry
47-
uses: snok/install-poetry@v1
48-
- name: Set up Python
49-
uses: actions/setup-python@v4
46+
- name: Install uv
47+
uses: astral-sh/setup-uv@v6
5048
with:
5149
python-version: '3.11'
52-
cache: 'poetry'
53-
- name: Poetry config
54-
run: poetry self add 'poethepoet[poetry_plugin]'
50+
activate-environment: true
51+
enable-cache: true
5552
- name: Install dependencies
56-
run: poetry install --with dev
53+
run: |
54+
uv lock --check
55+
uv sync --locked
5756
- name: Set the version
58-
run: poetry version ${{needs.semver-parser.outputs.semver}}
57+
run: make version-set VERSION=${{needs.semver-parser.outputs.semver}}
58+
- name: re-sync
59+
run: uv sync
5960
- name: Build
60-
run: poetry build
61+
run: uv build
6162
- name: Publish
62-
run: poetry publish
63+
run: |
64+
python -m pip install --upgrade twine
65+
twine upload --non-interactive -u __token__ -p "$PYPI_TOKEN_PYPI" dist/*
6366
env:
64-
POETRY_PYPI_TOKEN_PYPI: ${{ secrets.PYPI_TOKEN_PYPI }}
67+
PYPI_TOKEN_PYPI: ${{ secrets.PYPI_TOKEN_PYPI }}

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ docs/build
1717
.coverage*
1818
.Ds_Store
1919
.pylintrc
20+
2021
docs/source/generated
22+
# docs/source/_static/model_table
2123
**.orig
2224
.venv
25+

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ logits, activations = model.run_with_cache("Hello World")
5353
## Key Tutorials
5454

5555
* [Introduction to the Library and Mech
56-
Interp](https://arena3-chapter1-transformer-interp.streamlit.app/[1.2]_Intro_to_Mech_Interp)
56+
Interp](https://arena-chapter1-transformer-interp.streamlit.app/[1.2]_Intro_to_Mech_Interp)
5757
* [Demo of Main TransformerLens Features](https://neelnanda.io/transformer-lens-demo)
5858

5959
## Gallery

0 commit comments

Comments
 (0)