Understanding Capacity-Driven Scale-Out Neural Recommendation Inference

Consider that a deep learning model can be represented by a directed control and dataflow graph. A distributed model simply partitions a neural network into subnets, where each subnet can operate independently. Treating each subnet independently provides desirable flexibility in deployment, since all the infrastructure for serving the model already exists. An example of such partitioning is shown in Figure 2(b). This is traditionally referred to as model or layer parallelism.

Optimal model sharding is a challenging systems problem, due to the number of configurations and varying optimization objectives. While this problem has been studied in the context of training, the inference context presents new challenges. Training must contend with the accuracy implications of mini-batch sizing and parameter synchronization, and is targeted at maximizing throughput on a single instance that’s fed a large dataset [33]. Data-center inference, in contrast, must contend with varying request rates and sizes under strict latency constraints, and the impact of model replication to meet those requirements. Furthermore, in the context of recommendation, memory footprint is dominated by embedding tables which are computationally sparse. Model sharding for deep learning recommendation inference is motivated by the desire to enable huge models, and thus we consider this new challenge capacity-driven sharding.

Due to the number of shard configurations, an exhaustive search for an optimal sharding scheme is intractable. Instead, heuristics are used, which depend on the model architecture and include measurements or estimates of model capacity, compute resources, and communication costs. Naively, the heuristic should aim to minimize the number of shards due to the additional compute and network resources consumed with more shards. We study the resulting impact of this assumption in Section VI. The heuristics are guided by the following observations:

1. 1.

   Embedding tables used in sparse layers contribute to the vast majority ($>97\%$) of model size.

2. 2.

   Sparse layers are memory and communication bound, while dense layers are compute bound.

3. 3.

   Existing server infrastructure cannot support memory requirements of sparse layers but can support the compute and latency requirements of dense layers.

4. 4.

   Deep recommendation model architectures, as in Figure 2, can execute sparse operators in parallel, which feed successive dense layers.

Because the dense layers do not benefit from additional compute or parallelism provided by distributed inference, our sharding strategies are constrained to only move sparse layers to remote shards. This is specific to the architecture of deep recommendation models. Figure 2 shows that, after embedding table lookup and pooling, all of the resulting embedding vectors and dense features are combined in a series of feature interaction layers. Placing the interaction layer on its own shards results in an undesirable increase of communication and lessening of compute density. Placing this layer with existing sparse shards also increases communication and violates the stateless shard constraint. Alternate architectures may benefit from sharding interaction layers. Consider a model architecture that has dedicated feature interaction layers for specific sets of embedding tables. Such an architecture could shard those feature interaction layers with their respective sparse layers and indeed see performance benefits. While an intriguing model architecture, we are limited to models that are currently available, thus this work chooses to focus on the more traditional model in Figure 2(a).

Thus, sharding the sparse layers (1) directly addresses capacity concerns of the largest parameters, (2) effectively parallelizes sparse operators which otherwise execute sequentially, and (3) enables better resource provisioning by isolating the communication- and compute-bound portions of inference. Note that inference is not traditionally operator-parallel because operators do not typically produce enough work to offset scheduling overheads. Extra computing cores are instead utilized by increasing batch-level parallelism.

Three sharding strategies, using the above heuristic, are evaluated in this work (Table I). These strategies address the new and distinct challenges of huge deep learning recommendation models: the varied number and size of embedding tables, non-uniformity of sparsity, and the challenge of the real-time serving environment. Two trivial cases of (1) non-distributed inference, or singular, and (2) a single shard with all embedding tables are also presented in Table I.

We present an in-depth description of the customized open-source frameworks used in this work. This provides the reader with a concrete implementation to frame our results and provides guidance for implementation with other frameworks. For this work, distributed inference is built on top of a highly customized variant of Thrift and Caffe2, however the methodology is generalizable to any RPC or ML framework [30, 31]. While PyTorch has absorbed and succeeded Caffe2 as the state-of-the-art moving forward, the Caffe2 infrastructure used in this work is shared between both frameworks. Thrift serves ranking requests by loading models and appropriately splitting received requests to inference batches to the appropriate net. A modified version of Caffe2 is used that includes RPC operators that issues Thrift requests. The intermediate requests to sparse shards are routed via a universal service discovery protocol. All inter-server communication occurs through standard TCP/IP stack over Ethernet. The model is transformed for distributed inference after training. A custom partitioning tool employs a user-supplied configuration to group embedding tables and their operators, insert RPC operators, generate new Caffe2 nets, and then serialize the model to storage.

Distributed tracing is a proven technique for debugging and performance investigations across inter-dependent services[34, 35]. We leverage this technique to investigate the performance impact of the distributed inference infrastructure. Custom instrumentation was added at multiple layers in the production service stack to provide a complete view of costs for request/response serialization, RPC service boilerplate setup, and model execution for each request.

