File size: 3,822 Bytes
03903f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """
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)
|