Skip to content

Commit 5d86c0d

Browse files
author
Ronald Tse
committed
feat(distill): ara-diac-tiny spec — the browser/Worker tier probe
~25-45M-param student (d384, enc6/dec2, byte vocab), single-variable vs run-002: same r6 teacher, corpus, frozen labels. Adds config-only student init for widths with no pretrained backbone. Gate: no collapse; DER <= 5.0 windowed would ship as the tiny tier.
1 parent d0979f7 commit 5d86c0d

1 file changed

Lines changed: 44 additions & 1 deletion

File tree

src/gpu/modal_distill.py

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,40 @@ def _ensure_src_path() -> None:
141141
"mode": "sequence",
142142
"note": "r6 canonical (2.5793 DER); gate <= 3.07 windowed DER-CE",
143143
},
144+
"ara-diac-tiny": {
145+
# the browser/Worker tier: ~45M params -> ~45MB int8 zip. Same
146+
# teacher (r6), corpus, and frozen labels as run-002 — the ONLY
147+
# variable is student capacity. From-scratch risk is real (Thai
148+
# ablation: scratch byt5-small plateaued ~13% PER) but this is
149+
# dense CE on 30k paragraph units, not RL; gate: does not
150+
# collapse (finite output, DER bounded well below scratch) and
151+
# <= 5.0 windowed DER-CE would make it shippable.
152+
"teacher": "rababa_arabic_byt5/run-006-morph/best",
153+
"teacher_volume": "rababa",
154+
"student_config": {
155+
"vocab_size": 384,
156+
"d_model": 384,
157+
"d_ff": 1024,
158+
"d_kv": 64,
159+
"num_layers": 6,
160+
"num_decoder_layers": 2,
161+
"feed_forward_proj": "relu",
162+
"decoder_start_token_id": 0,
163+
"eos_token_id": 1,
164+
"pad_token_id": 0,
165+
},
166+
"train": "r5-units/domain.txt",
167+
"train_extra": ["r5-units/replay.txt"],
168+
"unit_limits": [24000, 6000],
169+
"max_len": 1450,
170+
"label_beams": "1",
171+
"out": "rababa_arabic_distill_tiny/run-001",
172+
"labels_file": "teacher_labels_v2.jsonl",
173+
"labels_complete": "true",
174+
"mode": "sequence",
175+
"note": "tiny-tier probe: r6 teacher, run-002 corpus/labels, "
176+
"~45M student; collapse check + DER gate <= 5.0",
177+
},
144178
"ara-diac-small-pkm": {
145179
# TODO.qwen-next/02 — the LongCat/Qwen capacity axis: keep the
146180
# ByT5-small compute, add product-key lookup memory (+~25M
@@ -367,7 +401,16 @@ def distill(spec_id: str, epochs: int = 3, alpha: float = 0.5, temperature: floa
367401
for p in teacher.parameters():
368402
p.requires_grad_(False)
369403

370-
student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]).to(device)
404+
if "student_config" in spec:
405+
# tiny tier: no pretrained backbone at this width — random init
406+
# from an explicit config (dense teacher-label supervision, see
407+
# the spec note on collapse risk)
408+
from transformers import T5Config, T5ForConditionalGeneration
409+
410+
cfg = T5Config(**spec["student_config"])
411+
student = T5ForConditionalGeneration(cfg).to(device)
412+
else:
413+
student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]).to(device)
371414
student.train()
372415

373416
class Pairs(Dataset):

0 commit comments

Comments
 (0)