From 24a8344afa6391d7c1dfc02803dde419f46b5ad9 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 22 Jan 2026 15:20:37 +0000 Subject: [PATCH] Refine compile threshold estimation for sentence transformers --- unsloth/models/sentence_transformer.py | 97 +++++++++++++++++++++++++- 1 file changed, 94 insertions(+), 3 deletions(-) diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 779a39b15a..fa1985fb10 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -927,6 +927,7 @@ class FastSentenceTransformer(FastModel): "mpnet", "bert", "distilbert", + "modernbert", "roberta", "xlm-roberta", "albert", @@ -934,7 +935,12 @@ class FastSentenceTransformer(FastModel): } @staticmethod - def _estimate_compile_threshold(model): + def _estimate_compile_threshold( + model, + batch_size = None, + grad_accum = None, + max_seq_length = None, + ): """ Estimate the minimum training steps needed for torch.compile to be beneficial. Returns the threshold with a 1.2x safety margin built in. @@ -942,6 +948,10 @@ class FastSentenceTransformer(FastModel): Based on empirical benchmarks: - Larger models have lower breakeven (more time saved per step) - Warmup time scales with model size but speedup also increases + + Optional inputs (batch_size, grad_accum, max_seq_length) allow + a coarse pre-run adjustment. These are intentionally conservative + and avoid any runtime measurements. """ # Get parameter count from inner model if hasattr(model, "__getitem__"): @@ -953,6 +963,15 @@ class FastSentenceTransformer(FastModel): else: params = sum(p.numel() for p in model.parameters()) + model_type = None + try: + if "inner" in locals(): + model_type = getattr(getattr(inner, "config", None), "model_type", None) + except Exception: + model_type = None + if isinstance(model_type, str): + model_type = model_type.lower() + params_m = params / 1e6 # Empirical formula based on benchmarks with batch_size=2, grad_accum=4 @@ -979,7 +998,59 @@ class FastSentenceTransformer(FastModel): breakeven = float("inf") # Return threshold with 1.2x safety margin - return int(breakeven * 1.2) + threshold = breakeven * 1.2 + + # Optional adjustment based on expected work per step. + # This uses only pre-run information (batch size, grad accum, seq length). + generic_scale = 1.0 + fast_scale = 1.0 + if batch_size is not None or grad_accum is not None or max_seq_length is not None: + try: + bs = int(batch_size) if batch_size is not None else 2 + ga = int(grad_accum) if grad_accum is not None else 4 + seq = int(max_seq_length) if max_seq_length is not None else 512 + except Exception: + bs, ga, seq = 2, 4, 512 + + bs = max(1, bs) + ga = max(1, ga) + # Guard against unbounded tokenizer.model_max_length + seq = max(64, min(seq, 8192)) + + ref_bs, ref_ga, ref_seq = 2, 4, 512 + + # Generic path: lighter scaling, less conservative than params-only. + ga_scale = (ref_ga / ga) ** 1.0 + bs_seq_scale = ((ref_bs * ref_seq) / (bs * seq)) ** 0.15 + generic_scale = 0.35 * ga_scale * bs_seq_scale + generic_scale = max(0.05, min(generic_scale, 5.0)) + + # Fast encoder path: stronger scaling based on observed behavior. + fast_ga_scale = (ref_ga / ga) ** 1.5 + fast_bs_seq_scale = ((ref_bs * ref_seq) / (bs * seq)) ** 0.25 + fast_scale = 0.2 * fast_ga_scale * fast_bs_seq_scale + fast_scale = max(0.05, min(fast_scale, 5.0)) + + # Conservative safety factors: generic is less conservative than fast. + generic_threshold = threshold * generic_scale * 1.25 + + is_fast_type = ( + isinstance(model_type, str) + and model_type in FastSentenceTransformer.ENCODER_MODEL_TYPES + ) + if is_fast_type: + fast_threshold = threshold * fast_scale * 1.5 + # Prefer the smaller (less conservative) of the two estimates. + final_threshold = min(generic_threshold, fast_threshold) + else: + final_threshold = generic_threshold + + # Reduce mpnet overestimation slightly. + if model_type == "mpnet": + final_threshold *= 0.7 + + # Lower bound to avoid compiling on extremely short runs. + return int(max(20, final_threshold)) @staticmethod def _apply_torch_compile(model, mode = "default"): @@ -1728,9 +1799,29 @@ def _patch_sentence_transformer_trainer(): and getattr(model, "_compile_pending", False) ): max_steps = getattr(training_args, "max_steps", -1) - threshold = getattr(model, "_compile_threshold", 0) compile_mode = getattr(model, "_compile_mode", "default") + # Re-estimate threshold now that training args are available + batch_size = getattr(training_args, "per_device_train_batch_size", None) + grad_accum = getattr(training_args, "gradient_accumulation_steps", None) + max_seq_length = getattr(model, "max_seq_length", None) + if max_seq_length is None and hasattr(model, "__getitem__"): + try: + max_seq_length = getattr(model[0], "max_seq_length", None) + except Exception: + max_seq_length = None + if max_seq_length is None: + tokenizer = getattr(model, "tokenizer", None) + max_seq_length = getattr(tokenizer, "model_max_length", None) if tokenizer is not None else None + + threshold = FastSentenceTransformer._estimate_compile_threshold( + model, + batch_size = batch_size, + grad_accum = grad_accum, + max_seq_length = max_seq_length, + ) + model._compile_threshold = threshold + if max_steps > 0 and max_steps >= threshold: print( f"Unsloth: Auto-compiling model ({max_steps} steps >= {threshold} threshold)"