GEM训练：Meta如何将其LLM规模广告基础模型的效率提升一倍
内容提要
Meta的GEM广告推荐模型通过软硬件协同设计，将端到端训练效率提升至20-25% MFU，训练FLOPs一年内增长4倍。核心创新包括定制内核库（如Jagged Flash Attention、GDPA、BlockAttention）和混合超低精度训练提升计算效率；拓扑感知的5D并行、SM-free通信、自动激活检查点和负载均衡优化扩展效率。这些方法解决了推荐系统与LLM混合架构的独特挑战。
延伸解读
为什么推荐模型训练不能直接套用LLM方案
GEM的混合架构（万亿级稀疏参数+数十亿稠密参数）和推荐场景的数据特性（如变长序列、非对称注意力）使其训练负载与典型LLM差异显著。文章指出，为LLM优化的内核、并行策略和低精度方案无法直接迁移，必须针对推荐负载重新设计。这提醒我们，AI基础设施的优化需要紧密结合具体工作负载，不能简单复用通用方案。
低精度训练的关键不只是精度，还有开销
虽然FP8/FP4能带来更高的Tensor Core吞吐，但量化本身会产生额外开销（如缩放因子计算、数据转换）。Meta通过预量化FSDP分片、融合量化到前序算子、以及随机舍入和Hadamard变换等技巧，既降低了开销又保证了数值稳定性。这表明低精度训练的成功依赖于系统级的精细设计，而非单纯降低位宽。
扩展效率的瓶颈往往在通信和负载均衡
文章强调，在数千GPU规模下，通信开销、SM占用和负载不均会严重侵蚀扩展效率。Meta通过拓扑感知的5D并行、SM-free通信和基于序列长度的负载均衡（BBS）来应对。特别是BBS方法，通过排序和交错子批次，在不增加通信的前提下实现了接近最优的负载均衡，这为处理数据倾斜提供了实用思路。
Q&A
Meta的GEM广告推荐模型在训练效率上取得了哪些具体成果？
Meta的GEM模型通过软硬件协同设计，将端到端训练效率提升至20-25%的模型FLOPs利用率（MFU），并在12个月内将训练FLOPs扩大了4倍。
GEM模型在训练中面临哪些独特挑战？
GEM模型面临两大挑战：一是计算效率方面，包括输入序列长度不一（jagged inputs）导致填充浪费、注意力模式不对称、内存受限操作以及数值敏感性问题；二是扩展效率方面，包括海量稀疏参数和密集参数带来的通信开销、架构多样性导致的重叠窗口不均、长序列导致的内存压力以及序列长度不均导致的负载不均衡。
Meta如何优化GEM模型的计算效率？
Meta通过定制内核库和混合超低精度训练来优化计算效率。定制内核包括Jagged Flash Attention（JFA）处理变长序列、BlockAttention降低长序列自注意力复杂度、GDPA统一并加速非对称注意力模块；超低精度训练采用MXFP8注意力与MLP，并配合数值稳定性增强技术，如随机Hadamard变换、随机舍入和混合精度策略。
GEM模型如何实现高效的大规模分布式训练？
Meta采用拓扑感知的5D并行策略，包括对密集参数使用2D FSDP加专家并行（EP），对稀疏参数使用全分片2D模型并行。同时，通过SM-free通信（如NCCLX）减少通信对计算资源的占用，使用自动激活检查点和激活量化降低内存压力，并通过序列长度感知的负载均衡（如Base Batch Shuffling）解决数据驱动的负载不均问题。
Jagged Flash Attention (JFA) 解决了什么问题？
JFA解决了推荐模型中用户序列长度不一（jagged inputs）导致的填充浪费问题。传统FlashAttention假设序列长度固定，而JFA直接处理变长张量，消除了填充开销，并支持自定义注意力偏置、非对称查询/键值长度和高效反向传播。JFA经过四代演进，最终在最新GPU上达到或超过SOTA性能。
GEM模型如何应对低精度训练带来的数值稳定性问题？
Meta采用了多种技术来保证低精度训练的数值稳定性：使用随机Hadamard变换分散异常值，采用随机舍入消除确定性舍入偏差，对权重梯度选择跳过或使用更高精度，以及混合精度策略（在敏感层使用BF16）。这些方法在提升训练速度的同时，避免了CTR/CVR等目标的精度回归。
GEM模型在负载均衡方面采用了什么创新方法？
Meta开发了Base Batch Shuffling（BBS）技术，通过分布式读取器生成小批次（128样本），按序列长度排序并交错合并（最重与最轻配对）形成完整训练批次，在不引入跨秩通信的情况下捕获了接近理论最优的负载均衡，实现了4%的效率提升（包括4%的QPS提升和4%的峰值内存降低）。
