From 0286d5de3914b9de16c8f7225c9de73d080816e3 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 22 Aug 2026 13:05:25 +0800 Subject: [PATCH 1/2] fix(distill): train from in-memory labels; never re-read the just-written file Arabic tiny labeled all 11,790 srcs, then the post-label re-read of teacher_labels.jsonl served a stale volume replica (2 visible pairs) and the run trained on 2 pairs, saving garbage. The just-generated labels are now kept in memory and used directly; the file is only read on resume runs. The torn-view guard is unconditional, and the label loop gets a final volume commit (the modulo could skip the tail). --- src/gpu/modal_distill.py | 43 +++++++++++++++++++++++++++++----------- 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 0decb60..ace1980 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -763,6 +763,7 @@ def collate(batch): 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] + fresh_rows: list[tuple[str, str]] = [] if todo: print(f"[{spec_id}] labeling {len(todo)} remaining...", flush=True) @@ -828,13 +829,15 @@ def label_batch(batch, max_len: int = 0): preds = label_batch(batch) for (src, _), pred in zip(batch, preds, strict=True): if pred is not None: + text = pred.strip() fh.write( json.dumps( - {"src": src, "teacher": pred.strip()}, + {"src": src, "teacher": text}, ensure_ascii=False, ) + "\n" ) + fresh_rows.append((src, text)) labeled += len(batch) if labeled <= 200 * 16 or labeled % 3200 < len(batch): mem = torch.cuda.memory_allocated() / 2**30 @@ -844,6 +847,12 @@ def label_batch(batch, max_len: int = 0): ) if labeled % 3200 < len(batch): SECRYST_CHECKPOINTS.commit() + # the modulo can leave the tail uncommitted + { + "secryst": SECRYST_CHECKPOINTS, + "rababa": CHECKPOINTS, + "persian": PERSIAN_CHECKPOINTS, + }.get(spec.get("out_volume", teacher_vol), SECRYST_CHECKPOINTS).commit() else: print(f"[{spec_id}] teacher labels already complete", flush=True) @@ -855,22 +864,32 @@ def label_batch(batch, max_len: int = 0): student.gradient_checkpointing_enable() teacher_labels = [] seen_labels: set[str] = set() - for line in teacher_labels_path.read_text(encoding="utf-8", errors="ignore").splitlines(): - if not line.strip(): - continue - try: - row = json.loads(line) - except json.JSONDecodeError: - continue # torn line from a volume replication race - label = (row.get("teacher") or "").strip() - src = (row.get("src") or "").strip() + + def accept_label(src: str, label: str) -> None: + src, label = src.strip(), label.strip() if src and src not in seen_labels and label and len(label.encode()) <= 384: seen_labels.add(src) teacher_labels.append((src, label)) + + if fresh_rows: + # this run generated the labels: use them directly. The volume + # replica can serve a stale view of the just-written file (the + # rababa/secrets tear: 2 visible pairs after 11,790 written). + for src, label in fresh_rows: + accept_label(src, label) + else: + for line in teacher_labels_path.read_text(encoding="utf-8", errors="ignore").splitlines(): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue # torn line from a volume replication race + accept_label(row.get("src") or "", row.get("teacher") or "") print(f"[{spec_id}] trainable label pairs: {len(teacher_labels)}", flush=True) - if spec.get("labels_complete") and len(teacher_labels) < 0.5 * len(train_ds.rows): + if len(teacher_labels) < 0.5 * len(train_ds.rows): raise RuntimeError( - f"labels file view is torn: {len(teacher_labels)} valid pairs for " + f"labels view is torn: {len(teacher_labels)} valid pairs for " f"{len(train_ds.rows)} srcs — volume replication race; relaunch" ) From 0231dd93cf49eef29caa172291353e8cbc4eb662 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 22 Aug 2026 15:54:22 +0800 Subject: [PATCH 2/2] fix(distill): regenerate labels on torn resume-read instead of raising MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The follow-on hole: when the Step-1 resume-read declares everything done, Step 2 re-reads the file — that read can itself serve a stale replica (6 valid pairs of 27,324 persisted), and raising there makes the watchdog relaunch into the identical state forever. On a torn read the run now regenerates all labels in-container and trains from the in-memory rows. Labeling extracted into a label_all() closure shared by the fresh/resume/regenerate paths. --- src/gpu/modal_distill.py | 116 +++++++++++++++++++++------------------ 1 file changed, 64 insertions(+), 52 deletions(-) diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index ace1980..368b352 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -764,54 +764,52 @@ def collate(batch): todo = [(s, t) for s, t in train_ds.rows if s not in done] fresh_rows: list[tuple[str, str]] = [] - if todo: - print(f"[{spec_id}] labeling {len(todo)} remaining...", flush=True) - - seq_max = int(spec.get("max_len", 384)) - - def label_batch(batch, max_len: int = 0): - # 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) - if not max_len: - max_len = seq_max - try: - enc = teacher_tok( - [s for s, _ in batch], - padding=True, - truncation=True, - max_length=max_len, - return_tensors="pt", - ).to("cuda") - with torch.inference_mode(): - out = teacher.generate( - # r5 contract: generation cap = 2x window bytes - # (diacritized output runs 1.4-1.6x input) - **enc, max_new_tokens=2 * max_len, num_beams=label_beams - ) - 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 + seq_max = int(spec.get("max_len", 384)) + + def label_batch(batch, max_len: int = 0): + # 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) + if not max_len: + max_len = seq_max + try: + enc = teacher_tok( + [s for s, _ in batch], + padding=True, + truncation=True, + max_length=max_len, + return_tensors="pt", + ).to("cuda") + with torch.inference_mode(): + out = teacher.generate( + # r5 contract: generation cap = 2x window bytes + # (diacritized output runs 1.4-1.6x input) + **enc, max_new_tokens=2 * max_len, num_beams=label_beams ) - + 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 + ) + + def label_all(pairs: list[tuple[str, str]]) -> list[tuple[str, str]]: # 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())) + pairs = sorted(pairs, key=lambda p: len(p[0].encode())) budget = 32 * max(200, seq_max) batches: list[list[tuple[str, str]]] = [] cur: list[tuple[str, str]] = [] cur_max = 0 - for pair in todo: + for pair in pairs: length = len(pair[0].encode()) new_max = max(cur_max, length) if cur and (len(cur) + 1) * new_max > budget: @@ -823,6 +821,7 @@ def label_batch(batch, max_len: int = 0): if cur: batches.append(cur) + rows: list[tuple[str, str]] = [] labeled = 0 with teacher_labels_path.open("a", encoding="utf-8") as fh: for batch in batches: @@ -837,12 +836,12 @@ def label_batch(batch, max_len: int = 0): ) + "\n" ) - fresh_rows.append((src, text)) + rows.append((src, text)) labeled += len(batch) if labeled <= 200 * 16 or labeled % 3200 < len(batch): mem = torch.cuda.memory_allocated() / 2**30 print( - f" labeled {labeled}/{len(todo)} (gpu {mem:.2f} GiB)", + f" labeled {labeled}/{len(pairs)} (gpu {mem:.2f} GiB)", flush=True, ) if labeled % 3200 < len(batch): @@ -853,6 +852,11 @@ def label_batch(batch, max_len: int = 0): "rababa": CHECKPOINTS, "persian": PERSIAN_CHECKPOINTS, }.get(spec.get("out_volume", teacher_vol), SECRYST_CHECKPOINTS).commit() + return rows + + if todo: + print(f"[{spec_id}] labeling {len(todo)} remaining...", flush=True) + fresh_rows = label_all(todo) else: print(f"[{spec_id}] teacher labels already complete", flush=True) @@ -871,13 +875,7 @@ def accept_label(src: str, label: str) -> None: seen_labels.add(src) teacher_labels.append((src, label)) - if fresh_rows: - # this run generated the labels: use them directly. The volume - # replica can serve a stale view of the just-written file (the - # rababa/secrets tear: 2 visible pairs after 11,790 written). - for src, label in fresh_rows: - accept_label(src, label) - else: + if not fresh_rows: for line in teacher_labels_path.read_text(encoding="utf-8", errors="ignore").splitlines(): if not line.strip(): continue @@ -886,11 +884,25 @@ def accept_label(src: str, label: str) -> None: except json.JSONDecodeError: continue # torn line from a volume replication race accept_label(row.get("src") or "", row.get("teacher") or "") + if len(teacher_labels) < 0.5 * len(train_ds.rows): + # stale replica of a complete file: regenerate rather than + # fail (relaunch-only loops forever on this path) + print( + f"[{spec_id}] labels view torn ({len(teacher_labels)} valid " + f"pairs); regenerating all labels", + flush=True, + ) + teacher_labels = [] + seen_labels = set() + fresh_rows = label_all(list(train_ds.rows)) + if fresh_rows: + for src, label in fresh_rows: + accept_label(src, label) print(f"[{spec_id}] trainable label pairs: {len(teacher_labels)}", flush=True) if len(teacher_labels) < 0.5 * len(train_ds.rows): raise RuntimeError( - f"labels view is torn: {len(teacher_labels)} valid pairs for " - f"{len(train_ds.rows)} srcs — volume replication race; relaunch" + f"labels view is torn even after regeneration: " + f"{len(teacher_labels)} valid pairs for {len(train_ds.rows)} srcs" ) class TeacherPairs(Dataset):