Skip to content

Check input + max_new_tokens against max_out_tokens in InferenceEngine.generate #8592

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
udsy19 wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Check input + max_new_tokens against max_out_tokens in InferenceEngine.generate #8592

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2137403
Check input + max_new_tokens against max_out_tokens in InferenceEngin…
udsy19 Sep 17, 2026
faeb85c
Also check max_length, and fix num_beams None crash blocking it
udsy19 Sep 19, 2026
b718dfd
Mirror transformers' max_new_tokens/max_length precedence in the leng…
udsy19 Sep 19, 2026
298f76f
Keep tensor_length in the max_length branch of the length guard
udsy19 Oct 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 24 additions & 4 deletions deepspeed/inference/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand All @@ -630,12 +630,32 @@ 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:
total_length = tensor_length if max_new_tokens is None else tensor_length + max_new_tokens
if max_length is not None:
total_length = max(total_length, max_length)

@ebarkhordar ebarkhordar Sep 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for picking both of these up so quickly. I ran the new commit in a clean container at faeb85cd (python:3.11-slim, torch 2.14.0+cpu, transformers 5.17.0), and the max_length half over-rejects when a caller sets both budgets.

max(total_length, max_length) takes the larger of the two, but transformers gives max_new_tokens precedence and ignores max_length when both are present, so the larger number is the one that will not be used. A call that fits now fails:

input=90  max_out_tokens=100  max_new_tokens=5  max_length=4000
  21374033 (before this commit) -> NO RAISE
  faeb85cd  -> RuntimeError: Input with size 90, max_new_tokens 5, and max_length 4000
               together exceed maximum length of 100
  transformers on the same inputs: "max_new_tokens will take precedence",
               out len 95, so 5 new tokens inside the budget of 100

The num_beams repair makes this reachable from a generation_config too, where it hit the TypeError first: GenerationConfig(max_length=4096) with max_new_tokens=50 against max_out_tokens=1024 raises, and that generation is really 140 tokens.

Mirroring the precedence instead of taking the max clears both, and your two new tests still pass on it:

if max_new_tokens is not None:
    total_length = tensor_length + max_new_tokens
elif max_length is not None:
    total_length = max_length
else:
    total_length = tensor_length

I ran this on CPU against a stub module, so I have not exercised a real GPU generate path.

@ebarkhordar ebarkhordar Sep 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for turning that around so fast. I ran b718dfd in a clean container (python:3.11-slim, torch 2.14.0+cpu, transformers 5.17.0) and the precedence part is right, but the snippet I handed you dropped tensor_length from the max_length branch, and that opens a gap I should have caught. max_length is a total, so when it is smaller than the input the guard now compares against a number below the input length and stops firing.

input=2000  max_out_tokens=1024  max_length=512
  21374033 -> RuntimeError
  faeb85cd -> RuntimeError
  b718dfdd -> NO RAISE, reaches generate()

same via generation_config=GenerationConfig(max_length=512): NO RAISE at b718dfdd
control, max_length=4000 on that input: still raises at all three

It is not a silent overflow on a real model: _validate_generated_length raises ValueError: Input length of input_ids is 2000, but max_length is set to 512 immediately after. What goes missing is the engine's own error, the one that names max_out_tokens.

Keeping the max inside the branch restores it:

elif max_length is not None:
    total_length = max(tensor_length, max_length)

On that, the two cases above raise again, max_length=900 over an 800 token input still passes through, and the test file's results are unchanged: 8 passed and 6 failed both with and without it, the 6 being a missing c++ for deepspeed_shm_comm plus the CUDA only cuda-graph test on this CPU box.

Sorry for the round trip, that one is on my suggestion.

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)

Expand Down
60 changes: 60 additions & 0 deletions tests/unit/inference/test_inference_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,66 @@ 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_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:
Expand Down
Loading