| """PC-SHO-DLM: LOCAL PC Training on A100 — Single mode, fast.""" |
| import json, os, sys, threading, time |
| import gradio as gr |
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader, Dataset |
|
|
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src")) |
| from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater, count_parameters |
|
|
| class CharDS(Dataset): |
| def __init__(self, text, seq_len): |
| self.seq_len = seq_len |
| self.data = torch.tensor([min(b+1,256) for b in text.encode("utf-8")], dtype=torch.long) |
| self.n = max(1, (len(self.data)-seq_len)//seq_len) |
| def __len__(self): return self.n |
| def __getitem__(self, i): |
| return {"input_ids": self.data[i*self.seq_len:(i+1)*self.seq_len]} |
|
|
| LOG = [] |
| STATUS = "Idle" |
|
|
| def train(): |
| global STATUS, LOG |
| LOG = ["Starting LOCAL PC training on A100..."] |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| LOG.append(f"Device: {device}") |
| if device == "cuda": |
| LOG.append(f"GPU: {torch.cuda.get_device_name()}") |
| LOG.append(f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB") |
|
|
| config = PCSHOConfig( |
| vocab_size=257, max_seq_len=512, d_model=512, n_heads=8, |
| n_layers=12, d_ff=2048, n_diffusion_steps=128, n_settling_steps=6, |
| mask_token_id=0, dropout=0.1, feedback_rank=128, |
| ) |
| LOG.append(f"Model: {count_parameters(PCSHODLM(config)):,} params") |
|
|
| STATUS = "Loading data..." |
| from datasets import load_dataset |
| ds = load_dataset("wikitext","wikitext-103-raw-v1",split="train") |
| dv = load_dataset("wikitext","wikitext-103-raw-v1",split="validation") |
| tt = "\n".join([r["text"] for r in ds if r["text"].strip()])[:100_000_000] |
| vt = "\n".join([r["text"] for r in dv if r["text"].strip()]) |
| train_ds, val_ds = CharDS(tt,config.max_seq_len), CharDS(vt,config.max_seq_len) |
| LOG.append(f"Data: {len(tt):,} train, {len(vt):,} val chars") |
|
|
| model = PCSHODLM(config).to(device) |
| model.train() |
| updater = LocalParameterUpdater(model, lr_forward=3e-4, lr_feedback=3e-4, lr_readout=3e-4, lr_precision=3e-5) |
| tl = DataLoader(train_ds, batch_size=64, shuffle=True, drop_last=True, num_workers=4, pin_memory=True) |
| vl = DataLoader(val_ds, batch_size=64, shuffle=False, num_workers=2, pin_memory=True) |
|
|
| STATUS = "Training LOCAL PC..." |
| LOG.append(f"\n{'='*60}\nMode: LOCAL PC | Params: {count_parameters(model):,}\n{'='*60}") |
| step, start, max_steps = 0, time.time(), 10000 |
| log_data = {"steps":[],"losses":[],"energies":[]} |
|
|
| while step < max_steps: |
| for batch in tl: |
| if step >= max_steps: break |
| x_0 = batch["input_ids"].to(device) |
| result = updater.step({"input_ids": x_0}) |
| step += 1 |
| loss = result.get("loss",0); energies = result.get("energies",[]) |
| if step % 100 == 0: |
| elapsed = time.time()-start |
| tps = step*64*config.max_seq_len/elapsed |
| e = f"{energies[-1]:.0f}" if energies else "N/A" |
| msg = f"[local] Step {step:6d} | Loss: {loss:.4f} | Energy: {e} | Tok/s: {tps:.0f} | {elapsed:.0f}s" |
| LOG.append(msg); log_data["steps"].append(step); log_data["losses"].append(loss); log_data["energies"].append(energies[-1] if energies else 0) |
| if step % 2000 == 0: |
| model.eval() |
| tl2,tt2=0.0,0 |
| with torch.no_grad(): |
| for vb in vl: |
| vx=vb["input_ids"].to(device); vo=model(vx); nm=vo["mask"].sum().item() |
| if nm>0: tl2+=vo["loss"].item()*nm; tt2+=nm |
| if tt2>100000: break |
| vl2=tl2/max(1,tt2); LOG.append(f" --> [local] Val loss: {vl2:.4f}"); log_data.setdefault("val",[]).append((step,vl2)) |
| model.train() |
|
|
| model.eval() |
| tl2,tt2=0.0,0 |
| with torch.no_grad(): |
| for vb in vl: |
| vx=vb["input_ids"].to(device); vo=model(vx); nm=vo["mask"].sum().item() |
| if nm>0: tl2+=vo["loss"].item()*nm; tt2+=nm |
| if tt2>200000: break |
| fv=tl2/max(1,tt2); log_data["final_val"]=fv |
| LOG.append(f"DONE: {step} steps in {time.time()-start:.0f}s | Final val: {fv:.4f}") |
| STATUS = "Complete!" |
|
|
| os.makedirs("results",exist_ok=True) |
| torch.save({"model":model.state_dict(),"config":config,"log":log_data},"results/local_pc_10k.pt") |
| with open("results/local_pc_log.json","w") as f: json.dump(log_data,f) |
| try: |
| from huggingface_hub import HfApi |
| HfApi().upload_folder(folder_path="results",repo_id="zotowata/pc-sho-dlm-local",repo_type="space",path_in_repo="results") |
| LOG.append("Results uploaded!") |
| except: pass |
|
|
| _thread = None |
| def start(): |
| global _thread |
| if _thread and _thread.is_alive(): return "Already running!" |
| _thread = threading.Thread(target=train, daemon=True); _thread.start() |
| return "LOCAL PC training started on A100!" |
|
|
| with gr.Blocks(title="PC-SHO-DLM Local PC") as demo: |
| gr.Markdown("# PC-SHO-DLM: LOCAL PC Training (A100)") |
| with gr.Row(): |
| btn = gr.Button("Start Training", variant="primary") |
| st = gr.Textbox(label="Status", value="Idle") |
| log = gr.Textbox(label="Log", lines=25, max_lines=50) |
| btn.click(start, outputs=st) |
| refresh = gr.Button("Refresh"); refresh.click(lambda: "\n".join(LOG[-50:]), outputs=log); refresh.click(lambda: STATUS, outputs=st) |
| timer = gr.Timer(5); timer.tick(lambda: "\n".join(LOG[-50:]), outputs=log); timer.tick(lambda: STATUS, outputs=st) |
|
|
| demo.launch() |
|
|
