| """ |
| clankerDiffusion — Modal L4 training. |
| |
| Builds a CUDA image, pulls the latest code from the HF code repo at runtime, |
| 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 |
| |
| Data is generated in-container by prep.py (fast HF egress), scaled by --scale, |
| so no large file transfer is needed. |
| |
| Launch (after `modal token new` on this machine): |
| modal run modal_train.py::train_on_l4 --hours 20 --size large --scale 3 |
| """ |
| import os |
| from modal import App, Image, Volume, Secret |
|
|
| 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") |
| .pip_install("numpy", "tokenizers", "datasets", "safetensors", |
| "huggingface_hub", "hf-transfer") |
| .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "TOKENIZERS_PARALLELISM": "false"}) |
| ) |
|
|
| 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), |
| |
| "08b": dict(d_model=1536, n_layers=20, n_heads=16, d_ff=4096, batch=12, |
| rope_scale=16.0), |
| } |
|
|
|
|
| @app.function( |
| image=image, |
| gpu="L4", |
| timeout=82800, |
| volumes={"/vol": volume}, |
| secrets=[Secret.from_dotenv(DOTENV)], |
| ) |
| def train_on_l4(hours: float = 20.0, ckpt_every: int = 250, |
| size: str = "large", batch: int = None, scale: float = 3.0): |
| import subprocess |
| import sys |
| import shutil |
|
|
| |
| from huggingface_hub import snapshot_download |
| os.makedirs("/root/clanker", exist_ok=True) |
| snapshot_download(repo_id=CODE_REPO, repo_type="model", |
| local_dir="/root/clanker", |
| token=os.environ.get("HF_TOKEN")) |
| 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) |
|
|
| |
| for name in ("tokenizer.json", "meta.json", "rag_corpus.txt"): |
| src = os.path.join("/root/clanker/data", name) |
| if os.path.exists(src): |
| shutil.copy(src, os.path.join(data_dir, name)) |
|
|
| os.makedirs(data_dir, exist_ok=True) |
| os.makedirs(ckpt_dir, exist_ok=True) |
|
|
| |
| |
| |
| 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): |
| print("[modal] found partial train.bin without meta.json; discarding.") |
| os.remove(stale) |
| print(f"[modal] generating data locally (scale={scale}) ...") |
| subprocess.run([sys.executable, "prep.py", "--out-dir", data_dir, |
| "--scale", str(scale)], check=True) |
|
|
| 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"]), |
| ] |
| if "rope_scale" in spec: |
| cmd += ["--rope-scale", str(spec["rope_scale"])] |
| 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, scale: float = 3.0): |
| train_on_l4.remote(hours=hours, ckpt_every=ckpt_every, size=size, |
| batch=batch, scale=scale) |
|
|
