Fix: rerotation with multi-gpus

[giulio98]

Fixes #240

This PR fixes a device mismatch in key re-rotation when models are split across multiple GPUs with device_map="auto".

KeyRerotationPress._rerotate_cos_sin() could receive x, selected_positions, and inv_freq on different CUDA devices. This caused re-rotation to fail with errors like:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

The fix uses x.device as the canonical device inside _rerotate_cos_sin() and moves selected_positions and inv_freq there before computing the rotary frequencies.

This fixes FinchPress and should also apply to other presses using KeyRerotationPress.rerotate_keys().

Tested locally with FinchPress and meta-llama/Llama-3.1-8B-Instruct using device_map="auto" on a multi-GPU setup.

Checklist

Before submitting a PR, please make sure:

• Tests are working (make test)
• Code is formatted correctly (make style, on errors try fix with make format)
• Copyright header is included
• All commits are signed-off using git commit -s
• (new press) mypress_press.py is in the presses directory
• (new press) MyPress is in __init__.py
• (new press) README.md is updated with a 1 liner about the new press in the Available presses section
• (new press) New press is in the default_presses list in tests/default_presses.py
• (new press) A docstring is provided that follows the same structure as the existing ones

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[SimJeg]

Thanks for the contribution @giulio98 !