Refine compile threshold estimation for sentence transformers

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)"