File size: 6,476 Bytes
0812823 ef82d47 0812823 ef82d47 0812823 ef82d47 0812823 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | # /// script
# requires-python = ">=3.10"
# dependencies = [
# "torch",
# "numpy",
# "datasets>=2.19",
# "huggingface_hub>=0.24",
# "safetensors>=0.4",
# ]
# ///
"""Long-context bounded-state test (the decisive experiment for the thesis).
Trident keeps a *fixed-size* recurrent state (independent of context length). The
question that matters: does that O(1) state actually carry useful long-range
information *beyond the training window*?
We stream a long byte sequence in fixed windows and score bits-per-byte two ways:
* CARRY - carry the recurrent state across windows (true streaming),
* RESET - reset the state every window (no cross-window memory).
If CARRY < RESET on later windows, the fixed state is demonstrably transporting
information across tens of thousands of bytes with constant memory. We also
report the constant state footprint vs a transformer KV cache that grows O(n).
"""
import json
import math
import os
import sys
from pathlib import Path
import torch
def env(k, d=None):
v = os.environ.get(k)
return v if v not in (None, "") else d
def log(m):
print(f"[longctx] {m}", flush=True)
def build_stream(dataset, name, split, text_field, total_bytes):
from datasets import load_dataset
ds = load_dataset(dataset, name=name, split=split, streaming=True)
buf = bytearray()
for row in ds:
t = row.get(text_field)
if not t:
continue
buf.extend(t.encode("utf-8"))
buf.extend(b"\n\n")
if len(buf) >= total_bytes:
break
return bytes(buf[:total_bytes])
def main():
device = "cuda" if torch.cuda.is_available() else "cpu"
code_repo = env("CODE_REPO", "farguney/trident")
ckpt_tag = env("CKPT_TAG", "final")
W = int(env("WINDOW", 2048))
B = int(env("STREAMS", 4))
num_windows = int(env("NUM_WINDOWS", 64))
total = B * num_windows * W
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
work = Path("/tmp/trident_code")
snapshot_download(repo_id=code_repo, repo_type="model",
allow_patterns=["src/**"], local_dir=str(work))
sys.path.insert(0, str(work / "src"))
from trident import Trident, TridentConfig
ckdir = Path("/tmp/ckpt")
ckpt_dir = env("CKPT_DIR", "checkpoints")
snapshot_download(repo_id=code_repo, repo_type="model",
allow_patterns=[f"{ckpt_dir}/{ckpt_tag}/*"], local_dir=str(ckdir))
cdir = ckdir / ckpt_dir / ckpt_tag
cfg = TridentConfig(**json.loads((cdir / "config.json").read_text()))
model = Trident(cfg).to(device).eval()
sd = load_file(str(cdir / "model.safetensors"))
model.load_state_dict(sd, strict=False)
nparams = sum(p.numel() for p in model.parameters())
log(f"device={device} ckpt={ckpt_tag} params={nparams/1e6:.1f}M "
f"window={W} streams={B} windows={num_windows} bytes={total}")
raw = build_stream(env("VAL_DATASET", "deepmind/pg19"), env("VAL_NAME"),
env("VAL_SPLIT", "test"), env("VAL_TEXT_FIELD", "text"), total)
data = torch.tensor(list(raw), dtype=torch.long)[:B * num_windows * W]
data = data.view(B, num_windows * W).to(device)
log(f"stream built: {tuple(data.shape)}")
amp = torch.bfloat16 if device == "cuda" else torch.float32
@torch.no_grad()
def run(carry: bool):
state = None
per_window = []
for w in range(num_windows):
win = data[:, w * W:(w + 1) * W]
with torch.autocast(device_type="cuda", dtype=amp, enabled=device == "cuda"):
out = model(win, return_logits=True, state0=state if carry else None,
return_state=True)
logits = out["logits"].float()
nll = torch.nn.functional.cross_entropy(
logits.reshape(-1, logits.shape[-1]), win.reshape(-1), reduction="sum")
per_window.append(nll.item() / (B * W) / math.log(2))
state = [s.detach() for s in out["state"]] if carry else None
return per_window
bpb_carry = run(True)
bpb_reset = run(False)
def block_mean(x, lo, hi):
seg = x[lo:hi]
return sum(seg) / len(seg) if seg else float("nan")
nb = num_windows
early_c, late_c = block_mean(bpb_carry, 0, 1), block_mean(bpb_carry, nb - 8, nb)
early_r, late_r = block_mean(bpb_reset, 0, 1), block_mean(bpb_reset, nb - 8, nb)
mean_c = sum(bpb_carry) / nb
mean_r = sum(bpb_reset) / nb
# constant state footprint vs transformer KV cache growth
state_bytes = B * cfg.n_blocks * cfg.n_heads * cfg.d_v * cfg.d_k * 4
ctx_bytes = num_windows * W
# a same-d_model,-n_layers transformer KV cache (fp16) at ctx_bytes tokens-ish
kv_at_ctx = 2 * cfg.n_blocks * cfg.d_model * ctx_bytes * 2 * B
log("==== LONG-CONTEXT BOUNDED-STATE RESULT (BPB, lower better) ====")
log(f" CARRY mean={mean_c:.4f} first_window={early_c:.4f} last8_windows={late_c:.4f}")
log(f" RESET mean={mean_r:.4f} first_window={early_r:.4f} last8_windows={late_r:.4f}")
log(f" state advantage (RESET-CARRY) mean={mean_r-mean_c:+.4f} late={late_r-late_c:+.4f}")
log(f" memory: Trident state={state_bytes/1e6:.2f} MB (constant) vs "
f"transformer KV@{ctx_bytes} ctx ~= {kv_at_ctx/1e9:.2f} GB (grows O(n))")
results = {
"ckpt": ckpt_tag, "params_M": nparams / 1e6, "window": W, "streams": B,
"num_windows": num_windows, "context_bytes": ctx_bytes,
"bpb_carry": bpb_carry, "bpb_reset": bpb_reset,
"carry_mean": mean_c, "reset_mean": mean_r,
"carry_first": early_c, "carry_last8": late_c,
"reset_first": early_r, "reset_last8": late_r,
"state_advantage_mean": mean_r - mean_c, "state_advantage_late": late_r - late_c,
"trident_state_bytes": state_bytes, "transformer_kv_bytes_at_ctx": kv_at_ctx,
}
from huggingface_hub import HfApi
out = Path("/tmp/results"); out.mkdir(exist_ok=True)
result_name = env("RESULT_NAME", "longctx.json")
(out / result_name).write_text(json.dumps(results, indent=2))
try:
HfApi().upload_folder(folder_path=str(out), repo_id=code_repo, repo_type="model",
path_in_repo="results", commit_message="long-context bounded-state result")
log("pushed results/longctx.json")
except Exception as e: # noqa
log(f"upload failed: {e}")
if __name__ == "__main__":
main()
|