Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[Olmo3] different RoPE per layer type - #46911

Merged
zucchini-nlp merged 11 commits into
huggingface:mainfrom
zucchini-nlp:olmo3-rope
Jun 29, 2026
Merged

[Olmo3] different RoPE per layer type#46911
zucchini-nlp merged 11 commits into
huggingface:mainfrom
zucchini-nlp:olmo3-rope

Conversation

@zucchini-nlp

Copy link
Copy Markdown
Member

What does this PR do?

Reverts back per layer-type RoPE in Olmo which was removed in #39847

Comment on lines +222 to +228
def test_real_model_7b_greedy_generation(self):
expectations = Expectations(
{
("cuda", None): 'system\nYou are a helpful function-calling AI assistant. You do not currently have access to any functions. <functions></functions>\nuser\nWho would win in a fight - a dinosaur or a cow named Moo Moo?\nassistant\nThis is a fun and imaginative question! Let’s break it down:\n\n### 1. **A Dinosaur (General Case)**\nDinosaurs were a huge and diverse group, spanning from tiny feathered raptors to massive sauropods like *Brachiosaurus* or *Tyrannosaurus rex',
}
) # fmt: skip

@zucchini-nlp zucchini-nlp Jun 26, 2026

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.

there were no integration tests with official ckpt, somehow it redirects to someone's personal repo 🫠

I added tests with ckpt that are supposed to use rope scaling, and we should see a difference now . Adding test for long seq beyond sliding window in a sec

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.

i didn't delete existing slow tests, not sure if that repo is supposed to be tested. LMK if you think we can delete everything to not waste resources on running them

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.

imo, best to keep but move/copy to internal testing. personal repos is not so nice

no need to save on resources, rather have something non broken

@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.

Comment thread tests/causal_lm_tester.py
@zucchini-nlp

Copy link
Copy Markdown
Member Author

run-slow: olmo3, olmo_hybrid

Comment on lines +961 to 964
# Released ckpt don't use any ROPE and have it set to `None`
self.rotary_emb = (
OlmoHybridRotaryEmbedding(config=config)
if getattr(config, "rope_parameters", None) is not None

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.

ig we can't delete the module, since some users might have added rope in non-official ckpts

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.

yep rather keep it now, that's why it's hard when official ckpts release after integration :/

@huggingface huggingface deleted a comment from github-actions Bot Jun 26, 2026
@huggingface huggingface deleted a comment from github-actions Bot Jun 26, 2026
@huggingface huggingface deleted a comment from github-actions Bot Jun 26, 2026
@zucchini-nlp
zucchini-nlp requested a review from vasqu June 26, 2026 10:58
@zucchini-nlp

Copy link
Copy Markdown
Member Author

run-slow: olmo3, olmo_hybrid

@github-actions

Copy link
Copy Markdown
Contributor

Workflow Run ⚙️

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

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

@github-actions

Copy link
Copy Markdown
Contributor

CI Results

Workflow Run ⚙️

Commit Info

Context Commit Description
RUN de5cfa59 workflow commit (merge commit)
PR a6dfe8d0 branch commit (from PR)
main ed7d6c8d 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.

I'm only a bit hesitant re that non downcasting on the rope of olmo3 - any source for that?

But other than that agree with most points, just nits/smaller comments

Comment thread tests/causal_lm_tester.py
Comment thread tests/models/olmo3/test_modeling_olmo3.py
Comment on lines +222 to +228
def test_real_model_7b_greedy_generation(self):
expectations = Expectations(
{
("cuda", None): 'system\nYou are a helpful function-calling AI assistant. You do not currently have access to any functions. <functions></functions>\nuser\nWho would win in a fight - a dinosaur or a cow named Moo Moo?\nassistant\nThis is a fun and imaginative question! Let’s break it down:\n\n### 1. **A Dinosaur (General Case)**\nDinosaurs were a huge and diverse group, spanning from tiny feathered raptors to massive sauropods like *Brachiosaurus* or *Tyrannosaurus rex',
}
) # fmt: skip

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.

imo, best to keep but move/copy to internal testing. personal repos is not so nice

no need to save on resources, rather have something non broken

Comment thread tests/models/olmo3/test_modeling_olmo3.py Outdated
Comment thread tests/models/olmo3/test_modeling_olmo3.py Outdated
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
@is_tensor_parallel_test
def test_tp_generation_quantized(self):
# If model uses rope-theta 50k (default value), the test fails

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.

that is surprising, wondering whether it also affects other models 👀

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.

ikr, very weird but I didn't want to dig yet. Most models init with 10k by default except for really weird ones, I will check a bit later and open an issue/another PR if necessary

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.

hmm seems like it is an error accumulation from QK norm over the whole hidden-dim and higher theta, I'll just leave it there as there are more models skipping this particular test with todo: investigate
maybe we'd need to raise tols for certain cases

Comment thread src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
Comment on lines +961 to 964
# Released ckpt don't use any ROPE and have it set to `None`
self.rotary_emb = (
OlmoHybridRotaryEmbedding(config=config)
if getattr(config, "rope_parameters", None) is not None

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.

yep rather keep it now, that's why it's hard when official ckpts release after integration :/

Comment thread src/transformers/models/olmo3/modeling_olmo3.py
self.num_key_value_heads = self.num_attention_heads
super().__post_init__(**kwargs)

def convert_rope_params_to_dict(self, **kwargs):

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.

not on you but we should avoid it whenever we can

@github-actions

Copy link
Copy Markdown
Contributor

CI Dashboard: View test results in Grafana

@zucchini-nlp
zucchini-nlp added this pull request to the merge queue Jun 29, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jun 29, 2026
@zucchini-nlp
zucchini-nlp added this pull request to the merge queue Jun 29, 2026
@github-actions

Copy link
Copy Markdown
Contributor

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

run-slow: olmo3, olmo_hybrid

Merged via the queue into huggingface:main with commit 326683d Jun 29, 2026
5 checks passed
@zucchini-nlp
zucchini-nlp deleted the olmo3-rope branch June 29, 2026 08:52
@Cyrilvallez

Copy link
Copy Markdown
Member

cc @zucchini-nlp @vasqu why are the cos/sin not downcasted at the end here like any other Rope module? Any reason? This means any later ops will be upcasted in the model...

@zucchini-nlp

Copy link
Copy Markdown
Member Author

@Cyrilvallez yeas, it is done on purpose to follow official ref impl. QK are then recasted back to model dtype after rotating

q_type, k_type = q.dtype, k.dtype
cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim)
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed.to(q_type), k_embed.to(k_type)

@Cyrilvallez

Copy link
Copy Markdown
Member

All right, thanks for confirming!

stevhliu pushed a commit to stevhliu/transformers that referenced this pull request Jul 30, 2026
* fix maybe, needs checking

* fix

* oops

* fix repo

* fix rope tests

* why was it added if ckpt has no rope?

* fix the TP test

* adjust expectations with runners

* Update tests/models/olmo3/test_modeling_olmo3.py

Co-authored-by: Anton Vlasjuk <[email protected]>

* spit out personal repo test in a new class

---------

Co-authored-by: Anton Vlasjuk <[email protected]>
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.

4 participants