File size: 4,610 Bytes
d995873 | 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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """RelayPatch2D — the certified per-block residual relay for diffusion
token streams (verbatim port of the r2 bed, pod2/aleph_diffusion_core.py;
certified relay-favorable 3-for-3 substrates: Lune flow / SD15-core eps /
Anima DiT flow — exp001/exp006/exp004).
Differences from the text-side RelayPatchwork (amoe.core.adapter), each a
paid-for law:
* trailing-dim generalization: forward works on any [..., d] stream;
* zero-init output weight AND bias (exp006 bias-leak law) — a fresh
relay is EXACTLY inert, P-INIT bit-exact;
* `enabled=False` returns x via CODE-PATH SKIP (`return x`, `is x`) —
the toggle law is bit-exact at every dtype;
* dims-from-state-dict reconstruction (never trust constructor defaults).
"""
from __future__ import annotations
import torch
import torch.nn as nn
from ...core.address import AlephAddress
from ...core.adapter import SquaredReLU
class RelayPatch2D(nn.Module):
"""forward(x[..., d]) = x + sigmoid(gate) * consume(m_hat(proj(x) slots))."""
# Hosts dispatch on this: relays ignore the sigma axis, multiband
# modules require band windows. Marked on both so a caller never has
# to guess from the class name.
needs_bands = False
def __init__(self, d: int, n_slots: int = 16, K: int = 64,
tau: float = 0.1, hidden: int = 178):
super().__init__()
self.d, self.n_slots, self.hidden = d, n_slots, hidden
self.proj = nn.Linear(d, n_slots * 4, bias=False)
nn.init.orthogonal_(self.proj.weight)
self.addr = AlephAddress(K, 4, tau)
self.consume = nn.Sequential(
nn.Linear(n_slots * 4, hidden), SquaredReLU(),
nn.LayerNorm(hidden), nn.Linear(hidden, d))
nn.init.zeros_(self.consume[-1].weight) # zero-init out: weight
nn.init.zeros_(self.consume[-1].bias) # AND bias (exp006 law)
self.gate = nn.Parameter(torch.tensor(-3.0))
self.enabled = True # plain attr, not a buffer
self.telemetry = False # read-only ratio stash
self.last_delta_ratio = None
def assert_zero_init(self):
w, b = self.consume[-1].weight, self.consume[-1].bias
assert w.abs().max().item() == 0.0, "out weight not zero-init"
assert b.abs().max().item() == 0.0, "out BIAS not zero-init (leak law)"
def forward(self, x: torch.Tensor) -> torch.Tensor:
if not self.enabled:
return x
lead = x.shape[:-1]
slots = self.proj(x).view(*lead, self.n_slots, 4)
feats = self.addr.m_hat(slots).reshape(*lead, self.n_slots * 4)
delta = torch.sigmoid(self.gate) * self.consume(feats)
if self.telemetry:
with torch.no_grad():
self.last_delta_ratio = (
delta.norm() / x.norm().clamp_min(1e-12)).item()
return x + delta
@classmethod
def from_state_dict(cls, sd: dict) -> "RelayPatch2D":
"""Rebuild dims FROM the state dict, then strict-load. Pre-amoe
fork saves lack the addr.home drift-gauge buffer — reconstruct
it from the codebook, LOUDLY (the gauge starts here)."""
sd = dict(sd)
d = sd["proj.weight"].shape[1]
n_slots = sd["proj.weight"].shape[0] // 4
hidden = sd["consume.0.weight"].shape[0]
K = sd["addr.codebook"].shape[0]
if "addr.home" not in sd:
import torch.nn.functional as F
sd["addr.home"] = F.normalize(
sd["addr.codebook"].detach().float(), dim=-1)
print("amoe.diffusion: addr.home reconstructed from codebook "
"(pre-amoe save; drift gauge starts here)")
m = cls(d, n_slots=n_slots, K=K, hidden=hidden)
m.load_state_dict(sd, strict=True)
return m
class BlockWithRelay(nn.Module):
"""Wrap one trunk block; the relay reads its hidden-states output.
Tuple outputs pass through (hybrid-safe, as amoe.core BlockWithAdapter).
Carries a `w_bands` slot for wrap-uniformity with multiband sites —
the relay ignores it."""
def __init__(self, block: nn.Module, relay: nn.Module):
super().__init__()
self.block = block
self.mod = relay
self.w_bands = None # unused by relays; uniform wrap API
def forward(self, *args, **kwargs):
out = self.block(*args, **kwargs)
if isinstance(out, tuple):
return (self.mod(out[0]),) + out[1:]
return self.mod(out)
|