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