rotary.py

"""Adjacent-pair (GPT-J style) rotary position embeddings.

Convention, which matters when exporting a trained checkpoint: this module
rotates *adjacent* channel pairs ``(x0, x1), (x2, x3), ...``. Llama and most
Hugging Face models instead rotate *half-split* pairs ``(x0, x_{d/2}), ...``
("NeoX style"). The two are related by a permutation of the query/key rows,
so a checkpoint trained here is NOT drop-in loadable as a Llama checkpoint
without permuting ``qkv_proj``.

Both conventions are first-class in the common inference runtimes -- select
GPT-J/``NORM``-style rotary rather than ``NEOX`` when converting. Concretely:
``llama.cpp`` ``rope_type=NORM``, vLLM ``is_neox_style=False``.

``cos``/``sin`` here have shape ``(..., sequence_length, head_dim / 2)``,
half the width of the Hugging Face convention, because adjacent-pair rotation
needs one angle per pair rather than a duplicated pair of angles. That makes
this form measurably cheaper than the half-split ``rotate_half`` formulation,
which needs full-width tables and a concatenation.

Regression coverage for the convention itself lives in
``tests/test_rotary.py::manual_adjacent_pair_rotation``.
"""

from __future__ import annotations

import math

import torch
from torch import Tensor, nn

__all__ = [
    "RotaryEmbedding",
    "apply_rotary_pos_emb",
]


_INTEGER_DTYPES = {
    torch.uint8,
    torch.int8,
    torch.int16,
    torch.int32,
    torch.int64,
}


class RotaryEmbedding(nn.Module):
    def __init__(
        self,
        head_dim: int,
        base: float = 10_000.0,
        *,
        device: torch.device | str | None = None,
    ) -> None:
        super().__init__()

        if type(head_dim) is not int or head_dim <= 0:
            raise ValueError(f"head_dim must be a positive integer, got {head_dim!r}")

        if head_dim % 2 != 0:
            raise ValueError(f"head_dim must be even, got head_dim={head_dim}")

        if (
            isinstance(base, bool)
            or not isinstance(base, (int, float))
            or not math.isfinite(float(base))
            or base <= 0.0
        ):
            raise ValueError(f"base must be a positive finite number, got {base!r}")

        self.head_dim = head_dim
        self.base = float(base)

        self.register_buffer(
            "inv_freq",
            torch.empty(
                head_dim // 2,
                dtype=torch.float32,
                device=device,
            ),
            persistent=False,
        )
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Reconstruct inverse frequencies on the buffer's current device."""

        frequency_indices = torch.arange(
            start=0,
            end=self.head_dim,
            step=2,
            dtype=torch.float32,
            device=self.inv_freq.device,
        )

        inv_freq = self.base ** (-frequency_indices / self.head_dim)

        # Assignment preserves the registered, non-persistent buffer while
        # also replacing storage allocated by Transformers' meta-device
        # loading path.
        self.inv_freq = inv_freq

    @torch.no_grad()
    def forward(
        self,
        hidden_states: Tensor,
        position_ids: Tensor | None = None,
    ) -> tuple[Tensor, Tensor]:

        if hidden_states.ndim < 2:
            raise ValueError(
                "hidden_states must have at least two dimensions, "
                f"got shape={tuple(hidden_states.shape)}"
            )

        if not hidden_states.is_floating_point():
            raise TypeError(
                "hidden_states must be a floating-point tensor, "
                f"got dtype={hidden_states.dtype}"
            )

        sequence_length = hidden_states.shape[-2]

        if position_ids is None:
            position_ids = torch.arange(
                sequence_length,
                device=hidden_states.device,
                dtype=torch.long,
            )
        else:
            if position_ids.ndim not in {1, 2}:
                raise ValueError(
                    "position_ids must have shape "
                    "(sequence_length,) or "
                    "(batch_size, sequence_length), "
                    f"got shape={tuple(position_ids.shape)}"
                )

            if position_ids.shape[-1] != sequence_length:
                raise ValueError(
                    "The final position_ids dimension must equal the "
                    f"sequence length {sequence_length}, "
                    f"got {position_ids.shape[-1]}"
                )

            if position_ids.dtype not in _INTEGER_DTYPES:
                raise TypeError(
                    "position_ids must contain integers, "
                    f"got dtype={position_ids.dtype}"
                )

            position_ids = position_ids.to(
                device=hidden_states.device,
            )

        # Compute frequencies in float32 even when the model is running in
        # float16 or bfloat16. Cast only the final cosine/sine tensors.
        inv_freq = self.inv_freq.to(
            device=hidden_states.device,
            dtype=torch.float32,
        )

        positions = position_ids.to(dtype=torch.float32)

        angles = positions.unsqueeze(-1) * inv_freq

        cos = angles.cos()
        sin = angles.sin()

        return (
            cos.to(dtype=hidden_states.dtype),
            sin.to(dtype=hidden_states.dtype),
        )

    def extra_repr(self) -> str:
        return f"head_dim={self.head_dim}, base={self.base}"


