RecShard: statistical feature-based memory optimization for industry-scale neural recommendation

Geet Sethi, Bilge Acun, Niket Agarwal, Christos Kozyrakis, Caroline Trippel and Carole-Jean Wu

Recommendation models hide terabytes of embedding tables behind a few hot rows — RecShard profiles the training data's statistics and solves one optimization problem that splits every table between fast GPU memory and slow host memory at exactly the right row, keeping 99.5% of lookups on the fast tier.

The argument in one glance

Problem

Recommendation models are memory monsters, and the memory is running out

Deep learning recommendation models (DLRMs) dominate datacenter ML — over 50% of training and 80% of inference demand at Facebook. Their embedding tables, which turn categorical features into vectors, hold over 99% of model capacity: terabytes. Between 2017 and 2021 model memory needs grew 16× and per-sample bandwidth demand almost 30×, while GPU memory grew less than 6×.
What people did

Spread whole tables across GPUs and spill the rest to host memory

Training systems shard embedding tables across GPUs (model parallelism) and, when GPU memory runs out, place overflow tables in CPU DRAM reached via UVM. Which table goes where is decided by greedy heuristics using a single cost per table: its size, or its lookup rate.
Why it fell short

Whole tables are the wrong unit

Access skew lives inside tables: most features follow power laws, so a sliver of rows serves most lookups. Bandwidth demand differs 10× between features (pooling factors from a few to ~200), usage frequency ranges from under 1% to 100% of samples, and — because tables are hash tables — the birthday paradox leaves large fractions of rows never touched at all. A per-table cost cannot see any of this, so heuristics misplace hot rows on slow memory and load-balance poorly.
This paper's idea

Profile the data, split every table at the statistically right row

RecShard samples about 1% of training data to estimate each feature’s access distribution, pooling factor, and coverage, then solves a mixed integer linear program that partitions every table between GPU HBM and host DRAM and balances all GPUs in one shot. Capacity-constrained models train their embeddings over 6× faster on average, with 12× better load balance — 61% of rows live on the slow tier, yet 99.5% of accesses hit the fast one.

Why this problem is worth solving

deep learning recommendation model — the architecture behind feeds, ads, and search ranking, predicting the probability that a user interacts with an item (the click-through rate)

is two very different networks stapled together. Small dense multilayer perceptrons do the arithmetic. And embedding tables do the remembering: each categorical (“sparse”) feature — user country, pages recently viewed, video IDs watched — has a table whose rows are learned vectors, one per category. A training sample activates a handful of categories per feature; their rows are fetched and summed (“pooled”) into a dense representation.

The tables are where the scale lives — over 99% of model capacity, terabytes in production — and the workload they generate is unlike anything dense networks produce: scatter-gather reads of 64-float rows, with almost no compute per byte. At Facebook, DLRMs account for more than half of all training demand. The hardware trend is unforgiving: in the paper’s four-year window, model memory requirements grew 16× and per-sample bandwidth demand nearly 30×, while GPU high-bandwidth memory capacity improved less than 6×. So production training systems became tiered: each GPU’s

high-bandwidth memory — the fast on-package DRAM of the GPU, about 1.5 TB/s on an A100

holds what fits, and the rest lives in host DRAM reached through

Unified Virtual Memory — a CUDA mechanism letting GPU kernels transparently read host memory over PCIe, at roughly two orders of magnitude lower bandwidth than HBM

. With a 100× cliff between tiers, what goes where becomes the training-throughput question.

What people did before, and where it broke

The state of the art placed whole tables using a greedy heuristic: assign each table a scalar cost, sort descending, then repeatedly give the next table to the least-loaded GPU, spilling to UVM when HBM fills. Two cost functions were standard — size-based (hash size × vector dimension) and lookup-based (average lookups × dimension) — plus combinations.

The paper’s characterization of hundreds of production features shows why this granularity fails.

Skew lives inside tables. Most features’ category frequencies follow power laws (country popularity mirrors world population; content IDs follow popularity curves). A small set of rows sources the vast majority of a table’s accesses — and the skew strength differs per feature, so no single rule of thumb fits all.

Bandwidth demand is a per-feature property. A feature’s

the number of categories active in one sample — hence the number of rows fetched and summed per lookup

averages near 1 for some features (user location) and around 200 for others (recent-activity lists): an order-of-magnitude spread in bytes per sample that a size-based cost never sees.

Usage frequency varies 100×. A feature’s coverage — the fraction of training samples in which it appears at all — ranges from under 1% to 100%. A huge table consulted in 1% of samples deserves different placement than a small one consulted in every sample.

Hash tables waste memory by mathematics. Raw category spaces are huge, so inputs are hashed into fixed-size tables. Hashing N distinct values into N slots leaves about 1/e — 37% — of slots empty (the birthday paradox), and production tables show it: one measured table was 26% unused from data sparsity plus another 22% from collisions. Growing a table to reduce collisions grows the empty fraction too. Whole-table placement drags all this dead weight into precious HBM.

