Upload source/files/in2_gen.py with huggingface_hub

source/files/in2_gen.py

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