model/rotary_embedding.py · bobbydeveloper2014/forge-500m-chat at main

forge-500m-chat / model /rotary_embedding.py

bobbydeveloper2014

Upload folder using huggingface_hub

4d74557 verified 2 days ago

6.37 kB

"""
RoPE (Rotary Positional Embedding) with YaRN support.
GPU-optimized via cached cos/sin buffers and optional Triton kernel.
"""

from __future__ import annotations
import math
import torch
import torch.nn as nn
from typing import Optional, Tuple


class RotaryEmbedding(nn.Module):
    """
    Rotary Positional Embedding (RoPE) with:
    - Long-context theta (500,000 base)
    - YaRN scaling for context extension
    - GPU-resident cos/sin cache (no recomputation per step)
    - BF16-native computation
    """
    
    def __init__(
        self,
        head_dim: int,
        max_seq_len: int = 4096,
        base: float = 500_000.0,
        scaling_type: Optional[str] = None,
        scaling_factor: float = 1.0,
        device: Optional[torch.device] = None,
    ):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.base = base
        self.scaling_type = scaling_type
        self.scaling_factor = scaling_factor
        
        # Compute inverse frequencies
        inv_freq = self._build_inv_freq(head_dim, base, scaling_type, scaling_factor)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        
        # Pre-compute and cache cos/sin for all positions
        self._cos_cache: Optional[torch.Tensor] = None
        self._sin_cache: Optional[torch.Tensor] = None
        self._cache_seq_len: int = 0
        
        # Build initial cache
        self._build_cache(max_seq_len, device)
    
    def _build_inv_freq(
        self,
        head_dim: int,
        base: float,
        scaling_type: Optional[str],
        scaling_factor: float,
    ) -> torch.Tensor:
        """Compute inverse frequencies with optional YaRN scaling."""
        dim = head_dim
        # Standard RoPE frequencies
        freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        
        if scaling_type == "yarn":
            # YaRN: scale only the low-frequency dimensions
            # High-frequency dims (short-range) are left unchanged
            # Low-frequency dims (long-range) are scaled down by 1/factor
            mscale = 0.1 * math.log(scaling_factor) + 1.0  # YaRN mscale
            
            # Threshold: dims where wavelength > original_max_len
            low_freq_factor = 1.0
            high_freq_factor = 4.0
            original_max_len = self.max_seq_len
            new_base = base * (
                (scaling_factor * original_max_len / (2 * math.pi)) -
                (original_max_len / (2 * math.pi))
            ) ** (dim / (dim - 2))
            
            # Apply NTK-by-parts
            d = torch.arange(0, dim, 2).float()
            wavelength = 2 * math.pi / freqs
            
            # Smooth interpolation between no-scale and full-scale
            scale = torch.where(
                wavelength < original_max_len / high_freq_factor,
                torch.ones_like(freqs),
                torch.where(
                    wavelength > original_max_len / low_freq_factor,
                    torch.full_like(freqs, 1.0 / scaling_factor),
                    (original_max_len / wavelength / low_freq_factor - 1.0) / 
                    (high_freq_factor / low_freq_factor - 1.0) *
                    (1.0 - 1.0 / scaling_factor) + 1.0 / scaling_factor
                )
            )
            freqs = freqs * scale
        
        elif scaling_type == "linear":
            freqs = freqs / scaling_factor
        
        return freqs
    
    @torch.no_grad()
    def _build_cache(self, seq_len: int, device=None):
        """Build/extend the cos/sin cache up to seq_len."""
        if seq_len <= self._cache_seq_len:
            return
        
        device = device or (self.inv_freq.device if hasattr(self, 'inv_freq') 
                            else torch.device('cpu'))
        
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        freqs = torch.outer(t, self.inv_freq.to(device))
        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, head_dim)
        
        # Store in BF16 for fast loading
        self._cos_cache = emb.cos().to(torch.bfloat16)
        self._sin_cache = emb.sin().to(torch.bfloat16)
        self._cache_seq_len = seq_len
    
    def get_cos_sin(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return cached cos/sin for positions 0..seq_len-1."""
        if seq_len > self._cache_seq_len:
            self._build_cache(seq_len * 2, device)  # extend with buffer
        
        cos = self._cos_cache[:seq_len].to(device)
        sin = self._sin_cache[:seq_len].to(device)
        return cos, sin  # each (seq_len, head_dim)
    
    def apply_rotary(
        self,
        q: torch.Tensor,   # (B, n_heads, T, head_dim)
        k: torch.Tensor,   # (B, n_kv_heads, T, head_dim)
        position_ids: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Apply RoPE to Q and K."""
        T = q.shape[2]
        device = q.device
        
        cos, sin = self.get_cos_sin(T, device)  # (T, head_dim)
        
        if position_ids is not None:
            # Custom positions (for packed sequences or position bias)
            cos = cos[position_ids]   # (B, T, head_dim) or (T, head_dim)
            sin = sin[position_ids]
            if position_ids.dim() == 2:
                cos = cos.unsqueeze(1)  # (B, 1, T, head_dim)
                sin = sin.unsqueeze(1)
            else:
                cos = cos.unsqueeze(0).unsqueeze(0)  # (1, 1, T, head_dim)
                sin = sin.unsqueeze(0).unsqueeze(0)
        else:
            # Reshape for broadcasting: (1, 1, T, head_dim)
            cos = cos.unsqueeze(0).unsqueeze(0)
            sin = sin.unsqueeze(0).unsqueeze(0)
        
        q_out = _rotate_half_apply(q.float(), cos, sin).to(q.dtype)
        k_out = _rotate_half_apply(k.float(), cos, sin).to(k.dtype)
        
        return q_out, k_out
    
    def forward(self, q, k, position_ids=None):
        return self.apply_rotary(q, k, position_ids)


def _rotate_half_apply(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Apply rotary: x * cos + rotate_half(x) * sin"""
    d = x.shape[-1] // 2
    x1, x2 = x[..., :d], x[..., d:]
    rotated = torch.cat([-x2, x1], dim=-1)
    return x * cos + rotated * sin

	"""
	RoPE (Rotary Positional Embedding) with YaRN support.
	GPU-optimized via cached cos/sin buffers and optional Triton kernel.
	"""

	from __future__ import annotations
	import math
	import torch
	import torch.nn as nn
	from typing import Optional, Tuple


	class RotaryEmbedding(nn.Module):
	"""
	Rotary Positional Embedding (RoPE) with:
	- Long-context theta (500,000 base)
	- YaRN scaling for context extension
	- GPU-resident cos/sin cache (no recomputation per step)
	- BF16-native computation
	"""

	def __init__(
	self,
	head_dim: int,
	max_seq_len: int = 4096,
	base: float = 500_000.0,
	scaling_type: Optional[str] = None,
	scaling_factor: float = 1.0,
	device: Optional[torch.device] = None,
	):
	super().__init__()
	self.head_dim = head_dim
	self.max_seq_len = max_seq_len
	self.base = base
	self.scaling_type = scaling_type
	self.scaling_factor = scaling_factor

	# Compute inverse frequencies
	inv_freq = self._build_inv_freq(head_dim, base, scaling_type, scaling_factor)
	self.register_buffer("inv_freq", inv_freq, persistent=False)

	# Pre-compute and cache cos/sin for all positions
	self._cos_cache: Optional[torch.Tensor] = None
	self._sin_cache: Optional[torch.Tensor] = None
	self._cache_seq_len: int = 0

	# Build initial cache
	self._build_cache(max_seq_len, device)

	def _build_inv_freq(
	self,
	head_dim: int,
	base: float,
	scaling_type: Optional[str],
	scaling_factor: float,
	) -> torch.Tensor:
	"""Compute inverse frequencies with optional YaRN scaling."""
	dim = head_dim
	# Standard RoPE frequencies
	freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))

	if scaling_type == "yarn":
	# YaRN: scale only the low-frequency dimensions
	# High-frequency dims (short-range) are left unchanged
	# Low-frequency dims (long-range) are scaled down by 1/factor
	mscale = 0.1 * math.log(scaling_factor) + 1.0 # YaRN mscale

	# Threshold: dims where wavelength > original_max_len
	low_freq_factor = 1.0
	high_freq_factor = 4.0
	original_max_len = self.max_seq_len
	new_base = base * (
	(scaling_factor * original_max_len / (2 * math.pi)) -
	(original_max_len / (2 * math.pi))
	) ** (dim / (dim - 2))

	# Apply NTK-by-parts
	d = torch.arange(0, dim, 2).float()
	wavelength = 2 * math.pi / freqs

	# Smooth interpolation between no-scale and full-scale
	scale = torch.where(
	wavelength < original_max_len / high_freq_factor,
	torch.ones_like(freqs),
	torch.where(
	wavelength > original_max_len / low_freq_factor,
	torch.full_like(freqs, 1.0 / scaling_factor),
	(original_max_len / wavelength / low_freq_factor - 1.0) /
	(high_freq_factor / low_freq_factor - 1.0) *
	(1.0 - 1.0 / scaling_factor) + 1.0 / scaling_factor
	)
	)
	freqs = freqs * scale

	elif scaling_type == "linear":
	freqs = freqs / scaling_factor

	return freqs

	@torch.no_grad()
	def _build_cache(self, seq_len: int, device=None):
	"""Build/extend the cos/sin cache up to seq_len."""
	if seq_len <= self._cache_seq_len:
	return

	device = device or (self.inv_freq.device if hasattr(self, 'inv_freq')
	else torch.device('cpu'))

	t = torch.arange(seq_len, device=device, dtype=torch.float32)
	freqs = torch.outer(t, self.inv_freq.to(device))
	emb = torch.cat([freqs, freqs], dim=-1) # (seq_len, head_dim)

	# Store in BF16 for fast loading
	self._cos_cache = emb.cos().to(torch.bfloat16)
	self._sin_cache = emb.sin().to(torch.bfloat16)
	self._cache_seq_len = seq_len

	def get_cos_sin(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
	"""Return cached cos/sin for positions 0..seq_len-1."""
	if seq_len > self._cache_seq_len:
	self._build_cache(seq_len * 2, device) # extend with buffer

	cos = self._cos_cache[:seq_len].to(device)
	sin = self._sin_cache[:seq_len].to(device)
	return cos, sin # each (seq_len, head_dim)

	def apply_rotary(
	self,
	q: torch.Tensor, # (B, n_heads, T, head_dim)
	k: torch.Tensor, # (B, n_kv_heads, T, head_dim)
	position_ids: Optional[torch.Tensor] = None,
	) -> Tuple[torch.Tensor, torch.Tensor]:
	"""Apply RoPE to Q and K."""
	T = q.shape[2]
	device = q.device

	cos, sin = self.get_cos_sin(T, device) # (T, head_dim)

	if position_ids is not None:
	# Custom positions (for packed sequences or position bias)
	cos = cos[position_ids] # (B, T, head_dim) or (T, head_dim)
	sin = sin[position_ids]
	if position_ids.dim() == 2:
	cos = cos.unsqueeze(1) # (B, 1, T, head_dim)
	sin = sin.unsqueeze(1)
	else:
	cos = cos.unsqueeze(0).unsqueeze(0) # (1, 1, T, head_dim)
	sin = sin.unsqueeze(0).unsqueeze(0)
	else:
	# Reshape for broadcasting: (1, 1, T, head_dim)
	cos = cos.unsqueeze(0).unsqueeze(0)
	sin = sin.unsqueeze(0).unsqueeze(0)

	q_out = _rotate_half_apply(q.float(), cos, sin).to(q.dtype)
	k_out = _rotate_half_apply(k.float(), cos, sin).to(k.dtype)

	return q_out, k_out

	def forward(self, q, k, position_ids=None):
	return self.apply_rotary(q, k, position_ids)


	def _rotate_half_apply(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
	"""Apply rotary: x * cos + rotate_half(x) * sin"""
	d = x.shape[-1] // 2
	x1, x2 = x[..., :d], x[..., d:]
	rotated = torch.cat([-x2, x1], dim=-1)
	return x * cos + rotated * sin