File size: 14,004 Bytes

f5f5621

#!/usr/bin/env python3
"""
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()