this sample func seems like a bug...

[JoursBleu]

没有考虑topk topp？
而且至少应该算个softmax出来？

TriForce/utils/sampling.py

Line 64 in 164c8c0

idx_next = torch.multinomial(probs, num_samples=num_samples, replacement=True)

[preminstrel]

sample 之前有 apply 一个 norm_logits 函数。

TriForce/utils/sampling.py

Line 43 in 164c8c0

def norm_logits(logits : torch.Tensor, temperature=0.6, top_k=-1, top_p=0.9) -> torch.Tensor:

token = sample(norm_logits(x))

[JoursBleu]

okk，不过这个norm_logits没有考虑temperature为0的情况，跑temperature=0会nan

[preminstrel]

temperature=0 的时候直接 argmax 就可以了，不用这么复杂的操作了，更 efficient 一些。

[JoursBleu]

就是这一行会错
https://github.com/Infini-AI-Lab/TriForce/blob/main/utils/decoding.py#L23

[preminstrel]

或者你直接temperature=0.001，和 0 没什么区别其实。

[JoursBleu]

还有这里

TriForce/utils/decoding.py

Line 32 in 164c8c0

next_token = sample(norm_logits(logits[:,-1,:], temperature=temperature ,top_k=top_k, top_p=top_p))

和这里

TriForce/utils/decoding.py

Line 62 in 164c8c0

next_token = sample(norm_logits(logits[:,-1,:], temperature=temperature ,top_k=top_k, top_p=top_p))