From 14f020bdc318ba84144d76113741243dad3875be Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 29 Aug 2026 12:59:08 +0800 Subject: [PATCH 1/2] feat(distill): layer-drop bridge + ara-diac-small-layerdrop rung Both width-cut rungs failed (scratch 74.68, naive SVD stitch 82.96) - every matrix surgically damaged. Depth instead: encoder 12->6, width untouched, surviving layers copied verbatim; Muon per E3; int4 target ~75MB browser artifact. --- src/gpu/distill_specs.yaml | 29 ++++++++++++++++++++++++ src/gpu/modal_distill.py | 46 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 73 insertions(+), 2 deletions(-) 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..f0e9a7b 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -124,6 +124,47 @@ 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, tgt in narrow.items(): + 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, tgt in narrow.items(): + if name in out: + continue + src = pretrained.get(name) + out[name] = (src if src is not None else tgt).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 +174,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 From 42d3a6f76c06050cd12fb1fad1a8f6f4743d96ab Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 29 Aug 2026 13:00:00 +0800 Subject: [PATCH 2/2] style: drop unused loop vars in layer_drop_state --- src/gpu/modal_distill.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index f0e9a7b..51bbd6c 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -152,17 +152,19 @@ def src_depth(prefix): 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, tgt in narrow.items(): + 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, tgt in narrow.items(): + for name in narrow: if name in out: continue src = pretrained.get(name) - out[name] = (src if src is not None else tgt).clone() + if src is None: + raise KeyError(name) + out[name] = src.clone() return out def _maybe_stitch(spec_id: str, spec: dict, student) -> None: