prepare data on a CPU job, train from the prepared files

v2/train.py

@@ -106,14 +106,34 @@ def enron_text(batch):
     return {"text": texts}
 
 
-def load_sources(args) -> dict:
-    """Downloads whole parquet files and cleans them on every core."""
+# Byte-level BPE with an 8K vocabulary averages a little over 4 characters per token on English. Reading 40% more
+# text than that covers the variance, and anything short gets repeated (and logged) in flatten().
+CHARS_PER_TOKEN = 4.0
+TEXT_MARGIN = 1.4
+
+
+def enough_rows(ds, budget_tokens: int, lengths: np.ndarray, seed: int):
+    """A random subset of rows whose text adds up to roughly the token budget, so only that much gets cleaned."""
+    order = np.random.default_rng(seed).permutation(len(ds))
+    needed = np.searchsorted(np.cumsum(lengths[order]), budget_tokens * CHARS_PER_TOKEN * TEXT_MARGIN) + 1
+    return ds.select(np.sort(order[:needed]))
+
+
+def load_sources(args, budgets: dict) -> dict:
+    """Downloads parquet files, keeps only as many documents as each source's budget needs, and cleans them on
+    every core."""
+    import pyarrow.compute as pc
     from datasets import load_dataset
 
     procs = cpus()
     fineweb = load_dataset("parquet", data_files=[FINEWEB.format(i) for i in range(args.fineweb_files)], split="train")
-    fineweb = fineweb.select_columns(["text"]).map(lambda b: {"text": [to_ascii(t) for t in b["text"]]}, batched=True, num_proc=procs)
+    fineweb = fineweb.select_columns(["text"])
+    fineweb = enough_rows(fineweb, budgets["fineweb"], pc.utf8_length(fineweb.data.column("text")).to_numpy(), seed=0)
+    fineweb = fineweb.map(lambda b: {"text": [to_ascii(t) for t in b["text"]]}, batched=True, num_proc=procs)
     smoltalk = load_dataset("parquet", data_files=[SMOLTALK.format(i) for i in range(args.smoltalk_files)], split="train")
+    sample = smoltalk.shuffle(seed=0).select(range(min(2000, len(smoltalk))))
+    average = np.mean([sum(len(m["content"]) for m in row["messages"]) for row in sample])
+    smoltalk = enough_rows(smoltalk, budgets["smoltalk"], np.full(len(smoltalk), average), seed=1)
     smoltalk = smoltalk.map(smoltalk_text, batched=True, num_proc=procs, remove_columns=smoltalk.column_names)
     enron = load_dataset("parquet", data_files=[ENRON.format(i) for i in range(3)], split="train")
     enron = enron.map(enron_text, batched=True, num_proc=procs, remove_columns=enron.column_names)
@@ -449,6 +469,55 @@ def headline(scored: dict) -> dict:
     return out
 
 
