#!/usr/bin/env python3 """ in2_gen.py -- synthesise IN2-style TRAINING shards for bluscriptCP. IN2 (arXiv:2404.16811) teaches long-range retrieval by making the answer impossible to produce without a fact planted at a random distant position. We cannot reproduce its "loss on the answer only" masking (see the plan: TI has no maskable CE, `narrow` has no autograd node, and sentinel target ids are unsafe because the fwd/bwd kernels do not bounds-check). So we remove the NEED to mask: * our filler is CONTIGUOUS REAL TEXT, not IN2's shuffled-segment bag, so unmasked LM loss on it is ordinary harmless LM training; and * dilution is fixed by construction -- K triples per example raise the answer tokens from ~0.1% of positions to several percent. Output is what DataLoaderLite already eats: a HEADERLESS little-endian uint16 token stream in a file whose name contains "train"/"val", under a directory passed as CP_DATA_ROOT. No C++ training/loss/loader change is required. LAYOUT (one example = exactly T tokens) [ filler, with needles and decoys planted throughout ] ^needle_i (key_i + val_i) ^decoy (key + val, NO probe) ... and at each site p_i: [ key_i ][ val_i ] probe answer * SITES p_i are drawn from an ARM-INDEPENDENT rng stream, so arms B and C have bit-identical answer-position histograms and differ ONLY in needle distance. Without this we would difference NLLs measured at different points on a position-varying curve -- the error needle_gen.py's fixed-slot design exists to prevent. * Sites are INTERLEAVED over the sequence, never concatenated at the tail: tail-packing would confine every supervised position to the last K*SPAN offsets and train a position-specific circuit indistinguishable from retrieval. * PAIRING is shuffled and DECOY needles (no probe) are planted, so "the i-th probe pairs with the i-th needle" is not a winning heuristic. That shortcut is order/counting, not content retrieval, and would earn full reward at any distance. * KEY ALPHABET is capped (default 512, not the full 50304). Uniform draws from the vocab cost ln(50304)=10.83 nats IRREDUCIBLY -- pure noise that inflates the pre-clip gradient norm and, against the GLOBAL clip_grad_norm_(1.0), could silently leave the arms not LR-matched. ln(512)=6.24 nats gives the same discriminability (512^4 = 6.9e10 keys) with 42% less noise. ARMS a plain text, no injection -- drift control b needle immediately before its probe -- the objective change at distance ~0 c needle uniform in [0, p_i - SPAN) -- DISTANCE, the actual IN2 claim Stdlib only, matching needle_gen.py (the local pyenv python3 lacks numpy). Usage python3 in2_gen.py --arm c --T 16384 --examples 64 --split train --out-dir /tmp/in2c python3 in2_gen.py --arm b --T 16384 --examples 64 --split train --out-dir /tmp/in2b """ import argparse, array, json, os, random, sys VOCAB = 50304 # must match cfg.vocab_size; the C++ side bound-checks too PAD_TAIL = 4 # >=2: the loader's y-shift makes the last position of each # example target the NEXT example's first token. Keep the # answer clear of it. MIX = 1_000_003 # odd multiplier for deterministic seed derivation MAX_SHARD_TOKENS = 100_000_000 # DataLoaderLite is constructed with this cap; # anything past it in a shard is SILENTLY invisible def check_platform(): if sys.byteorder != "little": raise SystemExit("training shards must be uint16 LITTLE-endian; this host is big-endian") if array.array("H").itemsize != 2: raise SystemExit("unexpected array itemsize ('H' must be 2B)") def load_filler(path, need, rng): """One CONTIGUOUS run of real tokens per example, so filler keeps local textual statistics. Source shards are already uint16, so unlike needle_gen.py (which must widen to int32 for CP_EVAL_TOKENS_BIN) we keep the native width -- training shards are uint16.""" ntok = os.path.getsize(path) // 2 if ntok < need + 1: raise SystemExit(f"{path} has {ntok} tokens, need >= {need + 1}") start = rng.randrange(0, ntok - need) buf = array.array("H") with open(path, "rb") as f: f.seek(start * 2) buf.fromfile(f, need) lo, hi = min(buf), max(buf) if lo < 0 or hi >= VOCAB: raise SystemExit(f"{path} has ids outside [0,{VOCAB}): min={lo} max={hi}") return buf def draw_sites(T, n, span, rng, lo_margin): """n site starts in [lo_margin, T-PAD_TAIL-span], separated by >= 2*span. Drawn from an ARM-INDEPENDENT rng so every arm sees the same sites. Returned sorted; the pairing to needles is shuffled separately. Separation is 2*span, not span, and the reason is load-bearing: arm B plants its needle immediately before the probe, at p-span. At only span separation, two adjacent sites can sit exactly span apart and B's needle then lands on top of the earlier site, so B silently drops that triple -- which would make B's answer_start histogram DIFFER from C's and destroy the matched-position discipline the B-vs-C contrast depends on. The constraint is applied for all arms so sites stay arm-independent. """ hi = T - PAD_TAIL - span if hi <= lo_margin: raise SystemExit(f"T={T} too small for span={span}") sep = 2 * span chosen = [] guard = 0 while len(chosen) < n: guard += 1 if guard > 400 * n: raise SystemExit(f"could not place {n} sites in T={T} at separation {sep}; " f"lower --K (need roughly K*{sep} << {hi - lo_margin})") p = rng.randrange(lo_margin, hi + 1) if all(abs(p - q) >= sep for q in chosen): chosen.append(p) return sorted(chosen) def place_free(occupied, T, span, rng, hi_bound): """A start in [0, hi_bound-span] disjoint from `occupied`; None if it fails.""" for _ in range(400): if hi_bound - span < 0: return None s = rng.randrange(0, hi_bound - span + 1) if all(s + span <= a or s >= a + span for a, _ in occupied): return s return None def make_example(arm, T, K, key_len, val_len, alphabet, decoys, d_min, filler, site_rng, arm_rng): """Returns (tokens uint16 array of length T, list of records). site_rng -- ARM-INDEPENDENT: sites, keys, values, pairing. arm_rng -- arm-specific: needle placement only. """ span = key_len + val_len toks = filler # fresh copy per example recs = [] if arm == "a": # plain text, no injection return toks, recs, 0 sites = draw_sites(T, K, span, site_rng, lo_margin=span + d_min) # Distinct key/value pairs, then a shuffled site<->needle pairing so the # i-th probe does not correspond to the i-th needle by position. keys = [[site_rng.randrange(alphabet) for _ in range(key_len)] for _ in range(K)] vals = [[site_rng.randrange(alphabet) for _ in range(val_len)] for _ in range(K)] order = list(range(K)) site_rng.shuffle(order) occupied = [(p, "site") for p in sites] for si, p in enumerate(sites): j = order[si] # which key/value this site queries key, val = keys[j], vals[j] toks[p:p + key_len] = array.array("H", key) # probe toks[p + key_len:p + span] = array.array("H", val) # answer if arm == "b": ns = p - span # immediately before the probe # Guaranteed free by the 2*span site separation; assert rather than # skip, because a silent skip would desynchronise B's histogram from C's. if ns < 0 or any(ns + span > a and ns < a + span for a, _ in occupied if a != p): raise SystemExit(f"arm b: needle slot {ns} for site {p} is occupied -- " "site separation invariant violated") else: # arm c: uniform, far hi = p - d_min ns = place_free(occupied, T, span, arm_rng, hi) if ns is None: raise SystemExit(f"arm c: could not place a needle before site {p} " f"(d_min={d_min}); lower --K or --d-min") toks[ns:ns + span] = array.array("H", key + val) occupied.append((ns, "needle")) answer_start = p + key_len recs.append(dict(needle_start=ns, probe_start=p, answer_start=answer_start, distance=p - ns, csv_first=answer_start - 1)) # see off-by-one note below n_decoy = 0 # Decoy needles: no probe queries them, so needle count != probe count and # the i-th-of-K index carries no information. for _ in range(decoys): key = [site_rng.randrange(alphabet) for _ in range(key_len)] val = [site_rng.randrange(alphabet) for _ in range(val_len)] ds = place_free(occupied, T, span, arm_rng, T - PAD_TAIL) if ds is None: continue toks[ds:ds + span] = array.array("H", key + val) occupied.append((ds, "decoy")) n_decoy += 1 return toks, recs, n_decoy def main(): ap = argparse.ArgumentParser() ap.add_argument("--arm", required=True, choices=["a", "b", "c"]) ap.add_argument("--T", type=int, required=True) ap.add_argument("--examples", type=int, required=True) ap.add_argument("--split", default="train", choices=["train", "val"]) ap.add_argument("--K", type=int, default=64, help="triples per example") ap.add_argument("--key-len", type=int, default=4, help="probe/key tokens") ap.add_argument("--val-len", type=int, default=8, help="answer/value tokens") ap.add_argument("--alphabet", type=int, default=512, help="key/value alphabet size") ap.add_argument("--decoys", type=int, default=None, help="default K//4") ap.add_argument("--d-min", type=int, default=64, help="arm c: min needle->probe distance") ap.add_argument("--filler-shards", nargs="+", required=True) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--out-dir", required=True) ap.add_argument("--shards", type=int, default=1, help="split the examples across this many files. REQUIRED above " "100M tokens: DataLoaderLite is constructed with a " "max_tokens_per_shard of 1e8 and anything past it in a single " "shard is SILENTLY invisible. The example index stays GLOBAL " "across shards, so the arm-independent RNG stream (sites, keys, " "pairing, filler) is unchanged by how the split falls -- arms " "remain comparable.") args = ap.parse_args() check_platform() if args.alphabet > VOCAB: raise SystemExit(f"--alphabet {args.alphabet} exceeds VOCAB {VOCAB}") decoys = args.decoys if args.decoys is not None else args.K // 4 span = args.key_len + args.val_len for f in args.filler_shards: if not os.path.exists(f): raise SystemExit(f"filler shard not found: {f}") if args.split == "train" and "val" in os.path.basename(f): raise SystemExit(f"refusing to build TRAIN data from a val shard: {f}\n" " eval filler must stay disjoint from training data") os.makedirs(args.out_dir, exist_ok=True) total_tok = args.examples * args.T per_shard = (args.examples + args.shards - 1) // args.shards if per_shard * args.T > MAX_SHARD_TOKENS: raise SystemExit(f"{per_shard * args.T} tokens/shard exceeds the loader's " f"{MAX_SHARD_TOKENS} cap; raise --shards to at least " f"{(total_tok + MAX_SHARD_TOKENS - 1) // MAX_SHARD_TOKENS}") all_recs, hist, n_decoy_total = [], {}, 0 shard_paths = [] stream = array.array("H") shard_i = 0 for i in range(args.examples): # Site/key/value/pairing rng is ARM-INDEPENDENT -- identical across arms. site_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX + 7) # Needle placement is arm-specific. arm_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX + 1000 + ord(args.arm)) # Filler draw is arm-independent too, so arms see the same documents. fill_rng = random.Random(((args.seed * MIX + args.T) * MIX + i) * MIX + 99) shard = args.filler_shards[i % len(args.filler_shards)] filler = load_filler(shard, args.T, fill_rng) toks, recs, n_dec = make_example(args.arm, args.T, args.K, args.key_len, args.val_len, args.alphabet, decoys, args.d_min, filler, site_rng, arm_rng) if len(toks) != args.T: raise SystemExit(f"example {i} is {len(toks)} tokens, expected {args.T}") stream.extend(toks) for r in recs: hist[r["answer_start"]] = hist.get(r["answer_start"], 0) + 1 all_recs.append(recs) n_decoy_total += n_dec # Flush a shard whenever it is full, or on the last example. The global `i` # above drives every RNG, so the split point cannot change the data. if len(stream) // args.T >= per_shard or i == args.examples - 1: sp = os.path.join(args.out_dir, f"in2{args.arm}_{args.split}_{shard_i:05d}.bin") with open(sp, "wb") as f: stream.tofile(f) shard_paths.append(sp) stream = array.array("H") shard_i += 1 path = shard_paths[0] n_ans = sum(len(r) for r in all_recs) * args.val_len # Counted from ACTUAL placements, never derived from parameters: arm A plants # nothing, so a parameter-derived figure reported phantom decoy tokens for it. n_syn = sum(len(r) for r in all_recs) * (2 * span) + n_decoy_total * span man = dict(arm=args.arm, T=args.T, examples=args.examples, split=args.split, K=args.K, key_len=args.key_len, val_len=args.val_len, span=span, alphabet=args.alphabet, decoys=decoys, d_min=args.d_min, seed=args.seed, vocab=VOCAB, pad_tail=PAD_TAIL, filler_shards=args.filler_shards, bin=path, shards=shard_paths, total_tokens=total_tok, dtype="uint16", answer_tokens=n_ans, answer_frac=n_ans / total_tok, synthetic_tokens=n_syn, synthetic_frac=n_syn / total_tok, decoys_placed=n_decoy_total, answer_start_hist=dict(sorted(hist.items())), records=all_recs) with open(path.replace(".bin", ".json"), "w") as f: json.dump(man, f, indent=1) print(f"wrote {len(shard_paths)} shard(s) in {args.out_dir} {total_tok} tokens total") print(f" answer tokens {n_ans} ({n_ans/total_tok:.2%}) " f"synthetic {n_syn} ({n_syn/total_tok:.2%})") if all_recs and all_recs[0]: ds = [r["distance"] for rs in all_recs for r in rs] ds.sort() print(f" distance: min {ds[0]} median {ds[len(ds)//2]} max {ds[-1]} n={len(ds)}") if __name__ == "__main__": main()