Thrift and Caffe2 both provide instrumentation hooks to capture salient points in execution, leaving the service handler to be more intrusively modified, albeit still lightweight. Source instrumentation was chosen over other binary instrumentation approaches because it’s lighter weight and can interface with Thrift’s RequestContext abstraction to propagate contextual data for distributed tracing. At each trace point, metadata specific to the layer and a wall-clock timestamp are logged to a lock-free buffer and then asynchronously flushed to disk. Wall-clock time is desirable because its ordering helps achieve a useful trace visualization. Furthermore, most spans are small and sequential, enabling wall-clock time as a proxy for CPU time. Per-shard CPU time for each request is also logged to verify this claim. The trace points are then collected and post-processed offline for overhead analysis and to reconstruct a visualization of events, resembling Figure 3.

An example distributed trace is shown in Figure 3. Shards are separated by horizontal slices, with the main shard servicing requests located at the top. The example trace represents a single batch, however it’s typical for multiple batches to be executed in parallel depending on the number of ranking requests and the configurable batch size. Furthermore, the ordering of operators is typical for the recommendation models studied in this work. Operators are scheduled to execute sequentially–unless specifically asynchronous like the RPC ops–because other cores are utilized via request- and batch-level parallelism. The initial dense layers are first executed, followed by the sparse lookups, and then the later feature interaction and top dense layers.

As discussed in SectionIII-A, each sparse shard is a full RPC service for deployment flexibility, and as such incurs service-related overheads compared to local, in-line computation. Figure 3 demonstrates that latency overhead in each sparse shard is any time not spent executing SLS operators. In the non-distributed case, the SLS ops are directly executed in the main shard. The extra latency is attributed to the network link, extra (de)serialization of inputs and outputs, and time spent preparing and scheduling the Caffe2 net. Despite such overheads, we also see that the asynchronous nature of the distributed computation enables more parallelization of the sparse operators. An important takeaway from this characterization is the trade-off between increased parallelization to reduce latency overhead and decreased parallelization to reduce compute and data-center resource overhead. Furthermore, this trade-off is specific to each model.

Compared to simple, operator-based counters for compute or end-to-end (E2E) latency, capturing a cross-layer trace provides a holistic view of compute and latency. The example flow in Figure 3 shows compute overheads from RPC serialization and deserialization and request/response processing of the software infrastructure. Additional compute is incurred in each thread’s scheduling and book-keeping of asynchronous RPC operators. Simple end-to-end latency overhead is easily measured at the main shard. However, because each sparse shard is executed in parallel, attribution of latency overheads is more complex and involves overlap between sparse shard requests. To simplify analysis, the slowest asynchronous sparse shard request, per main shard request, is used for latency breakdown. The network-, serialization-, and RPC service- latencies in the associated sparse shard are used. Because the clocks on disparate servers will be skewed, network latency is measured as the difference between the outstanding request measured at the main shard and the end-to-end service latency measured at the sparse shard.

Three DLRM-like models, referred to as DRM1, DRM2 and DRM3, were used that span a range of large model attributes, such as varying input features and embedding table characteristics. The D-prefix is to contrast the distributed models to the specific models discussed in recent deep learning recommendation works [2, 1, 3]. The DRM* models are a subset of many possible, evolving models and chosen as early candidates for distributed inference given their sparse feature characteristics. The goal of studying these models is to present a basis for evaluating overheads of distributed inference via our cross-layer distributed trace analysis, and they should not be interpreted as canonical benchmark models such as those included in the MLPerf benchmark suite [9].

Figure 4, Figure 5, and Table II demonstrate the variety within large neural recommendation models. DRM1 and DRM2 have two nets, each having their own respective sparse features, while DRM3 has a single net. Because real, sampled requests were used for each model, the inference request size and corresponding compute and latency also varied between models. All parameters were uncompressed as single-precision floating point. Section VII-D discusses the impact of compression. Per-operator group attributions for each model are shown in Figure 4, as a simple mean average across all sampled requests for the non-distributed model. DRM1 and DRM2 are the most similar architectures, reflected in Figure 4. Compared to DRM3, DRM1 and DRM2 have a more complex structure evidenced by additional tensor transform costs. More relevant to this work, sparse operators consume a much larger portion of all operators compared to DRM3. Specifically, sparse operators contribute to 9.7%, 9.6%, and 3.1% of all operator time, in DRM1, DRM2, and DRM3 respectively. Despite their low proportional compute, the sparse operators account for >97% of model capacity in DRM1 and DRM2, and >99.9% of capacity for DRM3. Embedding tables larger than a given threshold were scaled down by a proportional factor to fit the entire model on a single 256GB server. This provided a straight comparison of compute and latency overheads across all sharding strategies listed in Table I, including singular. The original data-center scale models are many times larger.

The distribution of embedding tables sizes within each model is shown in Figure 5. DRM1 was sized to 200GB with 257 embedding tables and the largest table is at 3.6GB. DRM2 was proportionately sized to 138GB with 133 embedding tables and the largest table at 6.7GB. Lastly, DRM3 was sized to 200GB with 39 embedding tables and the largest table at 178.8GB. DRM1 and DRM2 demonstrate a long tail of embedding table size compared to DRM3, explaining the additional sparse operator cost shown in Figure 4. Comparatively, DRM3 is dominated by a single large table. The embedding tables are representative versions of the tables discussed in prior deep learning recommendation work [3, 1, 9].

For DRM1 and DRM2, ten sharding configurations were evaluated. Sharding strategy results for DRM1 are described in Table II. For the load-balanced configuration, per-shard capacities varied up to 50% compared to capacity-balanced where each shard is the same total capacity. For the capacity-balanced configuration, per-shard estimated load varied up to 371%, between shards 4 and 3 with eight shards, compared to load-balanced where each shard had the same total pooling factor. Lastly, the net-specific bin-packing (NSBP) strategy restricted each shard’s tables to a single net. This is most noticeable at the 2-shard configuration, where each net is placed on its own shard. Shard 2 consumes 4.75$\times$ as much memory capacity as Shard 1, yet is estimated to perform just 6.3% of Shard 1’s compute work.

DRM3 is only sharded with NSBP, and not capacity-balanced or load-balanced strategy, due to existing technical challenges of sharding huge tables. Because it is dominated by a single large table, for each additional shard the largest table is further split, while the smaller tables remained grouped as one shard. The dominating table has a pooling factor of 1, thus only one of the shards spanning the table will be accessed. For example, given four sparse shards, the largest table is partitioned into three shards and the remaining tables grouped together into one shard. Each inference, only two shard would be accessed: one for the sharded large table and one for the smaller tables.

Two classes of servers, representative of the data-center environment, were used for characterization. SC-Large is representative of a typical large server in a data-center and has 256GB of DRAM and two 20-core Intel CPUs. SC-Small is representative of a typical, more efficient web server and has 64GB of DRAM and two, slower clocked, 18-core Intel CPUs, and less network bandwidth than SC-Large. Because of the limited DRAM available on SC-Small, only a subset of configurations could be tested on this platform. The majority of results discussed in Section VI were collected on SC-Large platforms for an apples-to-apples comparison with the non-distributed models. A discussion on the impact of server platform follows in Section VII.

All inference experiments were run on CPU platforms, as discussed in Section II. The reserved, bare-metal servers were located in the same data centers as production recommendation ranking and utilized the same intranet. Their locations within the data center was representative of a typical inference serving tier. Shards were assigned to unique servers and not co-located. Shard-to-server mappings were randomized across repeated trials. Recommendation ranking requests were sampled from production servers and replayed on the test infrastructure; a database of de-identified requests were sampled evenly across a five-day time period in order to capture any diurnal behavior within the requests. The production replayer then pre-processed and cached the requests before sending them to the inference servers using the same networking infrastructure used in online production ranking.

For the majority of experimental runs, batch sizes for inference were set to production defaults, where each batch represents a number of recommendation items to rank and is executed in parallel. In Section VI requests were sent serially, to isolate inherent overheads. In Section VII-A requests were sent asynchronously, to simulate a higher QPS rate more representative of the production environment.

The P50, P90, and P99 End-to-End (E2E) latency and aggregate compute time overheads, across all sharding strategies, for all models, for serial blocking requests, are shown in Figures 6 and 7. A description of each sharding strategy is located in Table I. Note that singular is the non-distributed case, and 1-shard is the worse-case with all embedding tables on 1 sparse shard. P90 and P99 latencies are typical metrics for inference serving because a less accurate fallback recommendation is returned if the inference request is not processed within SLA targets. P50 is also presented for completeness to show the median case, instead of the mean average, due to long tail latencies discussed in prior work [1].

Understanding compute overheads is vital to minimize additional resource requirements at the data-center scale. Sharding strategy is one method for the system designer to balance latency constraints, discussed in the previous section, and compute overheads which impact resource requirements.

A trade-off between increased latency or increased compute, as a result of shard count, was established in the previous section. Shard count is a straightforward knob for system designers to balance compute overheads with latency constraints. The sharding strategy, discussed in this section, provides another knob for designers, but the effects on latency and compute are more nuanced.

