Skip to content

Complete type checking for OLMo support (builds on #816) #1081

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
taziksh wants to merge 78 commits into
base: dev
Choose a base branch
from
Open

Complete type checking for OLMo support (builds on #816) #1081

taziksh wants to merge 78 commits into from

Conversation

@taziksh
Copy link

@taziksh taziksh commented Oct 12, 2025 •
edited
Loading

Description

This PR builds on #816 and resolves all remaining type checking issues, allowing the OLMo implementation to pass mypy validation.

Changes:

  • Added type assertions for OLMo2/OLMoE decoder layers to resolve union-attr errors
  • Bumped minimum Python version to 3.10 for jaxtyping mypy plugin compatibility (Python 3.8 reached EOL Oct 2024)
  • Fixed jaxtyping dimension annotation in create_alibi_multipliers (head_idx → n_heads)

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • My changes generate no new warnings
  • New and existing unit tests pass locally with my changes

jonasrohw and others added 27 commits December 12, 2024 09:58
@jonasrohw
added and tested: OLMo-1B,OLMo-7B
1fe4d04
@jonasrohw
fixed: numpy do not do a major upgrade!
0f3e3b3
@jonasrohw
fixed: dimensions of 7b to be correct
3a101f4
@jonasrohw
tested: Loading checkpoints & model variations
1b34ccd
@joelburget
Reimplement OLMoE changes.
f0a0a68 
Originally from TransformerLensOrg#718.
@joelburget
Implement TODO (norm_topk_prob)
8c094e5
@joelburget
Disable bos token for OLMoE.
7565c06
@joelburget
Add q and k norm.
04cd309
@joelburget
Correct normalization type for OLMoE.
68d6961
@jonasrohw
Merge pull request TransformerLensOrg#1 from joelburget/olmoe
9afd032 
Add OLMoE
@jonasrohw
Merge branch 'dev' into OLMo
96c1fbb
@jonasrohw
ran formatting
72fb903
@bryce13950
Merge branch 'dev' into OLMo
9d3a85e
@bryce13950
Merge branch 'dev' into OLMo
d4519b2
@Ja1Zhou @jonasrohw
tmp update for olmo2
064310f
@jonasrohw
Fix: Olmo2 uses normalization after the attention/mlp
b1fd04b
@bryce13950
Merge branch 'dev' into OLMo
871ba03
@bryce13950
ran format
7939e8d
@bryce13950
fixed some type issues
97fd1e7
@bryce13950
Merge branch 'dev' into OLMo
9032fe7
@jleechung
OLMo 2 RMS
39703c4
@jleechung
OLMo 2 RMS
1c283c1
@jleechung
Tested Instruct models
688a421
@jonasrohw
Merge pull request TransformerLensOrg#3 from jleechung/OLMo
9febc5c 
Fix to OLMo 2 normalization
@taziksh
fix: Olmo2DecoderLayer type issues
86b1fce
@taziksh
fix type assertions for attention
fa5c885
@taziksh
chore: bump min Python to 3.10 for jaxtyping mypy plugin compatibility
148df46
@taziksh taziksh mentioned this pull request Oct 12, 2025
Open
6 tasks
taziksh and others added 29 commits January 16, 2026 17:34
@taziksh @jlarson4
chore: bump min Python to 3.10 for jaxtyping mypy plugin compatibility
884aeb6
@taziksh @jlarson4
fix: sort imports in olmo2.py
d51ab7d
@taziksh @jlarson4
docs: update Colab notebook for OLMo models
aa6d3b8
@jlarson4
Adjust error message to improve testing
9458988
@jlarson4
rebasing to dev
41280fe
@jlarson4
conflict resolution
80b8835
@jlarson4
Updating lock
b2ac313
@jlarson4
Fixed formatting, update error messages to properly test
8c92fc8
@jlarson4
more formatting
fc3da3e
@jlarson4
fixing type error
d7e5523
@jlarson4
fix format error
c20859c
@jlarson4
Fix type issues
5ccbf68
@jlarson4
Fix type issues
6d3c870
@jlarson4
Fix format issues
7151270
@jlarson4
Fix format issues again
1d0aebb
@jlarson4
Fix format issues for black
4316e8b
@jlarson4
another attempt at black formatting
040b19b
@jlarson4
Fix format issues for black again
fb259ce
@jlarson4
Retyping the blocks in HookedTransformer and HookedEncoder
72521c7
@jlarson4
undo modulelist typing
f0ddc0e
@jlarson4
Improve type checking in test_detect_head_with_invalid_head_name
bdbd649
@jlarson4
removing unused import
0ec06b9
@jlarson4
Fixing Patchscopes_Generation_Demo.ipynb
09a9bdd
@jlarson4
Fixing the rest of the notebooks
7933afc
@jlarson4
Fixing the more notebooks
61347e0
@jlarson4
run_line_magic
d4e986a
@jlarson4
BERT ipynb fix
4fed25a
@jlarson4
Trying to fix the BERT set_grad cell
7b22ce4
@jlarson4
more set_grad cell fixes
0899f3c
@taziksh
Copy link
Author

taziksh commented Jan 19, 2026

@jlarson4 do you know when this might be merged?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

7 participants

@taziksh @jonasrohw @joelburget @bryce13950 @Ja1Zhou @jleechung @jlarson4