| |
| """ |
| mk_lenbucket_notebook.py -- executed .ipynb for the training-document-length |
| experiment: does the LENGTH of training documents change NIAH retrieval across |
| the context window? |
| |
| usage: mk_lenbucket_notebook.py [--out lenbucket_results.ipynb] |
| |
| THREE ARMS, ONE VARIABLE |
| armA_c flux plain text median 826, p90 2,721 tok |
| l1 b1_4096_8192 docs in [4096, 8192) |
| l2 b2_8192_16384 docs in [8192, 16384) |
| All: 400 steps, global batch 524,288, from run7' (run 6 step 750), |
| YARN_ORIG_MAXPOS=4096 matched, Muon+ZeRO-1 bf16, 4xH100 at CP_B=2. |
| |
| ONE SHARED HAYSTACK: b3_16384_32768's val shard (docs 16-32k), held out from |
| every arm. armA_c is RE-SCORED here rather than reusing its 2026-09-03 numbers, |
| which used the flux val shard as filler -- a different haystack is a different |
| test. That is why armA_c reads 33.0 at 16,360 here and 25.0 there: same weights, |
| different filler. Columns here compare with each other, NOT with that table. |
| |
| ALL NUMBERS READ FROM lenb_len_*/summary.csv and depth_summary.csv. |
| """ |
| import argparse, csv, json, math, os, sys |
|
|
| ARMS = [("armA_c", "flux plain text (median 826, p90 2,721 tok)", "#2a78d6", "o"), |
| ("l1", "b1: docs 4-8k", "#eb6834", "s"), |
| ("l2", "b2: docs 8-16k", "#1baf7a", "D")] |
| LENS = [1024, 2048, 4096, 8192, 16360] |
| TASKS = ["niah_single", "val_backward"] |
| STYLE = {"surface": "#fcfcfb", "ink": "#0b0b0b", "ink2": "#52514e", "grid": "#e4e4e1"} |
|
|
|
|
| def read(d, f): |
| p = os.path.join(d, f) |
| return list(csv.reader(open(p))) if os.path.exists(p) else [] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--out", default="lenbucket_results.ipynb") |
| args = ap.parse_args() |
|
|
| summ, depth = {}, {} |
| for L in LENS: |
| rows = read(f"lenb_len_{L}", "summary.csv") |
| if len(rows) > 1: |
| hdr = rows[0][1:] |
| summ[L] = {r[0]: dict(zip(hdr, [float(x) if x else None for x in r[1:]])) |
| for r in rows[1:]} |
| dr = read(f"lenb_len_{L}", "depth_summary.csv") |
| if len(dr) > 1: |
| depth[L] = [dict(zip(dr[0], r)) for r in dr[1:] if len(r) == len(dr[0])] |
| if not summ: |
| sys.exit("no lenb_len_*/summary.csv -- run lenbucket_experiment.sh first") |
| print(" lengths:", sorted(summ)) |
|
|
| losses = {} |
| for a, lg in (("l1", "lenbucket_logs/l1.log"), ("l2", "lenbucket_logs/l2.log")): |
| if os.path.exists(lg): |
| import re |
| m = re.findall(r'^step\s+\d+ \| loss: ([0-9.]+)', open(lg, errors='ignore').read(), re.M) |
| if m: losses[a] = float(m[-1]) |
|
|
| payload = json.dumps({"summary": {str(k): v for k, v in summ.items()}, |
| "depth": {str(k): v for k, v in depth.items()}, |
| "arms": ARMS, "lens": LENS, "tasks": TASKS, |
| "style": STYLE, "losses": losses}, indent=1) |
|
|
| import nbformat as nbf |
| from nbclient import NotebookClient |
| C = [] |
| src = lambda s: s.strip("\n") |
| code = lambda s: C.append(nbf.v4.new_code_cell(src(s))) |
| md = lambda s: C.append(nbf.v4.new_markdown_cell(src(s))) |
|
|
| md(""" |
| # Does training-document length change long-context retrieval? |
| |
| **1.06B — 36 layers, d_model 1536, 12q/4kv heads, head_dim 128, SwiGLU 4096, |
| T=16384, YaRN scale 4, `YARN_ORIG_MAXPOS=4096`.** |
| |
| The 2026-09-03 corrected-cache result left NIAH retrieval at 27% at the full |
| 16,384 window. One standing explanation for that ceiling is the *training corpus*, |
| not the model: flux documents have **median 826 tokens, p90 2,721, and only 0.27% |
| reach 16k** — so almost no training example ever required attending across the |
| window. This isolates that by changing **only** document length. |
| |
| | arm | training documents | tokens available | |
| |---|---|---| |
| | `armA_c` | flux plain text — median 826, p90 2,721 | control | |
| | `l1` | `b1_4096_8192` — docs ∈ [4096, 8192) | 13.98 B | |
| | `l2` | `b2_8192_16384` — docs ∈ [8192, 16384) | 6.99 B | |
| |
| Everything else is held fixed: 400 steps, global batch 524,288, same parent |
| (`run7'`, run 6 step 750), same matched YaRN cache, Muon + ZeRO-1 bf16, 4×H100 at |
| `CP_B=2` (68.8 GB of 81.6 GB — the H200 runs used `CP_B=4`/108 GB). |
| |
| **One shared haystack: `b3_16384_32768`'s val shard** (documents of 16-32k tokens), |
| held out from every arm. Using b1's or b2's own val shard would be in-distribution |
| for exactly one arm and would flatter it. |
| |
| > `armA_c` is **re-scored here**, not carried over from 2026-09-03. Those numbers |
| > used the flux val shard as filler, and a different haystack is a different test — |
| > which is why `armA_c` reads **33.0** at 16,360 here versus **25.0** there, on |
| > identical weights. The columns below compare with each other; they do **not** |
| > compare with that earlier table. |
| """) |
|
|
| code(f""" |
| import json, math |
| import matplotlib.pyplot as plt |
| from matplotlib.ticker import FuncFormatter |
| |
| D = json.loads(r'''{payload}''') |
| S, ARMS, LENS = D["style"], D["arms"], D["lens"] |
| |
| def wilson(h, n): |
| \"\"\"95% Wilson half-width in points -- honest at the 0% floor.\"\"\" |
| if not n: return float("nan") |
| z, p = 1.96, h / n |
| return 100.0 * z * math.sqrt(p*(1-p)/n + z*z/(4*n*n)) / (1 + z*z/n) |
| |
| N = 100 # prompts per (arm, task, length) cell |
| |
| def acc(L, task, arm): |
| v = D["summary"].get(str(L), {{}}).get(task, {{}}).get(arm) |
| return v |
| |
| def style(ax, xl, yl, title=None): |
| ax.set_facecolor(S["surface"]); ax.grid(True, color=S["grid"], lw=.8, zorder=0) |
| ax.set_axisbelow(True) |
| for k, sp in ax.spines.items(): |
| sp.set_visible(k in ("left","bottom")); sp.set_color(S["grid"]) |
| ax.tick_params(colors=S["ink2"], labelsize=9) |
| ax.set_xlabel(xl, color=S["ink2"], fontsize=10) |
| ax.set_ylabel(yl, color=S["ink2"], fontsize=10) |
| if title: ax.set_title(title, color=S["ink"], fontsize=12, loc="left", pad=12) |
| |
| klabel = lambda v, _=None: (f"{{round(v/1024)}}k" if v >= 1024 else str(int(v))) |
| print("lengths:", LENS) |
| print("arms :", [a[0] for a in ARMS]) |
| """) |
|
|
| md(""" |
| ## 1. Headline — retrieval vs context length, by training document length |
| |
| `niah_single`: recall a 7-digit number given its key. Answer space 10⁷, so the |
| guess floor is ~1e-7 and any non-zero rate is real retrieval. Error bars are 95% |
| Wilson half-widths. |
| """) |
|
|
| code(""" |
| task = "niah_single" |
| fig, ax = plt.subplots(figsize=(9.5, 5.4), facecolor=S["surface"]) |
| for arm, lab, col, mk in ARMS: |
| xs, ys, es = [], [], [] |
| for L in LENS: |
| v = acc(L, task, arm) |
| if v is None: continue |
| xs.append(L); ys.append(v); es.append(wilson(v*N/100.0, N)) |
| if not xs: continue |
| ax.errorbar(xs, ys, yerr=es, color=col, marker=mk, ms=6, lw=2.3, capsize=3, |
| label=lab, zorder=4) |
| ax.set_xscale("log", base=2); ax.set_xticks(LENS) |
| ax.xaxis.set_major_formatter(FuncFormatter(klabel)) |
| ax.set_ylim(-4, 104); ax.axhline(0, color=S["grid"], lw=1) |
| style(ax, "evaluation context length (tokens)", "niah_single accuracy (%)", |
| "Longer training documents lift retrieval, and the gap widens with context") |
| ax.legend(frameon=False, fontsize=9, labelcolor=S["ink2"], loc="lower left", |
| title="training documents", title_fontsize=9) |
| plt.tight_layout(); plt.show() |
| """) |
|
|
| md(""" |
| ## 2. The gap vs the control, and whether it is resolvable |
| |
| `resolvable` compares each gap against the two cells' Wilson bars added in |
| quadrature. A gap smaller than that is not distinguishable at N=100 and calling |
| it an effect would be reading noise. |
| """) |
|
|
| code(""" |
| hdr = f"{'context':>8} {'armA_c':>13} {'l1':>13} {'l2':>13} {'l1-A':>7} {'l2-A':>7} {'l2-l1':>7} resolv(l2-A)" |
| print(hdr); print("-"*len(hdr)) |
| for L in LENS: |
| a, b, c = acc(L,task,"armA_c"), acc(L,task,"l1"), acc(L,task,"l2") |
| if None in (a,b,c): continue |
| f = lambda v: f"{v:5.1f}+-{wilson(v*N/100.0,N):4.1f}" |
| comb = math.hypot(wilson(a*N/100.0,N), wilson(c*N/100.0,N)) |
| res = "YES" if abs(c-a) > comb else "no" |
| print(f"{klabel(L):>8} {f(a):>13} {f(b):>13} {f(c):>13} " |
| f"{b-a:+7.1f} {c-a:+7.1f} {c-b:+7.1f} {res} (bar +-{comb:.1f})") |
| |
| print() |
| print("ordering check (l2 > l1 > armA_c at each length):") |
| for L in LENS: |
| a, b, c = acc(L,task,"armA_c"), acc(L,task,"l1"), acc(L,task,"l2") |
| if None in (a,b,c): continue |
| print(f" {klabel(L):>6}: {'l2>l1>A' if c>=b>=a else 'mixed':<9} " |
| f"({a:.0f} / {b:.0f} / {c:.0f})") |
| """) |
|
|
| md(""" |
| ## 3. Reach |
| |
| **reach** = the context length at which accuracy crosses a threshold, interpolated |
| between measured lengths. Two thresholds, because one hides how steeply the curve |
| falls: **50%** is "works more often than not"; **10%** is "still clearly above the |
| ~1e-7 guess floor". |
| |
| Interpolated reach inherits the error bars of the cells it sits between, so treat |
| each value as bracketed by its two adjacent measured lengths rather than exact. |
| """) |
|
|
| code(""" |
| def reach(xs, ys, thr): |
| pts = sorted(zip(xs, ys)) |
| if not pts or pts[0][1] < thr: return None |
| last = pts[0][0] |
| for i in range(1, len(pts)): |
| (x0,y0),(x1,y1) = pts[i-1], pts[i] |
| if y1 >= thr: last = x1; continue |
| if y0 == y1: return x0 |
| return x0 + (x1-x0)*(y0-thr)/(y0-y1) |
| return last |
| |
| print(f"{'arm':<10} {'reach@50%':>14} {'reach@10%':>14} training docs") |
| print("-"*72) |
| for arm, lab, _, _ in ARMS: |
| xs = [L for L in LENS if acc(L,task,arm) is not None] |
| ys = [acc(L,task,arm) for L in xs] |
| r50, r10 = reach(xs,ys,50.0), reach(xs,ys,10.0) |
| g = lambda r: (f"{r:,.0f} tok" if r else "below thr") |
| print(f"{arm:<10} {g(r50):>14} {g(r10):>14} {lab}") |
| print() |
| print("full window = 16,384 tokens; longest measured prompt = 16,360") |
| """) |
|
|
| md(""" |
| ## 4. Depth profile at the full window |
| |
| `depth` is the needle's **position fraction**: 0.0 = the very start of the window, |
| farthest from the question; 1.0 = immediately before it. A U-shape |
| (lost-in-the-middle) would need a bump at *both* ends. Note the cells here are |
| n=20 per (arm, depth), so bars are wide — this panel shows shape, not magnitude. |
| """) |
|
|
| code(""" |
| fig, ax = plt.subplots(figsize=(9.5, 5.0), facecolor=S["surface"]) |
| L = 16360 |
| rows = D["depth"].get(str(L), []) |
| for arm, lab, col, mk in ARMS: |
| pts = sorted((float(r["depth"]), int(r["hits"]), int(r["n"])) |
| for r in rows if r["model"] == arm) |
| if not pts: continue |
| xs = [p[0] for p in pts] |
| ys = [100.0*p[1]/p[2] if p[2] else 0 for p in pts] |
| es = [wilson(p[1],p[2]) for p in pts] |
| ax.errorbar(xs, ys, yerr=es, color=col, marker=mk, ms=6, lw=2.0, capsize=3, |
| label=lab, zorder=4) |
| mono = all(ys[i] <= ys[i+1] + 1e-9 for i in range(len(ys)-1)) |
| print(f"{arm:<8} " + " ".join(f"d{p[0]:.2f}={100.0*p[1]/p[2]:5.1f}%({p[1]}/{p[2]})" for p in pts)) |
| print(f"{'':<8} monotone rising: {mono}") |
| ax.set_ylim(-4, 104) |
| style(ax, "needle depth (0 = window start, 1 = next to the question)", |
| "accuracy (%)", f"Depth profile at {L:,} tokens (n=20 per cell)") |
| ax.legend(frameon=False, fontsize=9, labelcolor=S["ink2"], loc="upper left", |
| title="training documents", title_fontsize=9) |
| plt.tight_layout(); plt.show() |
| """) |
|
|
| md(""" |
| ## 5. Backward retrieval — `val_backward` |
| |
| Given the *number*, name the *key*. The answer space is 35 words, so chance is |
| ~2.9% — a far easier task than recalling a 7-digit number. |
| """) |
|
|
| code(""" |
| t2 = "val_backward" |
| print(f"{'context':>8} " + " ".join(f"{a[0]:>14}" for a in ARMS)) |
| print("-"*54) |
| for L in LENS: |
| cells=[] |
| for arm,_,_,_ in ARMS: |
| v = acc(L,t2,arm) |
| cells.append(f"{v:5.1f}+-{wilson(v*N/100.0,N):4.1f}" if v is not None else " -") |
| print(f"{klabel(L):>8} " + " ".join(f"{c:>14}" for c in cells)) |
| print() |
| print("chance with a 35-word answer space is ~2.9%") |
| print("-> flat at 0-2% for every arm: document length does nothing for value->key.") |
| """) |
|
|
| md(""" |
| ## 6. Training losses — and why they must not be compared |
| |
| ``` |
| """ + "\n".join(f"{k}: {v:.6f}" for k, v in sorted(losses.items())) + """ |
| ``` |
| |
| **These are not comparable across arms.** Each is measured on its own data |
| distribution, and predicting tokens inside 8-16k documents is a strictly harder |
| problem than inside 4-8k ones, or inside flux's 826-token median. `l2` has the |
| *worst* loss and the *best* retrieval — the same loss/capability divergence seen |
| when the YaRN cache was corrected (run7′ went to loss 2.384 from run7's 1.283 |
| while gaining 27pp at 16k). Only the retrieval numbers compare. |
| |
| ## 7. What this establishes, and what it does not |
| |
| **Established:** with the cache correct and everything else held fixed, training |
| document length raises long-context retrieval, monotonically in length and |
| increasingly with evaluation context. The gap is not resolvable at 1-4k and is |
| resolvable at 8k and 16k. |
| |
| **Not established:** |
| |
| - *Saturation.* Whether the gradient continues past `l2` is untested. The |
| `b3_16384_32768` bucket (docs 16-32k, 3.03 B tokens) would answer it, and would |
| need a different haystack since b3 is currently the filler. |
| - *Token-budget parity.* All arms saw 210M tokens, but b1 holds 13.98 B available |
| tokens and b2 only 6.99 B, so the arms drew from pools of different size and |
| document counts (2.57 M vs 0.63 M docs). At 210M tokens neither pool was |
| exhausted, so this is unlikely to matter, but it is not controlled. |
| - *Anything about backward retrieval*, which is flat at chance throughout. |
| - *Comparison with the 2026-09-03 table*, which used a different haystack. |
| """) |
|
|
| nb = nbf.v4.new_notebook() |
| nb["cells"] = C |
| nb.metadata["kernelspec"] = {"name": "python3", "display_name": "Python 3", |
| "language": "python"} |
| print(" executing...") |
| NotebookClient(nb, timeout=300, kernel_name="python3", |
| resources={"metadata": {"path": os.getcwd()}}).execute() |
| nbf.write(nb, args.out) |
| print(f" wrote {args.out} ({os.path.getsize(args.out)} B, {len(C)} cells)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
