VLLM tensor-parallel and RegexLogitsProcessor

[BenoitHardier]

Describe the issue as clearly as possible:

Hi,
I recently tried to use the RegexLogitsProcessor with VLLM introduced by #481.

When using it with a "small" model like 7B one on a unique GPU it works fine but when I try with a big one, namely Mixtral, on multiple GPUs with the vllm engine argument tensor-parallel, I ran into several problems (monkey patching not working and fsm_state in the Processor not initialize). I suspect the multiple workers of Ray to be the cause (monkey patching may not be propagated to all the workers same for the fsm_states).

I could have missed some relevant information but It seems that #481 only checks without tensor-parallel.

Steps/code to reproduce the bug:

import vllm
import vllm.model_executor.layers.sampler as sampler
from pydantic import BaseModel

from outlines.serve.vllm import JSONLogitsProcessor, _patched_apply_logits_processors

# Patch the _apply_logits_processors so it is compatible with `JSONLogitsProcessor`
sampler._apply_logits_processors = _patched_apply_logits_processors


class User(BaseModel):
    id: int
    name: str

model = "mistralai/Mixtral-8X7B-Instruct-v0.1"
#model = "mistralai/Mistral-7B-Instruct-v0.2"
llm = vllm.LLM(model=model, dtype='float16', max_model_len=1024,  tensor_parallel_size=4, max_num_seqs=512, enforce_eager=True)
logits_processor = JSONLogitsProcessor(User, llm)
result = llm.generate(
    ["A prompt", "Another prompt"],
    sampling_params=vllm.SamplingParams(
        max_tokens=100, logits_processors=[logits_processor]
    ),
)
print(result)

Expected result:

result from vllm

Error message:

