| |
| """ |
| 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): |
| |
| 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() |
|
|
| |
| BT = a.window |
| 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() |
|
|
