Upload a100_train.py with huggingface_hub

a100_train.py

@@ -0,0 +1,94 @@
+"""
+clankerDiffusion — single 80GB A100 training launcher.
+
+Target: 0.8B model (~08b config), 8k train context with NTK RoPE scaling
+(rope_scale=16) so it extends to 128k at inference. Trains up to --hours
+(12h/session, interruptible) and RESUMES from the latest checkpoint, so you
+can chain several 12h sessions to reach the 20h budget.
+
+Run on the A100 box (after `pip install torch ... numpy tokenizers datasets
+safetensors huggingface_hub` and `export HF_TOKEN=...`):
+    python a100_train.py --hours 12 --size 08b --scale 4
+
+Data is regenerated in-container by prep.py (fast HF egress); the scale is
+chosen so the 0.8B model sees enough tokens within the budget.
+"""
+import os
+import subprocess
+import argparse
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATA_DIR = os.path.join(HERE, "data")
+CKPT_DIR = os.path.join(HERE, "checkpoints")
+
+# 0.8B config (also in modal_train.py SIZES["08b"]).
+SIZES = {
+    "08b": dict(d_model=1536, n_layers=20, n_heads=16, d_ff=4096, batch=12),
+    "base": dict(d_model=768, n_layers=12, n_heads=12, d_ff=2048, batch=24),
+    "large": dict(d_model=2048, n_layers=24, n_heads=16, d_ff=5504, batch=8),
+}
+
+
+def main(hours=12.0, size="08b", scale=4.0, batch=None, ckpt_every=250,
+         seq_len=8192, rope_scale=16.0):
+    os.makedirs(DATA_DIR, exist_ok=True)
+    os.makedirs(CKPT_DIR, exist_ok=True)
+
+    # pull latest code so cloud + local stay in sync (optional; needs HF_TOKEN)
+    tok_path = os.path.join(DATA_DIR, "tokenizer.json")
+    if not os.path.exists(tok_path) and os.environ.get("HF_TOKEN"):
+        try:
+            from huggingface_hub import hf_hub_download
+            hf_hub_download(repo_id="coderofpears/clankerDiffusion-base",
+                            filename="data/tokenizer.json", repo_type="model",
+                            local_dir=DATA_DIR, token=os.environ.get("HF_TOKEN"))
+        except Exception as e:
+            print(f"[a100] tokenizer fetch skipped: {e}")
+
+    # (re)generate data if incomplete
+    meta_path = os.path.join(DATA_DIR, "meta.json")
+    if not os.path.exists(meta_path):
+        stale = os.path.join(DATA_DIR, "train.bin")
+        if os.path.exists(stale):
+            os.remove(stale)
+        print(f"[a100] generating data (scale={scale}) ...")
+        subprocess.run(["python", "prep.py", "--out-dir", DATA_DIR,
+                         "--scale", str(scale)], check=True)
+
+    spec = dict(SIZES.get(size, SIZES["08b"]))
+    if batch:
+        spec["batch"] = batch
+
+    cmd = [
+        "python", "train.py",
+        "--hours", str(hours),
+        "--batch", str(spec["batch"]),
+        "--ckpt-every", str(ckpt_every),
+        "--data-dir", DATA_DIR,
+        "--ckpt-dir", CKPT_DIR,
+        "--hf-repo", "coderofpears/clankerDiffusion-checkpoints",
+        "--d-model", str(spec["d_model"]),
+        "--n-layers", str(spec["n_layers"]),
+        "--n-heads", str(spec["n_heads"]),
+        "--d-ff", str(spec["d_ff"]),
+        "--seq-len", str(seq_len),
+        "--rope-scale", str(rope_scale),
+    ]
+    print("[a100] launching:", " ".join(cmd), flush=True)
+    subprocess.run(cmd, check=True)
+    print("[a100] session finished (resume next session with same command)",
+          flush=True)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--hours", type=float, default=12.0)
+    ap.add_argument("--size", default="08b")
+    ap.add_argument("--scale", type=float, default=4.0)
+    ap.add_argument("--batch", type=int, default=None)
+    ap.add_argument("--ckpt-every", type=int, default=250)
+    ap.add_argument("--seq-len", type=int, default=8192)
+    ap.add_argument("--rope-scale", type=float, default=16.0)
+    a = ap.parse_args()
+    main(hours=a.hours, size=a.size, scale=a.scale, batch=a.batch,
+         ckpt_every=a.ckpt_every, seq_len=a.seq_len, rope_scale=a.rope_scale)