Skip to content

Benchamrk framework for torchrec #3072

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

lizhouyu
Copy link
Contributor

@lizhouyu lizhouyu commented Jun 9, 2025

Summary:
Benchmark framework for MPZCH

Rollback Plan:

Differential Revision: D76150895

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

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

lizhouyu added 3 commits June 16, 2025 10:50
Summary:
Pull Request resolved: pytorch#3017

### Major changes
- Copy the following files from `fb` to corresponding location in the `torchrec` repository
  - `fb/distributed/hash_mc_embedding.py → torchrec/distributed/hash_mc_embedding.py`
  - `fb/modules/hash_mc_evictions.py → torchrec/modules/hash_mc_evictions.py`
  - `fb/modules/hash_mc_metrics.py → torchrec/modules/hash_mc_metrics.py`
  - `fb/modules/hash_mc_modules.py → torchrec/modules/hash_mc_modules.py`
  - `fb/modules/tests/test_hash_mc_evictions.py → torchrec/modules/tests/test_hash_mc_evictions.py`
  - `fb/modules/tests/test_hash_mc_modules.py → torchrec/modules/tests/test_hash_mc_modules.py`
- Create a `test_hash_zch_mc.py` file in `torchrec/distributed/tests` folder following the `test_quant_mc_embedding.py` in `torchrec/fb/distributed/tests`.
  - trimmed quantization and inference codes, and only kept the training part.
  - rewire the related packages from `torchrec.fb` to `torchrec`
- Update `BUCK` files in related folders
- Update the affected repos to use `torchrec` modules instead of the modules in `torchrec.fb`
- Update `/modules/hash_mc_metrics.py`
  - Replace the tensorboard module with a local file logger in `hash_mc_metrics.py` module
- Update the license declaration headers for the four OSS files

### ToDos after landing this Diff
- Clean the duplicated `hash_mc_modules.py` file in the `fb` folder for safe landing.

Differential Revision: D76476676
Summary:
### Major changes
- Create a `mpzch` folder under the `torchrec/github/examples` folder
- Implement a simple SparseArch module with a flag to switch between original and MPZCH managed collision modules
- Profile the running time and QPS for model training(GPU)/inference(CPU)
- Create a notebook tutorial for ZCH basics and the use of ZCH modules in TorchRec

### ToDos for OSS
- When the internal torchrec MPZCH module is OSS
  - Remove the `BUCK` file
  - Replace all the `from torchrec.fb.modules` in `sparse_arch.py` to `from torchrec.modules`

### Potential improvement
- Add hash collision counter
- Show profiling results in the Readme file
- Add multi-batch profiling

Differential Revision: D75570684

Reviewed By: aporialiao
Summary:
# Benchmark framework for MPZCH

### Major changes
- Add a `benchmark prober`  in `torchrec/distributed/benchmark/benchmark_zch_utils.py` to collect and calculate the zero collision hash related metrics like hit count, insert count, and collision count.
- Implement a `benchmark_zch_dlrmv2` local testbed in `torchrec/distributed/benchmark/benchmark_zch_dlrmv2.py`, which allows to profile a DLRMv2 model with and without the MPZCH enabled, and record the metrics including ZCH-related metrics, QPS, NE, and AUROC.
- Add `mc_adapter` modules in `torchrec/modules/mc_adapter.py`. These modules enable seamless replacement of embedding collection and embedding bag collection modules with the managed collision version.
- Add two dictionaries `self.table_name_on_device_remapped_ids_dict` and `self.table_name_on_device_input_ids_dict` in the `HashZchManagedCollisionModule` module in `torchrec/modules/hash_mc_modules.py` to record the remapped identities and input feature values to the MPZCH module on current rank respectively after input mapping.
- Add `count_non_zch_collision.py` script to count the collision rate of non-zch modules after performing `murmur_hash3`.
- Add the criteo kaggle dataset data loader in `torchrec/distributed/benchmark/data` and revise the `hashes` attribute of data pipeline in the `_get_in_memory_dataloader` function in the `torchrec/distributed/benchmark/data/dlrm_dataloader.py` file to pre-hash the input feature values to the passed-in argument `input_hash_size` (defaultly as 100000).
- Note that we can change the `single_ttl` in the `HashZchSingleTtlScorer` module of `torchrec/modules/hash_mc_evictions.py` to change the eviticability of identities in each `HashZchManagedCollisionModule` module, since exiting benchmark workflow only takes several minutes on the subset of the criteo-kaggle dataset. By default the identities become evictable after one hour. This descrepency leads to non-eviction during the profiling process.

### Dataset
- [Criteo Kaggle Small](https://drive.google.com/file/d/1__rPcUSa45FHkmnBwivuM7K4nMYWD7b7/view?usp=sharing)
- [Criteo Kaggle](https://drive.google.com/file/d/1_lAbXTEOk5vlPGXd4UvTrxGV6sCPer_R/view?usp=drive_link)

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

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

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