mirror alephlm.py from geolip 0.7.0 @a63a29b (MIRROR law; n_const=1 bit-identical)
Browse files- alephlm.py +17 -1
alephlm.py
CHANGED
|
@@ -40,7 +40,8 @@ class Block(nn.Module):
|
|
| 40 |
self.n2 = nn.LayerNorm(d)
|
| 41 |
if self.is_hub:
|
| 42 |
self.attn = CausalSplatHUB(d, cfg.hub_K, cfg.hub_D, cfg.tau,
|
| 43 |
-
chunk=cfg.hub_chunk
|
|
|
|
| 44 |
else:
|
| 45 |
self.attn = CausalSDPA(d, cfg.n_heads)
|
| 46 |
self.bank = AnchoredBank(d, cfg.bank_experts, cfg.bank_ff, cfg.tau,
|
|
@@ -183,6 +184,21 @@ class AlephLM(nn.Module):
|
|
| 183 |
idx = torch.cat([idx, nxt], dim=1)
|
| 184 |
return idx
|
| 185 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 186 |
def param_count(self) -> int:
|
| 187 |
seen, total = set(), 0
|
| 188 |
for p in self.parameters():
|
|
|
|
| 40 |
self.n2 = nn.LayerNorm(d)
|
| 41 |
if self.is_hub:
|
| 42 |
self.attn = CausalSplatHUB(d, cfg.hub_K, cfg.hub_D, cfg.tau,
|
| 43 |
+
chunk=cfg.hub_chunk,
|
| 44 |
+
n_const=getattr(cfg, "hub_const", 1))
|
| 45 |
else:
|
| 46 |
self.attn = CausalSDPA(d, cfg.n_heads)
|
| 47 |
self.bank = AnchoredBank(d, cfg.bank_experts, cfg.bank_ff, cfg.tau,
|
|
|
|
| 184 |
idx = torch.cat([idx, nxt], dim=1)
|
| 185 |
return idx
|
| 186 |
|
| 187 |
+
def compile_hubs(self, **compile_kw):
|
| 188 |
+
"""Opt-in: torch.compile each hub's attention (measured 4.0x over
|
| 189 |
+
eager at ctx 2048 with the fused forward, parity 1.4e-06 vs the
|
| 190 |
+
naive oracle; MHA-parity wall-clock by ctx 8192). Call AFTER
|
| 191 |
+
loading weights, and do NOT save a checkpoint while compiled (the
|
| 192 |
+
OptimizedModule wrapper prefixes state_dict keys). TRAINING use is
|
| 193 |
+
gated on a grad-parity + throughput check on the training
|
| 194 |
+
hardware — the speed verdict was no-grad forward only."""
|
| 195 |
+
import torch as _torch
|
| 196 |
+
for wrap in self.blocks:
|
| 197 |
+
blk = getattr(wrap, "block", wrap)
|
| 198 |
+
if blk.is_hub:
|
| 199 |
+
blk.attn = _torch.compile(blk.attn, **compile_kw)
|
| 200 |
+
return self
|
| 201 |
+
|
| 202 |
def param_count(self) -> int:
|
| 203 |
seen, total = set(), 0
|
| 204 |
for p in self.parameters():
|