fix OOM bug: add torch.OutOfMemoryError exception

Author: hanhainebula

FlagEmbedding/inference/embedder/decoder_only/base.py

@@ -182,6 +182,8 @@ def encode_single_device(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         # encode
         all_embeddings = []

FlagEmbedding/inference/embedder/decoder_only/icl.py

@@ -280,6 +280,8 @@ def encode_queries_single_device(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         # encode
         all_embeddings = []
@@ -391,6 +393,8 @@ def encode_single_device(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         # encode
         all_embeddings = []

FlagEmbedding/inference/embedder/encoder_only/base.py

@@ -172,6 +172,8 @@ def encode_single_device(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         # encode
         all_embeddings = []

FlagEmbedding/inference/embedder/encoder_only/m3.py

@@ -306,6 +306,8 @@ 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.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         # encode
         all_dense_embeddings, all_lexical_weights, all_colbert_vecs = [], [], []

FlagEmbedding/inference/reranker/decoder_only/base.py

@@ -319,6 +319,8 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         dataset, dataloader = None, None
         if use_dataloader:

FlagEmbedding/inference/reranker/decoder_only/layerwise.py

@@ -227,6 +227,8 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         dataset, dataloader = None, None
         if use_dataloader:

FlagEmbedding/inference/reranker/decoder_only/lightweight.py

@@ -302,7 +302,9 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
-        
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
+
         all_scores = []
         for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):
             queries_inputs = all_queries_inputs_sorted[batch_start:batch_start+batch_size]

FlagEmbedding/inference/reranker/encoder_only/base.py

@@ -136,6 +136,8 @@ def compute_score_single_gpu(
                 flag = True
             except RuntimeError as e:
                 batch_size = batch_size * 3 // 4
+            except torch.OutOfMemoryError as e:
+                batch_size = batch_size * 3 // 4
 
         all_scores = []
         for start_index in tqdm(range(0, len(all_inputs_sorted), batch_size), desc="Compute Scores",