Skip to content

Add OLMo 2 support (1B, macOS + iOS) - #232

Merged
stikves merged 6 commits into
apple:mainfrom
stikves:sukru/olmo2-support
Sep 11, 2026
Merged

stikves merged 6 commits into
apple:mainfrom
stikves:sukru/olmo2-support

Conversation

@stikves

@stikves stikves commented Sep 4, 2026 •

Copy link
Copy Markdown
Contributor

Summary

  • Add OLMo 2 1B model (Allen AI, Apache 2.0) with macOS and iOS variants
  • Post-norm architecture (x + norm(attn(x))) and flat QK-norm (RMSNorm over all head dims jointly before reshape)
  • macOS: INT4 quantized, iOS: 6-bit palettized
  • MHA (16/16 heads), SwiGLU, standard RoPE, no attention bias

Evaluation

word_perplexity on WikiText-2, 4K context.

Model Compression BPW Platform Perplexity
1B none (float16) 16.00 macOS 17.06
1B 4-bit quantized 4.50 macOS 18.72
1B none (float16) 16.00 iOS 17.09
1B 6-bit palettized 6.00 iOS 17.19

Test plan

  • 8 unit tests (parity, shape, structural, incremental decode)
  • PPL within 0.01% of HF baseline
  • iOS parity at 0.00 max logit diff

@stikves stikves changed the title Add OLMo 2 support (1B, macOS + iOS, Apache 2.0) Add OLMo 2 support (1B, macOS + iOS) Sep 4, 2026
@stikves
stikves force-pushed the sukru/olmo2-support branch from 546c24c to de1b526 Compare September 10, 2026 00:41
@stikves
stikves marked this pull request as ready for review September 10, 2026 00:44
Comment thread models/olmo2/olmo2_1b_6bit.yaml Outdated
Comment thread models/olmo2/olmo2_1b_6bit.yaml Outdated
stikves and others added 4 commits September 10, 2026 12:49
Co-authored-by: Prathamesh <46148373+pkmandke@users.noreply.github.com>
group_size 8 was unnecessarily fine; a wikitext PPL sweep at the model's
native 4096 context shows 8/16/32 are equivalent within noise (17.18 / 17.09 /
17.23). group_size 16 halves the palettization LUT overhead versus 8 while
holding quality (17.09). Update the README iOS 6-bit perplexity accordingly.
@stikves
stikves merged commit f16a60c into apple:main Sep 11, 2026
3 checks passed
@stikves
stikves deleted the sukru/olmo2-support branch September 11, 2026 19:30
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.

3 participants