Wenyi5608 greedy sampling

[wenyi5608]

Greedy Greedy search added repeat_penalty top_k top_p and temperature function

[ilya-lavrenov]

Can you maybe create a generic function which performances sampling and adopts top_p, top_k, num_beams and temperature ?
Currently, we already have beam search with repetition penalty, but other parameters don't have effect here.

Ideally, we need a single method similar to HF's generate, which hides generation loop from user. Could you please try it?

[wenyi5608]

@ilya-lavrenov HuggingFace's group_beam_search does not implement sampling( top_p, top_k, and temperature) function. Only beam_search and greedy_search have sampling function. genai implements group_beam_search and greedy_search function.

[apaniukov]

You must take the values from logits for the remaining tokens here, not scores, and then do the softmax on them. The softmax already modified the scores in the top_p call, so the second softmax will skew the distribution towards the most probable token.

[wenyi5608]

I refer to chatglm cpp(https://github.com/li-plus/chatglm.cpp/blob/main/chatglm.cpp#L825) to implement this function.
Is this implementation different from HuggingFace.

[apaniukov]

I checked the softmax effects on the logits and the results for the llama model are that the second softmax makes the distribution nearly uniform and the third one makes the distribution exactly uniform. Here is how the probabilities of the most probable token are evolving:

Softmax 0 times: tensor([[ 0.0000,  8.6805,  9.4036,  8.4245, 10.8393]])
Softmax 1 times: tensor([[3.1250e-05, 2.7906e-01, 2.0102e-01, 7.2127e-02, 2.6164e-01]])
Softmax 2 times: tensor([[3.1250e-05, 4.1308e-05, 3.8207e-05, 3.3586e-05, 4.0595e-05]])
Softmax 3 times: tensor([[3.1250e-05, 3.1250e-05, 3.1250e-05, 3.1250e-05, 3.1250e-05]])
For uniform distribution: 3.125e-05

The second softmax makes model predictions useless.

Huggingface doesn't transform the logits during top_p, just calculates softmax and filters the original logits based on the result. So sampling_softmax_inplace is not needed inside the sampling_top_p function.

Code for the softmax probs:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


model_checkpoint = "JackFram/llama-68m"
hf_tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
hf_model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
texts = ["This is a test"]
tokenized = hf_tokenizer(texts, return_tensors="pt")

with torch.no_grad():
    t = hf_model(**tokenized).logits
    print("Softmax 0 times:", torch.max(t, dim=-1).values)
    for i in range(1, 4):
        t = torch.nn.functional.softmax(t, dim=-1,)
        print(f"Softmax {i} times:", torch.max(t, dim=-1).values)

print(f"For uniform distribution: {1 / hf_tokenizer.vocab_size}")

[wenyi5608]

@apaniukov Thanks a lot!

[wenyi5608]

I re-implemented the sampling_top_p function, the implementation is similar to the implementation of huggingface.

[as-suvorov]

It seems tokens accumulation needed for applying sampling_repetition_penalty. This ids needs to preserved between get_out_token calls.

[wenyi5608]

Do you mean that the logits info remains the same before and after get_out_token calls?

[as-suvorov]

@wenyi5608 , random sampling decoding based on your PR was merged to generate api: pavel-esir#6. Generate API is preparing for the merge to main repository: #334. Thank you for your contribution!