fix(distill): label length cap follows spec max_len; teacher stays on GPU until labels are usable

The Step-2 label filter hard-rejected anything over 384 bytes — with
Arabic's 1450-byte windows that discarded essentially every label (8
of 11,793 survived). This, not volume replication, was the original
'2-6 valid pairs' disease. Cap is now 2x spec max_len. Regeneration
also ran after the teacher had been moved to CPU — the confirmation
step now precedes the free.

Author: Ronald Tse

src/gpu/modal_distill.py

@@ -898,18 +898,16 @@ def label_all(pairs: list[tuple[str, str]]) -> list[tuple[str, str]]:
     else:
         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()
+    # Step 2: student trains on teacher labels. The teacher stays on the
+    # GPU until the labels are confirmed usable — regeneration (below)
+    # still needs it.
+    label_cap = 2 * int(spec.get("max_len", 384))
     teacher_labels = []
     seen_labels: set[str] = set()
 
     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:
+        if src and src not in seen_labels and label and len(label.encode()) <= label_cap:
             seen_labels.add(src)
             teacher_labels.append((src, label))
 
@@ -923,10 +921,10 @@ def accept_label(src: str, label: str) -> None:
                 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)
+            # the file is unusable from this container: regenerate rather
+            # than fail (relaunch-only loops forever on this path)
             print(
-                f"[{spec_id}] labels view torn ({len(teacher_labels)} valid "
+                f"[{spec_id}] labels unusable ({len(teacher_labels)} valid "
                 f"pairs); regenerating all labels",
                 flush=True,
             )
@@ -939,10 +937,15 @@ def accept_label(src: str, label: str) -> None:
     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 even after regeneration: "
+            f"labels unusable even after regeneration: "
             f"{len(teacher_labels)} valid pairs for {len(train_ds.rows)} srcs"
         )
 
+    teacher.to("cpu")
+    torch.cuda.empty_cache()
+    student.to("cuda")
+    student.gradient_checkpointing_enable()
+
     class TeacherPairs(Dataset):
         def __len__(self):
             return len(teacher_labels)