diff --git a/src/gpu/distill_specs.yaml b/src/gpu/distill_specs.yaml index 87ae3f7..99384a1 100644 --- a/src/gpu/distill_specs.yaml +++ b/src/gpu/distill_specs.yaml @@ -134,6 +134,35 @@ ara-diac-small-2: labels_file: teacher_labels_r7.jsonl mode: sequence note: 'r7 teacher + Muon; E4 gate: beat the shipped 8.259 by >= 2pp' +ara-diac-small-layerdrop: + teacher: rababa_arabic_byt5/run-006-morph/best + teacher_volume: rababa + out_volume: secryst + student_init: google/byt5-small + layer_drop: 'true' + student_config: + d_model: 1472 + d_ff: 3584 + num_heads: 6 + enc_layers: 6 + dec_layers: 4 + feed_forward_proj: gated-gelu + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 24000 + - 6000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_small/run-006-layerdrop + labels_file: teacher_labels_v2.jsonl + labels_complete: 'true' + mode: sequence + note: depth-cut rung - encoder 12->6, width untouched, layers copied + verbatim (both width-cut rungs failed); int4 target ~75MB; Muon + optimizer per E3 + optimizer: muon ara-diac-tiny-stitched: teacher: rababa_arabic_byt5/run-006-morph/best teacher_volume: rababa diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index c597308..51bbd6c 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -124,6 +124,49 @@ def svd_stitch_state(wide: dict, narrow: dict) -> dict: return out + +def layer_drop_state(pretrained: dict, narrow: dict, keep: int = 2) -> dict: + """Depth-reduction bridge: keep every `keep`-th encoder layer (and + all decoder layers), copying weights VERBATIM - no projection. The + width-cut failures (run-005/006) damaged every matrix; this touches + none of the survivors. Remaps block indices after the drop.""" + out = {} + import re + + def depth(prefix): + idxs = [ + int(m.group(1)) + for k in narrow + if (m := re.match(rf"{prefix}block\.(\d+)\.", k)) + ] + return (max(idxs) + 1) if idxs else 0 + + def src_depth(prefix): + idxs = [ + int(m.group(1)) + for k in pretrained + if (m := re.match(rf"{prefix}block\.(\d+)\.", k)) + ] + return (max(idxs) + 1) if idxs else 0 + + for prefix in ("encoder.", "decoder."): + d_dst, d_src = depth(prefix), src_depth(prefix) + kept = list(range(0, d_src, keep))[:d_dst] if d_dst < d_src else list(range(d_src)) + for name in narrow: + if not name.startswith(f"{prefix}block."): + continue + m = re.match(rf"{prefix}block\.(\d+)\.(.*)", name) + src_name = f"{prefix}block.{kept[int(m.group(1))]}.{m.group(2)}" + out[name] = pretrained[src_name].clone() + for name in narrow: + if name in out: + continue + src = pretrained.get(name) + if src is None: + raise KeyError(name) + out[name] = src.clone() + return out + 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 @@ -133,8 +176,9 @@ def _maybe_stitch(spec_id: str, spec: dict, student) -> None: from transformers import AutoModelForSeq2SeqLM pretrained = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]) - print(f"[{spec_id}] svd-stitching from {spec['student_init']}", flush=True) - student.load_state_dict(svd_stitch_state(pretrained.state_dict(), student.state_dict())) + bridge = layer_drop_state if spec.get("layer_drop") else svd_stitch_state + print(f"[{spec_id}] {bridge.__name__} from {spec['student_init']}", flush=True) + student.load_state_dict(bridge(pretrained.state_dict(), student.state_dict())) del pretrained