fix(distill): logit-KD student vocab matches teacher

src/gpu/distill_specs.yaml

@@ -538,6 +538,7 @@ heb-diac-small-s46-layerdrop:
   student_init: google/byt5-small
   layer_drop: 'true'
   student_config:
+    vocab_size: 384  # logit-KD: must match the ByT5 teacher's vocab
     d_model: 1472
     d_kv: 64
     d_ff: 3584

src/gpu/modal_distill.py

@@ -177,7 +177,7 @@ def student_t5_config(cfg: dict):
     from transformers import T5Config
 
     return T5Config(
-        vocab_size=259,
+        vocab_size=cfg.get("vocab_size", 259),
         d_model=cfg.get("d_model", 384),
         d_ff=cfg.get("d_ff", 1536),
         d_kv=cfg.get("d_kv", cfg.get("d_model", 384) // cfg.get("num_heads", 6)),

tests/test_student_config.py

@@ -42,3 +42,14 @@ def test_byte_model_defaults() -> None:
     cfg = student_t5_config({})
     assert cfg.num_layers == 8 and cfg.num_decoder_layers == 8
     assert cfg.d_model == 384
+
+
+def test_vocab_size_override_matches_teacher() -> None:
+    # logit-KD computes teacher-vs-student KL: the student vocab must
+    # match the teacher's (ByT5 = 384), while the sequence path's
+    # byte-table default (259) stays intact
+    from gpu.modal_distill import student_t5_config
+
+    cfg = student_t5_config({"vocab_size": 384, "enc_layers": 6, "dec_layers": 4})
+    assert cfg.vocab_size == 384
+    assert student_t5_config({"enc_layers": 6, "dec_layers": 4}).vocab_size == 259