Add RawIdTrackerWrapper within TBE to access tracked ids and raw ids #3506
+565
−24
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
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
Summary: internal psutil was added as a new dependency in FBGEMM in this diff D85604160 and is causing the OSS build to fail. Fix: Added psutil in requirements Reviewed By: chouxi Differential Revision: D86103237
Summary: pre-commit fialing due to formatting issue in utils.py https://github.com/meta-pytorch/torchrec/actions/runs/19036463336/job/54361995037?fbclid=IwY2xjawN1ptZleHRuA2FlbQIxMQBicmlkETFnNTYxN1B1aTNIdFR4ZGNpAR5-KnjfEgM8ItENTpyyuDfASD_rKCs13J0CzkWJ3ADHJi0KkiX2llWT44OKjA_aem_hPBvWb3vyCZK4L1L9M4IlA Reviewed By: chouxi Differential Revision: D86102057
Summary: Adding post lookup tracker function within MMC module to allow tracking of hash_zch_identities with delta tracker. 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: D84920121
…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 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 Differential Revision: D84925177
aliafzal
force-pushed
the
export-D84925177
branch
from
b8f750a to
c30c1d0
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