def _reshape_frequencies_for_broadcast(
    frequencies: Tensor,
    target: Tensor,
) -> Tensor:

    extra_dimensions = target.ndim - frequencies.ndim

    if extra_dimensions < 0:
        raise ValueError(
            "Rotary frequencies have too many dimensions for the target: "
            f"frequencies.ndim={frequencies.ndim}, "
            f"target.ndim={target.ndim}"
        )

    broadcast_shape = (
        *frequencies.shape[:-2],
        *((1,) * extra_dimensions),
        *frequencies.shape[-2:],
    )

    return frequencies.reshape(broadcast_shape)


def _apply_rotary(
    hidden_states: Tensor,
    cos: Tensor,
    sin: Tensor,
) -> Tensor:

    if hidden_states.shape[-1] % 2 != 0:
        raise ValueError(
            f"The final hidden dimension must be even, got {hidden_states.shape[-1]}"
        )

    even_states = hidden_states[..., 0::2]
    odd_states = hidden_states[..., 1::2]

    cos = _reshape_frequencies_for_broadcast(
        cos,
        even_states,
    )
    sin = _reshape_frequencies_for_broadcast(
        sin,
        even_states,
    )

    rotated_even = even_states * cos - odd_states * sin
    rotated_odd = even_states * sin + odd_states * cos

    return torch.stack(
        (rotated_even, rotated_odd),
        dim=-1,
    ).flatten(start_dim=-2)


def apply_rotary_pos_emb(
    query: Tensor,
    key: Tensor,
    cos: Tensor,
    sin: Tensor,
) -> tuple[Tensor, Tensor]:

    if query.ndim < 2 or key.ndim < 2:
        raise ValueError("query and key must each have at least two dimensions")

    if query.shape[-2] != key.shape[-2]:
        raise ValueError(
            "query and key sequence lengths must match, "
            f"got {query.shape[-2]} and {key.shape[-2]}"
        )

    if query.shape[-1] != key.shape[-1]:
        raise ValueError(
            "query and key head dimensions must match, "
            f"got {query.shape[-1]} and {key.shape[-1]}"
        )

    if query.shape[-1] % 2 != 0:
        raise ValueError(
            f"The query/key head dimension must be even, got {query.shape[-1]}"
        )

    if query.device != key.device:
        raise ValueError(
            "query and key must be on the same device, "
            f"got {query.device} and {key.device}"
        )

    if query.dtype != key.dtype:
        raise ValueError(
            f"query and key must have the same dtype, got {query.dtype} and {key.dtype}"
        )

    if cos.shape != sin.shape:
        raise ValueError(
            "cos and sin must have identical shapes, "
            f"got {tuple(cos.shape)} and {tuple(sin.shape)}"
        )

    expected_frequency_shape = (
        query.shape[-2],
        query.shape[-1] // 2,
    )

    if cos.shape[-2:] != expected_frequency_shape:
        raise ValueError(
            "The final cosine/sine dimensions must be "
            "(sequence_length, head_dim / 2), "
            f"expected {expected_frequency_shape}, "
            f"got {tuple(cos.shape[-2:])}"
        )

    if cos.device != query.device or sin.device != query.device:
        raise ValueError("query, key, cos, and sin must be on the same device")

    # PATCHED (see scripts/prepare_neuronai_5b_base.py): align cos/sin with
    # the query dtype instead of rejecting the pair. Under mixed precision the
    # qkv projections emit bf16 while hidden_states -- and therefore cos/sin --
    # stay fp32, which is normal and which upstream HF models handle by
    # implicit type promotion.
    if cos.dtype != query.dtype:
        cos = cos.to(dtype=query.dtype)
    if sin.dtype != query.dtype:
        sin = sin.to(dtype=query.dtype)

    return (
        _apply_rotary(query, sin=sin, cos=cos),
        _apply_rotary(key, sin=sin, cos=cos),
    )