diff options
| -rw-r--r-- | EMAIL_BEN_DRAFT2.md | 19 | ||||
| -rw-r--r-- | docs/campaign/CASCADE_ABLATION_PLAN.md | 24 | ||||
| -rw-r--r-- | docs/campaign/COST_MODEL.md | 47 | ||||
| -rw-r--r-- | docs/hardware/COMPONENT_HW_MAP.md | 10 | ||||
| -rw-r--r-- | ep_run/etier_probe.py | 227 |
5 files changed, 314 insertions, 13 deletions
diff --git a/EMAIL_BEN_DRAFT2.md b/EMAIL_BEN_DRAFT2.md index f367b12..0b6fc4a 100644 --- a/EMAIL_BEN_DRAFT2.md +++ b/EMAIL_BEN_DRAFT2.md @@ -16,17 +16,17 @@ The results promised yesterday, concretely: - **The gap ledger.** At matched hyperparameters, 4k steps: EP statistically indistinguishable from its tuned backprop control (three seeds each). Full epoch: 0.050 nats behind, with the mechanism identified — the late-training decline of estimator signal-to-noise as the true gradient shrinks against a fixed precision floor (the quantified version of your β remark) — and a β-schedule that removes most of it. Muon's improvement over AdamW transfers to EP gradients at full magnitude. - **Solver cost.** Our nudged phase runs the full EP step at 3.2× the FLOPs of a backprop step (measured; single-sided, K=3 full-depth sweeps per nudge). For calibration against your Appendix B: that ratio is what turns an 18-day run into a ~4-day one at equal hardware. -On costs, anchored on measured throughput (TF32 verified bit-exact): +On costs — calibrated against a measured datapoint rather than list pricing (a 1B × 10B-token backprop run we executed on rented H200s cost ≈$1k all-in), combined with the measured 3.2× EP/BP FLOP ratio: -| run | cost, incl. 1.4× failure allowance | +| run | cost at market H100/H200 rates, incl. 1.3× failure allowance | |---|---| -| 300M × 6B tokens — recipe validation on a standard corpus | ~$1k | -| **1B × 20B tokens, Chinchilla-optimal** | **~$12k** | -| 1B backprop control | ~$4k | +| 300M × 6B tokens — recipe validation on a standard corpus | ~$400 | +| **1B × 20B tokens, Chinchilla-optimal** | **~$4k** (≤$8k even at the uncalibrated worst case) | +| 1B backprop control, same tokens | ~$1.5k | -Within $50k this leaves room for a second seed or a first 1.5B segment. The allowance is modest because the failure modes are characterized rather than assumed, and training checkpoints every 5k steps — a failure loses only the segment since the last checkpoint. A one-day calibration on your instances (~$500) would reduce the uncertainty from ±30% to ±10–15%; we would be glad to run that first. +Three qualifications. The numbers assume the TF32 path we have verified bit-exact against fp32; a standard mixed-precision variant (bf16 matrix products, fp32 states and energy accumulation) is not yet validated for EP and would roughly halve them if it passes our gradient-fidelity gate. AWS list pricing runs 1.5–2.5× market rental per GPU-hour, with spot or capacity-block reservations closing most of that gap. And checkpoints every 5k steps bound any failure to the segment since the last one. A one-day calibration on your instances (~$500) would replace these estimates with measurements; we would be glad to run that first. -The question we are weighing beyond that: **whether 3B at full Chinchilla, or 7–8B at reduced token count, is feasible** (~$25–30k and ~$50k respectively at current measured throughput). We are not requesting this now — the intention is to complete the runs above and revisit the larger scale on the basis of those results. +This pricing changes the scale question. Within $50k, the 1B milestone and its control leave room for one of: **3B at full Chinchilla scaling (~$30–40k at market rates), or 7–8B at a GPT-3-style token budget (~$25–35k)** — AWS on-demand pricing adds roughly half on top, and the mixed-precision variant, if it validates, subtracts it back. Our intention is unchanged — complete the 1B result first and let it argue for the larger run — but the envelope does reach one of the two. Separately, and possibly of direct interest to Rain: your paper leaves open whether F_PCN can be realized efficiently in hardware — **that question is the one our recipe was designed around.** Every trained operation was chosen to have a known analog implementation: divisive (current-mode) normalization, fixed rotations for the position encoding, translinear multipliers for the gated MLP, the subthreshold-exponential/KCL circuit for attention's softmax, and EP's two relaxation phases for the learning rule. A component-by-component hardware mapping is written up; we can share it. @@ -39,5 +39,8 @@ Yuren ## 发送前 checklist（内部,不进邮件） 1. [LINK] 换成真实 Overleaf 链接（zip 已备好: overleaf_dynamics_v3.zip,含两对象改写+KHS26引用+floss重调数字） 2. 附 a1_demo.ipynb（token 已内置） -3. gap 数字核对: f3e3cont 终值若 <1.2808 → "0.050 nats" 段落改为新 gap + "a β-schedule removes most of it" 升级为具体数字 +3. ~~gap 数字核对~~ RESOLVED: f3e3cont 终值 1.2883 > 1.2808 → "0.050 nats" 保持不变 4. 3.2× 的对照句("18-day → ~4-day")口吻是否合适你再定——它来自他们自己的 Appendix B,准确但直接 +5. 成本表 2026-07-12 已按用户实测锚点重校准（1B×10B BP H200 租赁 ≈$1k → $16.7k/EF 全包上限; + 市场价 H100/H200 TF32 路径 $6-9k/EF)。旧表锚 AWS-A100-savings,每 FLOP 贵 2.5-3×,已废。 + 完整推导见 docs/campaign/COST_MODEL.md diff --git a/docs/campaign/CASCADE_ABLATION_PLAN.md b/docs/campaign/CASCADE_ABLATION_PLAN.md index 61860f6..8b52369 100644 --- a/docs/campaign/CASCADE_ABLATION_PLAN.md +++ b/docs/campaign/CASCADE_ABLATION_PLAN.md @@ -538,3 +538,27 @@ s1 alone beat the BP mean). Ledger of the residual-0.05 epoch gap after three pr Carry `--bsign_rand` and a future centered mode as Stage-2 A/B flags; the gap question re-opens at 300M/real-corpus where it means something. Effort pivots to: (1) Stage-2 data pipeline (FineWeb-Edu + 32k tokenizer), (2) E-tier tolerance suite on the idle farm (hardware track / UIUC outreach feed). + +### RESULT 14 (2026-07-12): E-TIER WAVE-1 — full analog-fault tolerance ledger at stage1b s55000. +`etier_probe.py`, farm GPUs 2/3/7 (shards A/B/C), stage1b_ep_muon_s55000.pt (clean valCE 1.2678), +B=8 eval batches; metrics = faulted valCE (Δ vs clean), cos(EP_faulted, BP_faulted) [self-consistency +of the learning signal under fault], cos(EP_faulted, BP_clean) [direction vs the ideal update]. + +| fault (component) | mild | medium | severe | verdict | +|---|---|---|---|---| +| wq — weight quant (crossbar #3) | 8b: +0.004 / 0.956 | 6b: +0.051 / 0.812 | 4b: +2.23 / 0.05 | **8b FREE, 6b marginal, 4b dead → ≥7b effective is the binding spec** | +| fnoise — fwd additive state noise (softmax/relax #5) | 1e-3: 0.000 / 0.975 | 3e-3: 0.000 / 0.974 | 1e-2: +0.001 / 0.970 | **FREE at 1% — looped-era 1e-3 cliff does NOT transfer to cascade** | +| divmis — divisive-norm mismatch (#4/#7) | 1%: 0.000 / 0.971 | 3%: +0.003 / 0.948 | 10%: +0.042 / 0.785 | 3% (routine matching) FREE; 10% marginal | +| rope — phase error rad (#2) | 0.01: 0.000 / 0.973 | 0.03: +0.001 / 0.967 | 0.1: +0.012 / 0.923 | 0.03 rad FREE; ~2° I/Q accuracy suffices | +| gilbert — gate gain error (#6) | 1%: 0.000 / 0.974 | 3%: 0.000 / 0.971 | 10%: +0.006 / 0.945 | **FREE at 10%** — translinear practice is comfortably inside | +| fbnoise — nudge/error-channel noise | 1e-2: 0.969 / 0.975 | **1e-1: 0.951 / 0.957** | 3e-1: 0.764 / 0.768 | **10% relative noise on the ERROR CHANNEL is FREE** (cos 0.95) — the r-indifference/large-nudge gift, now measured on cascade | + +Reading: (a) the only hard constraint is crossbar weight precision (≥7b effective — inside standard +SRAM-CIM capability; 6b rescue = wave-2 quant-aware co-training); (b) everything dynamic — forward +noise 1%, error-channel noise 10%, gate/divider/phase mismatch at routine device tolerances — is +FREE at this scale. cos(EP,BP_faulted) stays ~0.97 under every non-fatal fault: the EP estimator +tracks whatever network the faults define, i.e. learning co-adapts to the fault (the analog-training +thesis in one number). CAVEAT: static probes at a trained checkpoint (eval CE + one-step gradient +direction), not training-under-fault; wave-2 = co-training with faults injected from step 0 +(expectation from the literature and from (c): tolerances IMPROVE). Feeds COMPONENT_HW_MAP.md +(per-row status updated) + UIUC outreach dossier. diff --git a/docs/campaign/COST_MODEL.md b/docs/campaign/COST_MODEL.md new file mode 100644 index 0000000..2b9d8b6 --- /dev/null +++ b/docs/campaign/COST_MODEL.md @@ -0,0 +1,47 @@ +# Cost model v2 (2026-07-12) — recalibrated on a measured rental datapoint + +Supersedes the v1 table in EMAIL_BEN_DRAFT2.md (drafts ≤5) and the numbers quoted in the +2026-07-11 scale discussion. **v1's error: anchored $/FLOP on AWS-A100-savings TF32 +(~$19.7k/EF) — the wrong reference class for 2026.** Market H100/H200 rentals are 2.5–3× +cheaper per FLOP. + +## Anchors + +| anchor | value | source | +|---|---|---| +| **Empirical all-in ceiling** | **$16.7k/EF** | user's real run: 1B × 10B tok BP on rented H200s ≈ $1k all-in; 6ND = 0.06 EF. Implied throughput only 42–58 TF/s eff (4–6% of bf16 peak) — an UN-optimized run, so this is a ceiling, not a target | +| Market rental rate | $2–3.5/GPU·h | H100/H200 marketplace (Vast/RunPod class), 2026 | +| Our stack's utilization | ~19% of TF32 peak | measured on A6000 (14 TF eff / 75 TF TF32 peak), eager TF32 | +| → projected H200 throughput | 75–125 TF/s eff | 15–25% × 494 TF TF32 peak | +| → **working rate (TF32 path)** | **$6–9k/EF** (mid $7.5k) | $2.5–3/h ÷ 75–125 TF/s | +| EP/BP FLOP ratio | **3.2×** (measured) | EP ≈ 19.2N FLOPs/tok vs BP 6N; single-sided, K=3 | +| Failure allowance | 1.3× | failure modes characterized; ckpt every 5k steps | +| AWS multiplier | 1.5–2.5× on-demand; 1.2–1.5× spot/capacity-block | vs marketplace $/GPU·h | +| Mixed-precision upside | ÷~2 | proper autocast (bf16 matmul, fp32 states + E-accum) — UNVALIDATED for EP; naive-cast is dead (cos ≤0.67) | + +Cross-check: v1's per-EF rate ($19.7k) ≈ the empirical ceiling ($16.7k) — v1 rows were not +order-of-magnitude wrong, they were "un-optimized-run" priced. The row deltas people remember +($1k vs $12k) decompose as: ×2 tokens (20B vs 10B) × 3.2 EP × 1.4 buffer ≈ ×9. + +## Table (market rates, TF32 path, ×1.3 buffer; AWS on-demand ≈ +50–150%) + +| run | EP FLOPs (19.2ND) | cost | note | +|---|---|---|---| +| 300M × 6B (recipe validation) | 0.035 EF | **~$400** | | +| **1B × 20B Chinchilla EP** | 0.384 EF | **~$4k** | ≤$8k at the empirical-ceiling rate | +| 1B × 20B BP control | 0.12 EF (6ND) | **~$1.5k** | user's datapoint directly: 10B = $1k | +| **3B × 60B full-Chinchilla EP** | 3.46 EF | **~$30–40k** | inside $50k alone (market); AWS on-demand $45–60k → needs spot or bf16 | +| **7B × 20B (GPT-3-style) EP** | 2.69 EF | **~$25–35k** | inside envelope | +| 7B × 140B full-Chinchilla EP | 18.8 EF | $110–170k | out of scope; dedicated funding | + +## Consequences + +1. **$50k envelope reaches ONE of {3B full-Chinchilla, 7–8B reduced-token} after the 1B milestone** — + the scale conversation upgrades from speculative to budgetable. +2. The **calibration day (~$500 on Ben's AWS instances)** now resolves a 2× spread + (working rate vs empirical ceiling, AWS vs market) — worth more than before. +3. **bf16 mixed-precision validation is the highest-leverage engineering item**: ÷2 across the + table brings BOTH large options comfortably inside AWS on-demand pricing. Distinct from the + dead naive-cast: fp32 states + fp32 energy accumulation, bf16 matmuls only, gate = cos(EP,BPTT). +4. Throughput assumptions are conservative (eager, no compile/flash in the nudged phase); + the speed-profile levers (--holofast, --sdpa ≈ 1.5×) are not priced in. diff --git a/docs/hardware/COMPONENT_HW_MAP.md b/docs/hardware/COMPONENT_HW_MAP.md index 95c5689..ff47123 100644 --- a/docs/hardware/COMPONENT_HW_MAP.md +++ b/docs/hardware/COMPONENT_HW_MAP.md @@ -11,11 +11,11 @@ Companion docs: `HW_RESEARCH_FINDINGS.md` (softmax dossier, OLMo2 analog audit, | # | Component | Computes | Analog primitive | Reuse source (no tapeout) | Feedback path (nudged phase) | Tolerance status | |---|---|---|---|---|---|---| | 1 | Token embedding | table lookup | **digital** (SRAM lookup + DAC line drivers) | any MCU/FPGA + DAC array | none (input boundary) | n/a | -| 2 | RoPE | fixed per-position 2×2 rotations on q,k | **I/Q quadrature mixer**: cos/sin DDS + multiplying DACs | century-old RF practice; COTS mixer/DDS chips | transpose = rotation by −θ — **same mixer, negated sin** | phase error ≈ position blur — expect very tolerant; **E-tier item (queued)** | -| 3 | qkv / attn-proj / SwiGLU w1,w3,w2 / head | matrix–vector products | **crossbar arrays** (the bulk of all compute) | Demo-0/1: **SRAM-CIM (Shanbhag DIMA class)**; alternatives: Mythic flash CIM, IBM HERMES PCM eval | J^T reads = **bidirectional crossbar access** (the known PAR price; standard in every analog-EP proposal) | wq8 safe / wq6 marginal (looped-era static-tolerance); **re-run on cascade in E-tier** | -| 4 | QK-norm (full-width RMS) | q,k ← q/rms(q)·g | **divisive normalization**: square-law devices + KCL current sum + divider/AGC | same primitive class as softmax normalization periphery | Jacobian symmetric — self-transpose, no extra hardware | divider mismatch/offset — **E-tier item (queued)** | -| 5 | Softmax (causal) | exp + row-normalize | **Elfadel–Wyatt "entropic resistor"**: subthreshold-exp devices + KCL normalization (log-sum-exp co-content) | classic analog-VLSI circuit family; small arrays demonstrated | softmax Jacobian diag(p)−pp^T is symmetric; the QKV coupling around it is the non-reciprocal part | ±1% device mismatch → ±1% softmax error (linear); QK-norm bounds its input range. Dynamic-noise cliff (fnoise ≥1e-3, looped-era) **re-test on cascade = E-tier** | -| 6 | SwiGLU gate | silu(w1x) ⊙ (w3x) | sigmoid = differential pair (native); ⊙ = **Gilbert/translinear multiplier** per unit | textbook standard cells | transpose of ⊙ = multiply by the partner signal — cheap | per-multiply mismatch/noise, 1408 wide — **E-tier item (queued)** | +| 2 | RoPE | fixed per-position 2×2 rotations on q,k | **I/Q quadrature mixer**: cos/sin DDS + multiplying DACs | century-old RF practice; COTS mixer/DDS chips | transpose = rotation by −θ — **same mixer, negated sin** | **E-tier w1: 0.03 rad FREE; 0.1 rad ΔCE +0.012, cos_clean 0.92** — ~2° I/Q phase accuracy suffices (easy) | +| 3 | qkv / attn-proj / SwiGLU w1,w3,w2 / head | matrix–vector products | **crossbar arrays** (the bulk of all compute) | Demo-0/1: **SRAM-CIM (Shanbhag DIMA class)**; alternatives: Mythic flash CIM, IBM HERMES PCM eval | J^T reads = **bidirectional crossbar access** (the known PAR price; standard in every analog-EP proposal) | **E-tier w1 (cascade): wq8 FREE (ΔCE +0.004, cos_clean 0.96); wq6 MARGINAL (+0.051, 0.81); wq4 DEAD (+2.23)** — ≥7b effective weight precision is the binding spec; wave-2 lever = quant-aware co-training | +| 4 | QK-norm (full-width RMS) | q,k ← q/rms(q)·g | **divisive normalization**: square-law devices + KCL current sum + divider/AGC | same primitive class as softmax normalization periphery | Jacobian symmetric — self-transpose, no extra hardware | **E-tier w1: 3% mismatch FREE (ΔCE +0.003, cos_clean 0.95); 10% MARGINAL (+0.042, 0.78)** — 3% device matching is routine | +| 5 | Softmax (causal) | exp + row-normalize | **Elfadel–Wyatt "entropic resistor"**: subthreshold-exp devices + KCL normalization (log-sum-exp co-content) | classic analog-VLSI circuit family; small arrays demonstrated | softmax Jacobian diag(p)−pp^T is symmetric; the QKV coupling around it is the non-reciprocal part | ±1% device mismatch → ±1% softmax error (linear); QK-norm bounds its input range. **E-tier w1: forward additive state noise σ=1e-2 FREE (ΔCE +0.001, cos 0.97) — the looped-era 1e-3 cliff does NOT transfer to cascade** | +| 6 | SwiGLU gate | silu(w1x) ⊙ (w3x) | sigmoid = differential pair (native); ⊙ = **Gilbert/translinear multiplier** per unit | textbook standard cells | transpose of ⊙ = multiply by the partner signal — cheap | **E-tier w1: 10% gain error ΔCE +0.006, cos_clean 0.945 — FREE** at translinear-practice tolerances | | 7 | RMSNorm ×2 per block (norm-after-sublayer) | bound each sublayer output | divisive normalization (as #4), **no mean-subtraction path** (cheaper than LayerNorm) | same as #4 | symmetric Jacobian | as #4; ~4L+1 normalizers total = the largest NEW periphery count | | 8 | Residual stream | z + branch outputs | **current-summing bus (KCL node)** | wires | pass-through | norm-after bounds every injection — the arch change is itself the tolerance fix | | 9 | Final RMSNorm | bound pre-readout state | as #4 | as #4 | symmetric | bounds the **ADC** dynamic range at the digital boundary | diff --git a/ep_run/etier_probe.py b/ep_run/etier_probe.py new file mode 100644 index 0000000..343b77a --- /dev/null +++ b/ep_run/etier_probe.py @@ -0,0 +1,227 @@ +"""E-tier wave-1: device-fault tolerance probes on the trained OLMo2 cascade (stage1b ckpt). +For each fault x severity: (a) faulted free-forward val CE (inference survival), +(b) cos(EP_faulted, BP_faulted) — estimator robustness on the faulted system, +(c) cos(EP_faulted, BP_clean) — direction vs the clean-system gradient. +Faults: wq (weight quant bits) | fnoise (dynamic block-output noise, mult) | +divmis (RMSNorm divider mismatch, fixed per-channel) | gilbert (SwiGLU gate gain mismatch) +| rope (phase error) | fbnoise (additive error-channel noise in the nudged feedback). +Usage: etier_probe.py --shard {A,B,C} +""" +import argparse, math, pickle, copy +import numpy as np, torch, torch.nn as nn, torch.nn.functional as F +from pathlib import Path + +ap = argparse.ArgumentParser() +ap.add_argument('--shard', required=True, choices=['A', 'B', 'C']) +ap.add_argument('--ckpt', default='runs/stage1b_ep_muon_s55000.pt') +ap.add_argument('--beta', type=float, default=1e-3) +ap.add_argument('--K', type=int, default=3) +ap.add_argument('--nb', type=int, default=4) +a = ap.parse_args() +dev = 'cuda' if torch.cuda.is_available() else 'cpu' +torch.manual_seed(7) + +DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe') +vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size'] +B, T = 8, 256 + +def get_batch(): + data = np.memmap(DD / 'val.bin', dtype=np.uint16, mode='r') + ix = torch.randint(len(data) - T - 1, (B,)) + x = torch.stack([torch.from_numpy(data[i:i + T].astype(np.int64)) for i in ix]) + y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + T].astype(np.int64)) for i in ix]) + return x.to(dev), y.to(dev) + +# ---- model (OLMo2 cascade, matches trainer) with fault hooks ---- +FAULT = {'fnoise': 0.0, 'fbnoise': 0.0} + +class RMSNorm(nn.Module): + def __init__(self, C, eps=1e-6): + super().__init__(); self.g = nn.Parameter(torch.ones(C)); self.eps = eps + def forward(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.g + +class SwiGLU(nn.Module): + def __init__(self, C): + super().__init__() + h = ((8 * C // 3) + 63) // 64 * 64 + self.w1 = nn.Linear(C, h, bias=False); self.w3 = nn.Linear(C, h, bias=False) + self.w2 = nn.Linear(h, C, bias=False) + self.register_buffer('ggain', torch.ones(h), persistent=False) # gilbert mismatch + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x) * self.ggain) + +class Olmo2Attn(nn.Module): + def __init__(self, C, H, T): + super().__init__() + self.H, self.hd = H, C // H + self.qkv = nn.Linear(C, 3 * C, bias=False); self.proj = nn.Linear(C, C, bias=False) + self.qn, self.kn = RMSNorm(C), RMSNorm(C) + inv = 1.0 / (500000.0 ** (torch.arange(0, self.hd, 2).float() / self.hd)) + fr = torch.outer(torch.arange(T).float(), inv) + self.register_buffer('rc', fr.cos(), persistent=False) + self.register_buffer('rs', fr.sin(), persistent=False) + self.register_buffer('fr', fr, persistent=False) + def rope(self, x): + Tn = x.shape[2] + x1, x2 = x[..., ::2], x[..., 1::2] + c, s = self.rc[None, None, :Tn], self.rs[None, None, :Tn] + return torch.stack((x1 * c - x2 * s, x1 * s + x2 * c), dim=-1).flatten(-2) + def forward(self, x): + Bn, Tn, C = x.shape + q, k, v = self.qkv(x).split(C, dim=2) + q, k = self.qn(q), self.kn(k) + q = self.rope(q.view(Bn, Tn, self.H, self.hd).transpose(1, 2)) + k = self.rope(k.view(Bn, Tn, self.H, self.hd).transpose(1, 2)) + v = v.view(Bn, Tn, self.H, self.hd).transpose(1, 2) + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + return self.proj(y.transpose(1, 2).contiguous().view(Bn, Tn, C)) + +class Olmo2Block(nn.Module): + def __init__(self, C, H, T): + super().__init__() + self.attn = Olmo2Attn(C, H, T); self.ff = SwiGLU(C) + self.na, self.nf = RMSNorm(C), RMSNorm(C) + def forward(self, z): + z = z + self.na(self.attn(z)) + out = z + self.nf(self.ff(z)) + if FAULT['fnoise'] > 0: + out = out * (1.0 + FAULT['fnoise'] * torch.randn_like(out)) + return out + +ck = torch.load(a.ckpt, map_location=dev, weights_only=False) +cfg = ck['config']; C, H, L = cfg['C'], cfg['H'], cfg['L'] +tok = nn.Embedding(vocab, C).to(dev); tok.load_state_dict(ck['tok']) +blocks0 = nn.ModuleList([Olmo2Block(C, H, T) for _ in range(L)]).to(dev) +blocks0.load_state_dict(ck['blocks'], strict=False) # ggain buffers are non-persistent extras +W_out0 = ck['wout'].to(dev) +ln_f0 = RMSNorm(C).to(dev); ln_f0.load_state_dict(ck['lnf']) +NBT = B * T + +def apply_fault(kind, sev): + """return (blocks, W_out, ln_f) with the fault applied; also sets FAULT dict.""" + FAULT['fnoise'] = 0.0; FAULT['fbnoise'] = 0.0 + bl = copy.deepcopy(blocks0); Wo = W_out0.clone(); lf = copy.deepcopy(ln_f0) + g = torch.Generator(device='cpu').manual_seed(11) + if kind == 'clean': + pass + elif kind == 'wq': + bits = sev + with torch.no_grad(): + for m in bl.modules(): + if isinstance(m, nn.Linear): + s = m.weight.abs().max() / (2 ** (bits - 1) - 1) + m.weight.copy_(torch.round(m.weight / s) * s) + s = Wo.abs().max() / (2 ** (bits - 1) - 1) + Wo = torch.round(Wo / s) * s + elif kind == 'divmis': + with torch.no_grad(): + for m in bl.modules(): + if isinstance(m, RMSNorm): + m.g.mul_(1.0 + sev * torch.randn(m.g.shape, generator=g).to(dev)) + lf.g.mul_(1.0 + sev * torch.randn(lf.g.shape, generator=g).to(dev)) + elif kind == 'gilbert': + with torch.no_grad(): + for m in bl.modules(): + if isinstance(m, SwiGLU): + m.ggain.copy_(1.0 + sev * torch.randn(m.ggain.shape, generator=g).to(dev)) + elif kind == 'rope': + with torch.no_grad(): + for m in bl.modules(): + if isinstance(m, Olmo2Attn): + d = sev * torch.randn(m.fr.shape, generator=g).to(dev) + m.rc.copy_((m.fr + d).cos()); m.rs.copy_((m.fr + d).sin()) + elif kind == 'fnoise': + FAULT['fnoise'] = sev + elif kind == 'fbnoise': + FAULT['fbnoise'] = sev + return bl, Wo, lf + +def readout(z, Wo, lf): return lf(z) @ Wo.t() + +def ep_grad(bl, Wo, lf, x, y): + """single-sided fb EP grad on blocks params (matches trainer structure, K rounds).""" + with torch.no_grad(): + z = tok(x) + zs_free = [] + for b in bl: z = b(z); zs_free.append(z.clone()) + # graphed free pass + ins, outs, zs = [], [], [] + prev = zs_free[0] * 0 + tok(x).detach() + prev = tok(x).detach() + for b in bl: + i = prev.detach().requires_grad_(True) + o = b(i) + ins.append(i); outs.append(o); zs.append(o.detach()) + prev = zs[-1] + d = [None] * L + for k in range(a.K): + zc = zs[L - 1].detach().requires_grad_(True) + ce = F.cross_entropy(readout(zc, Wo, lf).reshape(-1, vocab), y.reshape(-1)) + d[L - 1] = (-a.beta * NBT * torch.autograd.grad(ce, zc)[0]).detach() + for l in range(L - 2, -1, -1): + d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1])[0].detach() + if FAULT['fbnoise'] > 0: + for l in range(L): + d[l] = d[l] + FAULT['fbnoise'] * d[l].norm() / math.sqrt(d[l].numel()) * torch.randn_like(d[l]) + last = (k + 1 == a.K) + prev = tok(x).detach() + n_ins, n_outs = [], [] + for l in range(L): + i = prev.detach().requires_grad_(True) + o = bl[l](i) + n_ins.append(i); n_outs.append(o) + zs[l] = (o + d[l]).detach() + prev = zs[l] + ins, outs = n_ins, n_outs + E = 0.0 + for z, o in zip(zs, outs): E = E + 0.5 * ((z.detach().float() - o.float()) ** 2).sum() + obj = E / (NBT * a.beta) + params = [p for p in bl.parameters()] + gs = torch.autograd.grad(obj, params, allow_unused=True) + return [g if g is not None else torch.zeros(1, device=dev) for g in gs] + +def bp_grad(bl, Wo, lf, x, y): + z = tok(x) + for b in bl: z = b(z) + ce = F.cross_entropy(readout(z, Wo, lf).reshape(-1, vocab), y.reshape(-1)) + params = [p for p in bl.parameters()] + gs = torch.autograd.grad(ce, params, allow_unused=True) + return [g if g is not None else torch.zeros(1, device=dev) for g in gs] + +def val_ce(bl, Wo, lf, x, y): + with torch.no_grad(): + z = tok(x) + for b in bl: z = b(z) + return float(F.cross_entropy(readout(z, Wo, lf).reshape(-1, vocab), y.reshape(-1))) + +def cos(ga, gb): + va = torch.cat([g.reshape(-1) for g in ga]); vb = torch.cat([g.reshape(-1) for g in gb]) + return float((va @ vb) / (va.norm() * vb.norm() + 1e-30)) + +SHARDS = { + 'A': [('wq', 8), ('wq', 6), ('wq', 4), ('divmis', 0.01), ('divmis', 0.03), ('divmis', 0.10)], + 'B': [('fnoise', 1e-3), ('fnoise', 3e-3), ('fnoise', 1e-2), ('rope', 0.01), ('rope', 0.03), ('rope', 0.10)], + 'C': [('gilbert', 0.01), ('gilbert', 0.03), ('gilbert', 0.10), ('fbnoise', 1e-2), ('fbnoise', 1e-1), ('fbnoise', 3e-1)], +} + +batches = [get_batch() for _ in range(a.nb)] +blc, Woc, lfc = apply_fault('clean', 0) +clean_ce = sum(val_ce(blc, Woc, lfc, x, y) for x, y in batches) / a.nb +gclean = [[g.cpu() for g in bp_grad(blc, Woc, lfc, x, y)] for x, y in batches] +torch.cuda.empty_cache() +print(f"[clean] val CE {clean_ce:.4f}", flush=True) + +for kind, sev in SHARDS[a.shard]: + torch.cuda.empty_cache() + bl, Wo, lf = apply_fault(kind, sev) + ce = sum(val_ce(bl, Wo, lf, x, y) for x, y in batches) / a.nb + c_self, c_clean = [], [] + for i, (x, y) in enumerate(batches): + ge = ep_grad(bl, Wo, lf, x, y) + gbf = bp_grad(bl, Wo, lf, x, y) + c_self.append(cos(ge, gbf)) + c_clean.append(cos([g.cpu() for g in ge], gclean[i])) + del ge, gbf; torch.cuda.empty_cache() + print(f"[{kind}={sev}] valCE {ce:.4f} (Δ{ce-clean_ce:+.4f}) | cos(EP,BP_faulted) {sum(c_self)/a.nb:.4f} | cos(EP,BP_clean) {sum(c_clean)/a.nb:.4f}", flush=True) +print(f"DONE_{a.shard}", flush=True) |