Check input + max_new_tokens against max_out_tokens in InferenceEngine.generate

deepspeed/inference/engine.py

@@ -621,7 +621,7 @@ def _generate(self, *inputs, **kwargs):
         num_beams = 1
         if "generation_config" in kwargs:
             gen_config = kwargs["generation_config"]
-            num_beams = getattr(gen_config, "num_beams", 1)
+            num_beams = getattr(gen_config, "num_beams", 1) or 1
         if "num_beams" in kwargs:
             num_beams = kwargs["num_beams"]
 
@@ -630,12 +630,41 @@ def _generate(self, *inputs, **kwargs):
                                       "add your request to: https://github.com/deepspeedai/DeepSpeed/issues/2506")
 
         if ("input_ids" in kwargs) and (kwargs["input_ids"].dim() == 2):
+            # max_out_tokens bounds the KV-cache workspace for the whole generation
+            # (input + newly generated tokens). max_length is already a total
+            # (input + new tokens) budget in its own right, so compare against
+            # whichever of max_new_tokens/max_length the caller actually set.
+            max_new_tokens = kwargs.get("max_new_tokens", None)
+            max_length = kwargs.get("max_length", None)
+            if "generation_config" in kwargs:
+                gen_config = kwargs["generation_config"]
+                if max_new_tokens is None:
+                    max_new_tokens = getattr(gen_config, "max_new_tokens", None)
+                if max_length is None:
+                    max_length = getattr(gen_config, "max_length", None)
             for input_tensor in kwargs["input_ids"]:
                 tensor_length = input_tensor.shape[-1]
-                if tensor_length > self._config.max_out_tokens:
+                # transformers gives max_new_tokens precedence over max_length when both
+                # are set, so mirror that instead of taking the larger of the two. When only
+                # max_length is set it is a total (input + new tokens), but it can be smaller
+                # than the input; max(tensor_length, max_length) keeps the guard comparing
+                # against at least the prompt length so it still fires (and transformers'
+                # own _validate_generated_length also rejects max_length < input afterwards).
+                if max_new_tokens is not None:
+                    total_length = tensor_length + max_new_tokens
+                elif max_length is not None:
+                    total_length = max(tensor_length, max_length)
+                else:
+                    total_length = tensor_length
+                if total_length > self._config.max_out_tokens:
+                    if max_new_tokens is None and max_length is None:
+                        raise RuntimeError(f"Input with size {tensor_length} exceeds maximum length of "
+                                           f"{self._config.max_out_tokens}. Please increase max_tokens in the "
+                                           "DeepSpeed Inference Config.")
                     raise RuntimeError(
-                        f"Input with size {tensor_length} exceeds maximum length of {self._config.max_out_tokens}. Please increase max_tokens in the DeepSpeed Inference Config."
-                    )
+                        f"Input with size {tensor_length}, max_new_tokens {max_new_tokens}, and max_length "
+                        f"{max_length} together exceed maximum length of {self._config.max_out_tokens}. "
+                        "Please increase max_tokens in the DeepSpeed Inference Config.")
 
         return self.module.generate(*inputs, **kwargs)
 

tests/unit/inference/test_inference_config.py

@@ -21,6 +21,106 @@ def test_max_out_tokens_must_be_positive(field, value):
         DeepSpeedInferenceConfig(**{field: value})
 
 
+@pytest.mark.inference
+def test_generate_rejects_input_plus_max_new_tokens_over_budget():
+    # Regression test for https://github.com/deepspeedai/DeepSpeed/issues/3081.
+    # max_out_tokens sizes the KV-cache workspace for the whole generation (input
+    # tokens plus newly generated ones), but the length check in
+    # InferenceEngine._generate only compared the input length against it, so a
+    # request whose input alone fit could still push input + max_new_tokens past
+    # the workspace size and reach the wrapped module's generate() unguarded.
+    class GenerateStub(torch.nn.Module):
+
+        def generate(self, *args, **kwargs):
+            return "reached-generate"
+
+    engine = deepspeed.init_inference(GenerateStub(), config={"max_out_tokens": 100, "dtype": torch.float32})
+
+    with pytest.raises(RuntimeError, match="exceed"):
+        engine.generate(input_ids=torch.zeros((1, 90), dtype=torch.long), max_new_tokens=20)
+
+    # An input that fits alongside its requested new tokens must still be let through.
+    assert engine.generate(input_ids=torch.zeros((1, 58), dtype=torch.long), max_new_tokens=20) == "reached-generate"
+
+
+@pytest.mark.inference
+def test_generate_rejects_max_length_over_budget():
+    # max_length is already a total (input + new tokens) budget in its own right,
+    # so an overflow expressed that way must be caught even when max_new_tokens
+    # is never set.
+    class GenerateStub(torch.nn.Module):
+
+        def generate(self, *args, **kwargs):
+            return "reached-generate"
+
+    engine = deepspeed.init_inference(GenerateStub(), config={"max_out_tokens": 100, "dtype": torch.float32})
+
+    with pytest.raises(RuntimeError, match="exceed"):
+        engine.generate(input_ids=torch.zeros((1, 90), dtype=torch.long), max_length=4000)
+
+
+@pytest.mark.inference
+def test_generate_rejects_input_over_budget_when_max_length_is_smaller_than_input():
+    # max_length is a total, but a caller can set it below the input length. Comparing
+    # against max_length alone then uses a number below the prompt size, so an input that
+    # overflows the KV-cache workspace slips through whenever max_length < input. The guard
+    # must compare against max(tensor_length, max_length) so the prompt length still counts.
+    class GenerateStub(torch.nn.Module):
+
+        def generate(self, *args, **kwargs):
+            return "reached-generate"
+
+    engine = deepspeed.init_inference(GenerateStub(), config={"max_out_tokens": 1024, "dtype": torch.float32})
+
+    # input (2000) > max_out_tokens (1024) > max_length (512): must still raise.
+    with pytest.raises(RuntimeError, match="exceed"):
+        engine.generate(input_ids=torch.zeros((1, 2000), dtype=torch.long), max_length=512)
+
+    # A total that genuinely fits (both input and max_length under budget) still passes.
+    assert engine.generate(input_ids=torch.zeros((1, 800), dtype=torch.long),
+                           max_length=900) == "reached-generate"
+
+
+@pytest.mark.inference
+def test_generate_lets_fitting_call_through_when_both_budgets_set():
+    # transformers gives max_new_tokens precedence over max_length when both are
+    # set (see GenerationMixin._prepare_generated_length), so the guard must
+    # mirror that precedence rather than reject on whichever of the two is
+    # larger: a call whose *effective* total (input + max_new_tokens) fits must
+    # not be rejected just because the unused max_length would not have fit.
+    class GenerateStub(torch.nn.Module):
+
+        def generate(self, *args, **kwargs):
+            return "reached-generate"
+
+    engine = deepspeed.init_inference(GenerateStub(), config={"max_out_tokens": 100, "dtype": torch.float32})
+
+    assert engine.generate(input_ids=torch.zeros((1, 90), dtype=torch.long), max_new_tokens=5,
+                           max_length=4000) == "reached-generate"
+
+
+@pytest.mark.inference
+def test_generate_with_generation_config_does_not_crash_and_still_checks_budget():
+    # transformers' GenerationConfig defaults num_beams to None (not 1), so the
+    # pre-existing `getattr(gen_config, "num_beams", 1)` returned None instead of
+    # falling back, and `None > 1` raised TypeError before the length guard below
+    # it ever ran -- silently making the generation_config fallback unreachable.
+    from transformers import GenerationConfig
+
+    class GenerateStub(torch.nn.Module):
+
+        def generate(self, *args, **kwargs):
+            return "reached-generate"
+
+    engine = deepspeed.init_inference(GenerateStub(), config={"max_out_tokens": 100, "dtype": torch.float32})
+
+    gen_config = GenerationConfig(max_new_tokens=20)
+    assert gen_config.num_beams is None
+
+    with pytest.raises(RuntimeError, match="exceed"):
+        engine.generate(input_ids=torch.zeros((1, 90), dtype=torch.long), generation_config=gen_config)
+
+
 @pytest.mark.inference
 @pytest.mark.skipif(not get_accelerator().is_available(), reason="requires accelerator")
 class TestInferenceCudaGraphConfig: