Skip to content

Add missing fields to KJT's PyTree flatten/unflatten logic for VBE KJT #2952

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

Closed
wants to merge 1 commit into from

Conversation

jd7-tr
Copy link
Contributor

@jd7-tr jd7-tr commented May 7, 2025

Summary:

Context

  • Currently torchrec IR serializer does not support exporting variable batch KJT, because the stride_per_rank_per_rank and inverse_indices fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
  • The diff updates KJT's PyTree flatten/unflatten function to include stride_per_rank_per_rank and inverse_indices.

Ref

Differential Revision: D74295924

@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 May 7, 2025
@facebook-github-bot
Copy link
Contributor

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

@jd7-tr jd7-tr force-pushed the export-D74295924 branch from 20069fd to e7b6360 Compare June 20, 2025 22:46
jd7-tr added a commit to jd7-tr/torchrec that referenced this pull request Jun 20, 2025
pytorch#2952)

Summary:


# Context
* Currently torchrec IR serializer does not support exporting variable batch KJT, because the `stride_per_rank_per_rank` and `inverse_indices` fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
* The diff updates KJT's PyTree flatten/unflatten function to include `stride_per_rank_per_rank` and `inverse_indices`.
# Ref

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

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

@jd7-tr jd7-tr force-pushed the export-D74295924 branch from e7b6360 to a9fcb17 Compare June 21, 2025 22:36
jd7-tr added a commit to jd7-tr/torchrec that referenced this pull request Jun 21, 2025
pytorch#2952)

Summary:


# Context
* Currently torchrec IR serializer does not support exporting variable batch KJT, because the `stride_per_rank_per_rank` and `inverse_indices` fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
* The diff updates KJT's PyTree flatten/unflatten function to include `stride_per_rank_per_rank` and `inverse_indices`.
# Ref

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

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

@jd7-tr jd7-tr force-pushed the export-D74295924 branch from a9fcb17 to a65e0e3 Compare June 23, 2025 18:51
jd7-tr added a commit to jd7-tr/torchrec that referenced this pull request Jun 23, 2025
pytorch#2952)

Summary:


# Context
* Currently torchrec IR serializer does not support exporting variable batch KJT, because the `stride_per_rank_per_rank` and `inverse_indices` fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
* The diff updates KJT's PyTree flatten/unflatten function to include `stride_per_rank_per_rank` and `inverse_indices`.
# Ref

Reviewed By: TroyGarden

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

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

jd7-tr added a commit to jd7-tr/torchrec that referenced this pull request Jun 23, 2025
pytorch#2952)

Summary:
Pull Request resolved: pytorch#2952

# Context
* Currently torchrec IR serializer does not support exporting variable batch KJT, because the `stride_per_rank_per_rank` and `inverse_indices` fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
* The diff updates KJT's PyTree flatten/unflatten function to include `stride_per_rank_per_rank` and `inverse_indices`.
# Ref

Reviewed By: TroyGarden

Differential Revision: D74295924
@jd7-tr jd7-tr force-pushed the export-D74295924 branch 2 times, most recently from b4dfa13 to 1bc6ae0 Compare June 26, 2025 17:54
jd7-tr added a commit to jd7-tr/torchrec that referenced this pull request Jun 26, 2025
pytorch#2952)

Summary:


# Context
* Currently torchrec IR serializer does not support exporting variable batch KJT, because the `stride_per_rank_per_rank` and `inverse_indices` fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
* The diff updates KJT's PyTree flatten/unflatten function to include `stride_per_rank_per_rank` and `inverse_indices`.
# Ref

Reviewed By: TroyGarden

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

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

@jd7-tr jd7-tr force-pushed the export-D74295924 branch from 1bc6ae0 to 0836fe8 Compare July 9, 2025 00:15
@facebook-github-bot
Copy link
Contributor

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

jd7-tr added a commit to jd7-tr/torchrec that referenced this pull request Jul 9, 2025
pytorch#2952)

Summary:


# Context
* Currently torchrec IR serializer does not support exporting variable batch KJT, because the `stride_per_rank_per_rank` and `inverse_indices` fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
* The diff updates KJT's PyTree flatten/unflatten function to include `stride_per_rank_per_rank` and `inverse_indices`.
# Ref

Reviewed By: TroyGarden

Differential Revision: D74295924
@jd7-tr jd7-tr force-pushed the export-D74295924 branch from 0836fe8 to 0aba112 Compare July 9, 2025 20:27
@facebook-github-bot
Copy link
Contributor

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

pytorch#2952)

Summary:


# Context
* Currently torchrec IR serializer does not support exporting variable batch KJT, because the `stride_per_rank_per_rank` and `inverse_indices` fields are needed for deserializing VBE KJTs but they are included in the KJT's PyTree flatten/unflatten function.
* The diff updates KJT's PyTree flatten/unflatten function to include `stride_per_rank_per_rank` and `inverse_indices`.
# Ref

Reviewed By: TroyGarden

Differential Revision: D74295924
@jd7-tr jd7-tr force-pushed the export-D74295924 branch from 0aba112 to 818ea93 Compare July 10, 2025 20:38
@facebook-github-bot
Copy link
Contributor

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

@facebook-github-bot
Copy link
Contributor

This pull request has been reverted by d797031.

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 Reverted
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants