Upload folder using huggingface_hub

modal_train.py

@@ -0,0 +1,106 @@
+"""
+clankerDiffusion — Modal L4 training.
+
+Bakes the project code into the image, keeps data + checkpoints on a persistent
+Modal Volume (/vol), injects HF_TOKEN from .env, trains a big hybrid model on an
+L4 (24 GB), and pushes every checkpoint to the HF repo:
+
+    coderofpears/clankerDiffusion-checkpoints
+
+Launch (after `modal token new` on this machine):
+    modal run modal_train.py::train_on_l4 --hours 20 --size large
+
+Data: stage train.bin + tokenizer.json onto the volume once, e.g.
+    modal volume put clanker-vol C:\Users\User\CalcGPU\clankerDiffusion\data /data
+(or it is downloaded from the HF data repo if present there).
+"""
+import os
+from modal import App, Image, Volume, Secret, function
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+CODE_REPO = "coderofpears/clankerDiffusion-base"
+DATA_REPO = "coderofpears/clankerDiffusion-data"
+CKPT_REPO = "coderofpears/clankerDiffusion-checkpoints"
+DOTENV = os.path.join(HERE, ".env")
+
+image = (
+    Image.debian_slim()
+    .pip_install(
+        "torch==2.11.0", index_url="https://download.pytorch.org/whl/cu128",
+        "numpy", "tokenizers", "datasets", "safetensors",
+        "huggingface_hub", "accelerate", "hf-transfer",
+    )
+    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "TOKENIZERS_PARALLELISM": "false"})
+    .add_local_dir(
+        HERE, "/root/clanker",
+        ignore=["data", "checkpoints", ".venv", "__pycache__", "*.bin",
+                "*.pt", ".git", "*.log", "temp_up"],
+    )
+)
+
+app = App("clanker-diffusion", image=image)
+volume = Volume.from_name("clanker-vol", create_if_missing=True)
+
+SIZES = {
+    "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),
+    "xl":    dict(d_model=2560, n_layers=28, n_heads=20, d_ff=6912, batch=4),
+}
+
+
+@app.function(
+    image=image,
+    gpu="L4",
+    timeout=60 * 60 * 25,
+    volumes={"/vol": volume},
+    secrets=[Secret.from_dotenv(DOTENV)],
+    _allow_background_volume_commits=True,
+)
+def train_on_l4(hours: float = 20.0, ckpt_every: int = 250,
+                size: str = "large", batch: int = None):
+    import os
+    import subprocess
+    import sys
+
+    os.chdir("/root/clanker")
+    data_dir, ckpt_dir = "/vol/data", "/vol/checkpoints"
+    os.makedirs(data_dir, exist_ok=True)
+    os.makedirs(ckpt_dir, exist_ok=True)
+
+    # ---- ensure data ----
+    if not os.path.exists(os.path.join(data_dir, "train.bin")):
+        print("[modal] train.bin missing; trying HF download ...")
+        try:
+            subprocess.run(["hf", "download", DATA_REPO, "train.bin", "-d", data_dir],
+                           check=True)
+            subprocess.run(["hf", "download", DATA_REPO, "tokenizer.json", "-d", data_dir],
+                           check=True)
+        except Exception as e:
+            print(f"[modal] HF data download failed: {e}; bailing.")
+            raise
+
+    spec = dict(SIZES.get(size, SIZES["large"]))
+    if batch:
+        spec["batch"] = batch
+
+    cmd = [
+        sys.executable, "train.py",
+        "--hours", str(hours),
+        "--batch", str(spec["batch"]),
+        "--ckpt-every", str(ckpt_every),
+        "--data-dir", data_dir,
+        "--ckpt-dir", ckpt_dir,
+        "--hf-repo", CKPT_REPO,
+        "--d-model", str(spec["d_model"]),
+        "--n-layers", str(spec["n_layers"]),
+        "--n-heads", str(spec["n_heads"]),
+        "--d-ff", str(spec["d_ff"]),
+    ]
+    print("[modal] launching:", " ".join(cmd), flush=True)
+    subprocess.run(cmd, check=True)
+    print("[modal] training finished", flush=True)
+
+
+@app.local_entrypoint()
+def main(hours: float = 20.0, ckpt_every: int = 250, size: str = "large", batch: int = None):
+    train_on_l4.remote(hours=hours, ckpt_every=ckpt_every, size=size, batch=batch)

train.py

@@ -1,16 +1,19 @@
 """
-clankerDiffusion — local training loop (RTX 5060 Ti, ~5h budget).
+clankerDiffusion — training loop (hybrid AR / masked-diffusion).
 
 Hybrid objective (per step, mode chosen at random, p(AR)=0.5):
   AR   (mode 0): causal LM cross-entropy over the whole window.
-  DIFF (mode 1): MDLM absorbing-state masked diffusion — mask each token
+  DIFF (mode 1): MDLM absorbing-state masked diffusion -- mask each token
                  independently with ratio r~U(0,1); reconstruct masked tokens
-                 with bidirectional attention, conditioned on r via the time embed.
+                 with bidirectional attention, conditioned on r via time embed.
 
-Runs in bf16 (Blackwell), AdamW + cosine LR, grad-clip, checkpoints locally.
-Stops after --hours (default 5) of wall-clock time.
+Runs in bf16, AdamW + cosine LR, grad-clip, checkpoints locally and (optionally)
+pushes each checkpoint to a HuggingFace repo via the `hf` CLI.
+
+Used both locally and on Modal L4 (override --data-dir/--ckpt-dir/--hf-repo and
+the model dimensions for a bigger model).
 """
-import os, json, time, argparse
+import os, json, time, argparse, subprocess, threading
 import numpy as np
 import torch
 import torch.nn as nn
@@ -24,21 +27,38 @@ DATADIR = os.path.join(HERE, "data")
 CKPTDIR = os.path.join(HERE, "checkpoints")
 os.makedirs(CKPTDIR, exist_ok=True)
 
-# ---- architecture for 16 GB -------------------------------------------------
-CFG = dict(
+# default architecture tuned for 16 GB (RTX 5060 Ti)
+DEFAULT_CFG = dict(
     d_model=768, n_layers=12, n_heads=12, d_ff=2048,
     max_len=1024, vocab_size=32768,
 )
 
 
-def load_data():
-    meta = json.load(open(os.path.join(DATADIR, "meta.json")))
-    arr = np.memmap(os.path.join(DATADIR, "train.bin"), dtype=np.uint16, mode="r")
-    CFG["vocab_size"] = meta["vocab_size"]
-    CFG["max_len"] = meta["seq_len"]
-    print(f"[train] data n_tokens={meta['n_tokens']:,} seq_len={meta['seq_len']} "
-          f"vocab={meta['vocab_size']}")
-    return arr, meta["seq_len"]
+def build_cfg(args):
+    cfg = dict(DEFAULT_CFG)
+    for k in ("d_model", "n_layers", "n_heads", "d_ff", "vocab_size", "max_len"):
+        v = getattr(args, k, None)
+        if v is not None:
+            cfg[k] = v
+    return cfg
+
+
+def _push_hf(path, repo):
+    """Upload a single checkpoint file to HF via the `hf` CLI (background)."""
+    if not repo:
+        return
+    try:
+        subprocess.run(["hf", "upload", repo, path, "--repo-type", "model"],
+                       check=True, capture_output=True, timeout=600)
+        print(f"[hf] pushed {os.path.basename(path)} -> {repo}", flush=True)
+    except Exception as e:
+        print(f"[hf] push failed for {path}: {e}", flush=True)
+
+
+def load_data(data_dir):
+    meta = json.load(open(os.path.join(data_dir, "meta.json")))
+    arr = np.memmap(os.path.join(data_dir, "train.bin"), dtype=np.uint16, mode="r")
+    return arr, meta["seq_len"], meta["vocab_size"], meta["n_tokens"]
 
 
 def sample_batch(arr, seq_len, batch):
@@ -49,24 +69,33 @@ def sample_batch(arr, seq_len, batch):
 
 
 def train(args):
-    tok = YKTokenizer.load(os.path.join(DATADIR, "tokenizer.json"))
-    arr, seq_len = load_data()
-
-    model = YKDiff(CFG).cuda()
+    data_dir = args.data_dir or DATADIR
+    ckpt_dir = args.ckpt_dir or CKPTDIR
+    os.makedirs(ckpt_dir, exist_ok=True)
+    cfg = build_cfg(args)
+
+    tok = YKTokenizer.load(os.path.join(data_dir, "tokenizer.json"))
+    arr, seq_len, vocab, n_tokens = load_data(data_dir)
+    cfg["vocab_size"] = vocab
+    cfg["max_len"] = seq_len
+    print(f"[train] data n_tokens={n_tokens:,} seq_len={seq_len} vocab={vocab}")
+    print(f"[train] model params = {sum(p.numel() for p in YKDiff(cfg).parameters())/1e6:.1f}M")
+
+    model = YKDiff(cfg).cuda()
     n_params = sum(p.numel() for p in model.parameters())
-    print(f"[train] model params = {n_params/1e6:.1f}M")
+    print(f"[train] allocated params = {n_params/1e6:.1f}M")
 
     optim = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95),
                               weight_decay=0.1)
-    V = CFG["vocab_size"]
+    V = cfg["vocab_size"]
     pad_id = tok.pad_id
     mask_id = tok.mask_id
 
     # resume
     step0 = 0
-    ckpts = sorted([f for f in os.listdir(CKPTDIR) if f.endswith(".pt")])
+    ckpts = sorted([f for f in os.listdir(ckpt_dir) if f.endswith(".pt")])
     if ckpts and not args.fresh:
-        path = os.path.join(CKPTDIR, ckpts[-1])
+        path = os.path.join(ckpt_dir, ckpts[-1])
         sd = torch.load(path, map_location="cuda")
         model.load_state_dict(sd["model"]); optim.load_state_dict(sd["optim"])
         step0 = sd["step"]
@@ -104,7 +133,7 @@ def train(args):
                 masked = idx.clone(); masked[is_mask] = mask_id
                 logits = model(masked, m, t=r)
                 ce = F.cross_entropy(logits.reshape(-1, V), idx.reshape(-1),
-                                    reduction="none", ignore_index=-100)
+                                     reduction="none", ignore_index=-100)
                 ce = ce * is_mask.reshape(-1) * not_pad.reshape(-1)
                 denom = (is_mask & not_pad).reshape(-1).sum().clamp(min=1)
                 loss = ce.sum() / denom
@@ -121,16 +150,19 @@ def train(args):
                   f"t={(time.time()-t0)/60:.1f}m", flush=True)
 
         if step % args.ckpt_every == 0:
-            path = os.path.join(CKPTDIR, f"clanker_{step:07d}.pt")
+            path = os.path.join(ckpt_dir, f"clanker_{step:07d}.pt")
             torch.save({"model": model.state_dict(), "optim": optim.state_dict(),
-                        "step": step, "cfg": CFG, "vocab": V}, path)
+                        "step": step, "cfg": cfg, "vocab": V}, path)
             print(f"[train] checkpoint -> {path}", flush=True)
+            if args.hf_repo:
+                threading.Thread(target=_push_hf, args=(path, args.hf_repo),
+                                 daemon=True).start()
 
     # final save
-    path = os.path.join(CKPTDIR, f"clanker_{step:07d}_final.pt")
+    path = os.path.join(ckpt_dir, f"clanker_{step:07d}_final.pt")
     torch.save({"model": model.state_dict(), "optim": optim.state_dict(),
-                "step": step, "cfg": CFG, "vocab": V}, path)
-    json.dump(CFG, open(os.path.join(CKPTDIR, "config.json"), "w"))
+                "step": step, "cfg": cfg, "vocab": V}, path)
+    json.dump(cfg, open(os.path.join(ckpt_dir, "config.json"), "w"))
     print(f"[train] DONE final={path} steps={step}")
 
 
@@ -142,4 +174,15 @@ if __name__ == "__main__":
     ap.add_argument("--log-every", type=int, default=25)
     ap.add_argument("--ckpt-every", type=int, default=500)
     ap.add_argument("--fresh", action="store_true")
+    ap.add_argument("--data-dir", default=None)
+    ap.add_argument("--ckpt-dir", default=None)
+    ap.add_argument("--hf-repo", default=None,
+                    help="HuggingFace repo id to push checkpoints to (via `hf` CLI)")
+    # model overrides (for a bigger Modal model)
+    ap.add_argument("--d-model", type=int, default=None)
+    ap.add_argument("--n-layers", type=int, default=None)
+    ap.add_argument("--n-heads", type=int, default=None)
+    ap.add_argument("--d-ff", type=int, default=None)
+    ap.add_argument("--vocab-size", type=int, default=None)
+    ap.add_argument("--seq-len", type=int, default=None)
     train(ap.parse_args())