@@ -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+
127168def _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