Align OlmoHybrid to use a native cache in generate - #47604
Conversation
|
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. |
|
run-slow: olmo_hybrid |
|
This comment contains models: ["models/olmo_hybrid"] |
vasqu
left a comment
There was a problem hiding this comment.
Small nits but happy with this landing ❤️
| class OlmoHybridPreTrainedModel(Qwen3NextPreTrainedModel): | ||
| _is_stateful = True | ||
| # Uses a custom ``OlmoHybridDynamicCache``; StaticCache compatibility hasn't been wired up here. | ||
| _can_compile_fullgraph = False |
There was a problem hiding this comment.
Does fullgraph work now maybe?
There was a problem hiding this comment.
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!
|
Added an alternative with fused conv here #47630, there is no difference between fused and unfused conv weights. The results are identical |
|
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? |
|
Closing as all models are aligned now with proper cache support |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: olmo_hybrid |
CI recapDashboard: View test results in Grafana |
vasqu
left a comment
There was a problem hiding this comment.
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]) |
There was a problem hiding this comment.
I guess inkling had same shape for all convs?
There was a problem hiding this comment.
Nop, but it's not correctly tested...... 😭😭
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
ShortConvolutionmodule fromfla.modules, and we now have to keep weird if/else paths based on that because they do not operate with the same inputs.....