Deep learning recommendation models (DLRMs) are used across many business-critical services at Facebook and are the single largest AI application in terms of infrastructure demand in its data-centers. In this paper we discuss the SW/HW co-designed solution for high-performance distributed training of large-scale DLRMs. We introduce a high-performance scalable software stack based on PyTorch and pair it with the new evolution of Zion platform, namely ZionEX. We demonstrate the capability to train very large DLRMs with up to 12 Trillion parameters and show that we can attain 40X speedup in terms of time to solution over previous systems. We achieve this by (i) designing the ZionEX platform with dedicated scale-out network, provisioned with high bandwidth, optimal topology and efficient transport (ii) implementing an optimized PyTorch-based training stack supporting both model and data parallelism (iii) developing sharding algorithms capable of hierarchical partitioning of the embedding tables along row, column dimensions and load balancing them across multiple workers; (iv) adding high-performance core operators while retaining flexibility to support optimizers with fully deterministic updates (v) leveraging reduced precision communications, multi-level memory hierarchy (HBM+DDR+SSD) and pipelining. Furthermore, we develop and briefly comment on distributed data ingestion and other supporting services that are required for the robust and efficient end-to-end training in production environments.

翻译：深度学习推荐模型（DLRM）在Facebook多个关键业务场景中得到广泛应用，是数据中心中基础设施需求最大的单一AI应用。本文探讨了面向大规模DLRM高性能分布式训练的软硬件协同设计方案。我们基于PyTorch构建了高性能可扩展软件栈，并将其与新一代Zion平台（即ZionEX）相结合。实验证明，该方案能够训练参数量高达12万亿的超大规模DLRM，相比先前系统可获得40倍的求解加速比。这一成果通过以下途径实现：（i）设计具有专用扩展网络、高带宽、最优拓扑与高效传输协议的ZionEX平台；（ii）实现支持模型并行与数据并行的优化PyTorch训练栈；（iii）开发支持嵌入表沿行/列维度分层分区、并在多工作节点间实现负载均衡的分片算法；（iv）在保持灵活性的同时增加高性能核心算子，支持完全确定性更新的优化器；（v）利用低精度通信、多级存储层次（HBM+DDR+SSD）及流水线技术。此外，我们还开发并简要讨论了分布式数据摄取及其他支撑服务，这些服务是生产环境中高效鲁棒端到端训练的必备组件。