"""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), )