Improve outlines.processors, add integration tests to test_generate.py

[lapp0]

A lot of these fixes were intended for #966 however that's blocked until there's a new transformers release.

These improvements are general to all models and will enable PRs resolving #806 and #965

Structure of OutlinesLogitsProcessor

The goal is to create a base class which allows a logits processors to be implemented once and used for any outlines.models inference library.

To accomplish this we must normalize the input array. It must have a consistent type (torch.Tensor) and consistent dimensionality (2). We can normalize both of these simply, and without any copy operations.

mlx.core.array, numpy.array, and torch.Tensor all support pythons array standard __dlpack__. This standard allows for casting between array types without copying.

torch.Tensor is the only input type which cannot always be cast to any other type because torch tensors may live in GPU memory. Therefore, we cast all arrays to torch.Tensor, implement logits processors using torch methods, and convert back to the original array type in OutlinesLogitsProcessor. See docstring of OutlinesLogitsProcessor.__call__() for more details.

Detailed Changes

• Rename BaseLogitsProcessor to OutlinesLogitsProcessor
• Ensure OutlinesLogitsProcessor.process_logits() is always passed a 2D batch request with torch.Tensor logits and List input_ids. Also clean up code to be more readable in OutlinesLogitsProcessor__call__()
• Ensure FSMLogitsProcessor allows unstable sequence ordering (beam search in transformers and vLLM change the order of sequences)
• Update tests/generate/test_generate.py to cover more permutations of
  • regex / text
  • batch / single
  • greedy / multinomial / beam search
  • stream() / generate()
• Ensure performance stability with difference array libraries through benchmark_processors.py

[rlouf]

Why not call it LogitsProcessor?

[rlouf]

Actually, do we really need that class if all logit processors are ever going to be FSMLogitsProcessors

[lapp0]

My thinking here was if we want a new logits processor which isn't for structured generation, e.g. RepetitionPenaltyLogitsProcessor. We can easily implement this processor once and use it in every outlines.models without any changes specific to the model.

This provides us the flexiiblity to have the inference library handle the decoder pass and sampling, and outlines can handle all logits augmentation.

Another example: mlxlm doesn't support stop_strings and stop_strings is broken in the latest transformers version (4.41.2). We could implement a single StopStringsLogitsProcessor and make it available to all inference libraries. Not arguing that implementing this specific logits processor should be a priority, but I am arguing that having a separate base class which doesn't demand an FSM provides us flexibility and opportunities.

[rlouf]

What is the cost of this conversion? We need to profile this.

I would personally handle this by dispatching the core logic used to process the logits, process_logits depending on the type of input_ids or logits

[lapp0]

The idea is that we can cast numpy and mlx.core arrays to torch.Tensor using shared memory, not a copy operation :)

Then we can implement a single OutlinesLogitsProcessor subclass with a single process_logits() method and it works out of the box for any library. While the tests aren't in main yet, I've tested FSMLogitsProcessor, and RegexLogitsProcessor with nearly all outlines.models options (haven't tested against exllamav2) and they work with no additional changes.

This makes the processors easy to maintain and implement.

[rlouf]

cf above remark on dispatching the function depending on the type of the tensors.

[rlouf]

I would not convert back-and-forth to torch arrays and instead dispatch the logic depending on the type of input_ids and logits so as to not add a performance penalty during inference.

[lapp0]

I would not convert back-and-forth to torch arrays and instead dispatch the logic depending on the type of input_ids and logits so as to not add a performance penalty during inference.

Here are the benchmark results:

[50.00%] ··· Running (bench_processors.LogitsProcessorBenchmark.time_logits_processor--).
[100.00%] ··· bench_processors.LogitsProcessorBenchmark.time_logits_processor                                                                                                                                   ok
[100.00%] ··· ======== ==========
               param1            
              -------- ----------
               torch    150±20μs 
               numpy    161±7μs  
                mlx     195±20μs 
              ======== ==========

The performance penalty for casting (not copying) mlx.core.array -> torch.Tensor -> mlx.core.array is 45 microseconds. The benchmark has float logits with shape (4, 30000).

• This is around 1-3 orders of magnitude faster than applying a mask in FSMLogitsProcessor (varies a lot based on the size of the mask, huge masks can take ~20ms).
• this is 3-4 orders of magnitude faster than a decoder pass on a tiny model (mlx-community/Qwen1.5-1.8B-Chat-4bit)
• I expect the 45μs to be closer to numpys 11μs once mlx supports conversion directly from torch without a numpy intermediate.