Traceback (most recent call last):
  File "/workspace/./python_scripts/vllm_integration.py", line 19, in <module>
    result = llm.generate(
  File "/usr/local/lib/python3.10/dist-packages/vllm/entrypoints/llm.py", line 165, in generate
    return self._run_engine(use_tqdm)
  File "/usr/local/lib/python3.10/dist-packages/vllm/entrypoints/llm.py", line 185, in _run_engine
    step_outputs = self.llm_engine.step()
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/llm_engine.py", line 581, in step
    output = self._run_workers(
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/llm_engine.py", line 755, in _run_workers
    self._run_workers_in_batch(workers, method, *args, **kwargs))
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/llm_engine.py", line 732, in _run_workers_in_batch
    all_outputs = ray.get(all_outputs)
  File "/usr/local/lib/python3.10/dist-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/ray/_private/worker.py", line 2624, in get
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(TypeError): ray::RayWorkerVllm.execute_method() (pid=250471, ip=172.17.0.12, actor_id=6003ef308c0e22f0e05d73c901000000, repr=<vllm.engine.ray_utils.RayWorkerVllm object at 0x7f8c9c37a770>)
  File "/usr/local/lib/python3.10/dist-packages/vllm/engine/ray_utils.py", line 31, in execute_method
    return executor(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/vllm/worker/worker.py", line 159, in execute_model
    output = self.model_runner.execute_model(seq_group_metadata_list,
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/vllm/worker/model_runner.py", line 354, in execute_model
    output = self.model.sample(
  File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/models/mixtral.py", line 390, in sample
    next_tokens = self.sampler(self.lm_head.weight, hidden_states,
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/layers/sampler.py", line 52, in forward
    logits = _apply_logits_processors(logits, sampling_metadata)
  File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/layers/sampler.py", line 172, in _apply_logits_processors
    logits_row = logits_processor(token_ids, logits_row)
TypeError: RegexLogitsProcessor.__call__() missing 1 required positional argument: 'scores'

Outlines/Python version information:

Outlines 0.0.22
Python 3.10.12

Context for the issue:

No response

[viktor-ferenczi]

I have the same issue.

Check the changes in this vLLM PR (not merged yet):
vllm-project/vllm#2105

It integrates grammar into vLLM itself and contains code which works with Ray, which may give some ideas how to solve this.

This is the specific code:

    if request.grammar:
        if engine.worker_use_ray:
            grammar_logits_processor = RayRemoteGrammarLogitsProcessor(
                tokenizer=tokenizer, grammar=request.grammar)
        else:
            grammar_logits_processor = GrammarLogitsProcessor(
                tokenizer=tokenizer, grammar=request.grammar)
        logits_processors = [grammar_logits_processor]
    else:
        logits_processors = []

I guess we need a class similar to RayRemoteGrammarLogitsProcessor.

[lapp0]

I'm considering closing the vllm PR and moving the work over here since outlines has a more mature, fleshed out implementation.

I'm seeing a substantial performance degradation with RayRemoteGrammarLogitsProcessor. I haven't profiled, but my guess is that sending a tensor over ray is not performant.

I will experiment with the actor only taking the tokens as inputs, and returning a boolean mask. This would prevent the tensor going back and forth through ray.

https://numpy.org/doc/stable/reference/generated/numpy.ma.MaskedArray.tobytes.html

[viktor-ferenczi]

I agree with closing the vLLM ticket. I think vLLM needs another, much smaller PR. That should include only small changes to pass enough information to outlines, so no tricky patches are needed. Namely vLLM needs to pass the seq_id to the logits processors. It is an API change, however, but warranted.

Regarding Ray it seems to be a good idea, indeed. Try to pass numpy arrays, it may be faster. Also, the tokens may be largely redundant, so sending only the difference may also help decreasing the amount of information sent between processes.

[rlouf]

I agree with closing the vLLM ticket. I think vLLM needs another, much smaller PR. That should include only small changes to pass enough information to outlines, so no tricky patches are needed. Namely vLLM needs to pass the seq_id to the logits processors. It is an API change, however, but warranted.

We can think of other ways to track the FSM state, like using token_ids as keys as well. Someone mentioned it earlier but I can't find the comment.

[lapp0]

Regarding Ray it seems to be a good idea, indeed. Try to pass numpy arrays, it may be faster. Also, the tokens may be largely redundant, so sending only the difference may also help decreasing the amount of information sent between processes.

Logit processors aren't associated with one specific sequence so this is challenging.

_patched_apply_logits_processors does not work for generating multiple sequences at once because seq_id doesn't necessarily map to the same sequence between generations.

However patching also isn't necessary if you have correct cached-state lookups.

It's necessary to send a representation of the token sequence. Either the minimal (hash(prev_token_id), last_token_id) or just token_ids needs to be sent to the logits processing actor so it can look up the state.

The actor can send back bytes form of the mask.

We can think of other ways to track the FSM state, like using token_ids as keys as well. Someone mentioned it earlier but I can't find the comment.

#416 (comment)

This is necessary to get concurrent generation within a sequence group in vLLM working. Otherwise the generations tokens won't necessarily be added to the correct sequence. I got beam search working with outlines by using the token_ids as the FSM state key.

[viktor-ferenczi]

It explains why my attempt to use beam search with outlines regex failed yesterday. I like the approach of taking a hash of previous tokens.

[lapp0]

@viktor-ferenczi could you try this PR and tell me if it also works for you?

#539

[viktor-ferenczi]

Sure, I will try this tomorrow.

[viktor-ferenczi]

Tested both in regex and json_schema modes

Failed with errors:

AttributeError: 'RegexLogitsProcessor' object has no attribute 'fsm_state'
AttributeError: 'JSONLogitsProcessor' object has no attribute 'fsm_state'

Full tracebacks for regex and json_schema modes:
error.txt

The cause looks like that self.fsm_state is not set in RegexLogitsProcessor.__init__. This member is then accessed when the logits processor is first called, but without initializing it.

Appending this line to RegexLogitsProcessor.__init__ fixes the crash:

self.fsm_state: DefaultDict[int, int] = defaultdict(int)

But the constraint does not work as expected, it generates only 1-2 characters then stops.

GPUs: 2x4090 (2x24GB)

vLLM command:

python -O -u -m outlines.serve.serve \
  --model=TheBloke/deepseek-coder-33B-instruct-AWQ \
  --quantization=awq \
  --dtype=float16 \
  --host=0.0.0.0 \
  --port=8000 \
  --max-model-len=16384 \
  --max-num-seqs=16 \
  --tensor-parallel-size=2 \
  --swap-space=8 \
  --gpu-memory-utilization=0.95 \
  --enforce-eager \
  --disable-log-requests

vLLM request:

{
  "prompt": "You are an AI programming assistant, utilizing the Deepseek Coder model, developed by Deepseek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer.\n\nYou are a helpful AI assistant. You give concise answers. If you do not know something, then say so.\n### Instruction:\nWrite down the first 10 prime numbers as a comma separated list.\n\n### Response:\n",
  "n": 1,
  "best_of": 1,
  "presence_penalty": 0.0,
  "frequency_penalty": 0.0,
  "repetition_penalty": 1.0,
  "temperature": 0.0,
  "top_p": 1.0,
  "top_k": -1,
  "min_p": 0.0,
  "use_beam_search": false,
  "length_penalty": 1.0,
  "early_stopping": false,
  "stop": [],
  "stop_token_ids": [],
  "include_stop_str_in_output": false,
  "ignore_eos": false,
  "max_tokens": 50,
  "logprobs": null,
  "prompt_logprobs": null,
  "skip_special_tokens": true,
  "spaces_between_special_tokens": true,
  "regex": "\\d+(\\s*,\\s*\\d+)*\\s*"
}

Response:

{
  "text": [
    "You are an AI programming assistant, utilizing the Deepseek Coder model, developed by Deepseek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer.\n\nYou are a helpful AI assistant. You give concise answers. If you do not know something, then say so.\n### Instruction:\nWrite down the first 10 prime numbers as a comma separated list.\n\n### Response:\n1,"
  ]
}

[lapp0]

Thanks for running, I'll write some more test cases, I think it got messed up on a rebase, my mistake.

I believe the early termination is a result of the issue being fixed in #544