Skip to content
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 and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

#12729: add row-major split for BERT #12804

Merged
merged 7 commits into from
Sep 19, 2024
Merged

#12729: add row-major split for BERT #12804

merged 7 commits into from
Sep 19, 2024

Conversation

jaykru-tt
Copy link
Contributor

@jaykru-tt jaykru-tt commented Sep 18, 2024

Ticket

#12729

Problem description

The split operation didn't support row major tensors or even the shape [1,256,2] required by BERT when in tiled format.

What's changed

This commit introduces a row major support for the split operation, additionally allowing for n-way splits on any dimension and splits on tensors of rank != 4. I've also added a few test cases, but these only cover 2-way splits on shapes similar to the shape required by BERT.

Checklist

  • Post commit CI passes
  • Blackhole Post commit (if applicable)
  • Model regression CI testing passes (if applicable)
  • Device performance regression CI testing passes (if applicable)
  • New/Existing tests provide coverage for changes

@jaykru-tt
Copy link
Contributor Author

jaykru-tt commented Sep 18, 2024

If @sjameelTT's recent changes to slice make it in first, I'll need to update some of the slice calls here to use end-exclusive indexing.

Copy link
Contributor

@sjameelTT sjameelTT left a comment

Choose a reason for hiding this comment

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

Are you a codeowner under data_movement yet? I feel like you should add yourself there @tarafdarTT thoughts?

@ntarafdar
Copy link
Contributor

ntarafdar commented Sep 18, 2024

hey @jaykru-tt , I'm a little confused with the description.

I'm asking because I thought for Bert split we were going to use Tile layout but your description alludes to Bert requiring RM layout.
Furthermore in the issue we have the Bert shape, and if this PR addresses that issue we should add that shape in the unit test.

Copy link
Contributor

@ntarafdar ntarafdar left a comment

Choose a reason for hiding this comment

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

great stuff @jaykru-tt

@jaykru-tt
Copy link
Contributor Author

post-commit passed, lfg, merging this

@jaykru-tt jaykru-tt mentioned this pull request Sep 19, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants