AR / ar /model /rope.py

Upload AR-4B pretraining codebase, configs, and cluster scripts

2ae6691 verified 19 days ago

1.56 kB

import torch
import torch.nn as nn


def build_rope_cache(head_dim: int, max_seq_len: int, theta: float = 1_000_000.0):
    """Precompute the cos/sin rotation tables used by RoPE.
    Returns two tensors of shape (max_seq_len, head_dim), ready to index by position.
    """
    # each of the head_dim/2 pairs gets its own rotation speed (fast -> slow)
    idx = torch.arange(0, head_dim, 2, dtype=torch.float32)
    inv_freq = 1.0 / (theta ** (idx / head_dim))

    # outer product: rotation angle for every (position, pair) combination
    pos = torch.arange(max_seq_len, dtype=torch.float32)
    angles = torch.outer(pos, inv_freq)                      # (max_seq_len, head_dim/2)

    # duplicate each angle so it lines up with both halves of the vector
    angles = torch.cat([angles, angles], dim=-1)             # (max_seq_len, head_dim)
    return angles.cos(), angles.sin()


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    # turns [x1, x2] into [-x2, x1] — the 90-degree partner needed for rotation
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Rotate query or key vectors by their positions.
    x: (batch, heads, seq_len, head_dim); cos/sin: (seq_len, head_dim).
    """
    # broadcast cos/sin over batch and head dims
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    # the rotation identity: x*cos + rotate_half(x)*sin
    return x * cos + _rotate_half(x) * sin

	import torch
	import torch.nn as nn


	def build_rope_cache(head_dim: int, max_seq_len: int, theta: float = 1_000_000.0):
	"""Precompute the cos/sin rotation tables used by RoPE.

	Returns two tensors of shape (max_seq_len, head_dim), ready to index by position.
	"""
	# each of the head_dim/2 pairs gets its own rotation speed (fast -> slow)
	idx = torch.arange(0, head_dim, 2, dtype=torch.float32)
	inv_freq = 1.0 / (theta ** (idx / head_dim))

	# outer product: rotation angle for every (position, pair) combination
	pos = torch.arange(max_seq_len, dtype=torch.float32)
	angles = torch.outer(pos, inv_freq) # (max_seq_len, head_dim/2)

	# duplicate each angle so it lines up with both halves of the vector
	angles = torch.cat([angles, angles], dim=-1) # (max_seq_len, head_dim)
	return angles.cos(), angles.sin()


	def _rotate_half(x: torch.Tensor) -> torch.Tensor:
	# turns [x1, x2] into [-x2, x1] — the 90-degree partner needed for rotation
	half = x.shape[-1] // 2
	x1, x2 = x[..., :half], x[..., half:]
	return torch.cat([-x2, x1], dim=-1)


	def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
	"""Rotate query or key vectors by their positions.

	x: (batch, heads, seq_len, head_dim); cos/sin: (seq_len, head_dim).
	"""
	# broadcast cos/sin over batch and head dims
	cos = cos[None, None, :, :]
	sin = sin[None, None, :, :]
	# the rotation identity: xcos + rotate_half(x)sin
	return x * cos + _rotate_half(x) * sin