ModernBert: reuse GemmaRotaryEmbedding via modular + Integration tests #35459
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. |
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Super nice to add the tests
| class ModernBertRotaryEmbedding(GemmaRotaryEmbedding): | ||
| def __init__(self, config: ModernBertConfig, dim: int, base: float, device: Optional[torch.device] = None): | ||
| super().__init__(self, config=config, device=device) | ||
| self.rope_kwargs = {"dim": dim, "base": base} |
There was a problem hiding this comment.
I think a PR #35589 might have removed the kwargs! but happy to add them back if need or just not use the gemma one!
There was a problem hiding this comment.
I'll reintroduce the rope_kwargs in LlamaRotaryEmbedding.
…to modernbert/rotary_embeds
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Actually it makes more sense to juste inherit and overwrite the call to inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs) rather than reverting the commit!
…odular_modernbert" This reverts commit 11b44b9.
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Good to go! Thanks
#35459) * Introduce 5 integration tests for the 4 model classes + torch export * ModernBert: reuse GemmaRotaryEmbedding via modular * Revert #35589, keep rope_kwargs; rely on them in modular_modernbert * Revert "Revert #35589, keep rope_kwargs; rely on them in modular_modernbert" This reverts commit 11b44b9. * Don't set rope_kwargs; override 'self.rope_init_fn' call instead
4 participants
What does this PR do?
Note that I based the expected values from the integration tests on the
mainimplementation, not the RotaryEmbedding refactored implementation. In short: the test passes indicate that the refactor preserves the performance.This is the refactor that I promised earlier offline.
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
@ArthurZucker
cc @warner-benjamin