Skip to content

Commit ed11639

Browse files
committed
Validate positive inference min output tokens
min_out_tokens has no lower-bound validation, so a negative or zero value passes config construction and is only caught later inside GenWorkSpace's C++ workspace allocation. Add the same gt=0 constraint #8343 gave the sibling max_out_tokens field. Refs #8339 Signed-off-by: Amir Fathi <amirfathi.me@gmail.com>
1 parent c389bae commit ed11639

2 files changed

Lines changed: 11 additions & 1 deletion

File tree

deepspeed/inference/config.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,7 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel):
265265
to the required token-length required for your use-case.
266266
"""
267267

268-
min_out_tokens: int = Field(1, alias="min_tokens")
268+
min_out_tokens: int = Field(1, alias="min_tokens", gt=0)
269269
"""
270270
This argument communicates to the runtime the minimum number of tokens you
271271
expect you will need to generate. This will cause the runtime to error

tests/unit/inference/test_inference_config.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,20 @@
66
import pytest
77
import torch
88
import deepspeed
9+
from pydantic import ValidationError
10+
from deepspeed.inference.config import DeepSpeedInferenceConfig
911
from unit.common import DistributedTest
1012
from unit.simple_model import create_config_from_dict
1113

1214

15+
@pytest.mark.inference
16+
@pytest.mark.parametrize("field", ["min_out_tokens", "min_tokens"])
17+
@pytest.mark.parametrize("value", [-1, 0])
18+
def test_min_out_tokens_must_be_positive(field, value):
19+
with pytest.raises(ValidationError):
20+
DeepSpeedInferenceConfig(**{field: value})
21+
22+
1323
@pytest.mark.inference
1424
class TestInferenceConfig(DistributedTest):
1525
world_size = 1

0 commit comments

Comments
 (0)