基于深度强化学习的时间感知存储控制优化大规模分布式深度学习训练

《Scientific Reports》：Adaptive storage optimization for deep learning workloads in converged HPC-cloud environments: a reinforcement learning approach

　　

摘要

存储系统仍然是大规模深度学习训练中的重要瓶颈。大多数评估研究依赖于合成基准测试，而这些测试无法反映实际训练过程中存在的输入输出（I/O）异质性。自动化配置方法通常将参数选择视为无状态的，忽略了重复训练阶段的时间结构。为了解决这些问题，我们采用了一个集成了Lustre和S3兼容对象存储的高性能计算（HPC）-云计算（cloud）平台。我们评估了三种具有不同I/O特性的工作负载：在ImageNet-1K数据集上运行的ResNet-50、在医学分割挑战赛（Medical Segmentation Decathlon）中使用的3D U-Net，以及在LibriSpeech数据集上运行的Whisper Large。收敛时间被用作主要评估指标。我们提出了一种深度强化学习框架，该框架结合了近端策略优化（Proximal Policy Optimization）和 Soft Actor-Critic 算法，并通过基于Transformer的状态编码器来预测工作负载的阶段转换，而不仅仅是被动应对这些转换。通过这种方法，检查点写入延迟最多减少了74.6%，同时提升了所有三种工作负载的GPU利用率。该方法能够在不针对每种工作负载进行重新配置的情况下适应不同的I/O特性。研究结果表明，在异构的HPC-cloud环境中，具备时间感知能力的存储控制可以提高分布式训练流程的效率。