Repository: vllm-project/vllm

Issue:
Use of `logits_processors` has become very slow in v0.3.2
I am using vLLM together with [`outlines`](https://github.com/outlines-dev/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:

```python
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.


Discussion:
This seems to be due to the `deepcopy` of the `SamplingParams` on [this line in the `LLMEngine`](https://github.com/vllm-project/vllm/blame/e46fa5d52e02ee48d5fdd12b35e39993008b4bd6/vllm/engine/llm_engine.py#L471-L473), 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](https://github.com/vllm-project/vllm/pull/2881), and which is part of the [v0.3.2 release](https://github.com/vllm-project/vllm/releases/tag/v0.3.2). 

Tagging relevant people to that PR: @njhill @Yard1 
@Yard1 I think this should be fixed by next release, especially we are shipping #2819 
Hmm, I see, we should probably make it so that the logit processors are exempt from the deepcopy (unless #2819 already fixes that)
Ah, yes sorry about this. I can open a PR to do what @Yard1 suggests.

Resolving patch:
diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py
--- a/vllm/engine/llm_engine.py
+++ b/vllm/engine/llm_engine.py
@@ -483,8 +483,9 @@ def add_request(
             prompt_token_ids[:prefix_pos], lora_request.lora_int_id
             if lora_request else 0) if prefix_pos is not None else None
 
-        # Defensive copy of SamplingParams, which are used by the sampler
-        sampling_params = copy.deepcopy(sampling_params)
+        # Defensive copy of SamplingParams, which are used by the sampler,
+        # this doesn't deep-copy LogitsProcessor objects
+        sampling_params = sampling_params.clone()
 
         # Create the sequence group.
         seq_group = SequenceGroup(request_id, [seq], sampling_params,
diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py
--- a/vllm/sampling_params.py
+++ b/vllm/sampling_params.py
@@ -1,4 +1,5 @@
 """Sampling parameters for text generation."""
+import copy
 from enum import IntEnum
 from functools import cached_property
 from typing import Callable, List, Optional, Union
@@ -237,6 +238,20 @@ def sampling_type(self) -> SamplingType:
             return SamplingType.RANDOM_SEED
         return SamplingType.RANDOM
 
+    def clone(self) -> "SamplingParams":
+        """Deep copy excluding LogitsProcessor objects.
+
+        LogitsProcessor objects are excluded because they may contain an
+        arbitrary, nontrivial amount of data.
+        See https://github.com/vllm-project/vllm/issues/3087
+        """
+
+        logit_processor_refs = None if self.logits_processors is None else {
+            id(lp): lp
+            for lp in self.logits_processors
+        }
+        return copy.deepcopy(self, memo=logit_processor_refs)
+
     def __repr__(self) -> str:
         return (
             f"SamplingParams(n={self.n}, "


Test patch: