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

[TorchRec][IR] Add IR serializer for KTRegroupAsDict Module #1900

Closed
wants to merge 1 commit into from

Conversation

TroyGarden
Copy link
Contributor

@TroyGarden TroyGarden commented Apr 18, 2024

context

  • previously KTRegroupAsDict can't really supported by torch.export (IR) because this module has an intialization step as running the first batch.
  • during the export the KTRegroupAsDict module will be initialized by a fake_tensor which is wrong
  • if we initialize the module before torch.export, the device would be an issue.
  • another issue is that current torch.export can't support conditional logic in training, where initialization step only runs once.

torch.cond is a prototype feature in PyTorch. It has limited support for input and output types and doesn’t support training currently. Please look forward to a more stable implementation in a future version of PyTorch.

NOTE: this is more like a workaround solution, real solution needs support from pytorch compile for conditional logic

details

  • we treat the KTRegroupAsDict as another sparse_arch and do the model swap before and after torch.export.
  • more context: D59019375

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 18, 2024
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D56282744

TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Apr 19, 2024
…port (pytorch#1900)

Summary:


reference:
* D54009459

Differential Revision: D56282744
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D56282744

TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Apr 19, 2024
…port (pytorch#1900)

Summary:


reference:
* D54009459

Differential Revision: D56282744
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Apr 19, 2024
…port (pytorch#1900)

Summary:


reference:
* D54009459

Differential Revision: D56282744
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D56282744

TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Apr 19, 2024
…port (pytorch#1900)

Summary:


reference:
* D54009459

Differential Revision: D56282744
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D56282744

TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Apr 19, 2024
…port (pytorch#1900)

Summary:


reference:
* D54009459

Differential Revision: D56282744
TroyGarden added a commit to TroyGarden/torchrec that referenced this pull request Apr 19, 2024
…port (pytorch#1900)

Summary:


reference:
* D54009459

Differential Revision: D56282744
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D56282744

@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D56282744

@TroyGarden TroyGarden closed this Jul 15, 2024
@TroyGarden TroyGarden deleted the export-D56282744 branch July 15, 2024 22:35
@TroyGarden TroyGarden restored the export-D56282744 branch August 10, 2024 17:54
@TroyGarden TroyGarden reopened this Aug 10, 2024
Summary:
Pull Request resolved: pytorch#1900

# context
* previously `KTRegroupAsDict` can't really supported by torch.export (IR) because this module has an intialization step as running the first batch.
* during the export the `KTRegroupAsDict` module will be initialized by a fake_tensor which is wrong
* if we initialize the module before torch.export, the device would be an issue.
* another issue is that current torch.export [can't support conditional logic in training](https://pytorch.org/docs/stable/cond.html), where initialization step only runs once.
> torch.cond is a prototype feature in PyTorch. It has limited support for input and output types and doesn’t support training currently. Please look forward to a more stable implementation in a future version of PyTorch.

NOTE: this is more like a workaround solution, real solution needs support from pytorch compile for conditional logic

# details
* we treat the `KTRegroupAsDict` as another sparse_arch and do the model swap before and after torch.export.
* more context: D59019375

Differential Revision: D56282744
@TroyGarden TroyGarden changed the title register a custom op to keep the PEA module unflattened when torch.export [TorchRec][IR] Add IR serializer for KTRegroupAsDict Module Aug 10, 2024
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D56282744

PaulZhang12 pushed a commit that referenced this pull request Aug 19, 2024
Summary:
Pull Request resolved: #1900

# context
* previously `KTRegroupAsDict` can't really supported by torch.export (IR) because this module has an intialization step as running the first batch.
* during the export the `KTRegroupAsDict` module will be initialized by a fake_tensor which is wrong
* if we initialize the module before torch.export, the device would be an issue.
* another issue is that current torch.export [can't support conditional logic in training](https://pytorch.org/docs/stable/cond.html), where initialization step only runs once.
> torch.cond is a prototype feature in PyTorch. It has limited support for input and output types and doesn’t support training currently. Please look forward to a more stable implementation in a future version of PyTorch.

NOTE: this is more like a workaround solution, real solution needs support from pytorch compile for conditional logic

# details
* we treat the `KTRegroupAsDict` as another sparse_arch and do the model swap before and after torch.export.
* more context: D59019375

Reviewed By: PaulZhang12

Differential Revision: D56282744

fbshipit-source-id: b86f6eafa3d453735df6c9d00b33b16f70279dea
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. fb-exported
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants