Support for grammar

[7071eth]

It would be highly beneficial if the library could incorporate support for Grammar and GBNF files.
https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md

[jpeig]

Support this request. Would be massive to produce structured data at scale.

[viantirreau]

I will test whether Shopify/torch-grammar works around these lines.

Will be posting here if it works :)

[jordan-barrett-jm]

Any updates?

[viantirreau]

Hi! Sorry for the delay. I ended up implementing a minor speedup for constrained generation on llama.cpp by skipping "obvious" tokens and just decoding the ones that need an LLM. Although it managed to reach around double the original decoding performance, it is still not enough for my use case, so I will likely revisit this approach with vLLM in the future.
It is not the current priority for my personal project though, so it may take some weeks before I start experimenting again.

If anyone is eager to try an implementation, the idea I had in mind is, on a high level, as follows:

1. Instantiate the grammars, one for each sequence that will need to be decoded in parallel. I think it should be an attribute of the Sampler class:

   class Sampler(nn.Module):
   
   
       def __init__(self, vocab_size: int) -> None:
           super().__init__()
           self.vocab_size = vocab_size
           # Dict[int, GrammarSampler]. Maps a sequence_id to a fresh sampler
           self.grammars = defaultdict(lambda: GrammarSampler(input_text, "root", tokenizer))
   
   
       def forward(...):

   This follows the example from torch-grammar's Readme. As I'm writing this, I don't see how to know how many grammar samplers to create in advance, so the idea is to use a defaultdict that will create a new grammar as soon as an unseen sequence_id needs it.

2. Somehow, we should keep track of all the token IDs generated so far, because they are needed in the grammar sampler to set the logits of invalid continuation tokens to zero, and thus prevent those tokens from being randomly sampled and selected. A great candidate for this seems to be the output_tokens = _get_output_tokens(input_metadata) call a few lines into the forward function. This holds a list of tokens ids for each sequence id. It's enlightening to read the definition of _get_output_tokens towards the middle of the file to know the general shape of the returned list.

3. Just before computing the logprobs, we could call the logits_processor for each pair of (seq_id, grammar). The important thing here is to match up a few things: each sequence id to the corresponding grammar, that keeps track of the state of the generation and defines which tokens will be discarded now, and also to the current output for each of those sequence ids.

   def forward(...):
       # More sampling code above ...
   
       # We use float32 for probabilities and log probabilities.
       # Compute the probabilities.
       probs = torch.softmax(logits, dim=-1, dtype=torch.float)
       
       # I don't know how to vectorize this, so I unroll it in a simple for loop
       output_idx = 0
       for seq_ids, sampling_params in input_metadata.seq_groups:
           for seq_id in seq_ids:
               #  the defaultdict will create a clean grammar when a given seq_id doesn't exist yet
               grammar_sampler = self.grammars[seq_id]
               output_for_this_seq_id = output_tokens[output_idx]
               logits_for_this_seq_id = logits[output_idx]
               # call the logits_processor(ids, logits)
               logits[output_idx] = grammar_sampler.logits_processor(output_for_this_seq_id,
                                                                     logits_for_this_seq_id)
               # ^ this will zero out the logits that correspond to invalid token ids under the grammar
               output_idx += 1
               # Note that I deliberately don't use an `enumerate` to update the output_idx variable,
               # because the _get_output_tokens function appends a sequence once for each seq_id,
               # not on every seq_group!
   
       # Continue with the existing code
   
       # Compute the log probabilities.
       # Use log_softmax to ensure numerical stability.
       logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)

4. Accept the tokens that have been sampled to update the grammar state for each sequence id. Each sampling strategy (greedy, random, beam) uses the same interface for the return: List[Tuple[List[int], List[int]]]. Refer to the implementation of _greedy_sample, _random_sample or _beam_sample for more details, but from what I can understand by skimming through the code, those functions return, for each sequence group that requested that specific sampling strategy (this is the outer List of the return type), a tuple of the form (next_token_ids, parent_ids) to support generating multiple completions (one new token id for each sequence_id) for any single sequence group. Apparently, this tracks which tokens are selected (in parallel) as the next id and which ones were responsible for that generation (parent_ids).
   I know this is like the movie Inception, but hold on.

   def forward(...):
       # More sampling code above ...
   
       # Compute the log probabilities.
       # Use log_softmax to ensure numerical stability.
       logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)
       
       # Sample the next tokens.
       sample_results = _sample(probs, logprobs, input_metadata)
   
       # Now we accept the token for each sequence id inside every sequence group
       output_idx = 0
       for seq_group_id, (seq_ids, sampling_params) in enumerate(input_metadata.seq_groups):
           next_token_ids, parent_ids = sample_results[seq_group_id]
           for seq_id in seq_ids:
               grammar_sampler = self.grammars[seq_id]
               next_token_for_this_seq_id = next_token_ids[output_idx]
               grammar_sampler.accept_token(next_token_for_this_seq_id)
               output_idx += 1
   
       # Continue with the existing code 

As a final note, I'm not well familiarized with this project, and this plan is only a high level review of how I would go about implementing grammars. Also, note that the torch-grammar library is clearly unmaintained, but I think it should work well enough as a proof of concept. At least its API is comfortable.

[viantirreau]

Another update. I wasn't aware of the existence of logits_processors, as implemented on this PR. That direction seems way more promising and already has one example in the form of a unit test.

[jpeig]

@viantirreau

Take a look at https://github.com/noamgat/lm-format-enforcer

They're using logits_processors to enforce JSON schema's and regex = basically Grammar

[CarlosLannister]

@viantirreau There is a LM Format Enforcer Integration with vLLM colab tutorial on https://github.com/noamgat/lm-format-enforcer/blob/main/samples/colab_vllm_integration.ipynb

Perhaps this can be use as guideline for future integrations

[viktor-ferenczi]

What is the technical reason why we cannot port over GBNF? It could be a logits_processor, plus a grammar need to be passed via API requests (both vLLM and OpenAI).