diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 8d9eba0..397944d 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -170,6 +170,27 @@ def src_depth(prefix): out[name] = src.clone() return out +def student_t5_config(cfg: dict): + """Spec vocabulary -> T5Config, shared by every distill path. + The spec speaks enc_layers/dec_layers and byte-model constants; + T5Config speaks num_layers/num_decoder_layers. One translation.""" + from transformers import T5Config + + return T5Config( + 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)), + num_layers=cfg.get("enc_layers", 8), + num_decoder_layers=cfg.get("dec_layers", 8), + num_heads=cfg.get("num_heads", 6), + dropout_rate=0.1, + feed_forward_proj=cfg.get("feed_forward_proj", "relu"), + decoder_start_token_id=0, + relative_attention_max_distance=128, + ) + + def _maybe_stitch(spec_id: str, spec: dict, student) -> None: """When a custom-width student also names a pretrained init, bridge the pretrained weights down instead of random init (the capacity @@ -307,9 +328,9 @@ def distill(spec_id: str, epochs: int = 3, alpha: float = 0.5, temperature: floa # tiny tier: no pretrained backbone at this width — random init # from an explicit config (dense teacher-label supervision, see # the spec note on collapse risk) - from transformers import T5Config, T5ForConditionalGeneration + from transformers import T5ForConditionalGeneration - cfg = T5Config(**spec["student_config"]) + cfg = student_t5_config(spec["student_config"]) student = T5ForConditionalGeneration(cfg).to(device) if spec.get("layer_drop"): # depth-cut students keep the pretrained init: verbatim @@ -719,22 +740,9 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict: # teacher needs the GPU there; eviction-prone A10G headroom matters. student_tok = AutoTokenizer.from_pretrained("google/byt5-small") if spec.get("student_config"): - from transformers import T5Config, T5ForConditionalGeneration + from transformers import T5ForConditionalGeneration - cfg = spec["student_config"] - config = T5Config( - 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)), - num_layers=cfg.get("enc_layers", 8), - num_decoder_layers=cfg.get("dec_layers", 8), - num_heads=cfg.get("num_heads", 6), - dropout_rate=0.1, - feed_forward_proj=cfg.get("feed_forward_proj", "relu"), - decoder_start_token_id=0, - relative_attention_max_distance=128, - ) + config = student_t5_config(spec["student_config"]) student = T5ForConditionalGeneration(config) _maybe_stitch(spec_id, spec, student) n_params = sum(q.numel() for q in student.parameters()) / 1e6 diff --git a/tests/test_student_config.py b/tests/test_student_config.py new file mode 100644 index 0000000..f92fd56 --- /dev/null +++ b/tests/test_student_config.py @@ -0,0 +1,44 @@ +"""student_t5_config: one translation of spec vocabulary -> T5Config, +shared by the sequence and logit-KD paths (the logit path's raw +T5Config(**spec) silently took T5 defaults for depths and produced a +6/6 student from a 6/4 spec — the layer-copy then indexed out of +range).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +modal = pytest.importorskip("modal") + + +def test_layerdrop_spec_translates_depths() -> None: + from gpu.modal_distill import student_t5_config + + cfg = student_t5_config( + { + "d_model": 1472, + "d_kv": 64, + "d_ff": 3584, + "num_heads": 6, + "enc_layers": 6, + "dec_layers": 4, + "feed_forward_proj": "gated-gelu", + } + ) + assert cfg.num_layers == 6 + assert cfg.num_decoder_layers == 4 + assert cfg.vocab_size == 259 + assert cfg.decoder_start_token_id == 0 + + +def test_byte_model_defaults() -> None: + from gpu.modal_distill import student_t5_config + + cfg = student_t5_config({}) + assert cfg.num_layers == 8 and cfg.num_decoder_layers == 8 + assert cfg.d_model == 384