3fc9deac5b
- train_specialist_pt.py: pure PyTorch trainer replacing Fabricate binary. No SQLite/WAL — saves checkpoint.pt + history.json only. Fabricate was causing D-state IO blocking on the HDD due to continuous WAL writes. - eval_specialist_pt.py: PyTorch eval with confusion matrix, threshold analysis, guardrail assessment, and deployment verdict. - capacity_sweep.py: updated to call PyTorch scripts; TIERS now include heads field (4/6/8) since PyTorch supports multi-head unlike Fabricate. - prereq_op_binary_split.py: evaluate_combined rewritten in PyTorch. - Performance fix in train_specialist_pt.py: load training data onto GPU once and sample via torch.randint instead of DataLoader (eliminates Python/CPU overhead on tiny datasets). Fix applied, not yet timed. Next session: time the fix, then run sweep_all_gates.sh --pilot-only. See HANDOFF-2026-03-31.md for full context. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
329 lines
12 KiB
Python
329 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
train_specialist_pt.py
|
|
|
|
PyTorch replacement for run_grokking_until.py + whetstone_fabricate.
|
|
Trains a small transformer text classifier on TSV data.
|
|
No SQLite — saves checkpoint.pt and history.json only.
|
|
|
|
TSV format: label_idx<TAB>0<TAB>text
|
|
|
|
Usage:
|
|
python3 specialists/scripts/train_specialist_pt.py \
|
|
--dataset specialists/data/generated/worker_type_train.tsv \
|
|
--eval-dataset specialists/data/generated/worker_type_eval.tsv \
|
|
--out-dir /tmp/runs/worker_type_small \
|
|
--labels "implementer,reviewer,architect,qa" \
|
|
--hidden-dim 256 --layers 2 --heads 4 \
|
|
--lr 0.001 --weight-decay 0.01 \
|
|
--max-steps 3000 --batch-size 256 \
|
|
--history-out /tmp/runs/worker_type_small/history.json \
|
|
--reset
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import random
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
|
|
# ── Vocabulary ────────────────────────────────────────────────────────────────
|
|
|
|
class Vocab:
|
|
PAD, UNK = 0, 1
|
|
|
|
def __init__(self):
|
|
self.str_to_id = {"<pad>": 0, "<unk>": 1}
|
|
self.id_to_str = ["<pad>", "<unk>"]
|
|
|
|
def build(self, texts):
|
|
for text in texts:
|
|
for tok in _tokenize(text):
|
|
if tok not in self.str_to_id:
|
|
self.str_to_id[tok] = len(self.id_to_str)
|
|
self.id_to_str.append(tok)
|
|
|
|
def encode(self, text, seq_len):
|
|
ids = [self.str_to_id.get(t, self.UNK) for t in _tokenize(text)]
|
|
ids = ids[:seq_len]
|
|
ids += [self.PAD] * (seq_len - len(ids))
|
|
return ids
|
|
|
|
def __len__(self):
|
|
return len(self.id_to_str)
|
|
|
|
|
|
def _tokenize(text):
|
|
return text.lower().split()
|
|
|
|
|
|
# ── Dataset ───────────────────────────────────────────────────────────────────
|
|
|
|
class TSVDataset(Dataset):
|
|
def __init__(self, path, vocab, seq_len, label_list, build_vocab=False):
|
|
self.rows = []
|
|
texts = []
|
|
raw = []
|
|
with open(path) as f:
|
|
for line in f:
|
|
parts = line.strip().split("\t", 2)
|
|
if len(parts) < 3:
|
|
continue
|
|
raw.append((int(parts[0]), parts[2]))
|
|
texts.append(parts[2])
|
|
|
|
if build_vocab:
|
|
vocab.build(texts)
|
|
|
|
for label_idx, text in raw:
|
|
ids = vocab.encode(text, seq_len)
|
|
self.rows.append((torch.tensor(ids, dtype=torch.long), label_idx))
|
|
|
|
def __len__(self):
|
|
return len(self.rows)
|
|
|
|
def __getitem__(self, i):
|
|
return self.rows[i]
|
|
|
|
|
|
# ── Model ─────────────────────────────────────────────────────────────────────
|
|
|
|
class SpecialistModel(nn.Module):
|
|
def __init__(self, vocab_size, hidden_dim, layers, heads, n_labels, seq_len):
|
|
super().__init__()
|
|
self.seq_len = seq_len
|
|
self.tok_emb = nn.Embedding(vocab_size, hidden_dim, padding_idx=0)
|
|
self.pos_emb = nn.Embedding(seq_len, hidden_dim)
|
|
encoder_layer = nn.TransformerEncoderLayer(
|
|
d_model=hidden_dim,
|
|
nhead=heads,
|
|
dim_feedforward=hidden_dim * 4,
|
|
dropout=0.0,
|
|
batch_first=True,
|
|
)
|
|
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=layers)
|
|
self.head = nn.Linear(hidden_dim, n_labels)
|
|
self._init_weights()
|
|
|
|
def _init_weights(self):
|
|
for p in self.parameters():
|
|
if p.dim() > 1:
|
|
nn.init.xavier_uniform_(p)
|
|
|
|
def forward(self, x):
|
|
# x: (B, seq_len)
|
|
pad_mask = (x == 0) # True = ignore
|
|
pos = torch.arange(x.size(1), device=x.device).unsqueeze(0)
|
|
h = self.tok_emb(x) + self.pos_emb(pos)
|
|
h = self.encoder(h, src_key_padding_mask=pad_mask)
|
|
# Mean-pool over non-padding positions
|
|
lengths = (~pad_mask).float().sum(dim=1, keepdim=True).clamp(min=1)
|
|
h = (h * (~pad_mask).unsqueeze(-1).float()).sum(dim=1) / lengths
|
|
return self.head(h)
|
|
|
|
|
|
# ── Training ──────────────────────────────────────────────────────────────────
|
|
|
|
BLOCK = 1000 # steps per eval block (matches Fabricate convention)
|
|
|
|
|
|
def evaluate(model, loader, device):
|
|
model.eval()
|
|
correct = total = 0
|
|
loss_sum = 0.0
|
|
criterion = nn.CrossEntropyLoss()
|
|
with torch.no_grad():
|
|
for x, y in loader:
|
|
x, y = x.to(device), y.to(device)
|
|
logits = model(x)
|
|
loss_sum += criterion(logits, y).item() * len(y)
|
|
correct += (logits.argmax(1) == y).sum().item()
|
|
total += len(y)
|
|
return correct / total * 100.0, loss_sum / max(total, 1)
|
|
|
|
|
|
def train(args):
|
|
out_dir = Path(args.out_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
ckpt_path = out_dir / "checkpoint.pt"
|
|
history_path = Path(args.history_out) if args.history_out else out_dir / "history.json"
|
|
|
|
if args.reset:
|
|
for f in [ckpt_path, history_path]:
|
|
f.unlink(missing_ok=True)
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}", flush=True)
|
|
|
|
label_list = [l.strip() for l in args.labels.split(",")]
|
|
n_labels = len(label_list)
|
|
seq_len = args.seq_len
|
|
|
|
# ── Vocab + data ──
|
|
vocab = Vocab()
|
|
train_ds = TSVDataset(args.dataset, vocab, seq_len, label_list, build_vocab=True)
|
|
eval_ds = TSVDataset(args.eval_dataset, vocab, seq_len, label_list, build_vocab=False)
|
|
|
|
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
|
|
pin_memory=(device.type == "cuda"), num_workers=0)
|
|
eval_loader = DataLoader(eval_ds, batch_size=args.batch_size, shuffle=False,
|
|
pin_memory=(device.type == "cuda"), num_workers=0)
|
|
|
|
print(f"Train: {len(train_ds)} examples Eval: {len(eval_ds)} examples", flush=True)
|
|
print(f"Vocab: {len(vocab)} Labels: {label_list}", flush=True)
|
|
|
|
# ── Resume or fresh ──
|
|
start_step = 0
|
|
history_entries = []
|
|
best_acc = 0.0
|
|
|
|
model = SpecialistModel(
|
|
vocab_size=len(vocab),
|
|
hidden_dim=args.hidden_dim,
|
|
layers=args.layers,
|
|
heads=args.heads,
|
|
n_labels=n_labels,
|
|
seq_len=seq_len,
|
|
).to(device)
|
|
|
|
optimizer = torch.optim.AdamW(
|
|
model.parameters(), lr=args.lr, weight_decay=args.weight_decay
|
|
)
|
|
criterion = nn.CrossEntropyLoss()
|
|
|
|
if ckpt_path.exists() and not args.reset:
|
|
ckpt = torch.load(ckpt_path, map_location=device)
|
|
model.load_state_dict(ckpt["model"])
|
|
optimizer.load_state_dict(ckpt["optimizer"])
|
|
vocab = ckpt["vocab_obj"]
|
|
start_step = ckpt.get("step", 0)
|
|
if history_path.exists():
|
|
history_entries = json.loads(history_path.read_text()).get("history", [])
|
|
best_acc = max((e.get("overall", 0.0) for e in history_entries), default=0.0)
|
|
print(f"Resuming from step {start_step}, best_acc={best_acc:.1f}%", flush=True)
|
|
|
|
n_params = sum(p.numel() for p in model.parameters())
|
|
print(f"Model: hidden={args.hidden_dim} layers={args.layers} heads={args.heads} "
|
|
f"params={n_params:,}", flush=True)
|
|
|
|
# ── Load all training data onto GPU once (tiny datasets — avoids DataLoader overhead) ──
|
|
all_x = torch.stack([row[0] for row in train_ds]).to(device)
|
|
all_y = torch.tensor([row[1] for row in train_ds], dtype=torch.long, device=device)
|
|
n_train = len(all_x)
|
|
|
|
# ── Grokking detection state ──
|
|
grokked = False
|
|
blocks_since_grok = 0
|
|
prev_block_loss = None
|
|
prev_block_acc = 0.0
|
|
|
|
# ── Training loop ──
|
|
step = start_step
|
|
model.train()
|
|
block_start = time.perf_counter()
|
|
|
|
while step < args.max_steps:
|
|
idx = torch.randint(0, n_train, (args.batch_size,), device=device)
|
|
x, y = all_x[idx], all_y[idx]
|
|
optimizer.zero_grad()
|
|
loss = criterion(model(x), y)
|
|
loss.backward()
|
|
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
|
optimizer.step()
|
|
step += 1
|
|
|
|
# ── Block boundary ──
|
|
if step % BLOCK == 0 or step == args.max_steps:
|
|
eval_acc, eval_loss = evaluate(model, eval_loader, device)
|
|
elapsed = time.perf_counter() - block_start
|
|
|
|
entry = {
|
|
"step": step,
|
|
"overall": round(eval_acc, 2),
|
|
"loss": round(eval_loss, 4),
|
|
"elapsed_s": round(elapsed, 1),
|
|
}
|
|
history_entries.append(entry)
|
|
history_path.write_text(json.dumps({"history": history_entries}, indent=2))
|
|
|
|
# Save checkpoint
|
|
torch.save({
|
|
"model": model.state_dict(),
|
|
"optimizer": optimizer.state_dict(),
|
|
"vocab_obj": vocab,
|
|
"vocab": vocab.str_to_id,
|
|
"config": {
|
|
"hidden_dim": args.hidden_dim,
|
|
"layers": args.layers,
|
|
"heads": args.heads,
|
|
"n_labels": n_labels,
|
|
"seq_len": seq_len,
|
|
"vocab_size": len(vocab),
|
|
},
|
|
"labels": label_list,
|
|
"step": step,
|
|
}, ckpt_path)
|
|
|
|
if eval_acc > best_acc:
|
|
best_acc = eval_acc
|
|
|
|
print(f" step={step:6d} acc={eval_acc:5.1f}% loss={eval_loss:.4f} "
|
|
f"best={best_acc:.1f}% {elapsed:.0f}s", flush=True)
|
|
|
|
# Grokking detection
|
|
if not grokked:
|
|
loss_drop = (prev_block_loss is not None and
|
|
eval_loss < args.grok_loss_threshold)
|
|
acc_jump = (eval_acc - prev_block_acc) >= args.grok_acc_jump
|
|
if loss_drop or acc_jump:
|
|
grokked = True
|
|
print(f" *** GROKKING detected at step={step} "
|
|
f"({'loss_drop' if loss_drop else 'acc_jump'}) ***", flush=True)
|
|
|
|
if grokked:
|
|
blocks_since_grok += 1
|
|
if blocks_since_grok >= args.stop_after_grokking_blocks:
|
|
print(f" Stopping after {blocks_since_grok} post-grokking blocks.", flush=True)
|
|
break
|
|
|
|
prev_block_loss = eval_loss
|
|
prev_block_acc = eval_acc
|
|
block_start = time.perf_counter()
|
|
model.train()
|
|
|
|
print(f"\nDone. best_acc={best_acc:.1f}% steps={step} checkpoint={ckpt_path}", flush=True)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--dataset", required=True)
|
|
ap.add_argument("--eval-dataset", required=True)
|
|
ap.add_argument("--out-dir", required=True)
|
|
ap.add_argument("--labels", required=True)
|
|
ap.add_argument("--lr", type=float, default=0.001)
|
|
ap.add_argument("--weight-decay", type=float, default=0.01)
|
|
ap.add_argument("--max-steps", type=int, default=10000)
|
|
ap.add_argument("--hidden-dim", type=int, default=256)
|
|
ap.add_argument("--layers", type=int, default=2)
|
|
ap.add_argument("--heads", type=int, default=4)
|
|
ap.add_argument("--batch-size", type=int, default=256)
|
|
ap.add_argument("--seq-len", type=int, default=96)
|
|
ap.add_argument("--history-out", default=None)
|
|
ap.add_argument("--reset", action="store_true")
|
|
ap.add_argument("--grok-loss-threshold", type=float, default=0.05)
|
|
ap.add_argument("--grok-acc-jump", type=float, default=10.0)
|
|
ap.add_argument("--stop-after-grokking-blocks", type=int, default=4)
|
|
args = ap.parse_args()
|
|
train(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|