Machine learning driven system level heterogeneous memory management for high-performance computing

我们发现，利用机器学习的系统软件级异构内存管理解决方案，特别是基于非监督学习的方法，如强化学习，需要根据跨内存设备的数据布局快速估计执行运行时间，以探索不同的数据放置策略，这使得架构级模拟器对于此目的不切实际。我们提出了一种基于差分跟踪的方法，该方法使用通过基于高频采样的方法（例如，Intel的PEBS）在真实的硬件上使用不同的存储器设备。我们开发了一个运行时估计的基础上，这样的痕迹，提供了一个执行时间估计的数量级快于全系统模拟器。在一些HPC mini应用程序上，我们发现，与真实的硬件上的测量值相比，估计器预测运行时间的平均误差为4.4%。对于深度学习数据洗牌子主题，我们研究了在DL工作者之间划分数据集并在每个训练时期仅执行部分分布式样本交换的可行性。通过在ABCI的2048个GPU和Fugaku的4096个计算节点上进行广泛的实验，我们证明了在实践中，当仔细调整部分分布式交换时，可以保持全局洗牌的验证精度。我们在PyTorch中提供了一个实现，使用户能够控制拟议的数据交换方案。

Results have been achieved in two parallel efforts of the project.We found that system-software-level heterogeneous memory management solutions utilizing machine learning, in particular nonsupervised learning- based methods such as reinforcement learning, require rapid estimation of execution runtime as a function of the data layout across memory devices for exploring different data placement strategies, which renders architecture-level simulators impractical for this purpose. We proposed a differential tracing-based approach using memory access traces obtained by high-frequency sampling-based methods (e.g., Intel's PEBS) on real hardware using of different memory devices. We developed a runtime estimator based on such traces that provides an execution time estimate orders of magnitude faster than full-system simulators. On a number of HPC mini applications we showed that the estimator predicts runtime with an average error of 4.4% compared to measurements on real hardware.For the deep learning data shuffling subtopic, we investigated the viability of partitioning the dataset among DL workers and performing only a partial distributed exchange of samples in each training epoch. Through extensive experiments on up to 2048 GPUs of ABCI and 4096 compute nodes of Fugaku, we demonstrated that in practice validation accuracy of global shuffling can be maintained when carefully tuning the partial distributed exchange. We provided an implementation in PyTorch that enables users to control the proposed data exchange scheme.