面向深度学习推荐模型快速可扩展训练的软硬件协同设计
分布式、并行与集群计算
2023-02-28 v7 人工智能
机器学习
性能
摘要
深度学习推荐模型（DLRMs）被用于 Facebook 众多业务关键型服务中，并且是其数据中心基础设施需求方面规模最大的单一 AI 应用。本文中，我们讨论用于大规模 DLRM 高性能分布式训练的软硬件协同设计方案。我们引入基于 PyTorch 的高性能可扩展软件栈，并将其与新演进的 Zion 平台（即 ZionEX）配对。我们展示了训练参数量高达 12 万亿的超大 DLRM 的能力，并表明相较先前系统在求解时间上可实现 40 倍加速。我们通过以下方式实现：(i) 设计具备专用横向扩展网络、高带宽、最优拓扑与高效传输的 ZionEX 平台；(ii) 实现支持模型并行与数据并行的优化基于 PyTorch 的训练栈；(iii) 开发能够沿行、列维度对嵌入表进行分层划分并在多个工作节点间负载均衡的分片算法；(iv) 在保留支持具有完全确定性更新的优化器灵活性的同时，添加高性能核心算子；(v) 利用低精度通信、多级存储层次（HBM+DDR+SSD）与流水线。此外，我们开发并简要说明了在生产环境中稳健高效端到端训练所需的分布式数据摄取与其他支撑服务。
引用
@article{arxiv.2104.05158,
title = {Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models},
author = {Dheevatsa Mudigere and Yuchen Hao and Jianyu Huang and Zhihao Jia and Andrew Tulloch and Srinivas Sridharan and Xing Liu and Mustafa Ozdal and Jade Nie and Jongsoo Park and Liang Luo and Jie Amy Yang and Leon Gao and Dmytro Ivchenko and Aarti Basant and Yuxi Hu and Jiyan Yang and Ehsan K. Ardestani and Xiaodong Wang and Rakesh Komuravelli and Ching-Hsiang Chu and Serhat Yilmaz and Huayu Li and Jiyuan Qian and Zhuobo Feng and Yinbin Ma and Junjie Yang and Ellie Wen and Hong Li and Lin Yang and Chonglin Sun and Whitney Zhao and Dimitry Melts and Krishna Dhulipala and KR Kishore and Tyler Graf and Assaf Eisenman and Kiran Kumar Matam and Adi Gangidi and Guoqiang Jerry Chen and Manoj Krishnan and Avinash Nayak and Krishnakumar Nair and Bharath Muthiah and Mahmoud khorashadi and Pallab Bhattacharya and Petr Lapukhov and Maxim Naumov and Ajit Mathews and Lin Qiao and Mikhail Smelyanskiy and Bill Jia and Vijay Rao},
journal= {arXiv preprint arXiv:2104.05158},
year = {2023}
}