深度学习中的I/O需求分析与带宽评估
Source: Glenn K. Lockwood, I/O Requirements of AI, June 14, 2023
以下是我一直在积累的笔记和参考资料，旨在说明设计良好的存储子系统与购买大量GPU用于训练神经网络的重要性相当。
Ibrahim等人在《高性能计算环境中深度学习工作负载的架构需求》（Architectural Requirements for Deep Learning Workloads in HPC Environments，https://dx.doi.org/10.1109/PMBS54543.2021.00007）一文中，建立了一种很好的方法来计算在训练不同模型时保持GPU充分利用所需的存储带宽。他们评估了与科学研究最相关的相对较小的模型，并证明：
如果可以确定在训练过程中通过模型传递一份样本（前向和后向）所需的浮点运算次数（FLOPS），则可以使用平均样本大小计算出模型的每FLOP比特数（MiB/FLOP）比率。这相当于每FLOPS的MiB/s，可以将其乘以GPU（或整个系统）的FLOPS能力，从而得到每个GPU训练特定模型所需的带宽数量级估计。
他们表明CosmoFlow是一个计算成本较低的模型，每petaFLOP/s需要65GB/s的带宽。他们发现，在实践中，CosmoFlow只能利用每个NVIDIA V100 GPU的35-50 TFLOP/s进行训练，因此在单个V100上训练此模型需要2.275-3.250GB/s，而8路V100节点需要18.2GB/s-26GB/s。相比之下，典型的NFS客户端在TCP上的传输速度不超过3GB/s。
与之相比，在ImageNet上训练的ResNet-50的I/O密集度最低，每个V100最多只需要573MB/s，每个8路V100节点最多只需要3.9GB/s。
我创建了一个简单的工具（https://github.com/glennklockwood/atgtools/blob/master/ml-model-io-requirements.py），用于说明如何进行这种计算，并计算在不同GPU上训练模型所需的GB/s。这是一个非常粗略的模型，并通过假设每个GPU都具有足够的内存带宽、PCIe带宽、电源、冷却等，以其各自规格表上的额定性能进行训练的上限来估算所需的带宽。正如Ibrahim的论文所示（他们发现CosmoFlow以130 TFLOPS的理论性能中的35-50 TFLOPS进行训练），情况并非如此。
Kurth等人在《用于气候分析的超级计算深度学习》（Exascale deep learning for climate analytics，https://dl.acm.org/doi/10.5555/3291656.3291724）一文中直接计算了他们在气候数据集分割和分类的修改版“Tiramisu”网络所需的存储带宽。每个V100 GPU需要189MB/s，6路Power9 GPU节点需要1.14GB/s，整个训练作业所需的带宽为1.16TB/s。他们在V100 GPU上的FLOP/s/sample为4.188，在训练过程中实现了每GPU 20.93 TFLOP/s的性能。