MTrainS：利用异构内存提升 DLRM 训练效率
信息检索
2023-05-03 v1 机器学习
性能
摘要
推荐模型规模非常庞大，训练期间需要太字节（TB）级内存。为追求更优质量，模型规模与复杂度随时间增长，这需要额外的训练数据以避免过拟合。此种模型增长对数据中心资源需求量巨大。因此，训练效率对于维持数据中心功耗可控正变得愈发重要。在深度学习的推荐模型（DLRM）中，通过嵌入表捕获类别型输入的稀疏特征是模型规模的主要来源，且需要高内存带宽。本文研究了实际部署模型中嵌入表的带宽需求与局部性。我们观察到不同表的带宽需求并不均匀，且嵌入表表现出高时间局部性。随后我们设计了 MTrainS，其分层利用异构内存，包括字节与块可寻址的存储级内存（Storage Class Memory）用于 DLRM。MTrainS 允许每节点更高内存容量，并通过在内存容量受限用例中降低向外扩展至多主机的需要来提升训练效率。通过优化平台内存层次结构，我们将训练节点数减少 4-8 倍，在达到目标训练性能的同时节省了训练功耗与成本。
引用
@article{arxiv.2305.01515,
title = {MTrainS: Improving DLRM training efficiency using heterogeneous memories},
author = {Hiwot Tadese Kassa and Paul Johnson and Jason Akers and Mrinmoy Ghosh and Andrew Tulloch and Dheevatsa Mudigere and Jongsoo Park and Xing Liu and Ronald Dreslinski and Ehsan K. Ardestani},
journal= {arXiv preprint arXiv:2305.01515},
year = {2023}
}