We evaluate the performance of two EmbeddingBagCollection modules:

1. EmbeddingBagCollection (EBC) (code): a simple module backed by torch.nn.EmbeddingBag.

2. FusedEmbeddingBagCollection (Fused EBC) (code): a module backed by FBGEMM kernels which enables more efficient, high-performance operations on embedding tables. It is equipped with a fused optimizer, and UVM caching/management that makes much larger memory available for GPUs.

We chose the embedding tables (sparse arch) in ML Perf DLRM as the model to compare the performance difference between EBC and Fused EBC. Below are the settings on the embedding tables:

num_embeddings_per_feature = [45833188, 36746, 17245, 7413, 20243, 3, 7114, 1441, 62, 29275261, 1572176, 345138, 10, 2209, 11267, 128, 4, 974, 14, 48937457, 11316796, 40094537, 452104, 12606, 104, 35]
embedding_dim_size = 128

Other setup includes:

• Optimizer: Stochastic Gradient Descent (SGD)
• Dataset: Random dataset (code)
• CUDA 11.7, NCCL 2.11.4.
• AWS EC2 instance with 8 16GB NVIDIA Tesla V100

After the installation of Torchrec (see "Binary" in the "Installation" section, link), run the following command under the benchmark directory (/torchrec/torchrec/benchmarks):

python ebc_benchmarks.py [--mode MODE] [--cpu_only]

where MODE can be specified as ebc_comparison_dlrm (default) / fused_ebc_uvm / ebc_comparison_scaling to see different comparisons.

To ease the reading, we use "DLRM EMB" to abbreviate "DLRM embedding tables" from below. Since 1 GPU can't accommondate the full sized tables in DLRM, we need to reduce the embedding_dim of the 5 largest tables to some degree (see the "Note" column in the following tables for the reduction degree). For the metrics, we use the average training time over 100 epochs to represent the performance of each module. speedup (defined as training time using EBC divided by training time using Fused EBC) is also computed to demonstrate the degree of improvement from EBC to Fused EBC.

We see that Fused EBC has much faster training efficiency compared to EBC. The speedup from EBC to Fused EBC is 13X, 18X and 23X when the DLRM EMB is reduced by 128 times, 64 times and 32 times, respectively.

Here, we demonstrate the advantage of UVM/UVM-caching with Fused EBC. With UVM caching enabled, we can put larger sized tables in DLRM EMB in Fused EBC without significant sacrifice on the efficiency. With UVM enabled, we can allocate full sized DLRM EMB in UVM, with expected slower training performance because of the extra sync points between host and GPU (see this example for more UVM explanation/usage).

The above performance comparison is also put in a bar chart for better visualization.

Here, we study how the scaling on the embedding table affects the performance difference between EBC and Fused EBC. In doing so, we vary three parameters, num_tables, embedding_dim and num_embeddings, and present speedup from EBC to Fused EBC in the following tables. In each table, we observe that embedding_dim and num_embeddings do not have significant effect on speedup. However, as num_tables increases, the improvement from EBC to Fused EBC becomes higher (speedup increases), suggesting the benefit of Fused EBC when it is to deal with many embedding tables.

• num_tables = 10

• num_tables = 100

• num_tables = 1000