feat: support 1d ITensor offsets for embedding_bag converter by zewenli98 · Pull Request #2677 · pytorch/TensorRT (original) (raw)

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Conversation12 Commits14 Checks0 Files changed

Conversation

This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters

[ Show hidden characters]({{ revealButtonHref }})

zewenli98

Description

Support 1d ITensor offsets for embedding_bag converter.
Note that this is only for 1d inputs.

There's a bug that same offsets with different types (tensor or ITensor) when include_last_offset=True will give different results. I doubt this is a bug from PyTorch, I'm still investigating it.

Fixes #2345

Type of change

Checklist:

@zewenli98 zewenli98 changed the title[WIP] feat: support 1d ITensor offsets for embedding_bag converter feat: support 1d ITensor offsets for embedding_bag converter

Mar 12, 2024

gs-olive

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall a great improvement to embedding_bag; added a few small comments and suggestions

@zewenli98

Thanks for the review! I just refactored the embedding_bag with native TRT layers.

gs-olive

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall, looks good! Added one small comment, and pending CI validation

gs-olive

@zewenli98

This is currently blocked by a Myelin bug. Filed a NVBug internally.

@zewenli98

With the latest Pytorch, TRT-10 GA, and Torch-TRT main branch, the embedding_bag converter works on DLRM now.

laikhtewari pushed a commit that referenced this pull request

May 24, 2024

@zewenli98 @laikhtewari

Reviewers

@peri044 peri044 Awaiting requested review from peri044

@narendasan narendasan Awaiting requested review from narendasan

@chohk88 chohk88 Awaiting requested review from chohk88

@bowang007 bowang007 Awaiting requested review from bowang007

@gs-olive gs-olive Awaiting requested review from gs-olive