Characterizing and predicting the training performance of modern machine learning (ML) workloads on compute systems with compute and communication spread between CPUs, GPUs, and network devices is not only the key to optimization and planning but also a complex goal to achieve. The primary challenges include the complexity of synchronization and load balancing between CPUs and GPUs, the variance in input data distribution, and the use of different communication devices and topologies (e.g., NVLink, PCIe, network cards) that connect multiple compute devices, coupled with the desire for flexible training configurations. Built on top of our prior work for single-GPU platforms, we address these challenges and enable multi-GPU performance modeling by incorporating (1) data-distribution-aware performance models for embedding table lookup, and (2) data movement prediction of communication collectives, into our upgraded performance modeling pipeline equipped with inter-and intra-rank synchronization for ML workloads trained on multi-GPU platforms. Beyond accurately predicting the per-iteration training time of DLRM models with random configurations with a geomean error of 5.21% on two multi-GPU platforms, our prediction pipeline generalizes well to other types of ML workloads, such as Transformer-based NLP models with a geomean error of 3.00%. Moreover, even without actually running ML workloads like DLRMs on the hardware, it is capable of generating insights such as quickly selecting the fastest embedding table sharding configuration (with a success rate of 85%).

翻译：表征和预测现代机器学习（ML）工作负载在计算和通信分布于CPU、GPU及网络设备之间的计算系统上的训练性能，不仅是优化和规划的关键，也是一个复杂的目标。主要挑战包括CPU和GPU之间同步与负载均衡的复杂性、输入数据分布的差异性，以及连接多个计算设备的不同通信设备和拓扑结构（如NVLink、PCIe、网卡）的使用，同时还需要灵活的训练配置。基于我们先前在单GPU平台上的工作，我们通过将（1）嵌入表查找的数据分布感知性能模型和（2）通信集合的数据移动预测集成到升级的性能建模流水线中（该流水线具备面向多GPU平台训练的ML工作负载的秩间和秩内同步能力），解决了这些挑战，实现了多GPU性能建模。除了能够以5.21%的几何平均误差在两个多GPU平台上准确预测随机配置的DLRM模型每迭代训练时间外，我们的预测流水线还良好地泛化到其他类型的ML工作负载（例如基于Transformer的NLP模型，几何平均误差为3.00%）。此外，即使未在硬件上实际运行如DLRM等ML工作负载，它也能生成诸如快速选择最优嵌入表分片配置（成功率达85%）等洞察。