| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """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 |
|
|
| |
| state_bytes = B * cfg.n_blocks * cfg.n_heads * cfg.d_v * cfg.d_k * 4 |
| ctx_bytes = num_windows * W |
| |
| 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: |
| log(f"upload failed: {e}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
