Skip to content
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.

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

Support eos_token_id from generation_config.json #4182

Merged
merged 3 commits into from
Apr 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
19 changes: 17 additions & 2 deletions vllm/engine/llm_engine.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import time
from typing import Iterable, List, Optional, Type, Union

from transformers import PreTrainedTokenizer
from transformers import GenerationConfig, PreTrainedTokenizer

import vllm
from vllm.config import (CacheConfig, DecodingConfig, DeviceConfig, LoadConfig,
Expand Down Expand Up @@ -34,6 +34,17 @@
_LOCAL_LOGGING_INTERVAL_SEC = 5


def _load_generation_config_dict(model_config: ModelConfig):
try:
return GenerationConfig.from_pretrained(
model_config.model,
revision=model_config.revision,
).to_diff_dict()
except OSError:
# Not found.
return {}


class LLMEngine:
"""An LLM engine that receives requests and generates texts.

Expand Down Expand Up @@ -124,6 +135,8 @@ def __init__(
self._init_tokenizer()
self.detokenizer = Detokenizer(self.tokenizer)
self.seq_counter = Counter()
self.generation_config_fields = _load_generation_config_dict(
model_config)

self.model_executor = executor_class(
model_config=model_config,
Expand Down Expand Up @@ -391,6 +404,8 @@ def add_request(
# inject the eos token id into the sampling_params to support min_tokens
# processing
sampling_params.eos_token_id = seq.eos_token_id
sampling_params.update_from_generation_config(
self.generation_config_fields)

# Create the sequence group.
seq_group = SequenceGroup(request_id, [seq], sampling_params,
Expand Down Expand Up @@ -435,7 +450,7 @@ def _process_model_outputs(
scheduled_seq_groups: List[SequenceGroup],
ignored_seq_groups: List[SequenceGroup]) -> List[RequestOutput]:
"""Apply the model output to the sequences in the scheduled seq groups.

Returns RequestOutputs that can be returned to the client.
"""

Expand Down
14 changes: 13 additions & 1 deletion vllm/sampling_params.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import copy
from enum import IntEnum
from functools import cached_property
from typing import Callable, List, Optional, Union
from typing import Any, Callable, Dict, List, Optional, Union

import torch
from pydantic import Field
Expand Down Expand Up @@ -271,6 +271,18 @@ def _verify_greedy_sampling(self) -> None:
raise ValueError("best_of must be 1 when using greedy sampling."
f"Got {self.best_of}.")

def update_from_generation_config(
self, generation_config: Dict[str, Any]) -> None:
"""Update if there are non-default values from generation_config"""
# Update eos_token_id for generation
if eos_ids := generation_config.get("eos_token_id"):
# it can be either int or list of int
if isinstance(eos_ids, int):
eos_ids = [eos_ids]
original_stop_token_ids = set(self.stop_token_ids)
original_stop_token_ids.update(eos_ids)
self.stop_token_ids = list(original_stop_token_ids)

@cached_property
def sampling_type(self) -> SamplingType:
if self.use_beam_search:
Expand Down
Loading