fix long dtype in topk sampling #15049
Conversation
chujiezheng
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
|
|
||
| # Get with the logprob of the prompt or sampled token. | ||
| token_ids = token_ids.unsqueeze(-1) | ||
| token_ids = token_ids.unsqueeze(-1).to(torch.long) |
There was a problem hiding this comment.
Just curious: How is it different from .to(torch.long)?
There was a problem hiding this comment.
it's the same. my reason is long() is shorter, :-)
There was a problem hiding this comment.
|
Some context on why it's long tensor: https://discuss.pytorch.org/t/why-does-the-indices-tensor-have-to-be-long-dtype/139675 |
houseroad
left a comment
houseroad
left a comment
There was a problem hiding this comment.
Looks good to me. Wondering if we can add some unittest?
houseroad
left a comment
houseroad
left a comment
There was a problem hiding this comment.
Since it's a simple fix, I am okay with landing the fix first, then have a follow up PR to add the test.
|
|
||
| # Get with the logprob of the prompt or sampled token. | ||
| token_ids = token_ids.unsqueeze(-1) | ||
| token_ids = token_ids.unsqueeze(-1).to(torch.long) |
There was a problem hiding this comment.
Signed-off-by: Louis Ulmer <ulmerlouis@gmail.com>
Signed-off-by: Mu Huai <tianbowen.tbw@antgroup.com>
3 participants
Fix the long dtype in topk sampling, since gather requires the indices to be the long dtype (int64)