fix OOM bug: add torch.OutOfMemoryError exception

Author: hanhainebula

FlagEmbedding/inference/embedder/decoder_only/base.py

@@ -180,7 +180,9 @@ def encode_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = last_token_pool(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.OutofMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/decoder_only/icl.py

@@ -278,7 +278,9 @@ def encode_queries_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = last_token_pool(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.OutofMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode
@@ -389,7 +391,9 @@ def encode_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = last_token_pool(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.OutofMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/encoder_only/base.py

@@ -170,7 +170,9 @@ def encode_single_device(
                 last_hidden_state = self.model(**test_inputs_batch, return_dict=True).last_hidden_state
                 embeddings = self.pooling(last_hidden_state, test_inputs_batch['attention_mask'])
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.OutofMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/embedder/encoder_only/m3.py

@@ -304,7 +304,9 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
                     return_colbert_vecs=return_colbert_vecs
                 )
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.OutofMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         # encode

FlagEmbedding/inference/reranker/decoder_only/base.py

@@ -317,7 +317,9 @@ def compute_score_single_gpu(
 
                 self.model(**batch_inputs, output_hidden_states=True)
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            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

FlagEmbedding/inference/reranker/decoder_only/layerwise.py

@@ -225,7 +225,9 @@ def compute_score_single_gpu(
 
                 self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            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

FlagEmbedding/inference/reranker/decoder_only/lightweight.py

@@ -300,7 +300,9 @@ def compute_score_single_gpu(
                     cutoff_layers=cutoff_layers
                 )
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.OutofMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         all_scores = []

FlagEmbedding/inference/reranker/encoder_only/base.py

@@ -134,7 +134,9 @@ def compute_score_single_gpu(
                 ).to(device)
                 scores = self.model(**test_inputs_batch, return_dict=True).logits.view(-1, ).float()
                 flag = True
-            except (RuntimeError, torch.OutofMemoryError) as e:
+            except RuntimeError as e:
+                batch_size = batch_size * 3 // 4
+            except torch.OutofMemoryError as e:
                 batch_size = batch_size * 3 // 4
 
         all_scores = []