Upload scripts/mix_buckets.py with huggingface_hub

scripts/mix_buckets.py

@@ -0,0 +1,190 @@
+#!/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()