The result: on the paper’s models, heuristics differ from RecShard on 23–59% of row placements, their GPUs’ iteration times vary by up to 5× within one job, and their throughput collapses as models outgrow HBM.

The RecShard pipeline

Training data profiling new

Sample roughly 1% of the training dataset, apply each feature’s hash, and estimate three statistics per table: the access-frequency CDF over rows, the average pooling factor, and the coverage. The tail can be ignored safely — unseen tail values either collide into existing rows or land in rows that stay cold anyway.
↓ →
MILP partitioning and placement new

A mixed integer linear program chooses, for every table, a split point on its access CDF (hot rows → HBM, cold rows → UVM) and a GPU assignment — minimizing the worst per-GPU cost, where each table’s cost is its per-step bytes (pooling × dimension × batch) weighted by tier bandwidths and coverage. Gurobi solves the 47,276-variable program in 21–42 seconds.
↓ →
Remapping layer new

Hot rows are scattered through each table, so a per-table remapping table translates every embedding index to its new HBM or UVM location. Implemented as a PyTorch C++ operator inside the data loader, it runs in parallel with training — off the critical path.
↓ →
Training

The standard hybrid-parallel DLRM loop (data-parallel MLPs, model-parallel embeddings, all-to-all exchange) runs unmodified, now with nearly every lookup landing in HBM.

Solve time scales with the number of GPUs and CDF approximation steps — not with model size — and the strategy is recomputed cheaply each time a model trains, tracking shifts in the data distribution.

The idea

RecShard’s premise is that the training data already tells you the cost of every row before training starts — you just have to read the statistics and then optimize globally instead of greedily.

Rows, not tables, as the placement unit. Because access CDFs are steep, a table is really three populations: a hot head that belongs in HBM, a long cold tail that UVM can serve harmlessly, and (thanks to hashing) a block of never-touched rows whose placement cost is zero. RecShard splits each table at a chosen point of its CDF and places the pieces independently — reclaiming the birthday-paradox waste that whole-table schemes carry into HBM.

A cost model with three ingredients. A table’s per-training-step demand is pooling factor × vector dimension × bytes × batch size — the bytes it actually moves. That demand splits between tiers by the CDF at the split point: the fraction of accesses served from HBM pays 1/BW-HBM per byte, the remainder pays 1/BW-UVM. Finally the whole cost is weighted by coverage, converting per-table numbers into a fleet-comparable global currency. The ablation proves all three ingredients earn their place: using the CDF alone leaves 4.6× more accesses falling to UVM than the full formulation.

One shot, globally optimal, instead of greedy. The placement problem — hundreds of tables, 16 GPUs, two tiers, capacity limits per GPU — is encoded as a

mixed integer linear program — an optimization problem with linear costs and constraints where some variables must be integers (here, which GPU gets which table slice); off-the-shelf solvers find provably optimal solutions

whose objective is to minimize the most expensive GPU, since the slowest GPU sets the iteration time in synchronized training. Greedy heuristics assign tables one at a time and cannot undo early mistakes; the MILP considers all placements simultaneously. The paper has a telling internal result: a better cost function (lookup-based) fed into the greedy heuristic performed 46% worse than the size-based one on RM1 — evidence that the optimizer, not just the cost model, is where the win comes from.

Keep the machinery off the critical path. Profiling samples 1% of data; the solver runs under a minute; and index remapping happens inside the data-loading pipeline, overlapped with GPU work.

One embedding table's journey through RecShard

EMB = embedding table; row = one category's learned vector (dimension 64 here). CDF = cumulative distribution of accesses over rows, hottest first. HBM = GPU memory (~1.5 TB/s); UVM = host DRAM over PCIe (~100× slower). The feature and its numbers are an illustrative composite of the paper's characterization figures; system numbers (16 GPUs, batch 16,384, 24 GB HBM budget per GPU) are the paper's.

Step 1 of 6

Profile: 1% of the data reveals the shape

Take a feature like "pages recently viewed": RecShard samples training data, applies the feature’s hash, and builds its statistics. The access CDF is steep — say the hottest 5% of rows draw ~90% of accesses. Average pooling factor: ~40 rows summed per sample. Coverage: present in ~80% of samples. And a chunk of the hash space was never hit at all — data sparsity plus birthday-paradox collisions can leave tens of percent of rows empty.

per-EMB stats: access CDF, avg pooling, coverage
sampling: ~1% suffices (tail values collide or stay cold)

Why it works

Feature statistics are stable properties of the training data → a 1% sample measured before training predicts per-row access costs → placement can be proactive (decided once, optimally) instead of reactive caching machinery on the critical path.
Power-law CDFs mean a table’s value is concentrated in a small head → splitting tables at a CDF point puts the value in HBM and the bulk in UVM → capacity stops being the binding constraint long before bandwidth does.
Hashed tables contain provably-wasted space (birthday paradox: ~1/e empty at equal size) → RecShard prices untouched rows at zero and demotes them → HBM stops storing memory that mathematics guarantees is dead.
Synchronous training runs at the speed of the slowest GPU → the MILP minimizes the maximum per-GPU cost over all placements jointly → load imbalance (the heuristics’ 5× spread) collapses to a few percent, converting directly into throughput.
Pooling and coverage weight each table’s cost by real bytes and real usage → the solver spends HBM where bandwidth is actually consumed → removing them (CDF-only ablation) leaks 4.6× more accesses to the slow tier.

