feat: implement the min_tokens sampling parameter

[tjohnson31415]

Adds the min_tokens sampling parameter to ensure a minimum number of generated tokens.

The implementation here is meant to align with https://github.com/IBM/text-generation-inference (IBM's fork of HF TGI). In particular, we want to ignore stop sequences and penalize the EOS token until min_tokens have been generated. Stop sequence can be generated within min_tokens tokens but generation will not terminate. stop_token_ids are treated like the EOS token and penalized so that they are not generated until min_tokens tokens have been generated.

Related PR that stalled: #1945

[tjohnson31415]

@simon-mo I am unable to add you as a reviewer but tagging you RE: #1945 (comment).

I made the change to use a Logits Processor to penalize the tokens, but the changes are more than just the logits processor. The additional changes are to skip the stop sequences check if min_tokens have not yet been generated. Let me know if you had something else in mind for the Processor. I'm also looking for a good way to inject the MinNewTokensProcessor automatically when min_tokens is specified in the sampling params.

[njhill]

@tjohnson31415 IMHO since an explicit parameter is being introduced for this, it would be best for the eos token suppression to be tied to that without having to additionally pass a LogitsProcessor. There is already an ignore_eos parameter which is similar but I think is practically only useful for performance testing; I think min_tokens wouldn't have much utility unless used in conjunction with the MinTokensProcessor anyhow.

An additional advantage is that it would then be possible to have a vectorized implementation (not necessarily in this first PR).

One question in this case is whether other provided stop tokens (if any) should be suppressed in addition to eos, I'd lean towards yes but would be good to get input from others on that too.

[njhill]

This looks great @tjohnson31415. But I think technically the token_ids_to_penalize should be determined per seq_group (i.e. also within the loop) since they may be different per seq group. The indexing gets a bit tricker but I think it might be possible with scatter_ with src=-torch.inf. Or else could group the sequences that share the same list of tokens to pernalize.

[tjohnson31415]

Heh, yup. Thanks for pointing that out. I still need to write some tests for this 😅. I pushed a fix to build a list of coordinates to penalize within the loop so the stop ids are per seq_group.

I was trying to use scatter initially, but couldn't figure out how to get it to work. In particular, scatter uses a rectangular tensor and doesn't seem to have a way to "skip" rows where we don't want to scatter into. So I think a gather-modify-scatter (where we gather across all sequences and stop token ids) would work, but we'd still need to index into the gather'd tensor to set the -inf values.

[njhill]

Thanks @tjohnson31415, LGTM!

We should add a test for this too, though perhaps lets wait for confirmation from the maintainers that this would be accepted.

[simon-mo]

Thank you for the contribution, and thank Nick for the review.

My original intention is that min_tokens can be implemented using a built-in logits processor so the interface is cleaner. But current approach is fine as well.

There are some readability issue with _apply_min_tokens_penalty, please redo the list comprehension to make future devs easier to understand. And please add this to OpenAI compatible server as well (see protocols.py in the entrypoints/openai` directory).

[njhill]

@simon-mo this should be ready now!