diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py new file mode 100644 index 0000000..618033f --- /dev/null +++ b/src/gpu/modal_distill.py @@ -0,0 +1,301 @@ +"""Modal app: logit distillation of byte-level seq2seq models (WO07). + +Case study (work order): Hebrew ByT5-base s43 (teacher) -> ByT5-small +(student) on the same hebrew-v4 corpus the teacher was trained on. +Loss = alpha * T^2 * KL(teacher_soft || student_soft) + (1 - alpha) * CE. +The student initializes from google/byt5-small pretraining — a pretrained +backbone is essential (from-scratch ByT5-small plateaus ~13% PER in the +Thai ablations; mode collapse resists all fixes). + + modal run --detach src/gpu/modal_distill.py::main --spec heb-diac-small + until modal run --detach src/gpu/modal_distill.py::main --spec heb-diac-small; do sleep 60; done + +Checkpoints on rababa-checkpoints:/rababa_hebrew_distill_small/run-001 +(periodic save + auto-resume from latest — server evictions are expected). +GPU is A10G; never competes with A100 training runs. +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent + +IMAGE = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.12.1", + "transformers==5.14.1", + "pyyaml>=6.0", + "numpy>=1.26", + ) + .add_local_dir(str(REPO_ROOT), "/root/ml-models", copy=True) + .workdir("/root/ml-models") +) + +CHECKPOINTS = modal.Volume.from_name("rababa-checkpoints") +DATASETS = modal.Volume.from_name("rababa-datasets") + +SPECS: dict[str, dict[str, str]] = { + "heb-diac-small": { + "teacher": "rababa_hebrew_byt5_s43/run-001/best", + "student_init": "google/byt5-small", + "train": "hebrew-v4/train.jsonl", + "val": "hebrew-v4/val.jsonl", + "out": "rababa_hebrew_distill_small/run-001", + }, +} + +app = modal.App("interscript-ml-distill", image=IMAGE) + + +@app.function( + gpu="A10G", + cpu=8, + memory=32 * 1024, + timeout=6 * 3600, + volumes={"/datasets": DATASETS, "/checkpoints": CHECKPOINTS}, +) +def distill(spec_id: str, epochs: int = 3, alpha: float = 0.5, temperature: float = 2.0) -> dict: + import json + import math + + import torch + import torch.nn.functional as F + from torch.utils.data import DataLoader, Dataset + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_cosine_schedule_with_warmup + + spec = SPECS[spec_id] + device = "cuda" + out_root = Path("/checkpoints") / spec["out"] + out_root.mkdir(parents=True, exist_ok=True) + + tokenizer = AutoTokenizer.from_pretrained("google/byt5-small") + teacher = AutoModelForSeq2SeqLM.from_pretrained( + Path("/checkpoints") / spec["teacher"], attn_implementation="eager" + ).to(device, dtype=torch.float16).eval() + for p in teacher.parameters(): + p.requires_grad_(False) + + student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]).to(device) + student.train() + + class Pairs(Dataset): + def __init__(self, path: Path, max_len: int = 384): + self.rows = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue + s, t = (row.get("src") or "").strip(), (row.get("tgt") or "").strip() + if s and t and len(s.encode()) <= max_len and len(t.encode()) <= max_len: + self.rows.append((s, t)) + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + s, t = self.rows[i] + return s, t + + def collate(batch): + src = tokenizer([s for s, _ in batch], padding=True, return_tensors="pt") + labels = tokenizer([t for _, t in batch], padding=True, return_tensors="pt").input_ids + labels[labels == tokenizer.pad_token_id] = -100 + return src.input_ids, src.attention_mask, labels + + train_loader = DataLoader( + Pairs(Path("/datasets") / spec["train"]), + batch_size=8, + shuffle=True, + collate_fn=collate, + num_workers=2, + drop_last=True, + ) + steps_per_epoch = len(train_loader) + total_steps = steps_per_epoch * epochs + + # resume from the newest periodic checkpoint if present + start_step = 0 + ckpts = sorted(out_root.glob("step-*"), key=lambda p: int(p.name.split("-")[1])) + if ckpts: + state = torch.load(ckpts[-1] / "student.pt", map_location=device, weights_only=True) + student.load_state_dict(state) + opt_state = torch.load(ckpts[-1] / "optim.pt", map_location=device, weights_only=True) + start_step = int(ckpts[-1].name.split("-")[1]) + print(f"[resume] from {ckpts[-1].name}", flush=True) + optimizer = torch.optim.AdamW(student.parameters(), lr=1e-4) + if ckpts: + optimizer.load_state_dict(opt_state) + scheduler = get_cosine_schedule_with_warmup(optimizer, total_steps // 20, total_steps) + for _ in range(start_step): + scheduler.step() + + save_every = 500 + log_every = 50 + step = start_step + best_val = math.inf + + def val_loss() -> float: + student.eval() + total, n = 0.0, 0 + with torch.no_grad(): + for _, (ids, am, labels) in enumerate( + DataLoader( + Pairs(Path("/datasets") / spec["val"]), + batch_size=8, + collate_fn=collate, + ) + ): + ids, am, labels = ids.to(device), am.to(device), labels.to(device) + loss = student( + input_ids=ids, attention_mask=am, labels=labels + ).loss + total += float(loss) + n += 1 + if n >= 50: # sample the val split, it is only for model selection + break + student.train() + return total / max(n, 1) + + while step < total_steps: + for ids, am, labels in train_loader: + if step >= total_steps: + break + ids, am, labels = ids.to(device), am.to(device), labels.to(device) + with torch.no_grad(): + t_logits = teacher(input_ids=ids, attention_mask=am, labels=labels).logits + s_out = student(input_ids=ids, attention_mask=am, labels=labels) + mask = labels != -100 + kd = F.kl_div( + F.log_softmax(s_out.logits[mask] / temperature, dim=-1), + F.softmax(t_logits.float()[mask] / temperature, dim=-1), + reduction="batchmean", + ) * (temperature ** 2) + loss = alpha * kd + (1 - alpha) * s_out.loss + loss.backward() + torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step += 1 + + if step % log_every == 0: + print( + f"[step {step}/{total_steps}] loss={float(loss):.4f} " + f"ce={float(s_out.loss):.4f} kd={float(kd):.4f}", + flush=True, + ) + if step % save_every == 0: + ck = out_root / f"step-{step}" + ck.mkdir(exist_ok=True) + torch.save(student.state_dict(), ck / "student.pt") + torch.save(optimizer.state_dict(), ck / "optim.pt") + CHECKPOINTS.commit() + + vl = val_loss() + if vl < best_val: + best = out_root / "best" + best.mkdir(exist_ok=True) + student.save_pretrained(str(best)) + tokenizer.save_pretrained(str(best)) + with (best / "val_loss.txt").open("w") as fh: + fh.write(f"{vl}\n") + CHECKPOINTS.commit() + return {"spec": spec_id, "steps": step, "val_loss": vl} + + +NIKUD = None + + +def _nikud_only(text: str) -> list[str]: + import re + + return re.findall(r"[\u0591-\u05bd\u05bf\u05c1\u05c2\u05c4\u05c5\u05c7]", text) + + +def _edit_distance(a, b) -> int: + if a == b: + return 0 + if not a: + return len(b) + if not b: + return len(a) + prev = list(range(len(b) + 1)) + for i, ai in enumerate(a, 1): + curr = [i] + for j, bj in enumerate(b, 1): + curr.append(min(prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + (ai != bj))) + prev = curr + return prev[-1] + + +@app.function( + gpu="A10G", + cpu=8, + memory=32 * 1024, + timeout=2 * 3600, + volumes={"/datasets": DATASETS, "/checkpoints": CHECKPOINTS}, +) +def evaluate(spec_id: str = "heb-diac-small", limit: int = 0) -> dict: + """Greedy DER/CER of teacher and student on the same test pairs, one + harness — the before/after for RESULTS.md.""" + import json + from pathlib import Path + + import torch + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer + + spec = SPECS[spec_id] + tokenizer = AutoTokenizer.from_pretrained("google/byt5-small") + device = "cuda" + teacher = AutoModelForSeq2SeqLM.from_pretrained( + Path("/checkpoints") / spec["teacher"], attn_implementation="eager" + ).to(device, dtype=torch.float16).eval() + student = AutoModelForSeq2SeqLM.from_pretrained( + Path("/checkpoints") / spec["out"] / "best" + ).to(device, dtype=torch.float16).eval() + + pairs = [] + for line in (Path("/datasets") / "nakdimon" / "test-imf.jsonl").read_text( + encoding="utf-8" + ).splitlines(): + if line.strip(): + row = json.loads(line) + pairs.append((row["src"], row["tgt"])) + if limit: + pairs = pairs[:limit] + + def greedy(model, text: str, max_len: int = 256) -> str: + ids = tokenizer(text, return_tensors="pt").to(device) + with torch.no_grad(): + out = model.generate(**ids, max_new_tokens=max_len, num_beams=1) + return tokenizer.batch_decode(out, skip_special_tokens=True)[0].strip() + + def metrics(model) -> dict: + der_sum = cer_sum = n = 0.0 + for src, tgt in pairs: + pred = greedy(model, src) + gold_n, pred_n = _nikud_only(tgt), _nikud_only(pred) + der_sum += _edit_distance(pred_n, gold_n) / max(1, len(gold_n)) + cer_sum += _edit_distance(list(pred), list(tgt)) / max(1, len(tgt)) + n += 1 + return {"der": round(100 * der_sum / n, 2), "cer": round(100 * cer_sum / n, 2), "n": int(n)} + + return {"teacher": metrics(teacher), "student": metrics(student)} + + +@app.local_entrypoint() +def main(spec: str = "heb-diac-small", epochs: int = 3) -> None: + result = distill.remote(spec, epochs=epochs) + print(result) + + +@app.local_entrypoint() +def eval_main(spec: str = "heb-diac-small", limit: int = 0) -> None: + print(evaluate.remote(spec, limit))