Read rope_theta from rope_parameters in the Llama injection policy

[alanhuangyoo]

Fixes #8340.

What breaks

DS_LLAMAContainer.create_module resolves rope_theta as:

if hasattr(self.policy.client_module.self_attn, 'config'):
    _config.rope_theta = self.policy.client_module.self_attn.config.rope_theta
else:
    _config.rope_theta = self.policy.client_module.self_attn.rope_theta

transformers 5.0 folded the rotary settings into config.rope_parameters and dropped the attribute. LlamaAttention still has .config, so the first branch is taken and it raises.

This is not specific to the DeepSeek checkpoint in the issue — on transformers >= 5.0 it reproduces with a default LlamaConfig. Checked on 5.8.0:

transformers 5.8.0
LlamaConfig().rope_theta        -> AttributeError
LlamaConfig().rope_parameters   -> {'rope_theta': 10000.0, 'rope_type': 'default'}

and against a real module, both branches are dead:

self_attn has .config                 True
self_attn has .rope_theta             False
self_attn.config has rope_theta       False
self_attn.config has rope_parameters  True

current code -> AttributeError: 'LlamaConfig' object has no attribute 'rope_theta'

The fix

Try the old spellings first, then rope_parameters['rope_theta'], so pre-5.0 installs take exactly the path they take today and nothing changes for them.

This is the same drift #7443 adapted to. The num_heads accessor it fixed alongside still resolves on 5.8 (config.num_attention_heads is intact), so rope_theta is the only one that moved again.

requirements-dev.txt asks for transformers>=4.51.3 with no upper bound, so 5.x is in range.

Verified that rope_theta is the only accessor that moved: against a real LlamaDecoderLayer on 5.8.0, get_hidden_heads(), attention(), mlp() and layernorm() all resolve once this one is fixed.

I exercised the policy and container layer, not a full init_inference() run against a downloaded checkpoint on GPU.

Test

tests/unit/module_inject/test_llama_rope_theta.py covers the three layouts plus the precedence between them and the not-found case — CPU only, no model download:

6 passed

End to end against a real LlamaAttention on 5.8.0:

current code   AttributeError: 'LlamaConfig' object has no attribute 'rope_theta'
with this PR   10000.0

tests/unit/module_inject/    43 passed
yapf --diff / flake8         clean

I did not touch inference/v2. Six of its model implementations read self._config.rope_theta directly and exaone4 is the only one using a getattr default, so they likely have the same exposure — but that is a different engine and a different change, and I have not reproduced it.

[alanhuangyoo]

Opened #8345 for the inference/v2 exposure I flagged at the bottom of this description — separate engine, separate change, so I kept it out of here.

It turned out to be eight models, and one of them fails quietly rather than loudly: exaone4 reads the base through getattr(self._config, "rope_theta", 1000000.0), and Exaone4Config carries rope_parameters["rope_theta"] == 10000.0, so on 5.x it puts a 100x rotary base into RotateHalfConfig with nothing raised.

[ebarkhordar]

I ran the helper against a real LlamaConfig on both majors, in clean python:3.12-slim containers, at head ecda5f94. It resolves correctly on each, and on 5.8.0 a non-default theta survives the move, which is the part that matters for an actual checkpoint rather than the stock default:

transformers 4.56.2   hasattr(config,'rope_theta')=True   rope_parameters=None
                      LlamaConfig(rope_theta=5e5) -> 500000.0
transformers 5.8.0    hasattr(config,'rope_theta')=False  rope_parameters={'rope_theta': 10000.0, 'rope_type': 'default'}
                      LlamaConfig(rope_theta=5e5) -> 500000.0

The suggestion is about the tests rather than the fix. All six build the config out of SimpleNamespace, so the claim the PR rests on, that rope_theta is absent from a stock LlamaConfig on 5.x and present on 4.x, is the one thing nothing pins. If transformers reinstates the attribute as a deprecated property, every test here still passes and the injection path silently changes branch.

One test with the real class covers it, and it stays meaningful on either side of the boundary, taking the legacy branch on 4.x and rope_parameters on 5.x:

def test_reads_a_real_llama_config():
    from transformers.models.llama.configuration_llama import LlamaConfig

    assert _get_rope_theta(SimpleNamespace(config=LlamaConfig(rope_theta=500000.0))) == 500000.0

I checked it passes on both versions above before suggesting it. I lifted _get_rope_theta into a standalone module to run this without a torch install, so I exercised the helper and the config, not the container.

[alanhuangyoo]

Right, and the gap is worse than "nothing pins it" — every one of those six passes against the implementation this PR replaces. I checked:

old helper, real LlamaConfig(rope_theta=5e5), transformers 5.16.1
  -> AttributeError
old helper, all six SimpleNamespace cases
  -> pass

So the file had no test that could fail before the fix. Added in 3dd2716, along with a second one for the part your suggestion does not cover.

Your test pins the value. It does not catch the branch flip you describe: if transformers reinstates rope_theta as a deprecated property, it still passes and the injection path silently takes the legacy branch. That only matters if the two spellings can disagree, so the second test asserts they cannot:

parameters = getattr(config, "rope_parameters", None) or {}
carried = [v for v in (getattr(config, "rope_theta", None), parameters.get("rope_theta")) if v is not None]
assert carried, "neither spelling carries rope_theta on the installed transformers"
assert all(v == 500000.0 for v in carried), f"the spellings disagree: {carried}"

It holds on both sides — one entry on 4.x, one on 5.x, two agreeing entries in the reinstated case, and a failure naming the values if they ever diverge.

Measured here on 5.16.1, matching your 5.8.0 run:

hasattr(config, 'rope_theta')  False
rope_parameters                {'rope_theta': 500000.0, 'rope_type': 'default'}
_get_rope_theta                500000.0
LlamaConfig() stock default    10000.0

8 passing, yapf and flake8 clean. This is the same thing you found on #8345 and it was the right thing to find twice — I had fixed the synthetic-config problem there and left it standing here.

[ebarkhordar]

Agreed on the second test. Mine pins the value and would sit green through exactly the reinstatement I described, so it does not cover the case I raised it for.

Ran both real-config tests standalone, with _get_rope_theta parsed out of llama.py at 3dd2716 so this needed no torch install:

transformers 4.56.2   hasattr rope_theta True    rope_parameters None
transformers 5.8.0    hasattr rope_theta False   rope_parameters {'rope_theta': 500000.0, 'rope_type': 'default'}
transformers 5.16.1   hasattr rope_theta False   rope_parameters {'rope_theta': 500000.0, 'rope_type': 'default'}

2 passed on each

[alanhuangyoo]

Thanks for running it across three majors — 4.56.2 is the half I could not measure here, so that closes the gap the tests are meant to cover. Nothing further from me on this one.

[tohtana]

Hi @alanhuangyoo,
Thank you for submitting this PR! This is an important fix.

I found a correctness blocker for the DeepSeek-R1-Distill-Llama-8B configuration reported in #8340. It uses rope_type="llama3" together with factor, low_freq_factor, high_freq_factor, and original_max_position_embeddings.

This patch passes only rope_theta to the legacy DeepSpeed inference configuration, so the additional scaling parameters do not reach the downstream kernel. The kernel currently implements only standard scalar-theta RoPE. Supporting this configuration would require propagating the additional parameters and adding an implementation and dispatch path for Llama-3 scaled RoPE.

[alanhuangyoo]

You are right, and the consequence is worse than the parameters not arriving.

I checked what the injected path could do with them and the answer is nothing. Its rotary embedding is built from a scalar base:

InferenceContext.get_rotary(rotary_dim, rope_theta)

and DeepSpeedInferenceConfig carries rope_theta with no scaling fields — grepping the inference kernel path for rope_scaling, rope_type, low_freq_factor or scaling_factor returns nothing but MoE routing.

So with only the crash fix, DeepSeek-R1-Distill-Llama-8B would have started and run with unscaled positions, producing wrong output and exiting 0. That is worse than the AttributeError this PR exists to remove — the crash is at least visible. Thank you for catching it before it merged.

16c9b4c refuses instead:

DeepSpeed kernel injection cannot serve rope_type='llama3'. The injected attention
kernel builds its rotary embedding from rope_theta alone, so the scaling parameters
this configuration carries would be dropped and the model would run with unscaled
positions. Run this model without kernel injection (replace_with_kernel_inject=False).

Both spellings are covered — rope_parameters on 5.x and rope_scaling on 4.x — since the same request is written either way depending on the installed version. default and an absent type mean no scaling and still resolve, so the configuration #8340 reported is fixed as it was.

16 passing, including a parametrized refusal over llama3/linear/dynamic/yarn/longrope, the legacy-spelling case, and a real LlamaConfig still resolving. yapf and flake8 clean.

On the larger fix: propagating the parameters through DeepSpeedInferenceConfig and adding an inv_freq path plus dispatch in the kernel is a feature rather than a crash fix, and I would rather not smuggle it into this PR. Happy to open it as its own issue and take it if that is wanted — say the word and I will write it up with the same shape as #8420.

[alanhuangyoo]

Separate from the review: the red modal-torch-latest / DeepSpeedAI CI here is a job timeout, not a test failure. The API reports "conclusion": "cancelled", the job ran exactly 1h30m, and there is no FAILED line in the log. collect tests passes on the same commit.

Three of my PRs hit the same 1h30m cancellation today (#8341, #8356, #8438) while six went green, so it looks like the full suite sometimes runs past the job limit rather than anything about these trees. Mentioning it because a red X plus a changes-requested reads worse than the situation is.

[alanhuangyoo]

@tohtana — re-review request rather than a new argument. Your changes-requested is still the only thing on this one and I think it is addressed, but I would rather you confirm than assume.

You said the patch passes only rope_theta, so a rope_type="llama3" config's factor / low_freq_factor / high_freq_factor / original_max_position_embeddings never reach the kernel, and the kernel implements scalar-theta RoPE only.

Agreed, and I went and checked what the injected path could do with them: nothing. Its rotary embedding is built from a scalar theta_base, so there is no representation for a scaled variant even if the values arrived. Serving DeepSeek-R1-Distill-Llama-8B through it would need the scaling implemented in the kernel, which is a much bigger change than this PR.

So 8447987 refuses instead of silently serving the wrong rotary: _get_rope_theta raises for any rope_type other than None / default, naming the type it saw. That turns a config which currently produces quietly wrong outputs into one that says why it cannot be served. The original bug — rope_theta moving to rope_parameters in transformers 5.x, which broke the unscaled Llama path too — is still fixed.

If you would rather this PR implement llama3 scaling in the kernel rather than refuse it, say so and I will close this and open that instead; I did not want to bundle them.

The red CI here was the 90-minute job timeout, not a failure — no FAILED line in the log. I have merged current master in, which picks up #8404 raising the limit to 105, so it is re-running now.

[tohtana]

Thank you for the update, @alanhuangyoo! Looks good to me.

[alanhuangyoo]

@tohtana thanks — the formatting job was the last red and it is fixed in 4c687f1 (yapf reflowing one SimpleNamespace call in the test, no behaviour change). Everything is green now:

SUCCESS  formatting / formatting checks
SUCCESS  cpu-torch-latest / unit tests
SUCCESS  nv-pre-compile-ops / precompile ops
SUCCESS  python / install smoke (3.10, 3.11, 3.12)
SUCCESS  modal-torch-latest / collect tests
SKIPPED  modal-torch-latest / DeepSpeedAI CI   (runs from the merge queue since #8412)

Ready for the queue whenever suits. I did not touch the branch otherwise — the master merge on it is yours.