+# ---- Prepared data ---------------------------------------------------------------------------------------------
+
+
+def data_key(args) -> str:
+    """Prepared data depends only on these settings, so a training run can find the matching files."""
+    return f"{args.name}-v{args.vocab}-s{args.seq}-t{int(args.tokens / 1e6)}M-fim{args.fim_rate}-{args.mix}".replace(",", "_").replace("=", "")
+
+
+def prepare(args, weights: dict, output: Path):
+    budgets = {name: int((args.tokens + args.val_tokens) * weight) for name, weight in weights.items()}
+    started = time.time()
+    sources = load_sources(args, budgets)
+    print(f"sources ready in {time.time() - started:.0f}s: " + ", ".join(f"{n} {len(d):,} docs" for n, d in sources.items()), flush=True)
+    tokenizer = train_tokenizer(sources, weights, args.vocab, args.tokenizer_docs)
+    tokenizer.save(str(output / "tokenizer.json"))
+    rows, val_rows, counts = build_rows(sources, tokenizer, weights, args)
+    return tokenizer, rows, val_rows, counts
+
+
+def upload_prepared(args, output: Path, rows: np.ndarray, val_rows: np.ndarray, counts: dict, minutes: float):
+    from huggingface_hub import HfApi
+
+    folder = output / "data"
+    folder.mkdir(exist_ok=True)
+    rows.astype(np.uint16).tofile(folder / "train.bin")
+    val_rows.astype(np.uint16).tofile(folder / "val.bin")
+    shutil.copy(output / "tokenizer.json", folder / "tokenizer.json")
+    meta = {"seq": args.seq, "train_rows": len(rows), "val_rows": len(val_rows), "counts": counts, "minutes": round(minutes, 1), "args": vars(args)}
+    (folder / "meta.json").write_text(json.dumps(meta, indent=1))
+    api = HfApi(token=os.environ.get("HF_TOKEN"))
+    api.create_repo(args.data_repo, repo_type="dataset", private=True, exist_ok=True)
+    api.upload_folder(folder_path=str(folder), repo_id=args.data_repo, repo_type="dataset", path_in_repo=data_key(args),
+                      commit_message=f"prepared data: {data_key(args)}")
+    print(f"Prepared data at https://huggingface.co/datasets/{args.data_repo}/tree/main/{data_key(args)}", flush=True)
+
+
+def download_prepared(args, output: Path):
+    from huggingface_hub import snapshot_download
+    from tokenizers import Tokenizer
+
+    local = Path(snapshot_download(args.data_repo, repo_type="dataset", allow_patterns=[f"{data_key(args)}/*"], max_workers=16))
+    folder = local / data_key(args)
+    meta = json.loads((folder / "meta.json").read_text())
+    shutil.copy(folder / "tokenizer.json", output / "tokenizer.json")
+    rows = np.fromfile(folder / "train.bin", dtype=np.uint16).reshape(meta["train_rows"], meta["seq"] + 1)
+    val_rows = np.fromfile(folder / "val.bin", dtype=np.uint16).reshape(meta["val_rows"], meta["seq"] + 1)
+    return Tokenizer.from_file(str(output / "tokenizer.json")), rows, val_rows, meta["counts"]
+
+
 # ---- Main ------------------------------------------------------------------------------------------------------
 
 
@@ -476,6 +545,10 @@ def parse_args():
     parser.add_argument("--folder", default="v2")
     parser.add_argument("--name", default="tinynext-v2")
     parser.add_argument("--trackio-space", default="cfahlgren1/tinynext-trackio")
+    parser.add_argument("--stage", choices=("all", "prepare", "train"), default="all",
+                        help="prepare: build the tokenizer and token arrays on a CPU job and save them to --data-repo. "
+                             "train: download them and train on the GPU. all: both in one job.")
+    parser.add_argument("--data-repo", default="cfahlgren1/tinynext-data")
     parser.add_argument("--smoke", action="store_true", help="Tiny settings to check the pipeline end to end")
     parser.add_argument("--no-push", action="store_true")
     args = parser.parse_args()
@@ -518,13 +591,15 @@ def main():
     tracker = start_tracker(args)
 
     started = time.time()
-    sources = load_sources(args)
-    print(f"sources ready in {time.time() - started:.0f}s: " + ", ".join(f"{n} {len(d):,} docs" for n, d in sources.items()), flush=True)
-    tokenizer = train_tokenizer(sources, weights, args.vocab, args.tokenizer_docs)
-    tokenizer.save(str(output / "tokenizer.json"))
-    rows, val_rows, counts = build_rows(sources, tokenizer, weights, args)
+    if args.stage == "train":
+        tokenizer, rows, val_rows, counts = download_prepared(args, output)
+    else:
+        tokenizer, rows, val_rows, counts = prepare(args, weights, output)
     data_minutes = (time.time() - started) / 60
     print(f"data ready in {data_minutes:.1f} min: {rows.size / 1e6:.0f}M train tokens, {val_rows.size / 1e6:.1f}M val", flush=True)
+    if args.stage == "prepare":
+        upload_prepared(args, output, rows, val_rows, counts, data_minutes)
+        return
 
     # The whole token array lives on the GPU; each step just indexes into it.
     rows_gpu = torch.from_numpy(rows.astype(np.int32)).to(device)