| """ |
| 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") |
|
|
| |
| 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) |
|
|
| |
| 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}") |
|
|
| |
| 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) |
|
|
