Reduce tensor-scoring inference memory in Hugging Face test translation path

silnlp/nmt/hugging_face_config.py

@@ -1491,61 +1491,67 @@ def _translate_sentences(
         batch_size: int = self._config.infer["infer_batch_size"]
 
         dictionary = self._get_dictionary()
-        if vrefs is None or len(dictionary) == 0:
-            yield from self._translate_sentence_helper(
-                pipeline,
-                sentences,
-                batch_size,
-                return_tensors,
-                produce_multiple_translations=produce_multiple_translations,
-            )
-        else:
-            for batch, force_words in batch_sentences(sentences, vrefs, batch_size, dictionary):
-                if force_words is None:
-                    force_words_ids = None
-                else:
-                    force_words_ids = [[tokenizer.convert_tokens_to_ids(v) for v in vs] for vs in force_words]
-                    force_words_ids = prune_sublists(force_words_ids)
-
-                yield from self._translate_sentence_helper(
-                    pipeline,
-                    batch,
-                    batch_size,
-                    return_tensors,
-                    force_words_ids,
-                    produce_multiple_translations=produce_multiple_translations,
+        current_batch_size = batch_size
+        for batch, force_words in batch_sentences(sentences, vrefs, batch_size, dictionary):
+            if force_words is None:
+                force_words_ids = None
+            else:
+                force_words_ids = [[tokenizer.convert_tokens_to_ids(v) for v in vs] for vs in force_words]
+                force_words_ids = prune_sublists(force_words_ids)
+
+            batch_list = list(batch)
+            index = 0
+            while index < len(batch_list):
+                effective_size = min(current_batch_size, len(batch_list) - index)
+                sub_batch = batch_list[index : index + effective_size]
+                sub_force_words_ids = (
+                    force_words_ids[index : index + effective_size] if force_words_ids is not None else None
                 )
+                try:
+                    yield from self._translate_batch(
+                        pipeline,
+                        sub_batch,
+                        return_tensors,
+                        sub_force_words_ids,
+                        produce_multiple_translations=produce_multiple_translations,
+                    )
+                    index += effective_size
+                except RuntimeError as e:
+                    if not _should_reduce_batch_size(e) or current_batch_size <= 1:
+                        raise
+                    current_batch_size //= 2
+                    LOGGER.warning(
+                        "OOM during translation inference; reducing batch size to %d and retrying current batch.",
+                        current_batch_size,
+                    )
+                    gc.collect()
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
 
-    def _translate_sentence_helper(
+    def _translate_batch(
         self,
         pipeline: TranslationPipeline,
-        sentences: Iterable[TSent],
-        batch_size: int,
+        batch: List[TSent],
         return_tensors: bool,
         force_words_ids: List[List[List[int]]] = None,
         produce_multiple_translations: bool = False,
     ) -> Iterable[ModelOutputGroup]:
-
         num_drafts = self.get_num_drafts()
         if produce_multiple_translations and num_drafts > 1:
             multiple_translations_method: str = self._config.infer.get("multiple_translations_method")
 
-            sentences = list(sentences)
-
             if multiple_translations_method == "hybrid":
                 beam_search_results: List[dict] = self._translate_with_beam_search(
                     pipeline,
-                    sentences,
-                    batch_size,
+                    batch,
                     return_tensors,
                     num_return_sequences=1,
                     force_words_ids=force_words_ids,
                 )
 
                 sampling_results: List[dict] = self._translate_with_sampling(
                     pipeline,
-                    sentences,
-                    batch_size,
+                    batch,
                     return_tensors,
                     num_return_sequences=num_drafts - 1,
                     force_words_ids=force_words_ids,
@@ -1562,8 +1568,7 @@ def _translate_sentence_helper(
                     ModelOutputGroup(result)
                     for result in self._translate_with_sampling(
                         pipeline,
-                        sentences,
-                        batch_size,
+                        batch,
                         return_tensors,
                         num_return_sequences=num_drafts,
                         force_words_ids=force_words_ids,
@@ -1575,8 +1580,7 @@ def _translate_sentence_helper(
                     ModelOutputGroup(result)
                     for result in self._translate_with_beam_search(
                         pipeline,
-                        sentences,
-                        batch_size,
+                        batch,
                         return_tensors,
                         num_return_sequences=num_drafts,
                         force_words_ids=force_words_ids,
@@ -1588,8 +1592,7 @@ def _translate_sentence_helper(
                     ModelOutputGroup(result)
                     for result in self._translate_with_diverse_beam_search(
                         pipeline,
-                        sentences,
-                        batch_size,
+                        batch,
                         return_tensors,
                         num_return_sequences=num_drafts,
                         force_words_ids=force_words_ids,
@@ -1603,8 +1606,7 @@ def _translate_sentence_helper(
                 ModelOutputGroup([translated_sentence[0]])
                 for translated_sentence in self._translate_with_beam_search(
                     pipeline,
-                    sentences,
-                    batch_size,
+                    batch,
                     return_tensors,
                     num_return_sequences=1,
                     force_words_ids=force_words_ids,
@@ -1617,11 +1619,15 @@ def _translate_sentence_helper(
     def _flatten_tokenized_translations(self, pipeline_output) -> List[dict]:
         return [[i if isinstance(i, dict) else i[0] for i in translation] for translation in pipeline_output]
 
+    def _normalize_batch_translations(self, batch_translations, num_return_sequences: int) -> List[List[dict]]:
+        if num_return_sequences == 1:
+            batch_translations = [[t] for t in batch_translations]
+        return self._flatten_tokenized_translations(batch_translations)
+
     def _translate_with_beam_search(
         self,
         pipeline: TranslationPipeline,
-        sentences: Iterable[TSent],
-        batch_size: int,
+        batch: List[TSent],
         return_tensors: bool,
         num_return_sequences: int = 1,
         force_words_ids: List[List[List[int]]] = None,
@@ -1630,54 +1636,41 @@ def _translate_with_beam_search(
         if num_beams is None:
             num_beams = self._config.params.get("generation_num_beams")
 
-        translations = pipeline(
-            sentences,
+        batch_translations = pipeline(
+            batch,
             num_beams=num_beams,
             num_return_sequences=num_return_sequences,
             force_words_ids=force_words_ids,
-            batch_size=batch_size,
             return_text=not return_tensors,
             return_tensors=return_tensors,
         )
-
-        if num_return_sequences == 1:
-            translations = [[t] for t in translations]
-
-        return self._flatten_tokenized_translations(translations)
+        return self._normalize_batch_translations(batch_translations, num_return_sequences)
 
     def _translate_with_sampling(
         self,
         pipeline: TranslationPipeline,
-        sentences: Iterable[TSent],
-        batch_size: int,
+        batch: List[TSent],
         return_tensors: bool,
         num_return_sequences: int = 1,
         force_words_ids: List[List[List[int]]] = None,
     ) -> List[List[dict]]:
-
         temperature: Optional[int] = self._config.infer.get("temperature")
 
-        translations = pipeline(
-            sentences,
+        batch_translations = pipeline(
+            batch,
             do_sample=True,
             temperature=temperature,
             num_return_sequences=num_return_sequences,
             force_words_ids=force_words_ids,
-            batch_size=batch_size,
             return_text=not return_tensors,
             return_tensors=return_tensors,
         )
-
-        if num_return_sequences == 1:
-            translations = [[t] for t in translations]
-
-        return self._flatten_tokenized_translations(translations)
+        return self._normalize_batch_translations(batch_translations, num_return_sequences)
 
     def _translate_with_diverse_beam_search(
         self,
         pipeline: TranslationPipeline,
-        sentences: Iterable[TSent],
-        batch_size: int,
+        batch: List[TSent],
         return_tensors: bool,
         num_return_sequences: int = 1,
         force_words_ids: List[List[List[int]]] = None,
@@ -1687,22 +1680,17 @@ def _translate_with_diverse_beam_search(
             num_beams = self._config.params.get("generation_num_beams")
         diversity_penalty: Optional[float] = self._config.infer.get("diversity_penalty")
 
-        translations = pipeline(
-            sentences,
+        batch_translations = pipeline(
+            batch,
             num_beams=num_beams,
             num_beam_groups=num_beams,
             num_return_sequences=num_return_sequences,
             diversity_penalty=diversity_penalty,
             force_words_ids=force_words_ids,
-            batch_size=batch_size,
             return_text=not return_tensors,
             return_tensors=return_tensors,
         )
-
-        if num_return_sequences == 1:
-            translations = [[t] for t in translations]
-
-        return self._flatten_tokenized_translations(translations)
+        return self._normalize_batch_translations(batch_translations, num_return_sequences)
 
     def _create_inference_model(
         self,

tests/smoke_tests/test_experiment.py

@@ -1,13 +1,14 @@
 import shutil
 from pathlib import Path
 from typing import cast
+from unittest.mock import Mock
 
 import torch
 
 from silnlp.common.environment import SilNlpEnv
 from silnlp.nmt.config_utils import load_config_from_exp_dir
 from silnlp.nmt.experiment import SILExperiment
-from silnlp.nmt.hugging_face_config import HuggingFaceConfig
+from silnlp.nmt.hugging_face_config import HuggingFaceConfig, HuggingFaceNMTModel
 from tests.smoke_tests.mock_pretrained_model import (
     MockModelOutput,
     MockPreTrainedModelProviderFactory,
@@ -16,6 +17,7 @@
 
 TEST_MT_DIR = Path(__file__).parent
 EXPERIMENT_NAME = "test_experiment"
+OOM_ERROR_MESSAGE = "CUDA out of memory. Tried to allocate 7.00 GiB."
 
 
 def test_experiment_full_pipeline():
@@ -32,6 +34,67 @@ def test_experiment_full_pipeline():
     clean_experiment_directory(environment.get_mt_exp_dir(EXPERIMENT_NAME))
 
 
+def test_translate_sentences_does_not_cap_batch_size_with_tensors():
+    model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel))
+    model._config = cast(HuggingFaceConfig, Mock(infer={"infer_batch_size": 4}))
+    model._get_dictionary = Mock(return_value={})
+    captured_sub_batch_size = {}
+
+    def fake_translate_batch(
+        pipeline, batch, return_tensors, force_words_ids=None, produce_multiple_translations=False
+    ):
+        captured_sub_batch_size["value"] = len(batch)
+        return iter(())
+
+    model._translate_batch = fake_translate_batch
+
+    list(
+        HuggingFaceNMTModel._translate_sentences(
+            model,
+            tokenizer=Mock(),
+            pipeline=Mock(),
+            sentences=[["a"], ["b"], ["c"], ["d"]],
+            vrefs=None,
+            return_tensors=True,
+        )
+    )
+
+    assert captured_sub_batch_size["value"] == 4
+
+
+def test_translate_sentences_retries_with_smaller_batch():
+    model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel))
+    model._config = cast(
+        HuggingFaceConfig, Mock(infer={"infer_batch_size": 4, "num_beams": 2}, params={})
+    )
+    model._get_dictionary = Mock(return_value={})
+    call_sub_batch_sizes: list[int] = []
+
+    def fake_translate_batch(
+        pipeline, batch, return_tensors, force_words_ids=None, produce_multiple_translations=False
+    ):
+        call_sub_batch_sizes.append(len(batch))
+        if len(batch) > 2:
+            raise RuntimeError(OOM_ERROR_MESSAGE)
+        return iter(())
+
+    model._translate_batch = fake_translate_batch
+
+    list(
+        HuggingFaceNMTModel._translate_sentences(
+            model,
+            tokenizer=Mock(),
+            pipeline=Mock(),
+            sentences=[["a"], ["b"], ["c"], ["d"]],
+            vrefs=None,
+            return_tensors=False,
+        )
+    )
+
+    # First call is batch-of-4 (OOM), then two calls of batch-of-2 succeed
+    assert call_sub_batch_sizes == [4, 2, 2]
+
+
 def set_up_environment() -> SilNlpEnv:
     # To avoid keeping large numbers of files in the repository, the tests assume that there is an
     # active connection to the MinIO bucket and use file from the "Scripture" and "Paratext" directories.