Skip to content

[Model] Add FlexOlmo model implementation #24923

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

Merged
merged 17 commits into from
Oct 10, 2025
Merged

[Model] Add FlexOlmo model implementation #24923

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
a00b375
Add FlexOlmoConfig
2015aroras Sep 15, 2025
1e734fa
Add FlexOlmo model implementation
2015aroras Sep 15, 2025
f98dd62
Update supported models docs
2015aroras Sep 15, 2025
d1adc72
Move weight loading logic to FlexOlmoModel
2015aroras Sep 16, 2025
56b3cdf
Implement get_expert_mapping
2015aroras Sep 16, 2025
8942c65
Merge branch 'main' into shanea/flex-olmo-official
2015aroras Sep 19, 2025
5b3c09e
Make OlmoE suitable for inheritance by FlexOlmo
2015aroras Sep 19, 2025
2ba0a2d
Make FlexOlmo inherit from OlmoE
2015aroras Sep 19, 2025
a4a2bbb
Merge branch 'main' into shanea/flex-olmo-official
2015aroras Oct 9, 2025
bdec7d6
Remove deprecated(?) sampling logic
2015aroras Oct 9, 2025
6bdb8a6
Remove unneeded bias from gate
2015aroras Oct 9, 2025
d14a596
Update supported models
2015aroras Oct 9, 2025
c9705f8
Update test model in registry
2015aroras Oct 9, 2025
8fef613
Merge branch 'main' into shanea/flex-olmo-official
2015aroras Oct 9, 2025
a32b5f2
Fix formatting
2015aroras Oct 9, 2025
907a8d9
Merge branch 'main' into shanea/flex-olmo-official
2015aroras Oct 10, 2025
acab96d
Merge branch 'main' into shanea/flex-olmo-official
2015aroras Oct 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/models/supported_models.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -348,6 +348,7 @@ th {
| `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ | ✅︎ |
| `FalconMambaForCausalLM` | FalconMamba | `tiiuae/falcon-mamba-7b`, `tiiuae/falcon-mamba-7b-instruct`, etc. | | ✅︎ | ✅︎ |
| `FalconH1ForCausalLM` | Falcon-H1 | `tiiuae/Falcon-H1-34B-Base`, `tiiuae/Falcon-H1-34B-Instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
| `FlexOlmoForCausalLM` | FlexOlmo | TBA | | ✅︎ | ✅︎ |
| `GemmaForCausalLM` | Gemma | `google/gemma-2b`, `google/gemma-1.1-2b-it`, etc. | ✅︎ | ✅︎ | ✅︎ |
| `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ | ✅︎ |
| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ | ✅︎ |
Expand Down
1 change: 1 addition & 0 deletions tests/models/registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,7 @@ def check_available_online(
"Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"), # noqa: E501
"FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),
"FalconH1ForCausalLM":_HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"),
"FlexOlmoForCausalLM": _HfExamplesInfo("shanearora/Flex-reddit-2x7B-1T"),
Copy link
Collaborator

@jeejeelee jeejeelee Sep 16, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
"FlexOlmoForCausalLM": _HfExamplesInfo("shanearora/Flex-reddit-2x7B-1T"),
"FlexOlmoForCausalLM": _HfExamplesInfo(allenai/Flex-reddit-2x7B-1T"),

Ignore this, I misunderstood

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

c9705f8 Your suggestion is valid now that I've updated the official model.

"GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"),
"Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),
"Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),
Expand Down
Loading
Loading