Optimizing LLM-Scale Recommendation Model Training at Meta: A Co-Design Approach

Meta's Generative Ads Recommendation Model (GEM) represents a significant engineering challenge, blending large language model (LLM) scale with the distinct characteristics of recommendation systems. Unlike traditional LLMs, GEM's hybrid architecture and recommendation-domain data properties — such as highly variable sequence lengths (jagged inputs), diverse attention patterns, and memory-bound operations — necessitate specialized solutions to achieve high GPU utilization and efficient scaling.

Key Challenges in Recommendation Model Training

• Jagged Inputs: User activity histories vary wildly, leading to variable sequence lengths. Naive padding wastes up to 50% compute.
• Diverse Interaction Patterns: Asymmetric sequence shapes for different attention mechanisms (self-attention, cross-attention, pooled multi-head attention) hinder intra-kernel pipelining and compute saturation.
• Memory-Bound Operations: Small embedding dimensions and numerous normalizations lead to underutilized compute units.
• Numerical Sensitivity: Ads optimization tasks are highly sensitive to precision, making low-precision training risky for model quality.
• Scaling Inefficiency: Trillions of sparse and billions of dense parameters, coupled with diverse layer architectures and jagged data, create heavy communication, uneven overlap windows, and load skew, preventing near-linear scaling across thousands of GPUs.

Meta's Efficiency Framework: Compute vs. Scaling

Meta decomposes end-to-end Model FLOPs Utilization (MFU) into Local MFU (compute efficiency) and Scaling Ratio (scaling efficiency). This framework allows for targeted optimization:

• Local MFU (Compute Efficiency): Focuses on maximizing single-GPU utilization. Addressed through custom kernel design (e.g., Jagged Flash Attention, Generalized Dot-Product Attention) and ultra-low-precision training (e.g., MXFP8 attention + MLP) optimized for recommendation workloads and specific GPU architectures.
• Scaling Ratio (Scaling Efficiency): Focuses on minimizing performance degradation when distributing workloads across many GPUs. Addressed through advanced parallelism strategies (e.g., topology-aware 5D parallelism, 2D FSDP + Expert Parallelism, Fully Sharded 2D Model Parallelism), network hierarchy co-design, memory management, and load balancing.

Innovations in Compute Efficiency

• Jagged Flash Attention (JFA): A custom FlashAttention implementation that operates directly on variable-length jagged tensors, eliminating padding waste. Achieved through novel masking via subtraction schemes, optimized backward parallelization, warp specialization, and persistent kernels.
• Generalized Dot-Product Attention (GDPA): Unifies and accelerates diverse attention-like patterns in GEM (self-attention, PMA, cross-attention) that differ from typical LLM FlashAttention assumptions. Involves pipeline redesign for non-softmax activations and software-level tile scheduling for jagged tensors, closing performance gaps between real-world workloads and hardware capabilities.
• MXFP8 attention + MLP: Enables ultra-low-precision training to leverage Tensor Core throughput without regressing precision-sensitive objectives like CTR/CVR prediction.

Innovations in Scaling Efficiency

To scale efficiently across thousands of GPUs, Meta implemented a sophisticated 5D parallelism strategy co-designed with their multi-tiered network. This includes 2D FSDP (Fully Sharded Data Parallel) and Expert Parallelism for dense parameters, combined with Fully Sharded 2D Model Parallelism for sparse parameters, and SM-free collectives to minimize communication overhead. These techniques are crucial for handling the massive parameter count and diverse compute patterns of GEM, ensuring that communication doesn't bottleneck compute.