#!/usr/bin/env python3 """ mix_buckets.py -- build a mixed-length training set from the per-length buckets. python3 mix_buckets.py --out Data_Loader/Data/mix50 \ --src b5_65536_131072=0.50 --src b3_16384_32768=0.30 --src b1_4096_8192=0.20 \ --total-tokens 900e6 --window 65536 --grad-accum 8 WHY THIS EXISTS The 2026-09-12 run trained a full epoch on 100% b5 (every document >= 65,536 tokens) and lost long-context retrieval: niah_single went 67/65/52/29 -> 45/17/2/1 at 8k/16k/32k/64k while training loss FELL from 2.89 to 1.14. The literature is consistent that 100% long-context data causes severe degradation and that 25-66% short-context data should be mixed in. FORMAT (verified against the real shards, not assumed) Raw uint16, NO header, documents delimited by EOT = 50256. b5 shard 0 holds 1,059 documents of 65,573-130,893 tokens. DOCUMENT-LEVEL INTERLEAVING Shard-level mixing would give ~183 consecutive steps of one bucket at 96M tokens/shard -- long enough for the model to settle into a per-distribution strategy, which is the thing we are trying to prevent. This picks the next document from whichever bucket is furthest BELOW its target token share, so the mixture is right at every scale. A NOTE ON WHAT "SHORT DATA" MEANS AT A 64k WINDOW The training window stays 65,536. Short documents are PACKED, so a window drawn from b1 holds ~10 documents separated by EOT rather than one contiguous text. That is the standard recipe and it is the point: the model has to learn that not everything in context is relevant, which is exactly the signal that 100%-b5 training removed. It does NOT give the model short sequences -- that would need a variable T the trainer does not support mid-run. STEPS ARE COMPUTED CORRECTLY HERE The old sizing was STEPS = total_tokens // global_batch. That ignores the per-shard tail the loader discards on advance_shard(), and it came within one batch of silently wrapping on the last run. The manifest below emulates DataLoaderLite exactly and reports the true usable step count. """ import argparse, os, sys import numpy as np EOT = 50256 def doc_stream(paths, eot=EOT): """Yield documents (each ending with EOT) from a list of .bin shards. Streams with memmap: b5 shards are 192 MB each and there can be 139 of them. A document that straddles a shard boundary is stitched across the join. """ carry = np.empty(0, dtype=np.uint16) for p in paths: a = np.memmap(p, dtype=np.uint16, mode="r") idx = np.flatnonzero(np.asarray(a) == eot) prev = 0 for i in idx: seg = np.asarray(a[prev:i + 1]) if carry.size: seg = np.concatenate([carry, seg]); carry = np.empty(0, dtype=np.uint16) yield seg prev = i + 1 tail = np.asarray(a[prev:]) if tail.size: carry = np.concatenate([carry, tail]) if carry.size else tail.copy() if carry.size: yield carry def main(): ap = argparse.ArgumentParser() ap.add_argument("--out", required=True, help="output directory") ap.add_argument("--src", action="append", required=True, metavar="BUCKET=WEIGHT", help="repeatable, weights are TOKEN shares") ap.add_argument("--data-root", default="Data_Loader/Data") ap.add_argument("--total-tokens", type=float, required=True) ap.add_argument("--shard-tokens", type=float, default=96e6) ap.add_argument("--window", type=int, default=65536, help="CP_T") ap.add_argument("--grad-accum", type=int, default=8) ap.add_argument("--split", default="train") a = ap.parse_args() srcs = [] for s in a.src: name, _, w = s.partition("=") srcs.append((name, float(w))) tw = sum(w for _, w in srcs) srcs = [(n, w / tw) for n, w in srcs] total = int(a.total_tokens) shard_tokens = int(a.shard_tokens) os.makedirs(a.out, exist_ok=True) gens, want = {}, {} print(" sources:") for name, w in srcs: d = os.path.join(a.data_root, name) paths = sorted( os.path.join(d, f) for f in os.listdir(d) if f.endswith(".bin") and a.split in f) if not paths: sys.exit(f"no {a.split} shards in {d}") avail = sum(os.path.getsize(p) for p in paths) // 2 need = int(total * w) print(f" {name:<22} weight {w:5.1%} need {need/1e6:8.1f} M " f"have {avail/1e6:8.1f} M {'OK' if avail >= need else '*** SHORT ***'}") if avail < need: print(f" (will wrap within {name}; reduce its weight or --total-tokens)") gens[name] = doc_stream(paths) want[name] = w got = {n: 0 for n, _ in srcs} written = 0 shard_idx = 0 buf = [] buf_n = 0 exhausted = set() def flush(): nonlocal buf, buf_n, shard_idx if not buf: return arr = np.concatenate(buf) p = os.path.join(a.out, f"{a.split}_shard_{shard_idx:05d}.bin") arr.tofile(p) print(f" wrote {os.path.basename(p)} {arr.size/1e6:7.2f} M tokens") shard_idx += 1 buf = [] buf_n = 0 print("\n building:") while written < total and len(exhausted) < len(srcs): # the bucket furthest below its target share goes next pick, deficit = None, None for n, w in srcs: if n in exhausted: continue d = w - (got[n] / written if written else 0.0) if deficit is None or d > deficit: pick, deficit = n, d if pick is None: break try: doc = next(gens[pick]) except StopIteration: exhausted.add(pick) print(f" {pick} exhausted at {got[pick]/1e6:.1f} M tokens") continue buf.append(doc) buf_n += doc.size got[pick] += doc.size written += doc.size if buf_n >= shard_tokens: flush() flush() # --- emulate DataLoaderLite so STEPS is right --------------------------- BT = a.window # B=1 paths = sorted(os.path.join(a.out, f) for f in os.listdir(a.out) if f.endswith(".bin")) batches = 0 for p in paths: n = os.path.getsize(p) // 2 pos = 0 while pos + BT + 1 <= n: batches += 1 pos += BT if pos + BT + 1 > n: break steps = batches // a.grad_accum print("\n MIX SUMMARY") for n, w in srcs: print(f" {n:<22} target {w:5.1%} actual {got[n]/written:5.1%} " f"({got[n]/1e6:.1f} M tokens)") print(f" total {written/1e6:.1f} M tokens in {len(paths)} shards") print(f"\n USABLE STEPS at window {BT:,} x grad_accum {a.grad_accum}: {steps}") print(f" (naive total//global_batch would say " f"{written // (BT * a.grad_accum)} -- use {steps}, it accounts for the") print(f" per-shard tail the loader discards, and for the silent wrap)") with open(os.path.join(a.out, "MANIFEST.txt"), "w") as f: f.write(f"tokens {written}\nshards {len(paths)}\nbatches {batches}\n" f"window {BT}\ngrad_accum {a.grad_accum}\nusable_steps {steps}\n") for n, w in srcs: f.write(f"mix {n} target {w:.4f} actual {got[n]/written:.4f} tokens {got[n]}\n") print(f" wrote {a.out}/MANIFEST.txt") if __name__ == "__main__": main()