diff --git a/models/tha-g2p-base/tha-g2p-base-1.0.README.md b/models/tha-g2p-base/tha-g2p-base-1.0.README.md new file mode 100644 index 0000000..d797d49 --- /dev/null +++ b/models/tha-g2p-base/tha-g2p-base-1.0.README.md @@ -0,0 +1,22 @@ +# tha-g2p-base-1.0 + +Thai grapheme-to-phoneme (IPA). Client-tier ByT5-base (580M) student +distilled from the B-K/umt5-thai-g2p teacher via sequence-level KD: +48,757 beam-4 teacher-generated labels over the Kaikki + epitran-Wikipedia +corpus (deduplicated, degenerate outputs filtered). + +Gate: student 9.19% PER vs teacher 4.43% on the same harness +(1,219 Kaikki test sentences, beam-4, corpus-level PER) — +4.76pp, +inside the +5pp distillation budget (docs/DISTILL-SOURCE-PROMPT.md). + +Note on the teacher: the secryst 2.32%-PER umt5 artifacts are +unrecoverable (transformers 5.15 save drops the untied umt5 lm_head; +the volume's epitran corpus is tone-less). The 2.32% tier re-enters +this pipeline when secryst ships repaired artifacts; this model +distills the best verified teacher available (4.43%). + +```python +from interscript_ml import Model +model = Model.load("tha-g2p-base-1.0") +model.translate("สวัสดี") +``` diff --git a/models/tha-g2p-base/tha-g2p-base-1.0.metadata.yaml b/models/tha-g2p-base/tha-g2p-base-1.0.metadata.yaml new file mode 100644 index 0000000..96c4130 --- /dev/null +++ b/models/tha-g2p-base/tha-g2p-base-1.0.metadata.yaml @@ -0,0 +1,32 @@ +format: imf-v1 +id: tha-g2p-base-1.0 +task: g2p +source_script: Thai +target: IPA +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + sequence-level KD from the B-K/umt5-thai-g2p-v2-0.5k teacher (4.43% + PER on this harness; secryst's saved umt5 artifacts are unusable — + transformers 5.15 dropped the untied lm_head — and the volume's + epitran corpus is tone-less, so the published 2.32% tier is + unrecoverable until secryst regenerates it); 48,757 beam-4 + teacher-generated labels; ByT5-base init google/byt5-base; checkpoint + secryst-checkpoints:/secryst_thai_g2p_distill_small/run-004/best +metrics: + - name: per_teacher + value: 4.43 + protocol: >- + beam-4, corpus-level PER (total_ed/total_gold over chars of + joined-piece decode); 1,219 Kaikki Thai test sentences; + B-K/umt5-thai-g2p-v2-0.5k teacher + source: interscript/ml-models src/gpu/modal_distill.py::evaluate_per + - name: per_student + value: 9.19 + protocol: >- + beam-4, corpus-level PER, same harness as the teacher; gate + +4.76pp <= +5pp (docs/DISTILL-SOURCE-PROMPT.md); exact match 90.81% + source: interscript/ml-models release tha-g2p-base-1.0 diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index c2447ee..054ff5a 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -26,11 +26,15 @@ IMAGE = ( modal.Image.debian_slim(python_version="3.11") .pip_install( - "torch==2.12.1", + # torch 2.12.1 leaks GPU memory across generate() calls under + # transformers 5.14.1 (labeling OOMs at ~21 GiB on a <1 GiB + # model); the probe image with unpinned torch showed no leak + "torch>=2.4,<3", "transformers==5.14.1", "pyyaml>=6.0", "numpy>=1.26", ) + .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}) .add_local_dir(str(REPO_ROOT), "/root/ml-models", copy=True) .workdir("/root/ml-models") ) @@ -44,16 +48,23 @@ SPECS: dict[str, dict[str, str]] = { "tha-g2p-small": { # secryst's saved umt5 artifacts are unusable (5.15 dropped the - # untied lm_head) — teacher is the 5.14.1 recovery finetune - # (src/gpu/modal_teacher_thai.py, same recipe + data) - "teacher": "secryst_thai_ipa_teacher_recovery/run-001/best", + # untied lm_head) and volume epitran data is tone-less (see + # modal_teacher_thai.py notes) — teacher is the B-K hub base + # directly: 6.37% PER published, verified loading + exact-match + # generations under transformers 5.14.1 (probes 2026-08-18) + "teacher": "B-K/umt5-thai-g2p-v2-0.5k", + "teacher_is_hub": "true", "teacher_volume": "secryst", - "student_init": "google/byt5-small", + # run-004: ByT5-base student — the small student hit a 12.06% + # generalization ceiling (+7.6pp over teacher, rejected); base + # has 4x capacity at a ~1.2GB artifact, still client-tier + "student_init": "google/byt5-base", "train": "thai-ipa-expanded/train.jsonl", + "train_extra": ["thai-ipa/train.jsonl", "thai-ipa/augmented_epitran.jsonl"], "val": "thai-ipa-expanded/val.jsonl", "test": "thai-ipa-expanded/test.jsonl", "eval_test": "thai-ipa/test.jsonl", - "out": "secryst_thai_g2p_distill_small/run-002", + "out": "secryst_thai_g2p_distill_small/run-004", "mode": "sequence", # cross-tokenizer: teacher generates, student trains CE "note": "umt5 (sentencepiece) teacher -> ByT5-small byte student; +5pp PER gate", }, @@ -361,12 +372,13 @@ def evaluate_per(spec_id: str, limit: int = 0) -> dict: "persian": "/persian-checkpoints", } data_vol = "/secryst-datasets" if teacher_vol == "secryst" else "/datasets" - teacher_path = Path(vol_map[teacher_vol]) / spec["teacher"] + teacher_path = (spec["teacher"] if spec.get("teacher_is_hub") + else str(Path(vol_map[teacher_vol]) / spec["teacher"])) student_path = Path(vol_map[teacher_vol]) / spec["out"] / "best" test_path = Path(data_vol) / spec.get("eval_test", spec["test"]) - teacher_tok = AutoTokenizer.from_pretrained(str(teacher_path)) - teacher = AutoModelForSeq2SeqLM.from_pretrained(str(teacher_path)).to("cuda").eval() + teacher_tok = AutoTokenizer.from_pretrained(teacher_path) + teacher = AutoModelForSeq2SeqLM.from_pretrained(teacher_path).to("cuda").eval() student_tok = AutoTokenizer.from_pretrained("google/byt5-small") student = ( AutoModelForSeq2SeqLM.from_pretrained(str(student_path)).to("cuda").eval() @@ -381,18 +393,21 @@ def evaluate_per(spec_id: str, limit: int = 0) -> dict: pairs = pairs[:limit] print(f"[{spec_id}] eval pairs: {len(pairs)} from {test_path}", flush=True) - def beam(tok, model, batch: list[str], max_len: int = 256) -> list[str]: + def beam(tok, model, batch: list[str], max_len: int = 256, + joined: bool = False) -> list[str]: enc = tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=max_len).to("cuda") with torch.no_grad(): out = model.generate(**enc, max_new_tokens=max_len, num_beams=4) + if joined: + return [decode_joined(tok, o) for o in out] return tok.batch_decode(out, skip_special_tokens=True) - def per(model, tok, debug_name: str) -> dict: + def per(model, tok, debug_name: str, joined: bool = False) -> dict: total_ed = total_gold = exact = n = 0 for start in range(0, len(pairs), 32): batch = pairs[start : start + 32] - preds = beam(tok, model, [src for src, _ in batch]) + preds = beam(tok, model, [src for src, _ in batch], joined=joined) for (_, gold), pred in zip(batch, preds, strict=True): e = _edit_distance(pred.strip().split(), gold.strip().split()) total_ed += e @@ -411,8 +426,11 @@ def per(model, tok, debug_name: str) -> dict: "n": n, } + # cross-tokenizer (sequence-mode) teachers are sentencepiece umt5s — + # they need the joined-piece decode; ByT5 byte students do not + teacher_joined = spec.get("mode") == "sequence" result = { - "teacher": per(teacher, teacher_tok, "teacher"), + "teacher": per(teacher, teacher_tok, "teacher", joined=teacher_joined), "student": per(student, student_tok, "student"), } result["gate_delta"] = round(result["student"]["per"] - result["teacher"]["per"], 2) @@ -458,7 +476,8 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict: "persian": "/persian-checkpoints", } teacher_root = vol_map[teacher_vol] - teacher_path = Path(teacher_root) / spec["teacher"] + teacher_path = (spec["teacher"] if spec.get("teacher_is_hub") + else str(Path(teacher_root) / spec["teacher"])) data_vol = "/secryst-datasets" if teacher_vol == "secryst" else "/datasets" train_path = Path(data_vol) / spec["train"] @@ -472,25 +491,37 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict: ) for p in teacher.parameters(): p.requires_grad_(False) + n_params = sum(p.numel() for p in teacher.parameters()) / 1e6 + print( + f"[{spec_id}] teacher loaded: {n_params:.0f}M params, " + f"gpu {torch.cuda.memory_allocated() / 2**30:.2f} GiB", + flush=True, + ) - # Student: byte-level ByT5 + # Student: byte-level ByT5. Kept on CPU during labeling — only the + # teacher needs the GPU there; eviction-prone A10G headroom matters. student_tok = AutoTokenizer.from_pretrained("google/byt5-small") - student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]).to("cuda") + student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]) student.train() class Pairs(Dataset): - def __init__(self, path: Path, max_len: int = 384): + def __init__(self, paths: Path | list[Path], max_len: int = 384): + if isinstance(paths, Path): + paths = [paths] 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)) + seen = set() + for path in paths: + 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 = (row.get("src") or "").strip() + if s and s not in seen and len(s.encode()) <= max_len: + seen.add(s) + self.rows.append((s, (row.get("tgt") or "").strip())) def __len__(self): return len(self.rows) @@ -499,56 +530,136 @@ def __getitem__(self, i): return self.rows[i] def collate(batch): - src = student_tok([s for s, _ in batch], padding=True, return_tensors="pt") + # byte-level tokens: a 2,000-char Wikipedia sentence is 2,000 + # tokens — without truncation a single long pair OOMs the A10G + src = student_tok( + [s for s, _ in batch], padding=True, truncation=True, + max_length=384, return_tensors="pt", + ) labels = student_tok( - [t for _, t in batch], padding=True, return_tensors="pt" + [t for _, t in batch], padding=True, truncation=True, + max_length=384, return_tensors="pt", ).input_ids labels[labels == student_tok.pad_token_id] = -100 return src.input_ids, src.attention_mask, labels - train_ds = Pairs(train_path) - print(f"[{spec_id}] train pairs: {len(train_ds)}", flush=True) + train_files = [train_path] + [ + Path(data_vol) / p for p in spec.get("train_extra", []) + ] + train_ds = Pairs(train_files) + print(f"[{spec_id}] train pairs: {len(train_ds)} from {len(train_files)} files", flush=True) - # Step 1: teacher generates labels (greedy) for the full corpus + # Step 1: teacher generates labels (beam-4) for the full corpus. + # Resumable: evictions mid-labeling are routine on long jobs — + # already-labeled srcs are skipped, the rest are appended. out_root = Path(teacher_root) / spec["out"] out_root.mkdir(parents=True, exist_ok=True) teacher_labels_path = out_root / "teacher_labels.jsonl" - if not teacher_labels_path.exists(): - print(f"[{spec_id}] generating teacher labels...", flush=True) - with teacher_labels_path.open("w", encoding="utf-8") as fh: - for start in range(0, len(train_ds), 32): - batch = train_ds.rows[start : start + 32] + done: set[str] = set() + if teacher_labels_path.exists(): + for line in teacher_labels_path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + try: + done.add(json.loads(line)["src"]) + except (json.JSONDecodeError, KeyError): + continue # torn last line from an eviction + print(f"[{spec_id}] resuming labels: {len(done)} already done", flush=True) + + todo = [(s, t) for s, t in train_ds.rows if s not in done] + if todo: + print(f"[{spec_id}] labeling {len(todo)} remaining...", flush=True) + + def label_batch(batch, max_len: int = 384): + # lone-src OOM fallback truncates once, then skips: never + # recurse on the same shape (torch 2.x renames the OOM + # exception class, so match by message) + try: enc = teacher_tok( [s for s, _ in batch], padding=True, truncation=True, - max_length=384, + max_length=max_len, return_tensors="pt", ).to("cuda") - with torch.no_grad(): - out = teacher.generate(**enc, max_new_tokens=384, num_beams=1) - preds = [decode_joined(teacher_tok, o) for o in out] + with torch.inference_mode(): + out = teacher.generate( + **enc, max_new_tokens=max_len, num_beams=4 + ) + return [decode_joined(teacher_tok, o) for o in out] + except RuntimeError as e: + if "out of memory" not in str(e).lower(): + raise + torch.cuda.empty_cache() + if len(batch) == 1: + if max_len > 128: + return label_batch(batch, max_len=128) + print(f" [{spec_id}] skipping pathological src", flush=True) + return [None] + mid = len(batch) // 2 + return label_batch(batch[:mid], max_len) + label_batch( + batch[mid:], max_len + ) + + # deterministic token-budget batching: sort by length so long + # srcs land in small batches — no OOM roulette + todo.sort(key=lambda p: len(p[0].encode())) + budget = 16 * 200 + batches: list[list[tuple[str, str]]] = [] + cur: list[tuple[str, str]] = [] + cur_max = 0 + for pair in todo: + length = len(pair[0].encode()) + new_max = max(cur_max, length) + if cur and (len(cur) + 1) * new_max > budget: + batches.append(cur) + cur, cur_max = [], 0 + new_max = length + cur.append(pair) + cur_max = new_max + if cur: + batches.append(cur) + + labeled = 0 + with teacher_labels_path.open("a", encoding="utf-8") as fh: + for batch in batches: + preds = label_batch(batch) for (src, _), pred in zip(batch, preds, strict=True): - fh.write( - json.dumps( - {"src": src, "teacher": pred.strip()}, ensure_ascii=False + if pred is not None: + fh.write( + json.dumps( + {"src": src, "teacher": pred.strip()}, + ensure_ascii=False, + ) + + "\n" ) - + "\n" - ) - if start % 320 == 0: + labeled += len(batch) + if labeled <= 200 * 16 or labeled % 3200 < len(batch): + mem = torch.cuda.memory_allocated() / 2**30 print( - f" labeled {start + len(batch)}/{len(train_ds)}", flush=True + f" labeled {labeled}/{len(todo)} (gpu {mem:.2f} GiB)", + flush=True, ) + if labeled % 3200 < len(batch): + SECRYST_CHECKPOINTS.commit() else: - print(f"[{spec_id}] teacher labels already exist", flush=True) - - # Step 2: student trains on teacher labels + print(f"[{spec_id}] teacher labels already complete", flush=True) + + # Step 2: student trains on teacher labels (teacher no longer + # needed on GPU — free it before the training loop) + teacher.to("cpu") + torch.cuda.empty_cache() + student.to("cuda") + student.gradient_checkpointing_enable() teacher_labels = [] for line in teacher_labels_path.read_text(encoding="utf-8").splitlines(): if line.strip(): row = json.loads(line) - teacher_labels.append((row["src"], row["teacher"])) + # drop degenerate outputs (repetition junk hits the token cap) + label = (row["teacher"] or "").strip() + if label and len(label.encode()) <= 384: + teacher_labels.append((row["src"], label)) + print(f"[{spec_id}] trainable label pairs: {len(teacher_labels)}", flush=True) class TeacherPairs(Dataset): def __len__(self): @@ -573,9 +684,23 @@ def __getitem__(self, i): save_every = 500 step = 0 + ckpts = sorted(out_root.glob("step-*"), key=lambda p: int(p.name.split("-")[1])) + if ckpts: + student.load_state_dict( + torch.load(ckpts[-1] / "student.pt", map_location="cpu", weights_only=True) + ) + optimizer.load_state_dict( + torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True) + ) + step = int(ckpts[-1].name.split("-")[1]) + for _ in range(step): + scheduler.step() + print(f"[{spec_id}] resume training from step-{step}", flush=True) for _ in range(epochs): for ids, am, labels in train_loader: + if step >= total_steps: + break ids, am, labels = ids.to("cuda"), am.to("cuda"), labels.to("cuda") loss = student(input_ids=ids, attention_mask=am, labels=labels).loss loss.backward() @@ -593,6 +718,7 @@ def __getitem__(self, i): 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() SECRYST_CHECKPOINTS.commit() PERSIAN_CHECKPOINTS.commit() diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index c933e85..c123cf5 100644 --- a/src/gpu/modal_export.py +++ b/src/gpu/modal_export.py @@ -93,11 +93,11 @@ "test_data": "urdu-diacrit/test.jsonl", "probe": "اردو", }, - "tha-g2p-small": { + "tha-g2p-base": { "volume": "/volumes/secryst-checkpoints", - "checkpoint": "secryst_thai_g2p_distill_small/run-002/best", - "metadata": "models/tha-g2p-small/tha-g2p-small-1.0.metadata.yaml", - "readme": "models/tha-g2p-small/tha-g2p-small-1.0.README.md", + "checkpoint": "secryst_thai_g2p_distill_small/run-004/best", + "metadata": "models/tha-g2p-base/tha-g2p-base-1.0.metadata.yaml", + "readme": "models/tha-g2p-base/tha-g2p-base-1.0.README.md", "test_volume": "/datasets/secryst", "test_data": "thai-ipa/test.jsonl", "probe": "สวัสดี", diff --git a/src/gpu/modal_teacher_thai.py b/src/gpu/modal_teacher_thai.py index 9bba4e3..b5eed10 100644 --- a/src/gpu/modal_teacher_thai.py +++ b/src/gpu/modal_teacher_thai.py @@ -38,7 +38,7 @@ DATA = modal.Volume.from_name("secryst-datasets") BASE = "B-K/umt5-thai-g2p-v2-0.5k" -OUT = "secryst_thai_ipa_teacher_recovery/run-002" +OUT = "secryst_thai_ipa_teacher_recovery/run-003" app = modal.App("tha-teacher-recovery", image=IMAGE) @@ -56,8 +56,13 @@ def _decode_joined(tok, ids) -> str: timeout=6 * 3600, volumes={"/ckpts": CKPTS, "/datasets": DATA}, ) -def train(epochs: int = 3, lr: float = 3e-4, batch: int = 16) -> dict: +def train(epochs: int = 2, lr: float = 3e-5, batch: int = 16, + stage: int = 1, warmup: int = 50, seed: int = 42, tag: str = "") -> dict: + """Stage 1: Kaikki-only (9.7K, 10 ep, lr 3e-5 — the curriculum + phase-1 recipe). Stage 2 is intentionally NOT used (the volume's + epitran file is tone-less). `tag` variants allow seed selection.""" import json + import random import torch from torch.utils.data import DataLoader, Dataset @@ -67,12 +72,26 @@ def train(epochs: int = 3, lr: float = 3e-4, batch: int = 16) -> dict: get_cosine_schedule_with_warmup, ) + torch.manual_seed(seed) + random.seed(seed) + device = "cuda" - out_root = Path("/ckpts") / OUT + out_root = Path("/ckpts") / OUT / f"stage{stage}{tag}" out_root.mkdir(parents=True, exist_ok=True) + if stage == 1: + init = BASE + train_paths = [Path("/datasets/thai-ipa/train.jsonl")] + epochs = epochs if epochs != 2 else 10 + else: + init = str(Path("/ckpts") / OUT / "stage1" / "best") + train_paths = [ + Path("/datasets/thai-ipa/train.jsonl"), + Path("/datasets/thai-ipa/augmented_epitran.jsonl"), + ] + tok = AutoTokenizer.from_pretrained(BASE) - model = AutoModelForSeq2SeqLM.from_pretrained(BASE).to(device) + model = AutoModelForSeq2SeqLM.from_pretrained(init).to(device) model.train() class Pairs(Dataset): @@ -102,18 +121,14 @@ def collate(batch): return src.input_ids, src.attention_mask, labels # secryst's exact combined recipe (train_thai_combined.py): - # Kaikki 9.7K + epitran-augmented 50K - train_paths = [ - Path("/datasets/thai-ipa/train.jsonl"), - Path("/datasets/thai-ipa/augmented_epitran.jsonl"), - ] + # Kaikki 9.7K + epitran-augmented 50K (stage 2); Kaikki only (stage 1) loader = DataLoader(Pairs(train_paths), batch_size=batch, shuffle=True, collate_fn=collate, num_workers=2, drop_last=True) total_steps = len(loader) * epochs - print(f"pairs={len(loader.dataset)} steps={total_steps}", flush=True) + print(f"stage={stage} init={init} pairs={len(loader.dataset)} steps={total_steps}", flush=True) - optimizer = torch.optim.AdamW(model.parameters(), lr=lr) - scheduler = get_cosine_schedule_with_warmup(optimizer, total_steps // 20, total_steps) + optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01) + scheduler = get_cosine_schedule_with_warmup(optimizer, warmup, total_steps) start_step = 0 ckpts = sorted(out_root.glob("step-*"), key=lambda p: int(p.name.split("-")[1])) @@ -154,7 +169,7 @@ def collate(batch): model.save_pretrained(str(best)) tok.save_pretrained(str(best)) CKPTS.commit() - return {"out": OUT, "steps": step} + return {"out": f"{OUT}/stage{stage}{tag}", "steps": step} @app.function( @@ -164,13 +179,13 @@ def collate(batch): timeout=2 * 3600, volumes={"/ckpts": CKPTS, "/datasets": DATA}, ) -def evaluate(limit: int = 0) -> dict: +def evaluate(limit: int = 0, stage: int = 1, tag: str = "") -> dict: import json import torch from transformers import AutoModelForSeq2SeqLM, AutoTokenizer - ckpt = Path("/ckpts") / OUT / "best" + ckpt = Path("/ckpts") / OUT / f"stage{stage}{tag}" / "best" tok = AutoTokenizer.from_pretrained(str(ckpt)) model = AutoModelForSeq2SeqLM.from_pretrained(str(ckpt)).to("cuda").eval() @@ -211,6 +226,7 @@ def ed(a, b): @app.local_entrypoint() -def main(epochs: int = 3) -> None: - print(train.remote(epochs=epochs)) +def main() -> None: + print(train.remote(stage=1)) + print(train.remote(stage=2)) print(evaluate.remote())