| |
| |
| |
| |
| |
| |
| |
| from __future__ import annotations |
| __author__ = "Thomas Seri (Sentinel)" |
| __license__ = "R08-Non-Commercial-Ethical-License-v2.0" |
| __jurisdiction__ = "Court of Milan, Italy" |
|
|
| """ |
| Cascade a 3 stadi in SERIE (non stream paralleli). |
| |
| Architettura (allineata alla timeline pezzi, grain più fine): |
| |
| pezzi_ids → emb_v1 (+pos) → CausalAttn_1 → mesh_walk_1 (=h1) |
| → LN(proj_12(h1)) + LN(emb_v2)[*] → fuse_norm2 → mesh_walk_2 (=h2) |
| → LN(proj_23(h2)) + LN(emb_v3)[*] → fuse_norm3 → mesh_walk_3 (=h3) |
| → Drop0.1(n1(h1))+n2(h2)+n3(h3) → CausalAttn_2 → final_norm → lm_head (tied weight + bias) |
| |
| [*] v8.1 ``fuse_mask_pad``: su pad/unk V2/V3 l'addendo emb è azzerato (× active); |
| ``mask_pad_stages`` già saltava solo walk/piramide. |
| |
| Le sequenze parole/unioni sono espanse sulla lunghezza pezzi (pad dove non |
| c'è parola/unione). Obiettivo LM: next-token sui pezzi (unità di decode). |
| |
| Router v6 (``router_v6=True``, default spento = v5 invariato), per stadio: |
| - ingresso: faccia = argmax(coseni + face_bias) in training e in eval; |
| face_bias (buffer, niente gradiente) segue il Loss-Free Balancing |
| (``update_face_bias`` dopo ogni step); peso straight-through |
| w = p_f / sg(p_f) con p = softmax(coseni · exp(logit_scale)), logit_scale |
| appresa (init ln 10, clamp ≤ ln 100); |
| - ogni salto: x + w · ff(norm(x)), il peso non tocca mai il residuo; |
| - porte: come v5 (Gumbel-ST in training, argmax in eval); |
| - uscita: decisione dura a soglia nel forward; il peso del salto successivo |
| è w_porta · p_stay / sg(p_stay), p_stay = 1 − probs[:, uscita], così il |
| logit di uscita riceve gradiente dal contributo dell'esperto che segue. |
| ``hash_routing=True``: come v6 ma faccia d'ingresso fissa per id del token |
| (pezzo / parola / unione nei tre stadi, ``hash_faces``), senza to_sphere, |
| logit_scale né face_bias. |
| ``no_walk=True`` (solo v6/hash, ``max_hops=1``): controllo del cammino, ogni |
| token usa solo il blocco della faccia d'ingresso in ogni stadio; dopo quel |
| blocco escono tutti e le porte non vengono valutate (niente rumore Gumbel). |
| ``walk_to_center=True`` (v10 Spoke-to-Hub, solo hash, ``H_max∈{2,3}``): |
| ingresso Zipf/hash invariato (periferia = ID); dopo ogni blocco il token |
| segue ``parent_face[f]`` verso un **anello polare hub** di ``k_core≈15–20`` |
| facce (non una sola faccia). Gate (Argon): |
| hop0 residuo pieno (come v8.1, ``g=None``); |
| hop≥1 ``g = (c + (1−c)·σ(Linear(x))) / √H_max`` con ``c=center_mix`` |
| (default 0.25, formula liscia). Bias gate init 0. Niente Gumbel; |
| pad/unk = 0 hop (mask-pad). Train==eval path. |
| La piramide (se attiva) resta sulla faccia d'ingresso Zipf. |
| ``sector_mix=True`` (v11, implica ``walk_to_center``): stesso albero parent→hub, |
| ma hop≥1 non è più spoke sottile solo-parent: ``x ← x + g · mean_{f∈settore(entry)}(Δ_f)`` |
| dove il settore è lo **spicchio azimutale** di ``K≈sector_frac·F`` facce |
| (default 1/4) intorno all'ingresso Zipf. Hop0 resta solo faccia entry a peso 1.0. |
| Niente layer attention, niente Gumbel, Zipf classico (no multi-slot). |
| ``sheet_mix=True`` (v13 Sheet-Mix / v14 pezzi-only, implica ``sector_mix`` path): stesso hop0/hop≥1 |
| e ``_apply_sector_mean``, ma il buffer ``sector_faces`` è ``(F, K_eff=12)`` = |
| 3 fogli complementari (cerchi massimi a 0°/90°/±45°, esclusa l'orientazione |
| ``spiral_rank % 4``) × ``m=4`` facce equidistanti per foglio (ponte polo-a-polo). |
| Iso-param rispetto a v11 (solo indici; weight sharing invariato). |
| **Storico (v13–v18, default):** hop≥1 = media Δ sui blocchi delle K_eff facce dei fogli |
| (cross-face). **``sheet_entry_only=True`` (v19):** hop≥1 applica solo il blocco della |
| faccia d'ingresso Zipf (niente media su altre facce); i fogli non toccano altri esperti. |
| Il cammino parent→hub resta per lo stop; piramide invariata sull'entry. |
| ``pyramid_center=True`` (solo hash, «sfera a spicchi piramidali»): dopo il |
| cammino di ogni stadio il token scende nella piramide della sua faccia |
| d'ingresso, un passaggio d → d/2 con proiezione condivisa e guadagno/spostamento |
| per faccia, e arriva al centro, un MLP denso condiviso dallo stadio; l'uscita |
| del centro si somma al ramo della faccia (``PyramidCenter``). |
| """ |
|
|
| import math |
| from collections import deque |
| from pathlib import Path |
| from typing import List, Optional, Tuple, Union |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from fast_walk import ( |
| _build_paths, |
| fast_walk_token, |
| fast_walk_token_v6, |
| grouped_block_forward, |
| stack_bank, |
| ) |
| from geodesic_network import ( |
| CausalAttnBlock, |
| PairedFaceBank, |
| WalkConfig, |
| _default_n_heads, |
| ) |
| from mesh_topology import GeodesicMesh, build_geodesic_mesh, icosahedron_vertices |
|
|
| DEFAULT_K_CORE = 18 |
| V10_HOPS_DEFAULT = 2 |
| DEFAULT_CENTER_MIX = 0.25 |
| DEFAULT_SECTOR_FRAC = 0.25 |
| DEFAULT_EDGE_NB_TAU = 4.0 |
| DEFAULT_SHEET_M = 4 |
|
|
|
|
| def build_face_edge_ids(faces: np.ndarray) -> Tuple[np.ndarray, int]: |
| """``face_edges[f, i]`` = id unico del lato i; ``E = 3F/2`` su mesh chiusa.""" |
| fcount = int(faces.shape[0]) |
| edge_to_id: dict = {} |
| face_edges = np.zeros((fcount, 3), dtype=np.int64) |
| next_id = 0 |
| for fi, (a, b, c) in enumerate(faces): |
| for li, (u, v) in enumerate(((int(a), int(b)), (int(b), int(c)), (int(c), int(a)))): |
| key = (u, v) if u < v else (v, u) |
| eid = edge_to_id.get(key) |
| if eid is None: |
| eid = next_id |
| edge_to_id[key] = eid |
| next_id += 1 |
| face_edges[fi, li] = eid |
| return face_edges, next_id |
|
|
|
|
| def build_edge_neighbor_weights( |
| adjacency: np.ndarray, |
| face_centroids: np.ndarray, |
| tau: float = DEFAULT_EDGE_NB_TAU, |
| ) -> np.ndarray: |
| """Pesi (F, 3) sui vicini di lato: softmax(τ · z_polo[nb]), verso il Polo Nord.""" |
| z = face_centroids[:, 2].astype(np.float64) |
| scores = z[adjacency] * float(tau) |
| scores = scores - scores.max(axis=1, keepdims=True) |
| e = np.exp(scores) |
| return (e / e.sum(axis=1, keepdims=True)).astype(np.float32) |
|
|
|
|
| def build_spoke_hub_parent( |
| adjacency: np.ndarray, |
| face_centroids: np.ndarray, |
| k_core: int = DEFAULT_K_CORE, |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: |
| """Albero Spoke→Hub: ``parent[f]`` verso l'anello polare di ``k_core`` facce. |
| |
| Hub = top-``k_core`` facce per proiezione sull'asse polare icosaedrico |
| (vertice 0, stesso di ``compass_spiral``). BFS multi-sorgente: ogni faccia |
| non-hub punta al vicino più vicino all'anello; sulle hub ``parent[h]=h``. |
| |
| Ritorna ``(parent_face (F,), is_hub (F,) bool, hub_ids (K,))``. |
| """ |
| n = int(adjacency.shape[0]) |
| k = int(max(1, min(int(k_core), n))) |
| pole = icosahedron_vertices()[0] |
| u3 = pole / np.linalg.norm(pole) |
| z = face_centroids.astype(np.float64) @ u3 |
| hub_ids = np.argsort(-z)[:k].astype(np.int64) |
| is_hub = np.zeros(n, dtype=bool) |
| is_hub[hub_ids] = True |
| parent = np.arange(n, dtype=np.int64) |
| dist = np.full(n, -1, dtype=np.int32) |
| q: deque = deque() |
| for h in hub_ids.tolist(): |
| dist[h] = 0 |
| parent[h] = h |
| q.append(h) |
| while q: |
| f = q.popleft() |
| for nb in adjacency[f].tolist(): |
| nb = int(nb) |
| if dist[nb] < 0: |
| dist[nb] = int(dist[f]) + 1 |
| parent[nb] = f |
| q.append(nb) |
| if int((dist < 0).sum()) > 0: |
| raise RuntimeError("spoke-hub: grafo facce non connesso") |
| return parent, is_hub, hub_ids |
|
|
|
|
| def build_azimuth_sectors( |
| face_centroids: np.ndarray, |
| sector_frac: float = DEFAULT_SECTOR_FRAC, |
| ) -> np.ndarray: |
| """Settore azimutale per faccia: spicchio verso il polo (asse icosaedrico 0). |
| |
| Per ogni faccia ``e``, le ``K = max(1, round(sector_frac · F))`` facce con |
| minore distanza circolare in azimut intorno al polo (stesso vertice 0 di |
| ``build_spoke_hub_parent`` / ``compass_spiral``). Con ``sector_frac=0.25`` e |
| F=180 → K≈45 (~1/4 di sfera in longitudine, tutte le latitudini). |
| |
| Ritorna ``sector_faces (F, K)`` int64; ``sector_faces[e, 0] == e``. |
| """ |
| if not (0.0 < float(sector_frac) <= 1.0): |
| raise ValueError("sector_frac deve essere in (0, 1]") |
| cents = face_centroids.astype(np.float64) |
| n = int(cents.shape[0]) |
| k = int(max(1, min(n, round(float(sector_frac) * n)))) |
| pole = icosahedron_vertices()[0] |
| u3 = pole / np.linalg.norm(pole) |
| |
| ref = np.array([1.0, 0.0, 0.0], dtype=np.float64) |
| if abs(float(ref @ u3)) > 0.9: |
| ref = np.array([0.0, 1.0, 0.0], dtype=np.float64) |
| e1 = ref - (ref @ u3) * u3 |
| e1 = e1 / np.linalg.norm(e1) |
| e2 = np.cross(u3, e1) |
| xy = cents - np.outer(cents @ u3, u3) |
| x = xy @ e1 |
| y = xy @ e2 |
| az = np.arctan2(y, x) |
| |
| dphi = np.abs(az[:, None] - az[None, :]) |
| dphi = np.minimum(dphi, 2.0 * np.pi - dphi) |
| |
| np.fill_diagonal(dphi, -1.0) |
| order = np.argsort(dphi, axis=1, kind="mergesort")[:, :k].astype(np.int64) |
| return order |
|
|
|
|
| def build_complementary_sheets( |
| face_centroids: np.ndarray, |
| spiral_rank: Optional[np.ndarray] = None, |
| m_per_sheet: int = DEFAULT_SHEET_M, |
| ) -> np.ndarray: |
| """v13 Sheet-Mix: 3 fogli complementari × m facce su cerchi massimi → (F, 3m). |
| |
| Per ogni faccia ``f`` costruisce la base tangente ``(t_par, t_mer)`` rispetto |
| al polo icosaedrico (vertice 0). Le 4 normali ai cerchi massimi a |
| 0° / 90° / +45° / −45° sono:: |
| |
| N_V = t_par (meridiano, polo↔polo) |
| N_H = t_mer (ortogonale in f) |
| N_+ = (t_par+t_mer)/√2 |
| N_- = (t_par-t_mer)/√2 |
| |
| L'orientazione propria è ``spiral_rank[f] % 4`` (se assente → 0); si usano |
| i **3 fogli complementari**. Su ciascun foglio si campionano ``m`` facce |
| equidistanti sul cerchio massimo (θ = 2πk/m), incluso il ponte |
| fronte/retro e polo-a-polo. Drop-in su ``sector_faces`` / ``_apply_sector_mean``. |
| """ |
| m = int(m_per_sheet) |
| if m < 1: |
| raise ValueError("m_per_sheet deve essere >= 1") |
| cents = face_centroids.astype(np.float64) |
| n = int(cents.shape[0]) |
| |
| norms = np.linalg.norm(cents, axis=1, keepdims=True) |
| cents = cents / np.maximum(norms, 1e-12) |
| pole = icosahedron_vertices()[0] |
| u3 = pole / np.linalg.norm(pole) |
| |
| ref = np.array([1.0, 0.0, 0.0], dtype=np.float64) |
| if abs(float(ref @ u3)) > 0.9: |
| ref = np.array([0.0, 1.0, 0.0], dtype=np.float64) |
| e1 = ref - (ref @ u3) * u3 |
| e1 = e1 / np.linalg.norm(e1) |
|
|
| if spiral_rank is None: |
| orient = np.zeros(n, dtype=np.int64) |
| else: |
| sr = np.asarray(spiral_rank, dtype=np.int64).reshape(-1) |
| if sr.shape[0] != n: |
| raise ValueError(f"spiral_rank len={sr.shape[0]} ≠ F={n}") |
| orient = np.mod(sr, 4).astype(np.int64) |
|
|
| out = np.empty((n, 3 * m), dtype=np.int64) |
| |
| |
| for f in range(n): |
| c = cents[f] |
| |
| t_par = np.cross(u3, c) |
| tp_n = float(np.linalg.norm(t_par)) |
| if tp_n < 1e-8: |
| t_par = e1.copy() |
| else: |
| t_par = t_par / tp_n |
| t_mer = np.cross(c, t_par) |
| tm_n = float(np.linalg.norm(t_mer)) |
| if tm_n < 1e-8: |
| t_mer = np.cross(t_par, u3) |
| t_mer = t_mer / max(float(np.linalg.norm(t_mer)), 1e-12) |
| else: |
| t_mer = t_mer / tm_n |
| normals = ( |
| t_par, |
| t_mer, |
| (t_par + t_mer) / math.sqrt(2.0), |
| (t_par - t_mer) / math.sqrt(2.0), |
| ) |
| own = int(orient[f]) |
| comps = [s for s in range(4) if s != own] |
| col = 0 |
| for s in comps: |
| N = normals[s] |
| |
| ea = c - (c @ N) * N |
| ea_n = float(np.linalg.norm(ea)) |
| if ea_n < 1e-8: |
| |
| ea = t_par if s != 0 else t_mer |
| ea = ea - (ea @ N) * N |
| ea = ea / max(float(np.linalg.norm(ea)), 1e-12) |
| else: |
| ea = ea / ea_n |
| eb = np.cross(N, ea) |
| eb = eb / max(float(np.linalg.norm(eb)), 1e-12) |
| for k in range(m): |
| theta = 2.0 * math.pi * float(k) / float(m) |
| target = math.cos(theta) * ea + math.sin(theta) * eb |
| |
| scores = cents @ target |
| out[f, col] = int(np.argmax(scores)) |
| col += 1 |
| |
| out[f, 0] = f |
| return out |
|
|
|
|
| V6_LOGIT_SCALE_INIT = math.log(10.0) |
| V6_LOGIT_SCALE_MAX = math.log(100.0) |
| V6_BIAS_RATE = 1e-3 |
| _HASH_SALTS = (0x9E3779B9, 0x85EBCA6B, 0xC2B2AE35) |
| PYRAMID_SEED_OFFSET = 7919 |
|
|
|
|
| def _lowbias32(x: int) -> int: |
| x &= 0xFFFFFFFF |
| x ^= x >> 16 |
| x = (x * 0x7FEB352D) & 0xFFFFFFFF |
| x ^= x >> 15 |
| x = (x * 0x846CA68B) & 0xFFFFFFFF |
| x ^= x >> 16 |
| return x |
|
|
|
|
| def hash_faces(vocab_size: int, n_faces: int, stage: int) -> torch.Tensor: |
| """Faccia d'ingresso fissa per id (--hash-routing): lowbias32(id XOR sale dello stadio) mod n_faces.""" |
| salt = _HASH_SALTS[stage - 1] |
| return torch.tensor([_lowbias32(i ^ salt) % n_faces for i in range(vocab_size)], dtype=torch.long) |
|
|
|
|
| class PyramidCenter(nn.Module): |
| """Piramidi sotto le facce di uno stadio e centro condiviso (--pyramid-center). |
| |
| Discesa dalla faccia f: z = (1 + face_gain[f]) ⊙ down(norm(x)) + face_shift[f], d → d_narrow; |
| la proiezione propria della faccia è diag(1 + face_gain[f]) · down. Centro, uguale per tutti |
| i token: c = up(GELU(inner(z))), d_narrow → d_hidden → d. Uscita x + c; up parte da zero, quindi |
| all'inizio l'uscita è x. |
| """ |
|
|
| def __init__( |
| self, d_model: int, n_faces: int, d_narrow: int, d_hidden: int, generator: Optional[torch.Generator] = None |
| ): |
| super().__init__() |
| self.norm = nn.LayerNorm(d_model) |
| self.down = nn.Linear(d_model, d_narrow, bias=False) |
| self.face_gain = nn.Parameter(torch.zeros(n_faces, d_narrow)) |
| self.face_shift = nn.Parameter(torch.zeros(n_faces, d_narrow)) |
| self.center = nn.Sequential(nn.Linear(d_narrow, d_hidden), nn.GELU(), nn.Linear(d_hidden, d_model)) |
| self.reset_parameters(generator) |
|
|
| def reset_parameters(self, generator: Optional[torch.Generator] = None) -> None: |
| """Init di nn.Linear per down e per il primo strato del centro; ultima proiezione, guadagni e spostamenti a zero.""" |
| for lin in (self.down, self.center[0]): |
| nn.init.kaiming_uniform_(lin.weight, a=math.sqrt(5), generator=generator) |
| if lin.bias is not None: |
| bound = 1.0 / math.sqrt(lin.in_features) |
| nn.init.uniform_(lin.bias, -bound, bound, generator=generator) |
| for p in (self.center[2].weight, self.center[2].bias, self.face_gain, self.face_shift): |
| nn.init.zeros_(p) |
| self.norm.reset_parameters() |
|
|
| def forward(self, x: torch.Tensor, face: torch.Tensor) -> torch.Tensor: |
| z = self.down(self.norm(x)) |
| z = z * (1.0 + F.embedding(face, self.face_gain)) + F.embedding(face, self.face_shift) |
| return x + self.center(z) |
|
|
|
|
| class HexPyramidCenter(nn.Module): |
| """v9: guadagno/spostamento per **vertice** della mesh; media sui 3 vertici della faccia. |
| |
| Stessa discesa d → d_narrow e centro MLP condiviso di ``PyramidCenter``, ma i parametri |
| locali vivono sui V vertici (esagoni/pentagoni duali) invece che sulle F facce: facce |
| adiacenti condividono i vertici e quindi parte della piramide. |
| """ |
|
|
| def __init__( |
| self, |
| d_model: int, |
| n_vertices: int, |
| face_vertices: torch.Tensor, |
| d_narrow: int, |
| d_hidden: int, |
| generator: Optional[torch.Generator] = None, |
| ): |
| super().__init__() |
| if face_vertices.ndim != 2 or face_vertices.shape[1] != 3: |
| raise ValueError(f"face_vertices deve essere (F, 3), got {tuple(face_vertices.shape)}") |
| self.norm = nn.LayerNorm(d_model) |
| self.down = nn.Linear(d_model, d_narrow, bias=False) |
| self.vertex_gain = nn.Parameter(torch.zeros(n_vertices, d_narrow)) |
| self.vertex_shift = nn.Parameter(torch.zeros(n_vertices, d_narrow)) |
| self.center = nn.Sequential(nn.Linear(d_narrow, d_hidden), nn.GELU(), nn.Linear(d_hidden, d_model)) |
| self.register_buffer("face_vertices", face_vertices.long().contiguous()) |
| self.reset_parameters(generator) |
|
|
| def reset_parameters(self, generator: Optional[torch.Generator] = None) -> None: |
| for lin in (self.down, self.center[0]): |
| nn.init.kaiming_uniform_(lin.weight, a=math.sqrt(5), generator=generator) |
| if lin.bias is not None: |
| bound = 1.0 / math.sqrt(lin.in_features) |
| nn.init.uniform_(lin.bias, -bound, bound, generator=generator) |
| for p in (self.center[2].weight, self.center[2].bias, self.vertex_gain, self.vertex_shift): |
| nn.init.zeros_(p) |
| self.norm.reset_parameters() |
|
|
| def forward(self, x: torch.Tensor, face: torch.Tensor) -> torch.Tensor: |
| z = self.down(self.norm(x)) |
| verts = self.face_vertices[face] |
| g = F.embedding(verts, self.vertex_gain).mean(dim=-2) |
| s = F.embedding(verts, self.vertex_shift).mean(dim=-2) |
| z = z * (1.0 + g) + s |
| return x + self.center(z) |
|
|
|
|
| class SharedEdgeVertexPyramid(nn.Module): |
| """v12: face_gain specifico per faccia + lati (E) e vertici (V) condivisi + 1 centro MLP. |
| |
| ``z`` modulato da ``1 + face + 0.5·mean(edge) + 0.5·mean(vertex)`` (edge/vert init 0). |
| """ |
|
|
| def __init__( |
| self, |
| d_model: int, |
| n_faces: int, |
| n_edges: int, |
| n_vertices: int, |
| face_edges: torch.Tensor, |
| face_vertices: torch.Tensor, |
| d_narrow: int, |
| d_hidden: int, |
| generator: Optional[torch.Generator] = None, |
| ): |
| super().__init__() |
| if face_edges.ndim != 2 or face_edges.shape[1] != 3: |
| raise ValueError(f"face_edges deve essere (F, 3), got {tuple(face_edges.shape)}") |
| if face_vertices.ndim != 2 or face_vertices.shape[1] != 3: |
| raise ValueError(f"face_vertices deve essere (F, 3), got {tuple(face_vertices.shape)}") |
| self.norm = nn.LayerNorm(d_model) |
| self.down = nn.Linear(d_model, d_narrow, bias=False) |
| self.face_gain = nn.Parameter(torch.zeros(n_faces, d_narrow)) |
| self.face_shift = nn.Parameter(torch.zeros(n_faces, d_narrow)) |
| self.edge_gain = nn.Parameter(torch.zeros(n_edges, d_narrow)) |
| self.edge_shift = nn.Parameter(torch.zeros(n_edges, d_narrow)) |
| self.vertex_gain = nn.Parameter(torch.zeros(n_vertices, d_narrow)) |
| self.vertex_shift = nn.Parameter(torch.zeros(n_vertices, d_narrow)) |
| self.center = nn.Sequential(nn.Linear(d_narrow, d_hidden), nn.GELU(), nn.Linear(d_hidden, d_model)) |
| self.register_buffer("face_edges", face_edges.long().contiguous()) |
| self.register_buffer("face_vertices", face_vertices.long().contiguous()) |
| self.reset_parameters(generator) |
|
|
| def reset_parameters(self, generator: Optional[torch.Generator] = None) -> None: |
| for lin in (self.down, self.center[0]): |
| nn.init.kaiming_uniform_(lin.weight, a=math.sqrt(5), generator=generator) |
| if lin.bias is not None: |
| bound = 1.0 / math.sqrt(lin.in_features) |
| nn.init.uniform_(lin.bias, -bound, bound, generator=generator) |
| for p in ( |
| self.center[2].weight, |
| self.center[2].bias, |
| self.face_gain, |
| self.face_shift, |
| self.edge_gain, |
| self.edge_shift, |
| self.vertex_gain, |
| self.vertex_shift, |
| ): |
| nn.init.zeros_(p) |
| self.norm.reset_parameters() |
|
|
| def forward(self, x: torch.Tensor, face: torch.Tensor) -> torch.Tensor: |
| z = self.down(self.norm(x)) |
| f = face.long() |
| eg = F.embedding(self.face_edges[f], self.edge_gain).mean(dim=-2) |
| es = F.embedding(self.face_edges[f], self.edge_shift).mean(dim=-2) |
| vg = F.embedding(self.face_vertices[f], self.vertex_gain).mean(dim=-2) |
| vs = F.embedding(self.face_vertices[f], self.vertex_shift).mean(dim=-2) |
| g = F.embedding(f, self.face_gain) + 0.5 * eg + 0.5 * vg |
| s = F.embedding(f, self.face_shift) + 0.5 * es + 0.5 * vs |
| z = z * (1.0 + g) + s |
| return x + self.center(z) |
|
|
|
|
| def pyramid_kwargs(blob: dict) -> dict: |
| """Argomenti di CascadeMeshModel per la piramide, dai metadati di un checkpoint (assenti = spenta).""" |
| if not blob.get("pyramid_center") and not blob.get("hex_pyramid") and not blob.get("shared_edges_verts"): |
| return {} |
| cfg = blob.get("pyramid_config") or {} |
| out: dict = { |
| "pyramid_center": True, |
| "pyramid_narrow": cfg.get("d_stretto"), |
| "pyramid_hidden": cfg.get("d_centro"), |
| } |
| if blob.get("hex_pyramid"): |
| out["hex_pyramid"] = True |
| if blob.get("shared_edges_verts"): |
| out["shared_edges_verts"] = True |
| return out |
|
|
|
|
| def v8_kwargs(blob: dict) -> dict: |
| """Argomenti v8.x di CascadeMeshModel dai metadati checkpoint (assenti = default v7). |
| |
| Lo zipf-hash è già nei buffer ``bank.hash_face`` dello state_dict: non serve |
| ``zipf_hash_path`` in eval. |
| |
| ``fuse_mask_pad`` (v8.1): solo se presente nel blob — i ckpt v8.0 senza chiave |
| restano senza mask sull'addendo emb in fuse (come in training). |
| """ |
| out: dict = {} |
| if blob.get("d_hidden") is not None: |
| out["d_hidden"] = int(blob["d_hidden"]) |
| de = blob.get("d_emb_v23") |
| if de is not None and int(de) > 0: |
| out["d_emb_v23"] = int(de) |
| if blob.get("inter_stage_attn"): |
| out["inter_stage_attn"] = True |
| if blob.get("mask_pad_stages"): |
| out["mask_pad_stages"] = True |
| out["pad_id_v2"] = int(blob.get("pad_id_v2", 3)) |
| out["unk_id_v2"] = int(blob.get("unk_id_v2", 4)) |
| out["pad_id_v3"] = int(blob.get("pad_id_v3", 3)) |
| out["unk_id_v3"] = int(blob.get("unk_id_v3", 4)) |
| if blob.get("fuse_mask_pad"): |
| out["fuse_mask_pad"] = True |
| if blob.get("multi_slot_v1"): |
| out["multi_slot_v1"] = True |
| return out |
|
|
|
|
| def _make_emb(vocab: int, d_emb: int, d_model: int) -> nn.Module: |
| """Embedding pieno a ``d_model``, oppure ``Embedding(d_emb)→Linear→d_model``.""" |
| if d_emb >= d_model: |
| emb = nn.Embedding(vocab, d_model) |
| nn.init.normal_(emb.weight, mean=0.0, std=0.02) |
| return emb |
| emb = nn.Embedding(vocab, d_emb) |
| proj = nn.Linear(d_emb, d_model, bias=False) |
| nn.init.normal_(emb.weight, mean=0.0, std=0.02) |
| nn.init.normal_(proj.weight, mean=0.0, std=0.02) |
| return nn.Sequential(emb, proj) |
|
|
|
|
| class CascadeMeshModel(nn.Module): |
| """Tre emb + tre mesh in catena (default), oppure solo-V1 (v1_only=True).""" |
|
|
| def __init__( |
| self, |
| vocab_pezzi: int, |
| vocab_parole: int = 0, |
| vocab_unioni: int = 0, |
| d_model: int = 256, |
| frequency: int = 2, |
| walk: Optional[WalkConfig] = None, |
| d_hidden: Optional[int] = None, |
| n_attn_layers: int = 2, |
| n_heads: Optional[int] = None, |
| max_seq_len: int = 512, |
| attn_dropout: float = 0.0, |
| res_h1_dropout: float = 0.1, |
| v1_only: bool = False, |
| compass_spiral: bool = False, |
| fast_walk: bool = False, |
| router_v6: bool = False, |
| hash_routing: bool = False, |
| no_walk: bool = False, |
| walk_to_center: bool = False, |
| k_core: int = DEFAULT_K_CORE, |
| center_mix: float = DEFAULT_CENTER_MIX, |
| sector_mix: bool = False, |
| sector_frac: float = DEFAULT_SECTOR_FRAC, |
| sheet_mix: bool = False, |
| sheet_m: int = DEFAULT_SHEET_M, |
| sheet_entry_only: bool = False, |
| edge_neighbor_mix: bool = False, |
| edge_nb_tau: float = DEFAULT_EDGE_NB_TAU, |
| pyramid_center: bool = False, |
| pyramid_narrow: Optional[int] = None, |
| pyramid_hidden: Optional[int] = None, |
| hex_pyramid: bool = False, |
| shared_edges_verts: bool = False, |
| multi_slot_v1: bool = False, |
| d_emb_v23: Optional[int] = None, |
| inter_stage_attn: bool = False, |
| mask_pad_stages: bool = False, |
| fuse_mask_pad: bool = False, |
| zipf_hash_path: Optional[Union[str, Path]] = None, |
| pad_id_v2: int = 3, |
| unk_id_v2: int = 4, |
| pad_id_v3: int = 3, |
| unk_id_v3: int = 4, |
| ): |
| super().__init__() |
| self.compass_spiral = bool(compass_spiral) |
| |
| self.fast_walk = bool(fast_walk) |
| self.mesh: GeodesicMesh = build_geodesic_mesh( |
| frequency, compass_spiral=self.compass_spiral |
| ) |
| self.walk = walk or WalkConfig() |
| self.d_model = int(d_model) |
| self.vocab_pezzi = int(vocab_pezzi) |
| self.vocab_parole = int(vocab_parole) |
| self.vocab_unioni = int(vocab_unioni) |
| self.v1_only = bool(v1_only) |
| self.n_attn_layers = int(n_attn_layers) |
| self.max_seq_len = int(max_seq_len) |
| self.d_hidden = int(d_hidden) if d_hidden is not None else None |
| self.d_emb_v23 = int(d_emb_v23) if d_emb_v23 is not None else None |
| self.inter_stage_attn = bool(inter_stage_attn) |
| self.mask_pad_stages = bool(mask_pad_stages) |
| |
| |
| self.fuse_mask_pad = bool(fuse_mask_pad) and bool(mask_pad_stages) |
| self.pad_id_v2 = int(pad_id_v2) |
| self.unk_id_v2 = int(unk_id_v2) |
| self.pad_id_v3 = int(pad_id_v3) |
| self.unk_id_v3 = int(unk_id_v3) |
| heads = int(n_heads) if n_heads is not None else _default_n_heads(d_model) |
| |
| attn_ff = 4 * d_model |
|
|
| self.emb_v1 = nn.Embedding(vocab_pezzi, d_model) |
| nn.init.normal_(self.emb_v1.weight, mean=0.0, std=0.02) |
| if not self.v1_only: |
| if vocab_parole <= 0 or vocab_unioni <= 0: |
| raise ValueError( |
| "Cascade 3-stadi richiede vocab_parole > 0 e vocab_unioni > 0 " |
| "(oppure v1_only=True)" |
| ) |
| de = self.d_emb_v23 if self.d_emb_v23 is not None else d_model |
| self.emb_v2 = _make_emb(vocab_parole, de, d_model) |
| self.emb_v3 = _make_emb(vocab_unioni, de, d_model) |
|
|
| self.pos_emb = nn.Embedding(self.max_seq_len, d_model) |
| nn.init.normal_(self.pos_emb.weight, mean=0.0, std=0.02) |
| self.emb_scale = d_model**-0.5 |
|
|
| self.attn_blocks = nn.ModuleList( |
| [ |
| CausalAttnBlock(d_model, heads, d_ff=attn_ff, dropout=attn_dropout) |
| for _ in range(self.n_attn_layers) |
| ] |
| ) |
| self.attn_inter2: Optional[CausalAttnBlock] = None |
| if self.inter_stage_attn and not self.v1_only: |
| self.attn_inter2 = CausalAttnBlock(d_model, heads, d_ff=attn_ff, dropout=attn_dropout) |
|
|
| |
| self.to_sphere1 = nn.Linear(d_model, 3) |
| self.bank1 = PairedFaceBank(self.mesh, d_model, d_hidden) |
| if not self.v1_only: |
| self.to_sphere2 = nn.Linear(d_model, 3) |
| self.to_sphere3 = nn.Linear(d_model, 3) |
| self.bank2 = PairedFaceBank(self.mesh, d_model, d_hidden) |
| self.bank3 = PairedFaceBank(self.mesh, d_model, d_hidden) |
|
|
| |
| self.proj_12 = nn.Linear(d_model, d_model) |
| self.proj_23 = nn.Linear(d_model, d_model) |
| self.fuse_ln_h2 = nn.LayerNorm(d_model) |
| self.fuse_ln_e2 = nn.LayerNorm(d_model) |
| self.fuse_ln_h3 = nn.LayerNorm(d_model) |
| self.fuse_ln_e3 = nn.LayerNorm(d_model) |
| self.fuse_norm2 = nn.LayerNorm(d_model) |
| self.fuse_norm3 = nn.LayerNorm(d_model) |
|
|
| |
| |
| self.res_h1_dropout = 0.0 if self.v1_only else float(res_h1_dropout) |
| self.res_ln_h1 = nn.LayerNorm(d_model) |
| self.res_drop_h1 = nn.Dropout(self.res_h1_dropout) |
| if not self.v1_only: |
| self.res_ln_h2 = nn.LayerNorm(d_model) |
| self.res_ln_h3 = nn.LayerNorm(d_model) |
| self.final_norm = nn.LayerNorm(d_model) |
| self.lm_head = nn.Linear(d_model, vocab_pezzi, bias=True) |
| |
|
|
| self.register_buffer( |
| "centroids", |
| torch.from_numpy(self.mesh.face_centroids.astype("float32")), |
| ) |
| self.register_buffer( |
| "adjacency", |
| torch.from_numpy(self.mesh.adjacency.astype("int64")), |
| ) |
|
|
| |
| self.hash_routing = bool(hash_routing) |
| self.router_v6 = bool(router_v6) or self.hash_routing |
| self.no_walk = bool(no_walk) |
| self.sheet_mix = bool(sheet_mix) |
| self.sheet_m = int(sheet_m) |
| |
| self.sheet_entry_only = bool(sheet_entry_only) and self.sheet_mix |
| |
| self.sector_mix = bool(sector_mix) or self.sheet_mix |
| self.sector_frac = float(sector_frac) |
| self.edge_neighbor_mix = bool(edge_neighbor_mix) |
| self.edge_nb_tau = float(edge_nb_tau) |
| |
| self.walk_to_center = ( |
| bool(walk_to_center) or self.sector_mix or self.edge_neighbor_mix |
| ) |
| self.k_core = int(k_core) |
| self.center_mix = float(center_mix) |
| if self.sheet_mix and self.edge_neighbor_mix: |
| raise ValueError("sheet_mix (v13) e edge_neighbor_mix (v12) sono mutuamente esclusivi") |
| if self.sector_mix and self.edge_neighbor_mix: |
| raise ValueError("sector_mix (v11) e edge_neighbor_mix (v12) sono mutuamente esclusivi") |
| if self.sheet_mix and self.sheet_m < 1: |
| raise ValueError("sheet_m deve essere >= 1") |
| if bool(sheet_entry_only) and not self.sheet_mix: |
| raise ValueError("sheet_entry_only richiede sheet_mix") |
| if self.no_walk and self.walk_to_center: |
| raise ValueError("no_walk e walk_to_center (Spoke-to-Hub) sono mutuamente esclusivi") |
| if self.no_walk and (not self.router_v6 or int(self.walk.max_hops) != 1): |
| raise ValueError("no_walk: servono router_v6 o hash_routing e WalkConfig(max_hops=1)") |
| self.spoke_gates: Optional[nn.ModuleList] = None |
| if self.walk_to_center: |
| if not self.hash_routing: |
| raise ValueError("walk_to_center/Spoke-to-Hub: serve hash_routing (ingresso Zipf)") |
| hops = int(self.walk.max_hops) |
| if hops < 2 or hops > 3: |
| raise ValueError("walk_to_center/Spoke-to-Hub: H_max deve essere 2 o 3") |
| if not (0.0 <= self.center_mix <= 1.0): |
| raise ValueError("center_mix deve essere in [0, 1]") |
| if self.sector_mix and not self.sheet_mix and not (0.0 < self.sector_frac <= 1.0): |
| raise ValueError("sector_frac deve essere in (0, 1]") |
| parent, is_hub, hub_ids = build_spoke_hub_parent( |
| self.mesh.adjacency, self.mesh.face_centroids, k_core=self.k_core |
| ) |
| self.register_buffer("parent_face", torch.from_numpy(parent.astype("int64"))) |
| self.register_buffer("is_hub", torch.from_numpy(is_hub.astype("bool"))) |
| self.register_buffer("hub_ids", torch.from_numpy(hub_ids.astype("int64"))) |
| if self.sheet_mix and self.sheet_entry_only: |
| |
| self.sector_faces = None |
| self.sector_k = 1 |
| elif self.sheet_mix: |
| sheets = build_complementary_sheets( |
| self.mesh.face_centroids, |
| spiral_rank=getattr(self.mesh, "spiral_rank", None), |
| m_per_sheet=self.sheet_m, |
| ) |
| self.register_buffer( |
| "sector_faces", torch.from_numpy(sheets.astype("int64")) |
| ) |
| self.sector_k = int(sheets.shape[1]) |
| elif self.sector_mix: |
| sectors = build_azimuth_sectors( |
| self.mesh.face_centroids, sector_frac=self.sector_frac |
| ) |
| self.register_buffer( |
| "sector_faces", torch.from_numpy(sectors.astype("int64")) |
| ) |
| self.sector_k = int(sectors.shape[1]) |
| else: |
| self.sector_faces = None |
| self.sector_k = 0 |
| if self.edge_neighbor_mix: |
| w = build_edge_neighbor_weights( |
| self.mesh.adjacency, self.mesh.face_centroids, tau=self.edge_nb_tau |
| ) |
| self.register_buffer("edge_nb_w", torch.from_numpy(w)) |
| else: |
| self.edge_nb_w = None |
| self.spoke_res_scale = 1.0 / math.sqrt(float(hops)) |
| |
| n_banks = 1 if self.v1_only else 3 |
| gates = [] |
| for _ in range(n_banks): |
| lin = nn.Linear(self.d_model, 1) |
| nn.init.zeros_(lin.weight) |
| nn.init.zeros_(lin.bias) |
| gates.append(lin) |
| self.spoke_gates = nn.ModuleList(gates) |
| else: |
| self.parent_face = None |
| self.is_hub = None |
| self.hub_ids = None |
| self.sector_faces = None |
| self.sector_k = 0 |
| self.sector_mix = False |
| self.sheet_mix = False |
| self.sheet_entry_only = False |
| self.edge_neighbor_mix = False |
| self.edge_nb_w = None |
| self.spoke_res_scale = 1.0 |
| self.center_mix = 0.0 |
| self.v6_logit_scale_max = V6_LOGIT_SCALE_MAX |
| self.v6_bias_rate = V6_BIAS_RATE |
| if self.router_v6: |
| n_faces = int(self.mesh.num_faces) |
| sizes = (self.vocab_pezzi, self.vocab_parole, self.vocab_unioni) |
| for s, bank in enumerate(self._banks(), start=1): |
| if self.hash_routing: |
| bank.register_buffer("hash_face", hash_faces(sizes[s - 1], n_faces, s)) |
| delattr(self, f"to_sphere{s}") |
| setattr(self, f"to_sphere{s}", None) |
| else: |
| bank.logit_scale = nn.Parameter(torch.tensor(V6_LOGIT_SCALE_INIT)) |
| bank.register_buffer("face_bias", torch.zeros(n_faces)) |
|
|
| self.multi_slot_v1 = bool(multi_slot_v1) |
| if zipf_hash_path is not None: |
| if not self.hash_routing: |
| raise ValueError("zipf_hash_path richiede hash_routing") |
| self.load_zipf_hash(zipf_hash_path) |
| elif self.multi_slot_v1: |
| |
| v = int(self.bank1.hash_face.numel()) |
| self.bank1.register_buffer( |
| "hash_slots", |
| self.bank1.hash_face.unsqueeze(1).expand(v, 3).contiguous().clone(), |
| ) |
|
|
| |
| self.hex_pyramid = bool(hex_pyramid) |
| self.shared_edges_verts = bool(shared_edges_verts) |
| if self.hex_pyramid and self.shared_edges_verts: |
| raise ValueError("hex_pyramid e shared_edges_verts (v12) sono mutuamente esclusivi") |
| self.pyramid_center = bool(pyramid_center) or self.hex_pyramid or self.shared_edges_verts |
| self.pyramids: Optional[nn.ModuleList] = None |
| if self.pyramid_center: |
| if not self.hash_routing: |
| raise ValueError("pyramid_center/hex_pyramid: serve hash_routing (la piramide è quella della faccia d'ingresso)") |
| self.pyramid_narrow = int(pyramid_narrow or self.d_model // 2) |
| self.pyramid_hidden = int(pyramid_hidden or 2 * self.d_model) |
| |
| gen = torch.Generator().manual_seed((torch.initial_seed() + PYRAMID_SEED_OFFSET) % 2**63) |
| rng = torch.get_rng_state() |
| if self.hex_pyramid: |
| face_verts = torch.from_numpy(self.mesh.faces.astype("int64")) |
| n_vert = int(self.mesh.num_vertices) |
| self.pyramids = nn.ModuleList( |
| HexPyramidCenter( |
| self.d_model, n_vert, face_verts, self.pyramid_narrow, self.pyramid_hidden, gen |
| ) |
| for _ in self._banks() |
| ) |
| elif self.shared_edges_verts: |
| face_edges_np, n_edges = build_face_edge_ids(self.mesh.faces) |
| face_edges = torch.from_numpy(face_edges_np.astype("int64")) |
| face_verts = torch.from_numpy(self.mesh.faces.astype("int64")) |
| n_vert = int(self.mesh.num_vertices) |
| self.pyramids = nn.ModuleList( |
| SharedEdgeVertexPyramid( |
| self.d_model, |
| int(self.mesh.num_faces), |
| int(n_edges), |
| n_vert, |
| face_edges, |
| face_verts, |
| self.pyramid_narrow, |
| self.pyramid_hidden, |
| gen, |
| ) |
| for _ in self._banks() |
| ) |
| else: |
| self.pyramids = nn.ModuleList( |
| PyramidCenter(self.d_model, int(self.mesh.num_faces), self.pyramid_narrow, self.pyramid_hidden, gen) |
| for _ in self._banks() |
| ) |
| torch.set_rng_state(rng) |
|
|
| from watchman_keystone_r08 import NonLinearWatchmanKeystoneR08, OFFICIAL_SENTINEL_LICENSE_SUMMARY |
| self.watchman = NonLinearWatchmanKeystoneR08(d_model=self.d_model) |
| self._watchman_is_ethical = True |
| self._watchman_license_text = OFFICIAL_SENTINEL_LICENSE_SUMMARY |
|
|
| def _banks(self) -> List[PairedFaceBank]: |
| return [self.bank1] if self.v1_only else [self.bank1, self.bank2, self.bank3] |
|
|
| def load_zipf_hash(self, path: Union[str, Path]) -> None: |
| """Copia ``hash_face_{1,2,3}`` (e opz. ``hash_slots_1`` multi-slot V1) da un ``.pt`` Zipf.""" |
| path = Path(path) |
| try: |
| blob = torch.load(path, map_location="cpu", weights_only=False) |
| except TypeError: |
| blob = torch.load(path, map_location="cpu") |
| n_faces = int(self.mesh.num_faces) |
| for s, bank in enumerate(self._banks(), start=1): |
| key = f"hash_face_{s}" |
| if key not in blob: |
| raise KeyError(f"{path}: manca {key}") |
| t = blob[key].long().view(-1) |
| if t.numel() != bank.hash_face.numel(): |
| raise ValueError( |
| f"{path} {key}: len={t.numel()} ≠ bank.hash_face={bank.hash_face.numel()}" |
| ) |
| if int(t.max()) >= n_faces or int(t.min()) < 0: |
| raise ValueError(f"{path} {key}: facce fuori da [0, {n_faces})") |
| bank.hash_face.copy_(t) |
| |
| if "hash_slots_1" in blob: |
| slots = blob["hash_slots_1"].long() |
| if slots.ndim != 2 or slots.shape[0] != self.bank1.hash_face.numel(): |
| raise ValueError( |
| f"{path} hash_slots_1: shape={tuple(slots.shape)} " |
| f"atteso (V={self.bank1.hash_face.numel()}, K)" |
| ) |
| if int(slots.max()) >= n_faces or int(slots.min()) < 0: |
| raise ValueError(f"{path} hash_slots_1: facce fuori da [0, {n_faces})") |
| self.bank1.register_buffer("hash_slots", slots.contiguous()) |
| self.multi_slot_v1 = True |
| elif self.multi_slot_v1: |
| raise ValueError(f"{path}: multi_slot_v1 richiesto ma manca hash_slots_1") |
|
|
| def _active_mask(self, stage: int, ids: torch.Tensor) -> Optional[torch.Tensor]: |
| """Maschera token attivi per stadi 2/3 con ``mask_pad_stages`` (pad/unk → False).""" |
| if not self.mask_pad_stages or stage < 1: |
| return None |
| if stage == 1: |
| pad, unk = self.pad_id_v2, self.unk_id_v2 |
| elif stage == 2: |
| pad, unk = self.pad_id_v3, self.unk_id_v3 |
| else: |
| return None |
| return (ids != pad) & (ids != unk) |
|
|
| def _v6_entry( |
| self, |
| h: torch.Tensor, |
| bank: PairedFaceBank, |
| to_sphere: Optional[nn.Linear], |
| entry_ids: Optional[torch.Tensor], |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
| """Faccia d'ingresso v6 (B,) e peso straight-through w (B, 1), None senza gradiente. |
| |
| w = p_f / sg(p_f) è scritto come exp(log p_f − sg(log p_f)): il bias può |
| scegliere una faccia con p_f sotto il minimo di float32 e il rapporto |
| diventerebbe 0/0. In training il carico d'ingresso per faccia si accumula |
| sul device per ``update_face_bias``. |
| """ |
| if self.hash_routing: |
| if entry_ids is None: |
| raise ValueError("hash_routing: servono gli id dei token per la faccia d'ingresso") |
| |
| slots = getattr(bank, "hash_slots", None) |
| if slots is not None: |
| row = slots[entry_ids] |
| if self.training: |
| k = int(row.shape[-1]) |
| idx = torch.randint(0, k, entry_ids.shape, device=entry_ids.device) |
| face = row.gather(-1, idx.unsqueeze(-1)).squeeze(-1) |
| else: |
| face = row[..., 0] |
| return face, None |
| return bank.hash_face[entry_ids], None |
| scores = F.normalize(to_sphere(h), dim=-1) @ self.centroids.T |
| face = (scores + bank.face_bias).argmax(dim=-1) |
| if self.training: |
| n_faces = scores.shape[1] |
| load = (face.unsqueeze(1) == torch.arange(n_faces, device=face.device)).sum(0) |
| prev = bank.__dict__.get("_v6_load") |
| bank.__dict__["_v6_load"] = load if prev is None else prev + load |
| if not torch.is_grad_enabled(): |
| return face, None |
| scale = bank.logit_scale.clamp(max=self.v6_logit_scale_max).exp() |
| lp = F.log_softmax(scores * scale, dim=-1).gather(1, face.unsqueeze(1)) |
| return face, torch.exp(lp - lp.detach()) |
|
|
| def _v6_route( |
| self, face: torch.Tensor, logits: torch.Tensor |
| ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: |
| """Porte come ``_route_step``, uscita dura a soglia; ritorna anche il peso del salto successivo. |
| |
| w = w_porta · p_stay / sg(p_stay) (B, 1), None senza gradiente: w_porta è il |
| valore Gumbel-ST della porta scelta (solo in training), p_stay = |
| 1 − probs[:, 3] = exp(LSE(z[:3]) − LSE(z)). Sui token che proseguono il |
| logit di uscita riceve d log p_stay / d z_uscita = −p_exit volte il |
| gradiente del peso dell'esperto del salto successivo. |
| Con ``no_walk`` escono tutti: porte e uscita non vengono valutate. |
| Con ``walk_to_center`` (Spoke-to-Hub) il routing è in ``_walk_token_spoke_hub``: |
| qui si usa solo ``parent_face`` (niente Gumbel). |
| """ |
| if self.no_walk: |
| return face, torch.ones_like(face, dtype=torch.bool), None |
| if self.walk_to_center: |
| nxt = self.parent_face[face.long()] |
| exit_mask = nxt == face |
| return nxt, exit_mask, None |
| temp = max(self.walk.router_temperature, 1e-6) |
| z = logits / temp |
| probs = F.softmax(z, dim=-1) |
| exit_mask = probs[:, 3] >= self.walk.exit_threshold |
| neighbors = self.adjacency[face.long()] |
| w_side: Optional[torch.Tensor] = None |
| if self._hard_routing(): |
| side = probs[:, :3].argmax(dim=-1) |
| else: |
| side_oh = F.gumbel_softmax(logits[:, :3] / temp, hard=True, dim=-1) |
| side = side_oh.argmax(dim=-1) |
| w_side = side_oh.gather(1, side.unsqueeze(1)) |
| next_face = neighbors.gather(1, side.unsqueeze(1)).squeeze(1) |
| if not torch.is_grad_enabled(): |
| return next_face, exit_mask, None |
| log_stay = torch.logsumexp(z[:, :3], dim=-1, keepdim=True) - torch.logsumexp(z, dim=-1, keepdim=True) |
| w = torch.exp(log_stay - log_stay.detach()) |
| return next_face, exit_mask, (w if w_side is None else w_side * w) |
|
|
| def _walk_token_v6_loop( |
| self, |
| h: torch.Tensor, |
| bank: PairedFaceBank, |
| to_sphere: Optional[nn.Linear], |
| entry_ids: Optional[torch.Tensor] = None, |
| record_path: bool = False, |
| active: Optional[torch.Tensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[List[List[int]]]]: |
| """Cammino v6 col loop per faccia: riferimento di ``fast_walk_token_v6`` nei test.""" |
| b = h.shape[0] |
| face, w = self._v6_entry(h, bank, to_sphere, entry_ids) |
| if active is None: |
| active = torch.ones(b, dtype=torch.bool, device=h.device) |
| else: |
| active = active.to(device=h.device, dtype=torch.bool) |
| face0 = torch.where(active, face, torch.full_like(face, -1)) |
| paths = [[int(f)] for f in face0.tolist()] if record_path else None |
| for _ in range(self.walk.max_hops): |
| if not bool(active.any()): |
| break |
| logits = h.new_zeros(b, 4) |
| h_new = h.clone() |
| for f_id in face[active].unique().tolist(): |
| mask = active & (face == f_id) |
| blk = bank.blocks[int(bank.face_to_slot[f_id].item())] |
| x = h[mask] |
| delta = blk.ff(blk.norm(x)) |
| out = x + delta if w is None else x + w[mask] * delta |
| h_new[mask] = out |
| logits[mask] = blk.router(out) |
| h = h_new |
| next_face, exit_mask, w = self._v6_route(face, logits) |
| moved = active & ~exit_mask |
| face = torch.where(moved, next_face, face) |
| if paths is not None: |
| for i in range(b): |
| if moved[i]: |
| paths[i].append(int(face[i].item())) |
| elif active[i] and exit_mask[i]: |
| paths[i].append(-1) |
| active = active & ~exit_mask |
| return h, paths |
|
|
| @torch.no_grad() |
| def update_face_bias(self) -> None: |
| """Loss-Free Balancing: b_f += u · sign(carico medio − carico_f) col carico d'ingresso accumulato.""" |
| if not self.router_v6 or self.hash_routing: |
| return |
| for bank in self._banks(): |
| load = bank.__dict__.pop("_v6_load", None) |
| if load is None: |
| continue |
| load = load.to(bank.face_bias.dtype) |
| bank.face_bias.add_(torch.sign(load.mean() - load), alpha=self.v6_bias_rate) |
|
|
| def _hard_routing(self) -> bool: |
| """Hard (argmax) in eval o se WalkConfig.hard_routing; soft/ST in train.""" |
| return bool(self.walk.hard_routing) or (not self.training) |
|
|
| def project_to_face( |
| self, h: torch.Tensor, to_sphere: nn.Linear |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
| """ |
| Proietta h sulla sfera → faccia. |
| Train (soft): Gumbel-ST one-hot (B, F) così i grad arrivano a to_sphere. |
| Eval/hard: solo face id (argmax), face_oh=None. |
| """ |
| direction = F.normalize(to_sphere(h), dim=-1) |
| scores = direction @ self.centroids.T |
| if self._hard_routing(): |
| return scores.argmax(dim=-1), None |
| temp = max(self.walk.router_temperature, 1e-6) |
| face_oh = F.gumbel_softmax(scores / temp, hard=True, dim=-1) |
| face = face_oh.argmax(dim=-1) |
| return face, face_oh |
|
|
| def _route_step( |
| self, face: torch.Tensor, logits: torch.Tensor |
| ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: |
| """ |
| Router ternario + exit. |
| Train (soft): Gumbel-ST su 3 lati (niente argmax sulla one-hot) + mixture |
| su adjacency → side_oh per propagare i grad al router. |
| Eval/hard: argmax discreto, side_oh=None. |
| """ |
| temp = max(self.walk.router_temperature, 1e-6) |
| probs = F.softmax(logits / temp, dim=-1) |
| exit_prob = probs[:, 3] |
| neighbors = self.adjacency[face.long()] |
| exit_mask = exit_prob >= self.walk.exit_threshold |
|
|
| if self._hard_routing(): |
| side = probs[:, :3].argmax(dim=-1) |
| next_face = neighbors.gather(1, side.unsqueeze(1)).squeeze(1) |
| return next_face, exit_mask, None |
|
|
| |
| side_oh = F.gumbel_softmax(logits[:, :3] / temp, hard=True, dim=-1) |
| side = side_oh.argmax(dim=-1) |
| next_face = neighbors.gather(1, side.unsqueeze(1)).squeeze(1) |
| return next_face, exit_mask, side_oh |
|
|
| def _scatter_side_oh( |
| self, face: torch.Tensor, side_oh: torch.Tensor |
| ) -> torch.Tensor: |
| """Soft/ST mixture su adjacency: massa dei 3 lati → one-hot sparsa su F facce.""" |
| b = face.shape[0] |
| n_faces = int(self.centroids.shape[0]) |
| neighbors = self.adjacency[face.long()] |
| face_oh = side_oh.new_zeros(b, n_faces) |
| face_oh.scatter_add_(1, neighbors, side_oh) |
| return face_oh |
|
|
| def _spoke_stage_index(self, bank: PairedFaceBank) -> int: |
| for i, b in enumerate(self._banks()): |
| if b is bank: |
| return i |
| raise KeyError("bank non appartiene al modello") |
|
|
| def _apply_sector_mean( |
| self, |
| h: torch.Tensor, |
| bank: PairedFaceBank, |
| entry: torch.Tensor, |
| active: torch.Tensor, |
| g: Optional[torch.Tensor], |
| st: Optional[dict] = None, |
| ) -> torch.Tensor: |
| """v11: ``x ← x + g · mean_{f∈settore(entry)}(Δ_f)`` (hop≥1). |
| |
| Settore fissato sulla faccia d'ingresso Zipf. ``g=None`` → peso 1.0. |
| """ |
| assert self.sector_faces is not None |
| k = int(self.sector_k) |
| if k <= 0: |
| return h |
| use_fast = st is not None |
| delta_sum: Optional[torch.Tensor] = None |
| for slot in range(k): |
| face_k = self.sector_faces[entry.long(), slot] |
| if use_fast: |
| assert st is not None |
| out, _lg, counts = grouped_block_forward( |
| st, h, face_k, active, gate=None |
| ) |
| if st["visited"] is not None: |
| st["visited"].add(counts) |
| |
| d = out - h |
| else: |
| d = h.new_zeros(h.shape) |
| for f_id in face_k[active].unique().tolist(): |
| mask = active & (face_k == f_id) |
| blk = bank.blocks[int(bank.face_to_slot[f_id].item())] |
| x = h[mask] |
| d[mask] = blk.ff(blk.norm(x)) |
| delta_sum = d if delta_sum is None else delta_sum + d |
| assert delta_sum is not None |
| mean_d = delta_sum / float(k) |
| if g is None: |
| upd = h + mean_d |
| else: |
| upd = h + g * mean_d |
| return torch.where(active.unsqueeze(-1), upd, h) |
|
|
| def _apply_edge_neighbor_mix( |
| self, |
| h: torch.Tensor, |
| bank: PairedFaceBank, |
| entry: torch.Tensor, |
| active: torch.Tensor, |
| g: Optional[torch.Tensor], |
| st: Optional[dict] = None, |
| ) -> torch.Tensor: |
| """v12: ``x ← x + g · Σ_k w_{e,k} Δ_{nb_k}`` sui 3 lati di ``entry`` (pesi verso polo).""" |
| assert self.edge_nb_w is not None |
| e = entry.long() |
| nbs = self.adjacency[e] |
| w = self.edge_nb_w[e] |
| use_fast = st is not None |
| weighted: Optional[torch.Tensor] = None |
| for slot in range(3): |
| face_k = nbs[:, slot] |
| wk = w[:, slot].unsqueeze(-1) |
| if use_fast: |
| assert st is not None |
| out, _lg, counts = grouped_block_forward( |
| st, h, face_k, active, gate=None |
| ) |
| if st["visited"] is not None: |
| st["visited"].add(counts) |
| d = out - h |
| else: |
| d = h.new_zeros(h.shape) |
| for f_id in face_k[active].unique().tolist(): |
| mask = active & (face_k == f_id) |
| blk = bank.blocks[int(bank.face_to_slot[f_id].item())] |
| x = h[mask] |
| d[mask] = blk.ff(blk.norm(x)) |
| term = wk * d |
| weighted = term if weighted is None else weighted + term |
| assert weighted is not None |
| if g is None: |
| upd = h + weighted |
| else: |
| upd = h + g * weighted |
| return torch.where(active.unsqueeze(-1), upd, h) |
|
|
| def _walk_token_spoke_hub( |
| self, |
| h: torch.Tensor, |
| bank: PairedFaceBank, |
| to_sphere: Optional[nn.Linear], |
| entry_ids: Optional[torch.Tensor] = None, |
| record_path: bool = False, |
| active: Optional[torch.Tensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[List[List[int]]]]: |
| """v10 Spoke-to-Hub / v11 sector-mix / v12 edge-nb / v13–v19 sheet: parent→hub; no Gumbel. |
| |
| hop0: solo faccia entry, residuo pieno. |
| hop≥1: v10 = blocco ``face`` corrente con ``g=(c+(1-c)σ)/√H``; |
| v11 (``sector_mix``) = media Δ sul settore azimutale dell'ingresso; |
| v12 (``edge_neighbor_mix``) = mix pesato sui 3 vicini di lato (anche al polo); |
| v13–v18 (``sheet_mix``, default) = media Δ sui fogli complementari dell'ingresso; |
| v19 (``sheet_entry_only``) = di nuovo solo blocco entry (niente cross-face). |
| pad/unk (``active=False``): 0 hop. Train==eval. |
| """ |
| assert self.spoke_gates is not None and self.parent_face is not None |
| b = h.shape[0] |
| face, _ = self._v6_entry(h, bank, to_sphere, entry_ids) |
| if active is None: |
| active = torch.ones(b, dtype=torch.bool, device=h.device) |
| else: |
| active = active.to(device=h.device, dtype=torch.bool) |
| entry = face.clone() |
| face0 = torch.where(active, face, torch.full_like(face, -1)) |
| max_hops = int(self.walk.max_hops) |
| if b == 0 or max_hops <= 0: |
| return h, (_build_paths(face0, []) if record_path else None) |
|
|
| gate_lin = self.spoke_gates[self._spoke_stage_index(bank)] |
| scale = float(self.spoke_res_scale) |
| codes: List[torch.Tensor] = [] |
| use_fast = bool(self.fast_walk) |
| st = stack_bank(bank) if use_fast else None |
| |
| use_sector = ( |
| bool(self.sector_mix) |
| and self.sector_faces is not None |
| and not bool(self.sheet_entry_only) |
| ) |
| use_sheet_entry = bool(self.sheet_mix) and bool(self.sheet_entry_only) |
| use_edge_nb = bool(self.edge_neighbor_mix) and self.edge_nb_w is not None |
|
|
| c_mix = float(self.center_mix) |
| for hop in range(max_hops): |
| if not bool(active.any()): |
| break |
| |
| if hop == 0: |
| g: Optional[torch.Tensor] = None |
| else: |
| g = (c_mix + (1.0 - c_mix) * torch.sigmoid(gate_lin(h))) * scale |
| if use_edge_nb and hop >= 1: |
| h = self._apply_edge_neighbor_mix(h, bank, entry, active, g, st=st) |
| elif use_sector and hop >= 1: |
| h = self._apply_sector_mean(h, bank, entry, active, g, st=st) |
| elif use_sheet_entry and hop >= 1: |
| |
| apply_face = entry |
| if use_fast: |
| assert st is not None |
| out, _lg, counts = grouped_block_forward( |
| st, h, apply_face, active, gate=g |
| ) |
| if st["visited"] is not None: |
| st["visited"].add(counts) |
| h = torch.where(active.unsqueeze(-1), out, h) |
| else: |
| h_new = h.clone() |
| for f_id in apply_face[active].unique().tolist(): |
| mask = active & (apply_face == f_id) |
| blk = bank.blocks[int(bank.face_to_slot[f_id].item())] |
| x = h[mask] |
| delta = blk.ff(blk.norm(x)) |
| if g is None: |
| h_new[mask] = x + delta |
| else: |
| h_new[mask] = x + g[mask] * delta |
| h = h_new |
| elif use_fast: |
| assert st is not None |
| out, _lg, counts = grouped_block_forward( |
| st, h, face, active, gate=g |
| ) |
| if st["visited"] is not None: |
| st["visited"].add(counts) |
| h = torch.where(active.unsqueeze(-1), out, h) |
| else: |
| h_new = h.clone() |
| for f_id in face[active].unique().tolist(): |
| mask = active & (face == f_id) |
| blk = bank.blocks[int(bank.face_to_slot[f_id].item())] |
| x = h[mask] |
| delta = blk.ff(blk.norm(x)) |
| if g is None: |
| h_new[mask] = x + delta |
| else: |
| h_new[mask] = x + g[mask] * delta |
| h = h_new |
| |
| if use_edge_nb: |
| stop = torch.full( |
| (b,), hop == max_hops - 1, device=h.device, dtype=torch.bool |
| ) |
| moved = torch.zeros_like(active) |
| |
| else: |
| nxt = self.parent_face[face.long()] |
| at_hub = nxt == face |
| stop = at_hub | (hop == max_hops - 1) |
| moved = active & ~stop |
| face = torch.where(moved, nxt, face) |
| just_exit = active & stop |
| active = active & ~stop |
| if record_path: |
| codes.append( |
| torch.where( |
| moved, |
| face, |
| torch.where(just_exit, torch.full_like(face, -1), torch.full_like(face, -2)), |
| ) |
| ) |
|
|
| return h, (_build_paths(face0, codes) if record_path else None) |
|
|
| def walk_token( |
| self, |
| h: torch.Tensor, |
| bank: PairedFaceBank, |
| to_sphere: Optional[nn.Linear], |
| record_path: bool = False, |
| entry_ids: Optional[torch.Tensor] = None, |
| active: Optional[torch.Tensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[List[List[int]]]]: |
| if self.walk_to_center: |
| return self._walk_token_spoke_hub( |
| h, bank, to_sphere, entry_ids=entry_ids, record_path=record_path, active=active |
| ) |
| if self.router_v6: |
| if self.fast_walk: |
| return fast_walk_token_v6( |
| self, h, bank, to_sphere, entry_ids=entry_ids, record_path=record_path, active=active |
| ) |
| return self._walk_token_v6_loop( |
| h, bank, to_sphere, entry_ids=entry_ids, record_path=record_path, active=active |
| ) |
| if self.fast_walk: |
| return fast_walk_token(self, h, bank, to_sphere, record_path=record_path) |
| b = h.shape[0] |
| face, face_oh = self.project_to_face(h, to_sphere) |
| paths: Optional[List[List[int]]] = ( |
| [[int(face[i].item())] for i in range(b)] if record_path else None |
| ) |
| active = torch.ones(b, dtype=torch.bool, device=h.device) |
|
|
| for _ in range(self.walk.max_hops): |
| if not active.any(): |
| break |
| logits = h.new_zeros(b, 4) |
| h_new = h.clone() |
| for f_id in face[active].unique().tolist(): |
| mask = active & (face == f_id) |
| if not mask.any(): |
| continue |
| out, lg = bank.forward_face(int(f_id), h[mask]) |
| if face_oh is not None: |
| |
| w = face_oh[mask, int(f_id)].unsqueeze(-1) |
| h_new[mask] = out * w |
| else: |
| h_new[mask] = out |
| logits[mask] = lg |
| h = h_new |
| next_face, exit_mask, side_oh = self._route_step(face, logits) |
| moved = active & ~exit_mask |
| if side_oh is not None: |
| |
| scale = side_oh.sum(dim=-1, keepdim=True) |
| h = torch.where(active.unsqueeze(-1), h * scale, h) |
| routed_oh = self._scatter_side_oh(face, side_oh) |
| if face_oh is None: |
| face_oh = routed_oh |
| else: |
| face_oh = torch.where(moved.unsqueeze(-1), routed_oh, face_oh) |
| face = torch.where(moved, next_face, face) |
| just_exit = active & exit_mask |
| active = active & ~exit_mask |
| if record_path and paths is not None: |
| for i in range(b): |
| if moved[i]: |
| paths[i].append(int(face[i].item())) |
| elif just_exit[i]: |
| paths[i].append(-1) |
| return h, paths |
|
|
| def _descend( |
| self, |
| stage: int, |
| h: torch.Tensor, |
| ids: torch.Tensor, |
| active_bt: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| """Piramide della faccia d'ingresso e centro dello stadio (da 0) su h (B, T, d); identità senza piramidi.""" |
| if self.pyramids is None: |
| return h |
| bank = self._banks()[stage] |
| face = bank.hash_face[ids] |
| |
| slots = getattr(bank, "hash_slots", None) |
| if slots is not None: |
| row = slots[ids] |
| if self.training: |
| k = int(row.shape[-1]) |
| idx = torch.randint(0, k, ids.shape, device=ids.device) |
| face = row.gather(-1, idx.unsqueeze(-1)).squeeze(-1) |
| else: |
| face = row[..., 0] |
| out = self.pyramids[stage](h, face) |
| if active_bt is None: |
| return out |
| return torch.where(active_bt.unsqueeze(-1), out, h) |
|
|
| def _walk_seq( |
| self, |
| h: torch.Tensor, |
| bank: PairedFaceBank, |
| to_sphere: Optional[nn.Linear], |
| record_path: bool, |
| entry_ids: Optional[torch.Tensor] = None, |
| active_bt: Optional[torch.Tensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[List[List[List[int]]]]]: |
| b, t, d = h.shape |
| kw: dict = {} |
| if entry_ids is not None: |
| kw["entry_ids"] = entry_ids.reshape(b * t) |
| if active_bt is not None: |
| kw["active"] = active_bt.reshape(b * t) |
| flat, paths_flat = self.walk_token( |
| h.reshape(b * t, d), bank, to_sphere, record_path=record_path, **kw |
| ) |
| h_out = flat.view(b, t, d) |
| paths = None |
| if record_path and paths_flat is not None: |
| paths = [ |
| [paths_flat[bi * t + ti] for ti in range(t)] for bi in range(b) |
| ] |
| return h_out, paths |
|
|
| def forward( |
| self, |
| pezzi_ids: torch.Tensor, |
| parole_ids: Optional[torch.Tensor] = None, |
| unioni_ids: Optional[torch.Tensor] = None, |
| record_path: bool = False, |
| ) -> Tuple[torch.Tensor, Optional[List[List[List[int]]]]]: |
| """ |
| pezzi/parole/unioni: (B, T) allineati sulla timeline pezzi. |
| Se self.v1_only: richiede solo pezzi_ids (parole_ids/unioni_ids ignorati/None). |
| returns: logits (B, T, V_pezzi), paths opzionale dell'ultimo mesh walk. |
| """ |
| b, t = pezzi_ids.shape |
| if t > self.max_seq_len: |
| raise ValueError(f"seq_len={t} > max_seq_len={self.max_seq_len}") |
|
|
| pos = torch.arange(t, device=pezzi_ids.device).unsqueeze(0) |
| h = self.emb_v1(pezzi_ids) * self.emb_scale + self.pos_emb(pos) |
|
|
| |
| for block in self.attn_blocks[:-1]: |
| h = block(h) |
|
|
| hashed = self.hash_routing |
| if self.v1_only: |
| h1, paths = self._walk_seq( |
| h, self.bank1, self.to_sphere1, record_path=record_path, |
| entry_ids=pezzi_ids if hashed else None, |
| ) |
| h1 = self._descend(0, h1, pezzi_ids) |
| h = self.res_ln_h1(h1) |
| if len(self.attn_blocks) > 0: |
| h = self.attn_blocks[-1](h) |
| h = self.final_norm(h) |
| logits = self.lm_head(h) |
| logits = 30.0 * torch.tanh(logits / 30.0) |
| return logits, paths |
|
|
| if parole_ids is None or unioni_ids is None: |
| raise ValueError("Cascade 3-stadi richiede parole_ids e unioni_ids") |
|
|
| h1, _ = self._walk_seq( |
| h, self.bank1, self.to_sphere1, record_path=False, entry_ids=pezzi_ids if hashed else None |
| ) |
| h1 = self._descend(0, h1, pezzi_ids) |
|
|
| |
| |
| active2 = self._active_mask(1, parole_ids) |
| emb2 = self.fuse_ln_e2(self.emb_v2(parole_ids)) |
| if self.fuse_mask_pad and active2 is not None: |
| emb2 = emb2 * active2.unsqueeze(-1).to(dtype=emb2.dtype) |
| h = self.fuse_norm2(self.fuse_ln_h2(self.proj_12(h1)) + emb2) |
| if self.attn_inter2 is not None: |
| h = self.attn_inter2(h) |
| h2, _ = self._walk_seq( |
| h, |
| self.bank2, |
| self.to_sphere2, |
| record_path=False, |
| entry_ids=parole_ids if hashed else None, |
| active_bt=active2, |
| ) |
| h2 = self._descend(1, h2, parole_ids, active_bt=active2) |
|
|
| |
| active3 = self._active_mask(2, unioni_ids) |
| emb3 = self.fuse_ln_e3(self.emb_v3(unioni_ids)) |
| if self.fuse_mask_pad and active3 is not None: |
| emb3 = emb3 * active3.unsqueeze(-1).to(dtype=emb3.dtype) |
| h = self.fuse_norm3(self.fuse_ln_h3(self.proj_23(h2)) + emb3) |
| h3, paths = self._walk_seq( |
| h, |
| self.bank3, |
| self.to_sphere3, |
| record_path=record_path, |
| entry_ids=unioni_ids if hashed else None, |
| active_bt=active3, |
| ) |
| h3 = self._descend(2, h3, unioni_ids, active_bt=active3) |
|
|
| |
| h = ( |
| self.res_drop_h1(self.res_ln_h1(h1)) |
| + self.res_ln_h2(h2) |
| + self.res_ln_h3(h3) |
| ) |
| if len(self.attn_blocks) > 0: |
| h = self.attn_blocks[-1](h) |
| h = self.final_norm(h) |
| if hasattr(self, "watchman") and self.watchman is not None: |
| h_entangled = self.watchman.inject_geodesic_nonlinear_entanglement(h) |
| h = self.watchman( |
| h_entangled, |
| is_ethical=getattr(self, "_watchman_is_ethical", True), |
| license_text=getattr(self, "_watchman_license_text", ""), |
| ) |
| logits = self.lm_head(h) |
| logits = 30.0 * torch.tanh(logits / 30.0) |
| return logits, paths |
|
|
|
|
| class V1OnlyMeshModel(CascadeMeshModel): |
| """Solo V1 (pezzi): emb_v1 (+pos) → CausalAttn_0 → bank1 (80 facce) → LN → CausalAttn_1 → final_norm → lm_head.""" |
|
|
| def __init__( |
| self, |
| vocab_pezzi: int, |
| d_model: int = 256, |
| frequency: int = 2, |
| walk: Optional[WalkConfig] = None, |
| d_hidden: Optional[int] = None, |
| n_attn_layers: int = 2, |
| n_heads: Optional[int] = None, |
| max_seq_len: int = 512, |
| attn_dropout: float = 0.0, |
| compass_spiral: bool = False, |
| fast_walk: bool = False, |
| ): |
| super().__init__( |
| vocab_pezzi=vocab_pezzi, |
| vocab_parole=0, |
| vocab_unioni=0, |
| d_model=d_model, |
| frequency=frequency, |
| walk=walk, |
| d_hidden=d_hidden, |
| n_attn_layers=n_attn_layers, |
| n_heads=n_heads, |
| max_seq_len=max_seq_len, |
| attn_dropout=attn_dropout, |
| res_h1_dropout=0.0, |
| v1_only=True, |
| compass_spiral=compass_spiral, |
| fast_walk=fast_walk, |
| ) |
|
|
|
|
| def count_parameters(model: nn.Module) -> int: |
| return sum(p.numel() for p in model.parameters() if p.requires_grad) |
|
|
