Skip to content

feat: support 1d ITensor offsets for embedding_bag converter #2677

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.

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 14 commits into from
May 1, 2024
Merged

feat: support 1d ITensor offsets for embedding_bag converter #2677

merged 14 commits into from
May 1, 2024

Conversation

zewenli98
Copy link
Collaborator

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

  • New feature (non-breaking change which adds functionality)

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@github-actions github-actions bot added component: tests Issues re: Tests component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Mar 8, 2024
@github-actions github-actions bot requested a review from peri044 March 8, 2024 01:24
@zewenli98 zewenli98 self-assigned this Mar 8, 2024
@zewenli98 zewenli98 mentioned this pull request Mar 8, 2024
Closed
7 tasks
@github-actions github-actions bot added the component: lowering Issues re: The lowering / preprocessing passes label Mar 12, 2024
@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
gs-olive reviewed Mar 27, 2024
View reviewed changes
Copy link
Collaborator

@gs-olive gs-olive left a comment •
edited
Loading

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
Copy link
Collaborator Author

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

@zewenli98 zewenli98 requested a review from gs-olive April 4, 2024 01:05
gs-olive
gs-olive reviewed Apr 4, 2024
View reviewed changes
Copy link
Collaborator

@gs-olive gs-olive left a comment

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

@zewenli98 zewenli98 requested a review from gs-olive April 4, 2024 18:06
gs-olive
gs-olive reviewed Apr 5, 2024
View reviewed changes
@zewenli98
Copy link
Collaborator Author

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

@zewenli98 zewenli98 force-pushed the embedding_bag_1d_ITensor_offsets_dynamo_converter branch from 2d40091 to 1f58d47 Compare May 1, 2024 00:23
@zewenli98 zewenli98 requested a review from gs-olive May 1, 2024 00:27
@zewenli98
Copy link
Collaborator Author

zewenli98 commented May 1, 2024 •
edited
Loading

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

@zewenli98 zewenli98 merged commit de81be2 into pytorch:main May 1, 2024
laikhtewari pushed a commit that referenced this pull request May 24, 2024
@zewenli98 @laikhtewari
feat: support 1d ITensor offsets for embedding_bag converter (#2677)
efe24e5
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
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

Assignees

@zewenli98 zewenli98

Labels
cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Implement embedding bag convertor
3 participants
@zewenli98 @gs-olive @facebook-github-bot