Skip to content

Commit 163aedf

Browse files
Ronald Tseronaldtse
authored andcommitted
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.
1 parent 8470e22 commit 163aedf

2 files changed

Lines changed: 73 additions & 2 deletions

File tree

‎src/gpu/distill_specs.yaml‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,35 @@ ara-diac-small-2:
134134
labels_file: teacher_labels_r7.jsonl
135135
mode: sequence
136136
note: 'r7 teacher + Muon; E4 gate: beat the shipped 8.259 by >= 2pp'
137+
ara-diac-small-layerdrop:
138+
teacher: rababa_arabic_byt5/run-006-morph/best
139+
teacher_volume: rababa
140+
out_volume: secryst
141+
student_init: google/byt5-small
142+
layer_drop: 'true'
143+
student_config:
144+
d_model: 1472
145+
d_ff: 3584
146+
num_heads: 6
147+
enc_layers: 6
148+
dec_layers: 4
149+
feed_forward_proj: gated-gelu
150+
train: r5-units/domain.txt
151+
train_extra:
152+
- r5-units/replay.txt
153+
unit_limits:
154+
- 24000
155+
- 6000
156+
max_len: 1450
157+
label_beams: '1'
158+
out: rababa_arabic_distill_small/run-006-layerdrop
159+
labels_file: teacher_labels_v2.jsonl
160+
labels_complete: 'true'
161+
mode: sequence
162+
note: depth-cut rung - encoder 12->6, width untouched, layers copied
163+
verbatim (both width-cut rungs failed); int4 target ~75MB; Muon
164+
optimizer per E3
165+
optimizer: muon
137166
ara-diac-tiny-stitched:
138167
teacher: rababa_arabic_byt5/run-006-morph/best
139168
teacher_volume: rababa

‎src/gpu/modal_distill.py‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,47 @@ def svd_stitch_state(wide: dict, narrow: dict) -> dict:
124124
return out
125125

126126

127+
128+
def layer_drop_state(pretrained: dict, narrow: dict, keep: int = 2) -> dict:
129+
"""Depth-reduction bridge: keep every `keep`-th encoder layer (and
130+
all decoder layers), copying weights VERBATIM - no projection. The
131+
width-cut failures (run-005/006) damaged every matrix; this touches
132+
none of the survivors. Remaps block indices after the drop."""
133+
out = {}
134+
import re
135+
136+
def depth(prefix):
137+
idxs = [
138+
int(m.group(1))
139+
for k in narrow
140+
if (m := re.match(rf"{prefix}block\.(\d+)\.", k))
141+
]
142+
return (max(idxs) + 1) if idxs else 0
143+
144+
def src_depth(prefix):
145+
idxs = [
146+
int(m.group(1))
147+
for k in pretrained
148+
if (m := re.match(rf"{prefix}block\.(\d+)\.", k))
149+
]
150+
return (max(idxs) + 1) if idxs else 0
151+
152+
for prefix in ("encoder.", "decoder."):
153+
d_dst, d_src = depth(prefix), src_depth(prefix)
154+
kept = list(range(0, d_src, keep))[:d_dst] if d_dst < d_src else list(range(d_src))
155+
for name, tgt in narrow.items():
156+
if not name.startswith(f"{prefix}block."):
157+
continue
158+
m = re.match(rf"{prefix}block\.(\d+)\.(.*)", name)
159+
src_name = f"{prefix}block.{kept[int(m.group(1))]}.{m.group(2)}"
160+
out[name] = pretrained[src_name].clone()
161+
for name, tgt in narrow.items():
162+
if name in out:
163+
continue
164+
src = pretrained.get(name)
165+
out[name] = (src if src is not None else tgt).clone()
166+
return out
167+
127168
def _maybe_stitch(spec_id: str, spec: dict, student) -> None:
128169
"""When a custom-width student also names a pretrained init, bridge
129170
the pretrained weights down instead of random init (the capacity
@@ -133,8 +174,9 @@ def _maybe_stitch(spec_id: str, spec: dict, student) -> None:
133174
from transformers import AutoModelForSeq2SeqLM
134175

135176
pretrained = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"])
136-
print(f"[{spec_id}] svd-stitching from {spec['student_init']}", flush=True)
137-
student.load_state_dict(svd_stitch_state(pretrained.state_dict(), student.state_dict()))
177+
bridge = layer_drop_state if spec.get("layer_drop") else svd_stitch_state
178+
print(f"[{spec_id}] {bridge.__name__} from {spec['student_init']}", flush=True)
179+
student.load_state_dict(bridge(pretrained.state_dict(), student.state_dict()))
138180
del pretrained
139181

140182

0 commit comments

Comments
 (0)