fix(distill): train from in-memory labels; never re-read the just-written file

src/gpu/modal_distill.py

@@ -763,54 +763,53 @@ 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]
- 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
+ fresh_rows: list[tuple[str, str]] = []
+ 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:
@@ -822,28 +821,42 @@ 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:
  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"
  )
+ 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):
  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()
+ 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)
 
@@ -855,23 +868,41 @@ 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 not fresh_rows:
+ 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 "")
+ 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 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"{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):