Align with huggingface Top K sampling

[Abraham-Xu]

main modifications
all modifications are in vllm/model_executor/layers/sampler.py

1. Reverse the order of softmax and _apply_top_p_top_k function in forward function of class Sampler. So the _apply_top_p_top_k will use logits as input instead of probabilites.
2. In _apply_top_p_top_k function, it computes the temporary probabilities ( softmax of logits ) firstly for the top_p process to locate the indexes whose cumulative probability exceeds the probability p. ( ATTENTION: the output of _apply_top_p_top_k function is still logits instead of probabilities. This is the main difference from the original code. )
3. Change the top_p and top_k masked probability value from 0 to -float("Inf").
4. In _sample_from_prompt function, use torch.multinomial without parameter "replacement=True".

tested result
The input probabilites distribution of torch.multinomial for the first token is the same as for huggingface/transformers under the same weights and input sentence.

test code:
huggingface/transformers

PATH_TO_CONVERTED_WEIGHTS="/data/xutianci/llama_hf/"
PATH_TO_CONVERTED_TOKENIZER="/data/xutianci/vllm/llama-tokenizer/"

from transformers import AutoTokenizer,  AutoModelForCausalLM
import torch
import numpy as np

model = AutoModelForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)

prompt = "Hello, my name is"
inputs = tokenizer(prompt, return_tensors="pt")
print(f'inputs={inputs}')

# Generate
generate_ids = model.generate(inputs.input_ids, do_sample=True, max_new_tokens=1)
print(f'generate_ids={generate_ids}')
output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
print(f'output={output}')

vllm

from vllm import LLM, SamplingParams
import torch
import numpy as np

prompts = [
    "Hello, my name is",
]
sampling_params = SamplingParams(top_k=50, max_tokens=1)

llm = LLM(model="/data/xutianci/llama_hf/", tokenizer="/data/xutianci/vllm/llama-tokenizer/")

outputs = llm.generate(prompts, sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

related issue
https://github.com/vllm-project/vllm/issues/718

[zhuohan123]

LGTM! Thanks for the fix! Committed some code to make sure the format is correct and to rename some variables.

[Abraham-Xu]

LGTM! Thanks for the fix! Committed some code to make sure the format is correct and to rename some variables.

Glad to be merged. Thanks for checking and revising! The renaming made the algorithm more self-explanatory!