TensorSocket:面向深度学习训练的共享数据加载
机器学习
2025-08-04 v4 分布式、并行与集群计算
摘要
训练深度学习模型是一个重复且资源密集的过程。数据科学家在获得最高精度的参数集合(如超参数调优)和模型架构(如神经架构搜索)等之前,往往需要训练多个模型。这些训练任务的计算效率在很大程度上取决于训练数据供给训练过程的优劣。这些任务的重复性导致相同的数据处理管道被反复运行,加剧了对计算资源的需求及成本。在本文中,我们提出 TensorSocket,通过使同步训练进程共享同一数据加载器来减少深度学习训练的计算需求。在协同训练工作负载于 GPU 上具有高吞吐量、但受限于 CPU 端较低数据加载吞吐量的情况下,TensorSocket 缓解了 CPU 端的瓶颈。TensorSocket 通过减少协同训练进程间的冗余计算与数据重复,并利用现代 GPU-GPU 互连技术来实现这一点。在此过程中,TensorSocket 能够训练并平衡不同大小的模型,同时服务多种批大小,且在本质上与硬件和管道无关。我们的评估表明,TensorSocket 使原本在无数据共享下不可行的场景成为可能,将训练吞吐量提升了高达 100%,并在使用云实例时,通过减少 CPU 端的硬件资源需求实现了 50% 的成本节约。此外,TensorSocket 优于 CoorDL 和 Joader 等共享数据加载的 SOTA 解决方案;它更易于部署和维护,在需要更少 CPU 资源的同时,实现了更高或持平的吞吐量。
引用
@article{arxiv.2409.18749,
title = {TensorSocket: Shared Data Loading for Deep Learning Training},
author = {Ties Robroek and Neil Kim Nielsen and Pınar Tözün},
journal= {arXiv preprint arXiv:2409.18749},
year = {2025}
}