[CLUG] 中国人民大学鲁蔚征|深度学习负载IO模式和存储选择
发布时间:
2022-04-06 00:00
在2021年12月举行的中国Lustre用户峰会(China Lustre User Group,CLUG2021)上,来自中国人民大学校级计算平台技术负责人鲁蔚征老师分享了《深度学习负载IO模式和优化策略》,他不仅介绍了深度学习的基本概念,分析了训练任务的IO模式,还对存储选择提出了建议并呼吁超算中心和高校转变思维模式,在做超算的规划、建设以及运维的时候,务必要规划好 IO和存储需求。鲁老师毕业于北京大学,曾在小米大数据部等多家互联网头部公司工作,负责过多款单日活跃用户数级别为千万的APP,擅长大数据分析和机器学习,精通Flink、Spark、PyTorch等工具,具有丰富的工业界系统开发和运维实战经验。以下为鲁老师演讲的精彩摘要。

1深度学习基本概念
深度学习主要分为两部分,一部分是训练(即在数据集上训练一个模型),另一部分是推理(即在这个模型上对未知数据进行预测)。绝大多数的训练任务是在HPC场景下运行的。一个典型的训练过程首先是构建一个神经网络(由多种计算组成,比如卷积、全连接等),然后在训练数据集上找一些数据,经过前向传播,得到一个输出,然后使用损失函数与标注的数据进行对比,得到损失值。反向的过程就是根据损失值去更新模型中的参数。
与IO密切相关的一个重要概念是随机梯度下降算法(SGD算法)。由于处理器的内存有限,一次只能处理一个批大小(batch_size)的样本数据。但是,一个训练集可能很大(有成千上万的样本)。如何把成千上万的样本都迭代(iteration)完,就需要把将整个训练样本分成若干个批(batch),一次迭代(或一次step或一个batch)处理一个批大小的样本。把整个训练集的所有样本都处理完一次就叫一个Epoch,在一个Epoch中每个样本会被访问一次。还有一个重要的概念是shuffle,在一个Epoch中,我们需要进行data shuffle的过程。因为在随机梯度下降的过程中,我们希望数据不要总是保持同样的模式,而是经过一个随机化使模型学到更强的能力,不至于每次学到的都是固定的模式,从而让模型具有更强的泛化能力。
2深度学习应用
主流的深度学习应用分几大类,常见的有计算机视觉,自然语言处理,搜索/广告/推荐等。对于不同的应用来说,由于数据源不一样,因此IO也不太一样。计算机视觉应用主要是一些图片和视频段,自然语言处理应用绝大多数都是文本数据,对于搜索广告推荐来说,都是一些保存在企业数据仓库里的结构化数据。
如果我们原生地读写JPEG图片,对数据的IO有很大的压力,尤其是对10k左右的小图片并发地、高性能地读写的话,会有很大的IO压力。文本数据和图片数据在数据大小上有一定的区别,文本就是一些字符串,字符串的大小比图片视频小很多。启动深度学习训练任务前,需要先对文本数据进行预处理,处理后的数据以大文件形式保存的居多。
3 TensorFlow与PyTorch比较
深度学习离不开深度学习框架。深度学习框架主要分两大派系,一派是以TensorFlow为代表的静态图,另一派是PyTorch为代表的动态图。对于静态图,数据的读写与IO是静态图的一部分。对于以PyTorch为代表的动态图,数据的IO也很依赖Python社区提供的包。我们来具体比较一下TensorFlow和PyTorch。
对于数据源来说,首先TensorFlow需要进行多步的预处理,生成它预定义的TFRecord这种数据格式。TensorFlow的优势在于数据读取过程都是用C++写好的,集成到TensorFlow的一个静态图中的,数据读取已经在多线程、流水线等方面进行了优化,因此TensorFlow的读取有极大的性能优势。但它的缺点是它不方便调试,因为它基于定义好的TFRecord数据格式,我们看不到或者说不容易看到数据的真实值,当调试的时候就很麻烦。另外,它做了一个小文件打包成大文件的过程。由于小文件打包成大文件,对于随机化时的Shuffle阶段,它就不是绝对的随机化。

对于PyTorch来说,它没有自己定义的数据格式,只要是Python可读写文件系统上的数据都可以读出来。它是一种绝对的Random Access的过程,尤其在小文件的场景,PyTorch对存储和文件系统都有很大的读写压力,但它的优势是很方便调试,我们可以直接打印读出来的样本数据,所以一些科研机构需要快速进行迭代,快速实验的时候,很喜欢用PyTorch框架。而一些企业喜欢用TensorFlow框架,因为它在性能上有一定的优势。
4训练任务的IO的模式
HPC社区有一个工具叫做Darshan,它可以对IO进行一些分析。针对在TensorFlow上训练ImageNet的数据集,我们分析了它IO的模式。在下图(a)可以看到发生了哪些POSIX的读写等操作。用TensorFlow的话,主要是一些Read和Open。在训练任务中,绝大多数都是读训练样本数据,几乎没有写的过程,少量的写也只是生成一些log,生成Checkpoint。

上图(b)是细粒度去看每一次 POSIX的读写所访问的文件大小,横轴是文件大小,我们可以看到特别多的是0~100b左右的读操作。0~100b主要是一些元数据(Metadata)的读操作,所以这种小文件读写会对我们性能带来很大的瓶颈,尤其对Metadata读造成很大的性能瓶颈。小文件IO对于文件系统来说有很多次步骤,比如进行元数据的查询,比如一个并行的应用文件系统还要多次在网络上进行传输,所以当小文件IO以及元数据查询次数越多时,对文件系统的压力越大。

关于小文件IO的横向扩展性,以TensorFlow为例,当我们进行一个分布式的训练,我们可以看到在上方的柱状图中:蓝色部分是元数据读的过程,黄色部分是文件本身读,黄色部分的占比很小,几乎都看不到了。这说明在IO过程中,绝大多数都是读元数据,而文件本身的读,因为TensorFlow经过了高度的优化,读的过程已经很优化了,但是获取元数据这个过程还是有一定的压力。因为并行文件系统上元数据有一个专门的Metadata server,所以说元数据占用了大量的时间,成为一个瓶颈。上图横轴是不同的训练节点,分布式训练节点数越多,元数据的占比越大,而读文件本身一直都保持很小的占比,所以说元数据的读制约了吞吐量,也就是我们带宽的线性增长。
5存储的选择
根据前面的分析和数据,我们选择存储时需要考虑多个因素。
介质的选择 现在大家都选择闪存,甚至是内存,小文件读写的性能肯定是内存或者闪存NVMe SSD等比机械硬盘(HDD)的性能高很多。如果有深度学习的小文件场景,那么机械硬盘可能已经不太适合这样的场景了。我自己也在ImageNet上测试过端到端的训练,本地的SSD会比本地机械硬盘有1.2 ~ 2倍甚至更高的加速。
并行文件系统 并行文件系统一般是由Metadata和 Object Storage两部分组成的,在传统的HPC场景中,Object Storage是可以不断去横向扩展的,它可以带来很高的聚合带宽。但是,无论是单机还是分布式训练,它的瓶颈在于 Metadata上,所以说单纯去考虑聚合带宽,在深度学习训练场景上意义已经不太大了,一定要考虑好小文件的Metadata访问,比如我们在用并行文件系统时,Metadata这部分必须放在全闪上,或者说内存配的要更多一些。
6 思维模式的转变
过去,国内很多的超算中心和高校都很重视算力的提高,却相当轻视IO这块。从刚才的分析可以看到,IO其实是很影响 AI训练的。如果你的存储没有选好,可能我们的训练任务还不如用本地的SSD性能好。所以,我们在做超算的规划、建设以及运维的时候,务必要规划好 IO和存储需求,这是一种思维模式的转变。
7 优化IO的方向
优化方向1:改变深度学习框架IO模式
我们刚才分析了PyTorch框架有性能方面的问题,所以简单直观的方法就是把一些小文件打包成大文件,因为我们知道大文件的并行文件系统上带宽会很高,比如说webdataset就是一个PyTorch扩展库,它可以将小文件打包成一个大文件,然后集成到PyTorch的训练的一个任务里。
优化方向2:尽可能地利用本地的SSD
我们很多并行文件系统已经做了一些探索,比如说Lustre的PCC(Persistent Client Caching)*,它就把计算节点的 SSD尽可能地利用起来。因为深度学习训练绝大多数都是读,所以说不用担心写的问题,也不用担心一致性的问题,这就把本地的SSD当成一个高速的缓存就好了。
*备注:PCC功能在商业版Lustre文件系统(DDN EXAScaler6)中称为Hot Nodes功能。
优化方向3:内存缓存层
在内存缓存层方面已经有很多工作,但据我了解,截止到现在,还没有一个很可靠易用的开源软件。有一些大公司,比如微软做了一些公司内部的缓存层,发了一些论文。有一些科研机构也做了一些 IO的缓存层,但是科研机构的维护力量肯定不如企业的维护力量强,所以说它的开源软件的可用性还有待提高。
相关新闻