Add RawIdTrackerWrapper within TBE to access tracked ids and raw ids #3506
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
meta-cla
bot
added
the
CLA Signed
Contributor
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 3, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
b24f307 to
ef18305
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 3, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
2 times, most recently
from
ec35707 to
1d58979
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 3, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
1d58979 to
4f6d0a2
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 3, 2025
…eta-pytorch#3506) Summary: Pull Request resolved: meta-pytorch#3506 This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
4f6d0a2 to
b8f750a
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 3, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
b8f750a to
c30c1d0
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 3, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Differential Revision: D84925177
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 4, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
2 times, most recently
from
782ac51 to
f5fa809
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 5, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 7, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
2 times, most recently
from
8ce55cf to
fb8b404
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 7, 2025
…eta-pytorch#3506) Summary: Pull Request resolved: meta-pytorch#3506 This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 7, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
fb8b404 to
bb631f1
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 7, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
2 times, most recently
from
8e57084 to
23fe218
Compare
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 7, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
added a commit
to aliafzal/torchrec
that referenced
this pull request
Nov 7, 2025
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
23fe218 to
7fb9159
Compare
…eta-pytorch#3501) Summary: This diff introduces a new `RawIdTracker` class that extends TorchRec's model delta tracking infra to capture and track raw hash identities from MCC modules during training. This is specifically required for tracking raw ids for MPZCH tables. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: FriedCosey, chouxi Differential Revision: D84920167
Summary:
This diff introduces init_raw_id_tracker to initialize RawIdTracker if enabled through ModelTrackerConfig
Key Changes
-----------
1. **Added `TrackerType` enum** in [`**types.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fmaliafzal%2Ffbsource%2Ffbcode%2Ftorchrec%2Fdistributed%2Fmodel_tracker%2Ftypes.py%22%2Cnull%5D "/data/users/maliafzal/fbsource/fbcode/torchrec/distributed/model_tracker/types.py")
* Defines three tracker types: `NONE`, `TREC` (for EC/EBC), and `RAW_ID` (for MPZCH)
* Extended `ModelTrackerConfig` dataclass with `tracker_type` field (defaults to `TrackerType.NONE`)
2. **Enhanced DMP initialization** in [`**model_parallel.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fmaliafzal%2Ffbsource%2Ffbcode%2Ftorchrec%2Fdistributed%2Fmodel_parallel.py%22%2Cnull%5D "/data/users/maliafzal/fbsource/fbcode/torchrec/distributed/model_parallel.py")
* Added `init_raw_id_tracker()` method to create `RawIdTracker` instances
* Modified constructor to conditionally initialize trackers based on `tracker_type` configuration
internal
This is needed to support MPZCH modules for Raw embedding streaming.
Mode details :
https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld
Reviewed By: chouxi
Differential Revision: D84920233
…eta-pytorch#3506) Summary: This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP. internal This is needed to support MPZCH modules for Raw embedding streaming. Mode details : https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld Reviewed By: chouxi Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
7fb9159 to
dddc11b
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
This diff introduces RawIdTrackerWrapper, a wrapper class containing lookup and delete APIs registered during raw_ids_tracker initialization to access tracked ids and raw_ids. We needed to create a wrapper instead of passing in the tracker due to circular dependency issues since TBE is wrapped under DMP.
internal
This is needed to support MPZCH modules for Raw embedding streaming.
Mode details :
https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld
Differential Revision: D84925177