fix(test): stop passing removed device kwarg to Llama4 rotary embed… (#1516)

## Summary

`test/transformers/test_llama4_rope.py` constructs the rotary embedding
directly as
`Llama4TextRotaryEmbedding(config=config, device=device)`. Newer
HuggingFace `transformers`
deprecated the `device` kwarg on config-based
`*RotaryEmbedding.__init__`
(`@deprecate_kwarg("device", version="5.18")`) and **removes it at
5.18**. Because CI installs
the latest `transformers` (`setup.py` pins `transformers>=4.52.0`, no
upper bound), every test
in the file fails at setup with:

```
TypeError: Llama4TextRotaryEmbedding.__init__() got an unexpected keyword argument 'device'
```

This fixes all of the currently-failing `test_llama4_rope.py` cases by
dropping the removed kwarg.

## Why dropping `device=` is correct (not just silencing)

`Llama4TextRotaryEmbedding.forward` relocates the frequency buffer to
the input tensor's device:

```python
inv_freq_expanded = self.inv_freq[...].to(x.device)
```

and the test tensors are already created on the target device, so the
rotary module does not
need to be constructed on a specific device. Verified on
`transformers==5.15.0` that `cos`/`sin`
are **bitwise identical** (`torch.equal`) to the previous
`device=`-based construction.

## Relationship to #1514

#1514 fixed the same `device`-kwarg drift, but only for call sites
routed through
`transformers_version_dispatch`. `test_llama4_rope.py` constructs the
module **directly**, so it
was not covered. This PR is the direct-construction counterpart.

## Testing

- `pytest test/transformers/test_llama4_rope.py` → **19 passed**
(transformers 5.15.0, H100).
- Simulated transformers ≥5.18 (strict signature): the old `device=`
call raises the exact CI
  `TypeError`; the new call succeeds.
- `ruff check` / `ruff format --check` clean.

## Compatibility

| transformers | Before this PR | After this PR |
|---|---|---|
| < 4.48.0 | n/a (test path unchanged) | n/a |
| 4.48.0 – <5.18 | `device=` accepted (deprecation warning) | no
`device=`, identical result, no warning |
| ≥ 5.18 | **`TypeError` at construction** | works |

Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Vaibhav Jindal
2026-10-07 23:20:35 +00:00
committed by GitHub
co-authored by Copilot
parent 98f6371a9c
commit d5f281789a
+3 -3
View File
@@ -55,7 +55,7 @@ def test_correctness(bsz, seq_len, num_q_heads, num_kv_heads, head_dim, dtype, a
rope_theta=10000.0,
rope_scaling=None,
)
rotary_emb = Llama4TextRotaryEmbedding(config=config, device=device)
rotary_emb = Llama4TextRotaryEmbedding(config=config)
_tensor_q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).to(dtype)
_tensor_k = torch.randn((bsz, seq_len, num_kv_heads, head_dim), device=device).to(dtype)
@@ -109,7 +109,7 @@ def test_functional_correctness(bsz, seq_len, num_q_heads, num_kv_heads, head_di
rope_theta=10000.0,
rope_scaling=None,
)
rotary_emb = Llama4TextRotaryEmbedding(config=config, device=device)
rotary_emb = Llama4TextRotaryEmbedding(config=config)
_q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device, dtype=dtype)
_k = torch.randn((bsz, seq_len, num_kv_heads, head_dim), device=device, dtype=dtype)
@@ -200,7 +200,7 @@ def test_row_offset_does_not_wrap_int32(bsz, seq_len, num_q_heads, num_kv_heads,
rope_theta=10000.0,
rope_scaling=None,
)
rotary_emb = Llama4TextRotaryEmbedding(config=config, device=device)
rotary_emb = Llama4TextRotaryEmbedding(config=config)
pos_ids = torch.arange(seq_len, device=device).unsqueeze(0)
freqs_cis = rotary_emb(q, pos_ids)