GEM Training: How Meta Doubled the Efficiency of Its LLM-Scale Ads Foundation Model

來源原標題

GEM Training: How Meta Doubled the Efficiency of Its LLM-Scale Ads Foundation Model

作者

原始標籤

AI Research · Data Infrastructure · ML Applications

原始連結

https://engineering.fb.com/2026/08/03/ml-applications/training-gem-at-llm-scale-meta-ads-recommendation-foundation-model/

摘要原文

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. GEM is trained on ad content and user engagement data with two categories of features: sequence features (e.g., user activity history) and non-sequence features (e.g., 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)) Near-linear scaling requires four conditions: GEM’s workload threatens every one of these: 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: E2E MFU = Local MFU (compute efficiency) × Scaling Ratio (scaling efficiency) 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). Scaling Ratio (scaling efficiency) measures how much single-GPU performance is retained when distributing across thousands of GPUs. A scaling ratio of 1.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. 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. We evolved JFA through four generations, progressively closing the gap from being slower than padded SDPA (scaled dot-product attention) to matching SOTA CUDA/Cutlass performance on latest-generation GPUs: JFA v4 (TLX) achieves 40-140% TFLOPS improvement over JFA v2, which delivers consistent gains under production jagged distributions (sparsity 0.5), contributing to 18.5% relative local MFU gain and 12% QPS gain. 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. We observed a 2.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. With these optimizations, the optimized GDPA kernel achieves 2x forward speedup (1,145 BF16 TFLOPs, ~97% Tensor Core utilization) and 1.6x backward speedup over baseline. Under short K/V production settings, it achieves up to 3.5x forward speedup over Flash Attention 4 (FA4). Applied across the full model, these kernels deliver over 30% end-to-end training throughput improvement. For GEM self-attention, the core efficiency challenge was scaling long user sequences without paying the quadratic cost of full attention. We first moved the layer from full self-attention to sliding-window attention, limiting each token to nearby events and reducing complexity from O(L 2 ) to O(L * window). This made longer sequences practical. The Sliding Window Attention (SWA) kernel skipped off-window tiles in JFA and reduced long-sequence self-attention latency by up to 68% with neutral NE (normalized entropy, a model-quality metric) . We then pushed the structure further with block-aligned attention. Since GEM could safely use fixed 64-token blocks, each Q block only attends to its corresponding K/V block, turning attention into independent 64×64 problems. This removes the partial-window masking and multi-tile iteration still present in SWA, and lets a dedicated TLX kernel eliminate FlashAttention overheads such as online softmax correction, logsumexp HBM traffic, and separate Di preprocessing. Fusing RoPE backward into the attention epilogue removes another memory-bound kernel and keeps gradients in FP32 registers. Together, TLX block attention + fused rotary improves self-attention layer MFU by +30.6% over Triton block attention, or roughly +44% over the SWA baseline. On a GPU, lower precision directly translates to higher Tensor core throughput. For the latest generation GPU, FP8 delivers 2x peak FLOPS over FP16, and FP4 delivers 4x. We expect the peak FLOPS of low precision to increase faster in next-generation GPUs. This makes low-precision training increasingly attractive as hardware vendors scale low-precision FLOPS faster than FP16. However, making low-precision training work without quality regression — addressing both numerical stability and quantization overhead — remains an industry-wide challenge. We developed MXFP8 Attention and MLP with numerical stability enhancement, which addressed both training stability and quantization overhead. We extended the FA4 kernel with end-to-end MXFP8 blockscaled MMA for both forward and backward passes leveraging latest generation GPUs’ native support for low precision.