Model attributes also affect distributed inference performance. The number of nets, number of embedding tables, size distribution, and respective pooling factor are the model attributes most relevant attributes, and are described in Section V-A for the evaluated models. DRM3 has less total compute attributed to sparse operators and is dominated by a single large embedding table, compared to the long tail embedding tables in DRM1 and DRM2. Latency for DRM1 and DRM2 benefit from increased sharding, but DRM3 sees no such benefit shown in Figures 7 and 8.

Batch-sizing for inference splits a request into parallel tasks and is a careful balance between throughput and latency in order to meet SLA and QPS targets. To show its interaction with distributed inference, we set the batch size artificially large to perform one-batch per-request. Smaller batch sizes increase task-level parallelism per-request and can reduce latency, but consequently increase task-level overheads and reduce data-parallelism, which can reduce throughput. In contrast, larger batches increase sparse operator work and benefit from the parallelization of distributed inference. Batch-sizing for deep recommendation inference is an on-going research topic [2].

The request replayer was configured to send requests at 25QPS to the main server for DRM1, the most compute-heavy of all evaluated models. The overhead graph for each DRM1 configuration, compared to singular, is shown in Figure 16. All overheads in the 25QPS experiment are less than the same configuration when sent serially, which was shown in Figure 6. Across nearly all configurations, the P99 latency is improved compared to the singular configuration, shown in Figure 6. Furthermore, in some configurations, like 8-shard capacity- and load-balanced, P50 is either improved over singular, or has less than $\sim.05\%$ overhead. Distributed inference has better latency characteristics in higher QPS scenarios given a model with sufficient sparse operator compute.

To provide an apples-to-apples comparison of adding servers for distributed inference, the same hardware platform was used for the sparse shards. However, Figures 8 and 12 shows that the sparse shards have far less compute requirements compared to the main shard, as the sharding in this work is capacity-driven, not compute-driven. Thus, we run an additional configuration using lighter-weight SC-Small servers with DRM1, which again, is the most compute-intensive model. Figure 15 shows the per-shard operator latencies are nearly identical, when run on the original heavier-weight SC-Large and the lighter-weight SC-Small servers. Recall that the SC-Large server has more and faster cores, and $4\times$ memory capacity, resulting in increased energy footprint. This suggests the opportunity for coarse-grained platform specialization of sparse shards for increased serving- and energy-efficiency at the cluster-level.

In this section, we include a small discussion on how distributed inference can improve shard replication. Supporting singular models that have 100s of GBs of memory footprint requires servers that are inefficiently provisioned. For example, most of the memory capacity will be dedicated to the large embedding tables, but most cores will be dedicated to using the significantly smaller dense parameters. The majority of compute touches less than 3% of the model’s memory footprint for DRM1, DRM2, and DRM3. This inefficiency is scaled to support the QPS of millions or billions of users as inference servers are replicated dynamically. In such a case, the large load incurred by the dense layers, shown in Figure 4, will cause the entire model to be replicated to additional servers, including all embedding tables. Distributed inference alleviates this inefficiency by enabling compute resources to be allocated for dense layers and memory resources to be allocated for sparse layers. In other words, the memory requirements of replication are reduced. Our heuristic to shard the dense and sparse layers of the models into separate components simplifies this allocation. Lastly, sharding strategy plays an auxiliary role in replication. While the decision to shard sparse operators has the strongest impact on reducing memory requirements of replication, Table II shows that sharding strategy can further impact replication. Recall that DRM1 is comprised of two primary nets, where Net 2 consumes $4.75\times$ as much memory resources, but 6.3% of the compute resources as Net 1. The NSBP sharding strategy constrains each shard to hold either Net 1 or Net 2, and accordingly doesn’t parallelize work as well compared a net-agnostic sharding strategy. However, NSBP can further improve resource utilization by grouping the most compute-dense embedding tables together, even at higher shard counts necessary for larger models. This sharding strategy trade-off, between improved latency and improved resource utilization, will need to be made on a model-by-model basis and encourages further work in sharding automation.

Model compression is the traditional approach to constrain model capacity. We show the compressed model size for DRM1 in Table III as a point-of-reference. All compression techniques deployed on current data-center models were used to generate the compressed model, which utilize quantization and pruning[22]. All tables were row-wise linear quantized to at least 8-bits, and sufficiently large tables were quantized to 4-bits. Tables were manually pruned as specified by the model architect based on a threshold magnitude or training update frequency. Quantization and pruning options were chosen to preserve accuracy and latency. We note that latency and compute are marginally improved with compression, but because this was a focus of this work, we leave this analysis to future work. We speculate the cause is improved memory locality, relatively. More relevantly, Table III shows the compressed model is 5.56$\times$ smaller. While significant, even with these savings, large models will still not be able to fit on one, two, or even four commodity servers configured with $\sim$50GB of usable DRAM. Thus, compression alone is insufficient to enable emerging large deep recommendation systems.