Use of logits_processors has become very slow in v0.3.2

[saattrupdan]

I am using vLLM together with outlines for structured generation.

After having upgraded from v0.3.2, generation became very slow, and the RAM usage leads to OOM crashes now.

Here is a minimal example:

from vllm import LLM, SamplingParams
from outlines.serve.vllm import JSONLogitsProcessor
from pydantic import BaseModel, conlist
import datetime as dt

class Output(BaseModel):
    names: conlist(str, max_length=5)
    organizations: conlist(str, max_length=5)
    locations: conlist(str, max_length=5)
    miscellanous: conlist(str, max_length=5)

llm = LLM('mistralai/Mistral-7B-v0.1', max_model_len=10_000, gpu_memory_utilization=0.9)
logits_processor = JSONLogitsProcessor(schema=Output, llm=llm.llm_engine)
logits_processor.fsm.vocabulary = list(logits_processor.fsm.vocabulary)
prompt = """
Locate all the names, organizations, locations and other miscellaneous entities in the following sentence: 
"Charles went and saw Anna at the coffee shop Starbucks, which was based in a small town in Germany called Essen."
"""
sampling_params = SamplingParams(max_tokens=128, temperature=0, logits_processors=[logits_processor])

t0 = dt.datetime.now()
llm.generate([prompt] * 256, sampling_params=sampling_params)
time_elapsed = (dt.datetime.now() - t0).total_seconds()
print(f"Generation took {time_elapsed:,} seconds.")

When I run the above with vllm==0.3.1, the generation takes 58 seconds and use ~6GB memory, but if I upgrade vllm to v0.3.2 (and none of the other packages are changed), then suddenly the generation takes 418 seconds and spend ~18GB memory. Almost all of the time is spent stalling, not generating anything, but slowly using more and more memory, until it finally begins to generate.

I tried installing a forked version of outlines to see if the stalling was due to the internals of the JSONLogitsProcessor, but it is only called after the "stalling process" is done, so it seems like this is a vLLM issue.

[saattrupdan]

This seems to be due to the deepcopy of the SamplingParams on this line in the LLMEngine, which will thus also copy out the logits processors, which take up a considerable amount of memory in my case. This was added 2 weeks ago in this PR, and which is part of the v0.3.2 release.

Tagging relevant people to that PR: @njhill @Yard1

[simon-mo]

@Yard1 I think this should be fixed by next release, especially we are shipping #2819

[Yard1]

Hmm, I see, we should probably make it so that the logit processors are exempt from the deepcopy (unless #2819 already fixes that)

[njhill]

Ah, yes sorry about this. I can open a PR to do what @Yard1 suggests.

[njhill]

@Yard1 @simon-mo @saattrupdan fix is in #3099