File size: 1,954 Bytes
3e363e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """Small deterministic utilities used across GRM."""
from __future__ import annotations
import hashlib
import re
from collections.abc import Iterable
import torch
from torch import Tensor
def stable_hash(value: str, modulo: int) -> int:
if modulo <= 0:
raise ValueError("modulo must be positive")
digest = hashlib.blake2b(value.encode("utf-8"), digest_size=8).digest()
return int.from_bytes(digest, "big") % modulo
def tokenize(text: str | None) -> tuple[str, ...]:
if not text:
return ("<empty>",)
return tuple(re.findall(r"[A-Za-z0-9_]+|[^\w\s]", text.lower()))
def strings_to_ids(values: Iterable[str], size: int) -> Tensor:
return torch.tensor(
[stable_hash(value, size) for value in values], dtype=torch.long
)
def soft_cap(logits: Tensor, cap: float) -> Tensor:
"""``cap * tanh(logits / cap)`` in float32; the identity when ``cap <= 0``.
Bounds attention scores smoothly. Full-epoch runs blew up at the first
attention layer each time: query/key weights grew until the scores
saturated, and the pre-clip gradient norm went from single digits to
thousands. Near-identity for scores well below the cap, so a checkpoint
trained without it behaves almost the same with it.
"""
if cap <= 0:
return logits
return cap * torch.tanh(logits.float() / cap)
def qk_normalize(heads: Tensor) -> Tensor:
"""Scale each head vector to norm ``sqrt(head_dim)`` (QK normalisation).
Applied to queries and keys, the usual ``q.k / sqrt(d)`` score becomes
``sqrt(d) * cosine``: bounded however large the projection weights grow.
The soft cap bounded the scores but not that growth, and the first PAT
layer's query/key gradients still climbed (median 13 -> ~360 between 190K
and 245K steps).
"""
dim = heads.shape[-1]
unit = torch.nn.functional.normalize(heads.float(), dim=-1)
return (unit * dim**0.5).to(heads.dtype)
|