Skip to content

Align OlmoHybrid to use a native cache in generate - #47604

Merged
Cyrilvallez merged 12 commits into
mainfrom
fix-olmo-hybrid
Aug 5, 2026
Merged

Align OlmoHybrid to use a native cache in generate#47604
Cyrilvallez merged 12 commits into
mainfrom
fix-olmo-hybrid

Conversation

@Cyrilvallez

@Cyrilvallez Cyrilvallez commented Jul 28, 2026

Copy link
Copy Markdown
Member

CI

What does this PR do?

As per the title. Following the work in #47347, we have all the tools to make OlmoHybrid (almost) a good citizen. The weird part is that it was added with a ShortConvolution module from fla.modules, and we now have to keep weird if/else paths based on that because they do not operate with the same inputs.....

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@huggingface huggingface deleted a comment from github-actions Bot Jul 29, 2026
@huggingface huggingface deleted a comment from github-actions Bot Jul 29, 2026
@Cyrilvallez

Copy link
Copy Markdown
Member Author

run-slow: olmo_hybrid

@huggingface huggingface deleted a comment from github-actions Bot Jul 29, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/olmo_hybrid"]
quantizations: []

@github-actions

Copy link
Copy Markdown
Contributor

CI Results

Workflow Run ⚙️

Commit Info

Context Commit Description
RUN bb8c5c9f workflow commit (merge commit)
PR a77197dc branch commit (from PR)
main 3d7b75a8 base commit (on main)

✅ No failing test specific to this PR 🎉 👏 !

@vasqu vasqu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Small nits but happy with this landing ❤️

Comment thread src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py Outdated
Comment thread src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py Outdated
Comment thread src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py Outdated
class OlmoHybridPreTrainedModel(Qwen3NextPreTrainedModel):
_is_stateful = True
# Uses a custom ``OlmoHybridDynamicCache``; StaticCache compatibility hasn't been wired up here.
_can_compile_fullgraph = False

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Does fullgraph work now maybe?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Oh indeed it does of course! I added the Static versions for that with Inkling as well! Thanks for catching haha, we can start testing!

@vasqu

vasqu commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Added an alternative with fused conv here #47630, there is no difference between fused and unfused conv weights. The results are identical

@Cyrilvallez

Copy link
Copy Markdown
Member Author

Humm, I guess we could merge the weights if we want, but then we kind of break BC with regards to using the module from external package no?

@vasqu

vasqu commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Closing as all models are aligned now with proper cache support

@vasqu vasqu closed this Aug 5, 2026
@Cyrilvallez Cyrilvallez reopened this Aug 5, 2026
@github-actions

github-actions Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: olmo_hybrid

@Cyrilvallez Cyrilvallez changed the title Align OlmoHybrid to use a native cache Align OlmoHybrid to use a native cache in generate Aug 5, 2026
@github-actions

github-actions Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 31006155916:1
Result: failure | Jobs: 16 | Tests: 180,313 | Failures: 1 | Duration: 16h 2m

@vasqu vasqu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

One small question, i guess we could also revert the test change but i dont mind the change as its more general

self.assertEqual(layer.conv_states[i].shape[:-1], conv_shape[:-1])
self.assertTrue(layer.conv_states[i].shape[-1] >= conv_shape[-1])
self.assertEqual(layer.conv_states[i].shape[:-1], current_conv_shape[:-1])
self.assertTrue(layer.conv_states[i].shape[-1] >= current_conv_shape[-1])

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I guess inkling had same shape for all convs?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Nop, but it's not correctly tested...... 😭😭

@Cyrilvallez
Cyrilvallez merged commit 99fa125 into main Aug 5, 2026
111 of 113 checks passed
@Cyrilvallez
Cyrilvallez deleted the fix-olmo-hybrid branch August 5, 2026 15:13
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