Setting the file. One moment.
Train Sentence Transformer Distillation Example · Train Sentence Transformers · huggingface/skills · Skills Docs
Back to the top of the page Losses Cross Encoder
scripts/train_sentence_transformer_distillation_example.py
scripts/ train_sentence_transformer_distillation_example.py
Python · 303 lines · 12 KB
16
17 Three distillation patterns:
18 - Pattern 1 (this script): `(text, teacher_embedding)` + `MSELoss`. Cheapest,
19 most data-efficient. Use when student + teacher are both bi-encoders with the
20 same output dim and you have a pile of unlabeled text.
21 - Pattern 2: `(query, positive, negative, score_diff)` + `MarginMSELoss` from a
22 CrossEncoder teacher's score differences. Workhorse of ms-marco distillation.
23 See `../references/losses_sentence_transformer.md` (MarginMSELoss section).
24 - Pattern 3: `(query, positive, neg_1, ..., neg_n, labels)` + `DistillKLDivLoss`
25 to preserve the full teacher distribution. More data-hungry. Natural fit when
26 distilling from an ensemble of rerankers.
27
28 Mismatched dims: if the student's dim is smaller than the teacher's, MSELoss
29 fails. Add a PCA-init `Dense` projection so the teacher matches the student:
30
31 from sklearn.decomposition import PCA
32 from sentence_transformers.sentence_transformer.modules import Dense
33 pca = PCA(n_components=student.get_embedding_dimension())
34 pca.fit(teacher.encode(sentences[:20_000], convert_to_numpy=True))
35 dense = Dense(
36 in_features=teacher.get_embedding_dimension(),
37 out_features=student.get_embedding_dimension(),
38 bias=False, activation_function=torch.nn.Identity(),
39 )
40 dense.linear.weight = torch.nn.Parameter(torch.from_numpy(pca.components_).float())
41 teacher.add_module("dense", dense)
42
43 Distilling to a CrossEncoder student: construct with `activation_fn=nn.Identity()`
44 or eval ranking collapses silently. Every non-BCE CE loss expects raw logits
45 during training, but the model's `activation_fn` runs at eval time inside
46 `predict()`. Default `Sigmoid` (when `num_labels=1`) saturates raw logits >5 to
47 ~1.0, dropping nDCG from e.g. ~0.59 to ~0.14 with healthy-looking training loss.
48 Applies to all CE distillation / listwise / pairwise losses. See SKILL.md
49 Directive 7 ([CE]).
50
51 Layer pruning shortcut for Pattern 1: copy the teacher, delete layers (often
52 keeps 99%+ of quality at a fraction of the layers), then distill with MSELoss:
53
54 from copy import deepcopy
55 student = deepcopy(teacher)
56 layers = student.transformers_model.encoder.layer # BERT/MPNet/DistilBERT
57 student.transformers_model.encoder.layer = torch.nn.ModuleList(
58 [layers[0], layers[3], layers[6], layers[9]]
59 )
60 student.transformers_model.config.num_hidden_layers = 4
61
62 Tips: pre-compute teacher outputs once and cache (`dataset.save_to_disk`). LR
63 1e-4 (higher than the usual 2e-5: the target is dense regression). 1 epoch is
64 usually enough. The student inherits the teacher's weaknesses, so pick a
65 teacher strong on YOUR task. If the teacher expects an instruction prefix,
66 include it during teacher encoding so the student's target matches inference.
67
68 For multilingual student distillation (extend an English teacher to other
69 languages without in-language supervised data), see `train_sentence_transformer_make_multilingual_example.py`.
70 """
71
72 from __future__ import annotations
73
74 import argparse
75 import logging
76 import os
77 from contextlib import nullcontext
78
79 import torch
80 from datasets import Dataset, load_dataset
81
82 from sentence_transformers import (
83 SentenceTransformer,
84 SentenceTransformerModelCardData,
85 SentenceTransformerTrainer,
86 SentenceTransformerTrainingArguments,
87 )
88 from sentence_transformers.sentence_transformer.evaluation import EmbeddingSimilarityEvaluator
89 from sentence_transformers.sentence_transformer.losses import MSELoss
90 from sentence_transformers.sentence_transformer.modules import Normalize
91 from sentence_transformers.util.similarity import SimilarityFunction
92
93
94 def autocast_ctx ():
95 """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
96 if not torch.cuda.is_available():
97 return nullcontext()
98 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
99 return torch.autocast( "cuda" , dtype = dtype)
100
101
102 def log_trackio_dashboard ():
103 """Surface the Trackio dashboard URL so the user can watch training live."""
104 try :
105 from huggingface_hub import whoami
106
107 hf_user = whoami().get( "name" )
108 if hf_user:
109 logging.info(
110 f "Trackio dashboard (live training progress): https://huggingface.co/spaces/ { hf_user } /trackio"
111 )
112 except Exception :
113 pass
114
115
116 TEACHER_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
117 STUDENT_MODEL_NAME = "distilbert/distilbert-base-uncased"
118
119 CORPUS_DATASET = "sentence-transformers/all-nli"
120 CORPUS_SUBSET = "pair"
121
122 TRAIN_SIZE = 50_000
123 EVAL_SIZE = 1_000
124 OUTPUT_DIR = "models/distilbert-distilled-from-mpnet"
125 RUN_NAME = "distilbert-distill-from-mpnet"
126
127 TEACHER_ENCODE_BATCH_SIZE = 256
128 TRAIN_BATCH_SIZE = 128
129 SMOKE_TEST = os.environ.get( "SMOKE_TEST" ) == "1"
130
131
132 def setup_logging ():
133 """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
134 os.makedirs( "logs" , exist_ok = True )
135 logging.basicConfig(
136 format = " %(asctime)s - %(message)s " ,
137 datefmt = "%Y-%m- %d %H:%M:%S" ,
138 level = logging. INFO ,
139 handlers = [logging.StreamHandler(), logging.FileHandler( f "logs/ { RUN_NAME } .log" )],
140 force = True ,
141 )
142 for noisy in ( "httpx" , "httpcore" , "huggingface_hub" , "urllib3" , "filelock" , "fsspec" ):
143 logging.getLogger(noisy).setLevel(logging. WARNING )
144 if torch.cuda.is_available():
145 torch.set_float32_matmul_precision( "high" ) # TF32 on Ampere+, no quality loss
146
147
148 def main () -> None :
149 parser = argparse.ArgumentParser()
150 parser.add_argument(
151 "--eval-only" , type = str , default = None , help = "Skip training; load this saved model and run only the evaluator."
152 )
153 cli, _ = parser.parse_known_args()
154
155 setup_logging()
156
157 if cli.eval_only:
158 logging.info( f "Eval-only mode: loading model from { cli.eval_only } " )
159 model = SentenceTransformer(cli.eval_only)
160 stsb = load_dataset( "sentence-transformers/stsb" , split = "validation" )
161 evaluator = EmbeddingSimilarityEvaluator(
162 sentences1 = stsb[ "sentence1" ],
163 sentences2 = stsb[ "sentence2" ],
164 scores = stsb[ "score" ],
165 main_similarity = SimilarityFunction. COSINE ,
166 name = "sts-dev" ,
167 )
168 with autocast_ctx():
169 evaluator(model)
170 return
171
172 logging.info( f "Loading teacher: { TEACHER_MODEL_NAME } " )
173 teacher = SentenceTransformer( TEACHER_MODEL_NAME )
174
175 logging.info( f "Loading student: { STUDENT_MODEL_NAME } " )
176 student = SentenceTransformer(
177 STUDENT_MODEL_NAME ,
178 model_card_data = SentenceTransformerModelCardData(
179 language = "en" ,
180 license = "apache-2.0" ,
181 model_name = f " { STUDENT_MODEL_NAME .split( '/' )[ - 1 ] } distilled from { TEACHER_MODEL_NAME .split( '/' )[ - 1 ] } " ,
182 ),
183 )
184 # Match the teacher's final Normalize. MSELoss against unit-norm targets fights student
185 # outputs at norm ~5-10 and can silently regress
186 if any ( isinstance (m, Normalize) for m in teacher) and not any ( isinstance (m, Normalize) for m in student):
187 student.append(Normalize())
188
189 if student.get_embedding_dimension() != teacher.get_embedding_dimension():
190 raise SystemExit (
191 f "Student dim ( { student.get_embedding_dimension() } ) != teacher dim "
192 f "( { teacher.get_embedding_dimension() } ). Plain MSELoss requires matching dims. "
193 "Either pick a student with matching dim, or add a Dense projection layer "
194 "(PCA-initialized from teacher embeddings). See the 'MISMATCHED EMBEDDING DIMS' "
195 "section in this script's docstring."
196 )
197
198 logging.info( f "Loading corpus: { CORPUS_DATASET } ( { CORPUS_SUBSET } )" )
199 train_size = 50 if SMOKE_TEST else TRAIN_SIZE
200 eval_size = 20 if SMOKE_TEST else EVAL_SIZE
201 if SMOKE_TEST :
202 logging.info( "SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push" )
203 raw = load_dataset( CORPUS_DATASET , CORPUS_SUBSET , split = "train" )
204 sentences = list ( dict .fromkeys(s for row in raw for s in (row[ "anchor" ], row[ "positive" ]) if isinstance (s, str )))
205 sentences = sentences[: train_size + eval_size]
206 train_sentences = sentences[:train_size]
207 eval_sentences = sentences[train_size : train_size + eval_size]
208
209 logging.info( f "Encoding { len (train_sentences) :,} training sentences with the teacher (may take a while)" )
210 teacher_train = teacher.encode(
211 train_sentences, batch_size = TEACHER_ENCODE_BATCH_SIZE , convert_to_numpy = True , show_progress_bar = True
212 )
213
214 logging.info( f "Encoding { len (eval_sentences) :,} eval sentences with the teacher" )
215 teacher_eval = teacher.encode(
216 eval_sentences, batch_size = TEACHER_ENCODE_BATCH_SIZE , convert_to_numpy = True , show_progress_bar = True
217 )
218
219 train_dataset = Dataset.from_dict({ "sentence" : train_sentences, "label" : teacher_train.tolist()})
220 eval_dataset = Dataset.from_dict({ "sentence" : eval_sentences, "label" : teacher_eval.tolist()})
221
222 logging.info( f "Building training dataset ( { len (train_dataset) :,} ) and eval dataset ( { len (eval_dataset) :,} )" )
223
224 loss = MSELoss( model = student)
225
226 logging.info( "Setting up STS-B evaluator for quality tracking" )
227 stsb = load_dataset( "sentence-transformers/stsb" , split = "validation" )
228 evaluator = EmbeddingSimilarityEvaluator(
229 sentences1 = stsb[ "sentence1" ],
230 sentences2 = stsb[ "sentence2" ],
231 scores = stsb[ "score" ],
232 main_similarity = SimilarityFunction. COSINE ,
233 name = "sts-dev" ,
234 )
235 logging.info( "Teacher performance:" )
236 evaluator(teacher)
237 logging.info( "Student performance before distillation:" )
238 # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
239 baseline_eval = evaluator(student)[evaluator.primary_metric]
240 metric_key = f "eval_ { evaluator.primary_metric } "
241
242 args = SentenceTransformerTrainingArguments(
243 output_dir = OUTPUT_DIR ,
244 num_train_epochs = 1 ,
245 max_steps = 1 if SMOKE_TEST else - 1 ,
246 per_device_train_batch_size = TRAIN_BATCH_SIZE ,
247 per_device_eval_batch_size = TRAIN_BATCH_SIZE ,
248 learning_rate = 1e-4 ,
249 weight_decay = 0.01 ,
250 warmup_steps = 0.1 ,
251 bf16 = True ,
252 eval_strategy = "steps" ,
253 eval_steps = 0.1 ,
254 save_strategy = "steps" ,
255 save_steps = 0.1 ,
256 save_total_limit = 2 ,
257 logging_steps = 0.01 ,
258 logging_first_step = True ,
259 load_best_model_at_end = True ,
260 metric_for_best_model = metric_key,
261 greater_is_better = True ,
262 report_to = "none" if SMOKE_TEST else "trackio" ,
263 run_name = RUN_NAME ,
264 seed = 12 ,
265 )
266
267 trainer = SentenceTransformerTrainer(
268 model = student,
269 args = args,
270 train_dataset = train_dataset,
271 eval_dataset = eval_dataset,
272 loss = loss,
273 evaluator = evaluator,
274 )
275 if not SMOKE_TEST :
276 log_trackio_dashboard()
277 trainer.train()
278
279 logging.info( "Student performance after distillation:" )
280 score = evaluator(student)[evaluator.primary_metric]
281 delta = score - baseline_eval
282 verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
283 logging.info( f "VERDICT: { verdict } | score= { score :.4f} | baseline= { baseline_eval :.4f} | delta= { delta :+.4f} " )
284
285 final_dir = f " { OUTPUT_DIR } /final"
286 student.save_pretrained(final_dir)
287 logging.info( f "Saved to { final_dir } " )
288
289 if SMOKE_TEST :
290 logging.info( "SMOKE_TEST=1: skipping Hub push" )
291 return
292
293 try :
294 commit_url = student.push_to_hub( RUN_NAME )
295 logging.info( f "Pushed model to { commit_url.rsplit( '/commit/' , 1 )[ 0 ] } " )
296 except Exception :
297 import traceback
298
299 logging.error( f "Hub push failed: \n{ traceback.format_exc() } " )
300
301
302 if __name__ == "__main__" :
303 main()