transformers - 💡(How to fix) Fix `Qwen3_5MoeTextRotaryEmbedding.forward` is not compatible with CPU offload

import sys

import torch
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
    Qwen3_5MoeTextConfig,
    Qwen3_5MoeTextRotaryEmbedding,
)

if not torch.cuda.is_available():
    raise NotImplementedError("This reproducer needs at least one CUDA device")

# Default config; the bug is in the rotary forward, so the precise model
# dims don't matter. Using defaults avoids any checkpoint download
cfg = Qwen3_5MoeTextConfig()
rotary = Qwen3_5MoeTextRotaryEmbedding(cfg)
print(f"rotary.inv_freq.device = {rotary.inv_freq.device}")
assert rotary.inv_freq.device.type == "cpu", (
    "Expected inv_freq to land on CPU when no device is passed; got"
    f" {rotary.inv_freq.device}. Bug premise no longer holds."
)

# GPU activation, simulating what FSDP2 forward routes into the rotary call
bsz, seq_len = 1, 8
x = torch.randn(bsz, seq_len, cfg.hidden_size, device="cuda", dtype=torch.bfloat16)
position_ids = torch.arange(seq_len, device="cuda")[None, :]

try:
    cos, sin = rotary(x, position_ids)
except RuntimeError as exc:
    msg = str(exc)
    if (
        "Expected all tensors to be on the same device" in msg
        and "wrapper_CUDA_bmm" in msg
    ):
        print(f"Reproduced: {type(exc).__name__}: {msg}")
        sys.exit(0)
    raise NotImplementedError(
        f"Unexpected RuntimeError (not the device-mismatch we expect): {msg}"
    ) from exc

raise NotImplementedError(
    f"Did not reproduce: rotary forward returned {cos.device=} {sin.device=}."
)

---

RuntimeError: Expected all tensors to be on the same device, but got mat2 is on cuda:1, different from other tensors on cpu (when checking argument in method wrapper_CUDA_bmm)