Refactor and Generalize Benchmark Suite (#1632)

Summary:
Pull Request resolved: #1632

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. For benchmarking unsharded modules on the GPU, they have to be initialized on GPU. In order to keep the API more clean with taking in an input nn.Module instead of doing the initialization in the benchmark, benchmarking of unsharded modules have been removed from the test suite (for now).

Reviewed By: gnahzg

Differential Revision: D52677415

fbshipit-source-id: a9c7bd1a842d7b4812a6c7343478cae9b122904a

Author: PaulZhang12