What it costs

Profiling and solving, per training run. About 1% of data sampled, plus a 21–42 second MILP solve; both trivially amortized over training, but the pipeline must rerun when data distributions or the model change.
A remapping indirection. Every lookup passes through a per-table index remap; it is executed in the data loader off the critical path, at the price of CPU work and pipeline complexity.
Trust in the statistics. Placement is static per run: if access patterns drift sharply mid-training, the split points are stale until the next run (the paper argues distributions are stable at this horizon).
Whole-table-to-one-GPU constraint. The MILP splits tables across tiers, but each table’s row range stays on one GPU — a modeling simplification that bounds achievable balance.
Forward-pass modeling only. The MILP prices forward lookups and relies on backward-pass cost being proportional — true for embedding operators today, an assumption nonetheless.

The evaluation

The testbed is one production-class node: 16 NVIDIA A100 (40 GB) GPUs across two sockets, each GPU budgeted 24 GB of HBM for embeddings and 128 GB of host DRAM over PCIe 3.0 UVM, batch size 16,384. Three production-scale models share the same architecture — 397 sparse features, dimension-64 vectors — and differ by doubling every hash size: RM1 (318 GB, fits in the GPUs’ collective HBM), RM2 (635 GB), and RM3 (1.27 TB, needing UVM heavily). Baselines are the state-of-the-art greedy heuristic with size-based, lookup-based, and combined cost functions.

Sharding strategies on production-scale DLRMs (16 A100s)

Property	Size-based greedy	Lookup-based greedy (and combined)	RecShard
Placement unit and information used
Slowest-GPU embedding iteration time (RM1 / RM2 / RM3)
Load balance (std-dev of per-GPU time)
Behavior as models grow 4× (RM1 → RM3)

Modified open-source DLRM with FBGEMM embedding operators; per-GPU averages over training on production data, batch 16,384; 24 GB HBM per GPU reserved for embeddings, 128 GB host DRAM per GPU via PCIe 3.0 UVM. Min/max/mean/std-dev iteration times from the paper's Table 3; placement-difference and ablation numbers from Tables 4 and 6.Underlined values explain themselves when clicked.

The ablation cleanly separates the ingredients: on RM3, using only the access CDF in the MILP leaves 1.63 billion accesses per GPU falling to UVM; adding coverage cuts that to 881 million, adding pooling to 604 million, and the full formulation reaches 353 million — a 4.6× reduction attributable purely to the two “how much is this table actually used” statistics.

Where this stops working

What to keep in mind

All results come from Meta’s production models and data distributions; the power-law, coverage, and hashing phenomena are argued to be industry-typical, but the exact 6× depends on how steep your CDFs are. The comparison targets are per-table greedy heuristics — the deployed state of the art in 2022 — not caching-based alternatives, which the paper discusses qualitatively as reactive and critical-path-bound. The testbed is a single 16-GPU node with PCIe 3.0 UVM; faster host links (NVLink-C2C, CXL) narrow the tier gap that RecShard exploits, while larger models widen it. Training-side only: inference placement is out of scope. And the memory-tiering theme here is proactive and statistics-driven — an interesting contrast with reactive page-level tiering systems like Google’s TMTS (covered elsewhere on this site).

If you remember only three things

1. The right placement unit for embedding tables is the row range, not the table: power-law skew concentrates most accesses in a small hot head, pooling factors spread bandwidth demand across features by 10×, coverage varies 100×, and birthday-paradox hashing leaves ~1/e of slots empty — all invisible to per-table costs.
2. Training data statistics, sampled at 1%, price every row before training starts — so a 47,276-variable MILP can split each table at a CDF point and balance all 16 GPUs in one 21–42-second solve, with index remapping hidden in the data loader.
3. On production models, this turned capacity growth from a cliff into a slope: 2.6–7.4× faster embedding training (over 6× when UVM is needed), ~12× better GPU load balance, and 61% of rows on the slow tier while 99.5% of accesses stay on the fast one — a 4× model-size increase costs RecShard 1.2× versus over 3× for heuristics.

Check your understanding

1. A feature has exactly N distinct values, and its embedding table is given exactly N hashed slots. Predict roughly what fraction of the table is never used.
2. RecShard’s MILP weighs each table’s cost by pooling factor and coverage, not just its access CDF. The ablation shows why — what happens with the CDF alone?
3. Fed to the same greedy heuristic, the lookup-based cost function (more information) performed 46% WORSE than the size-based one on RM1. What does this reveal?
4. On RM3, RecShard places 61% of all embedding rows in slow UVM memory, yet throughput improves 7.41×. How can moving MORE data to the slow tier speed things up?