Repository: pytorch
Issue: Optimizing DLRM for CPU
Labels: module: performance,feature,module: cpu,triaged

## 🚀 Feature
A number of optimization and performance tuning for DLRM on CPU

## Motivation

Recommendation systems are one of the most common DL workloads in the cloud or enterprise server room. Very often the recommendation system burns most compute cycles in the data center among all DL workload.  DLRM is a state-of-the-art deep learning recommendation model which is composed of compute intensive MLP layers and memory intensive and capacity limited embedding layers.  Due to the memory capacity, people have found CPU's large DRAM capacity helpful training large DLRM models with very large embedding tables. 

Since the introdution of DLRM workload, Intel Pytorch team has been working to improve the performance using the configuration described in the DLRM paper. The work has generated ~90x performance improvement for DLRM trainning workload on CPU. This RFC summarizes the work we have done on the DLRM optimization and performance tuning, including a number code improvements and command option BKMs.   

1. Parallelize the Embedding_bag. The embedding_bag operation contains a table look up and reduction operation. We parallelize the index_select_add  and uses caffe2's Embedding look up function caffe2::EmbeddingLookupIdx(). https://github.com/pytorch/pytorch/pull/24385  Jianyu further improve this implementation by adding a new parameter into to EmbeddingBag API  https://github.com/pytorch/pytorch/pull/27477
2. Optimize the backward path of embedding_bag operation.  We parallelized the backward embedding_bag specialized for float-point tensor with SUM reduction operation with continuous stride.  This greatly improves the performance by avoid calling the original code which involves offset2bag(), cumsum(), index_select(), index_add(), and _embedding_bag_sparse_backward(). https://github.com/pytorch/pytorch/pull/27804
3. Optimize the gradient update. For the non-coalesced hybrid sparse tensor input, the add_out_dense_sparse() function originally does coalescing involves an expensive sort and merge operation. We parallelized gradient update by using fine-grain locked critical sections which adds the non-coalesced hybrid sparse tensor directly to the dense weight tensor. This bring more than 100x improvement on gradient update and 20x performance improvement for whole DLRM benchmark.https://github.com/pytorch/pytorch/pull/23057  
4. Parallelize the index_put. index_put is called by the backward of Zflat = Z[:, li, lj] in interaction. We optimize the accumulation path by using atomic add float to parallelize it. This improved about 10x on index_put accumulate operation in DLRM. https://github.com/pytorch/pytorch/pull/29705
5. Parallelize the cat operation. On the contiguous path, we pre-calculate the cat offset and then perform cat copy in parallel. This gives about 8x improvement on cat operation in DLRM.https://github.com/pytorch/pytorch/pull/30806
6. Parallelize Index_select. Index_select is used in the backward path of embedding_bag operation. The current implementation is in th module and only parallelize for contiguous input. We found that for DLRM, the input tensor could be non-contiguous, and the implementation go to a serialized version. We optimize the non-contiguous input and re-implement the operation in Aten library using parallel_for.  https://github.com/pytorch/pytorch/pull/30598
7. Uses all threads for all openmp tasks. We found that the best performance is achieve when all the threads running for all parallel_for tasks(). Having some light operation running at few threads doesn’t help the performance. https://github.com/pytorch/pytorch/issues/30803	

To further tune  the performance, you may consider the following command line options and software configuration. 
1. Export KMP_BLOCKTIME=1   KMP_BLOCKTIME sets the time that an OPENMP worker thread should wait ,after completing the execution of a parallel region, before sleeping. The default time is 200ms and can hurt performance. Setting to 0 is not good also since we don't want to put worker thread to sleep too often.  

2. Export KMP_AFFINITY="granularity=fine,compact,1,0"  
KMP_AFFINITY binds OpenMP threads to physical processing units.  In this setting, The OpenMP* thread n+1 is bound to a different physical core to OpenMP* thread n. 
	
3. Use je_malloc().  jemalloc is a general purpose malloc implementation that emphasizes fragmentation avoidance and scalable concurrency support. We found this helps performance, mostly due to the reduction of calling OS's zero_pages() and better cache locality. 
	




cc @VitalyFedyunin @ngimel @mruberry