Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Refactor and Generalize Benchmark Suite
Summary: Generalize the benchmarking suite to support general torchrec modules instead of just QuantizedEmbeddingBagCollection. Code was further refactored to be more clean and readable. Furthermore, since generalizing the module involves passing the torch.nn.Module in, it is required that the input module be on CPU, as when sharding on GPU, there could be a CUDA OOM if only one copy of the model can fit on the GPU. Therefore, unsharded modules are removed from testing (This can easily be replicated by world_size=1 as an argument) Differential Revision: D52677415
- Loading branch information