fix(inference): only retry batch-size probing on genuine OOM errors

FlagEmbedding/inference/embedder/decoder_only/base.py

@@ -7,6 +7,7 @@
 
 from FlagEmbedding.abc.inference import AbsEmbedder
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 
 
 # Pooling function for LLM-based embedding models
@@ -214,6 +215,8 @@ def encode_single_device(
         Returns:
             Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
         """
+        validate_batch_size(batch_size)
+
         if device is None:
             device = self.target_devices[0]
 
@@ -263,9 +266,11 @@ def encode_single_device(
                 embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         # encode
         all_embeddings = []

FlagEmbedding/inference/embedder/decoder_only/icl.py

@@ -11,6 +11,7 @@
 
 from FlagEmbedding.abc.inference import AbsEmbedder
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 
 
 # Pooling function for LLM-based embedding models
@@ -343,6 +344,8 @@ def encode_queries_single_device(
         Returns:
             Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
         """
+        validate_batch_size(batch_size)
+
         if device is None:
             device = self.target_devices[0]
 
@@ -418,9 +421,11 @@ def encode_queries_single_device(
                 embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         # encode
         all_embeddings = []
@@ -482,6 +487,8 @@ def encode_single_device(
         Returns:
             Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
         """
+        validate_batch_size(batch_size)
+
         if device is None:
             device = self.target_devices[0]
 
@@ -531,9 +538,11 @@ def encode_single_device(
                 embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         # encode
         all_embeddings = []

FlagEmbedding/inference/embedder/decoder_only/pseudo_moe.py

@@ -5,6 +5,7 @@
 
 from .base import BaseLLMEmbedder, last_token_pool
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 
 
 class PseudoMoELLMEmbedder(BaseLLMEmbedder):
@@ -95,6 +96,8 @@ def encode_single_device(
         device: Optional[str] = None,
         **kwargs: Any
     ):
+        validate_batch_size(batch_size)
+
         if device is None:
             device = self.target_devices[0]
 
@@ -153,10 +156,12 @@ def encode_single_device(
                     last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
                 _ = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
                 flag = True
-            except RuntimeError:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError:
-                batch_size = batch_size * 3 // 4
+            except RuntimeError as e:
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         # encode
         all_embeddings = []

FlagEmbedding/inference/embedder/encoder_only/base.py

@@ -7,6 +7,7 @@
 
 from FlagEmbedding.abc.inference import AbsEmbedder
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 
 
 class BaseEmbedder(AbsEmbedder):
@@ -195,6 +196,8 @@ def encode_single_device(
         Returns:
             Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
         """
+        validate_batch_size(batch_size)
+
         if device is None:
             device = self.target_devices[0]
 
@@ -244,9 +247,11 @@ def encode_single_device(
                 embeddings = self.pooling(last_hidden_state, inputs_batch['attention_mask'])
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         # encode
         all_embeddings = []

FlagEmbedding/inference/embedder/encoder_only/m3.py

@@ -11,6 +11,7 @@
 
 from FlagEmbedding.abc.inference import AbsEmbedder
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 from FlagEmbedding.finetune.embedder.encoder_only.m3 import (
     EncoderOnlyEmbedderM3ModelForInference, EncoderOnlyEmbedderM3Runner
 )
@@ -334,6 +335,8 @@ def encode_single_device(
         Returns:
             Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
         """
+        validate_batch_size(batch_size)
+
         # pop convert_to_numpy from kwargs
         kwargs.pop("convert_to_numpy", None)
 
@@ -412,9 +415,11 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
                 )
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         # encode
         all_dense_embeddings, all_lexical_weights, all_colbert_vecs = [], [], []

FlagEmbedding/inference/reranker/decoder_only/base.py

@@ -11,6 +11,7 @@
 from FlagEmbedding.abc.inference import AbsReranker
 from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat, prepare_for_model_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 
 
 def last_logit_pool(logits: Tensor,
@@ -289,6 +290,7 @@ def compute_score_single_gpu(
         """
         if prompt is None: prompt = self.prompt
         if batch_size is None: batch_size = self.batch_size
+        validate_batch_size(batch_size)
         if max_length is None: max_length = self.max_length
         if query_max_length is None:
             if self.query_max_length is not None:
@@ -417,9 +419,11 @@ def compute_score_single_gpu(
                 self.model(**batch_inputs, output_hidden_states=True)
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         dataset, dataloader = None, None
         if use_dataloader:

FlagEmbedding/inference/reranker/decoder_only/layerwise.py

@@ -11,6 +11,7 @@
 from FlagEmbedding.abc.inference import AbsReranker
 from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
 from FlagEmbedding.utils.tokenizer_compat import prepare_for_model_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 from FlagEmbedding.inference.reranker.decoder_only.base import DatasetForReranker, Collater
 
 from .models.modeling_minicpm_reranker import LayerWiseMiniCPMForCausalLM
@@ -168,6 +169,7 @@ def compute_score_single_gpu(
         if cutoff_layers is None: cutoff_layers = self.cutoff_layers
         if prompt is None: prompt = self.prompt
         if batch_size is None: batch_size = self.batch_size
+        validate_batch_size(batch_size)
         if max_length is None: max_length = self.max_length
         if query_max_length is None:
             if self.query_max_length is not None:
@@ -283,9 +285,11 @@ def compute_score_single_gpu(
                 self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         dataset, dataloader = None, None
         if use_dataloader:

FlagEmbedding/inference/reranker/decoder_only/lightweight.py

@@ -11,6 +11,7 @@
 from FlagEmbedding.abc.inference import AbsReranker
 from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat, prepare_for_model_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 
 
 def last_logit_pool_lightweight(logits: Tensor,
@@ -248,6 +249,7 @@ def compute_score_single_gpu(
         if compress_ratio is None: compress_ratio = self.compress_ratio
         if prompt is None: prompt = self.prompt
         if batch_size is None: batch_size = self.batch_size
+        validate_batch_size(batch_size)
         if max_length is None: max_length = self.max_length
         if query_max_length is None:
             if self.query_max_length is not None:
@@ -378,9 +380,11 @@ def compute_score_single_gpu(
                 )
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         all_scores = []
         for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):

FlagEmbedding/inference/reranker/encoder_only/base.py

@@ -6,6 +6,7 @@
 
 from FlagEmbedding.abc.inference import AbsReranker
 from FlagEmbedding.utils.tokenizer_compat import pad_with_compat, prepare_for_model_compat
+from FlagEmbedding.inference.utils import is_oom_error, shrink_batch_size_on_oom, validate_batch_size
 
 
 def sigmoid(x):
@@ -100,6 +101,7 @@ def compute_score_single_gpu(
             List[float]: Computed scores of queries and passages.
         """
         if batch_size is None: batch_size = self.batch_size
+        validate_batch_size(batch_size)
         if max_length is None: max_length = self.max_length
         if query_max_length is None:
             if self.query_max_length is not None:
@@ -172,9 +174,11 @@ def compute_score_single_gpu(
                 scores = self.model(**test_inputs_batch, return_dict=True).logits.view(-1, ).float()
                 flag = True
             except RuntimeError as e:
-                batch_size = batch_size * 3 // 4
-            except torch.cuda.OutOfMemoryError as e:
-                batch_size = batch_size * 3 // 4
+                # Only genuine OOM errors trigger a retry; any other error
+                # propagates immediately with its original traceback.
+                if not is_oom_error(e):
+                    raise
+                batch_size = shrink_batch_size_on_oom(batch_size, e)
 
         all_scores = []
         for start_index in tqdm(range(0, len(all_inputs_sorted), batch_size), desc="Compute Scores",

FlagEmbedding/inference/utils.py

@@ -0,0 +1,46 @@
+"""Shared helpers for single-device inference batch-size probing.
+
+Several inference paths probe for the largest workable batch size: they run a
+trial forward pass and, on out-of-memory (OOM) errors, shrink the batch size
+and retry. The helpers here implement the policy for that loop:
+
+* Only genuine OOM errors trigger a retry; any other ``RuntimeError``
+  propagates immediately with its original traceback.
+* Retrying stops (re-raising the OOM error) once the batch size would drop
+  below 1 -- a single sample cannot fit, so further shrinking is pointless.
+* The batch size therefore never drops below 1 inside the probing loop.
+"""
+
+from typing import Union
+
+
+def is_oom_error(error: BaseException) -> bool:
+    """Return True if *error* is an out-of-memory failure.
+
+    Matches ``torch.cuda.OutOfMemoryError`` (a ``RuntimeError`` subclass) via
+    its message, as well as allocator messages such as "CUDA out of memory"
+    or the CPU allocator's "not enough memory" errors.
+    """
+    message = str(error).lower()
+    return "out of memory" in message or "not enough memory" in message
+
+
+def shrink_batch_size_on_oom(batch_size: int, error: RuntimeError) -> int:
+    """Return the batch size to retry with after an OOM *error*.
+
+    Re-raises the original *error* when the reduced batch size would drop
+    below 1: a single sample does not fit, so retrying is pointless.
+    """
+    new_batch_size = batch_size * 3 // 4
+    if new_batch_size < 1:
+        raise error
+    return new_batch_size
+
+
+def validate_batch_size(batch_size: Union[int, None]) -> None:
+    """Reject non-positive caller-supplied batch sizes early.
+
+    ``None`` is allowed (callers fall back to the configured default).
+    """
+    if batch_size is not None and batch_size < 1:
+        raise ValueError(f"batch_size must be a positive integer, got {batch_size!r}")