diff --git a/src/gpu/distill_specs.yaml b/src/gpu/distill_specs.yaml index 50dc4ca..0fa9d30 100644 --- a/src/gpu/distill_specs.yaml +++ b/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 diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 397944d..2339614 100644 --- a/src/gpu/modal_distill.py +++ b/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)), diff --git a/tests/test_student_config.py b/tests/test_student_config.py index f92fd56..74932a6 100644 --- a/tests/test_student_config.py +++ b/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