GEM Training: How Meta Doubled the Efficiency of Its

What happened

GEM is the central recommendations foundation model behind Meta’s ads system. It has a hybrid architecture with trillions of sparse embedding parameters and billions of dense parameters. , user location, ad creative representation). Customized attention mechanisms are applied to each group independently, while also enabling cross-feature learning. The interplay between this hybrid architecture and rec-domain data properties is what makes GEM’s training uniquely challenging.

Today’s data center GPUs and their software stacks are mostly optimized for LLM workloads, whereas recommendation workloads have a fundamentally different profile due to unique data characteristics and rich user & ads signal interaction patterns that make it extremely difficult to achieve high GPU compute utilization for training a foundational recommendation model of GEM’s size.

Training GEM across thousands of GPUs with trillions of sparse embedding parameters and billions of dense parameters requires scaling efficiently, not just scaling up. Simply adding more GPUs does not translate to proportional speedup.

In distributed training, E2E latency per training step is determined by: E2E Latency = Max across GPU Rank (Max(Local Compute Time, Communication Time)) Given the challenges outlined above, we needed a framework that turned a sprawling co-design effort into a small number of technical levers. We measure training efficiency through E2E MFU, which decomposes into two factors: These factors describe two related but distinct optimization problems.

Local MFU (compute efficiency) measures how well a single GPU’s compute units are utilized — how close the workload runs to the hardware roofline. It is determined by kernel design, numerical precision, and how well the workload’s compute patterns (data dimensions, sequence lengths) map onto GPU architecture (Tensor cores, memory hierarchy, streaming multiprocessor scheduling).

To address the recommendations-system-specific challenges mentioned above and push up GPU FLOPS utilization, we built a custom kernel library and an ultra-low-precision training recipe custom-built and optimized for recommendation workloads on the latest GPU hardware. FlashAttention is designed for dense, fixed-length sequences common in LLMs.

In recommendation models, user sequences are inherently jagged — varying from hundreds to tens of thousands of tokens per sample — and padding to max length could waste up to 50% of compute. Standard FlashAttention implementations assume uniform sequence lengths for efficient tiling and parallelization; with jagged inputs, naive approaches either pad (wasting compute) or leave SMs idle when short sequences finish early.

We developed JFA, a custom FlashAttention implementation that operates directly on variable-length jagged tensors, eliminating padding overhead while supporting rec-specific features such as custom attention biases, asymmetric query/key-value lengths, and efficient backward passes. 5% relative local MFU gain and 12% QPS gain.

Why it matters

Scaling Ratio (scaling efficiency) measures how much single-GPU performance is retained when distributing across thousands of GPUs. 0 means perfect linear scaling; in practice, communication overhead, load imbalance, straggler effects, and activation recomputation from memory pressure all erode it. To isolate local MFU, we run model layers individually on a single GPU and compute a weighted average MFU without activation recomputation or communication exposure.

The scaling ratio is derived as the ratio between local and E2E MFU. This decomposition matters because it lets us treat compute efficiency and scaling efficiency as related but distinct optimization problems, each with its own dedicated set of techniques: Both must be addressed to maximize end-to-end MFU.

What to watch

GEM uses diverse attention-like interaction patterns — self-attention, PMA, and cross-attention — that share a common structure: two matrix multiplications with an element-wise activation in between, but replace softmax with activations like GELU or SiLU. We unify these modules under a single GDPA kernel optimized for production RecSys training workloads on latest generation GPUs.

Existing FlashAttention kernels are designed for LLM-style dense, long-sequence inputs and perform poorly under real production traffic. 6x forward performance gap and up to 4x worst-case gap between real-world workloads and synthetic benchmarks driven by short/asymmetric K/V sequences, jagged inputs, and large batch sizes that break pipeline occupancy assumptions. We redesigned the kernel pipeline, scheduling, and math to close the performance gap between real-world traffic and hardware roofline. 6x backward speedup over baseline.