关于nas:基于文件存储UFS的Pytorch训练IO优化实践

咱们在帮助某 AI 客户排查一个 UFS 文件存储的性能 case 时发现，其应用的 Pytorch 训练 IO 性能和硬件的 IO 能力有很大的差距（前面内容有具体性能比照数据）。

让咱们感到困惑的是：UFS 文件存储，咱们应用 fio 自测能够达到单实例最低 10Gbps 带宽、IOPS 也可达到 2w 以上。该 AI 客户在高 IOPS 要求的 AI 单机小模型训练场景下，或者之前应用 MXNet、TensorFlow 框架时，IO 都能跑到 UFS 实践性能，甚至在大型分布式训练场景中，UFS 也能够齐全胜任。

于是咱们开启了和客户的一次深度联结排查。

初步尝试优化

一、调整参数：

基于上述情况，首先思考是不是应用 Pytorch 的姿态不对？参考网上提到教训，客户调整 batch_size、Dataloader 等参数。

Batch_size

默认 batch_size 为 256，依据内存和显存配置尝试更改 batch_size 大小，让一次读取数据更多，发现理论对效率没有晋升。通过剖析是因为 batch_size 设置与数据读取逻辑没有间接关系，IO 始终会保留单队列与后端交互，不会升高网络交互上的整体延时（因为用的是 UFS 文件存储，前面会讲到为什么用）。

Pytorch Dataloader

Pytorch 框架 dataloader 的 worker 负责数据的读取和加载、调配。通过 batch_sampler 将 batch 数据调配给对应的 worker，由 worker 从磁盘读取数据并加载数据到内存，dataloader 从内存中读取相应 batch 做迭代训练。这里尝试调整了 worker_num 参数为 CPU 核数或倍数，发现晋升无限，反而内存和 CPU 的开销晋升了不少，整体减轻了训练设施的累赘，通过 worker 加载数据时的网络开销并不会升高，与本地 SSD 盘差距仍然存在。

这个也不难理解，前面用 strace 排查的时候，看到 CPU 更多的时候在期待。

所以：从目前信息来看，调整 Pytorch 框架参数对性能简直没有影响。

二、尝试不同存储产品

在客户调整参数的同时，咱们也应用了三种存储做验证，来看这里是否存在性能差别、差别到底有多大。在三种存储产品上放上同样的数据集：

1. 单张均匀大小 20KB 的小图片，总量 2w 张。
2. 以目录树形式存到三种存储下的雷同门路，应用 Pytorch 罕用的规范读图接口 CV2 和 PIL

测试后果，如下图：

注：SSHFS 基于 X86 物理机（32 核 /64G/480G SSD*6 raid10）搭建，网络 25Gbps

论断：通过对存储性能实测，UFS 文件存储较本地盘、单机 SSHFS 性能差距较大。

为什么会选用这两种存储（SSHFS 和本地 SSD）做 UFS 性能比照？

以后支流存储产品的选型上分为两类：自建 SSHFS/NFS 或采纳第三方 NAS 服务（相似 UFS 产品），个别场景中也会将须要的数据下载到本地 SSD 盘做训练。传统 SSD 本地盘领有极低的 IO 延时，一个 IO 申请解决根本会在 us 级别实现，针对越小的文件，IO 性能越显著。受限于单台物理机配置，无奈扩容，数据根本“即用即弃”。而数据是否平安也只能依赖磁盘的稳定性，一旦产生故障，数据恢复难度大。然而鉴于本地盘的劣势，个别也会用作一些较小模型的训练，单次训练任务在较短时间即可实现，即便硬件故障或者数据失落导致训练中断，对业务影响通常较小。

用户通常会应用 SSD 物理机自建 SSHFS/NFS 共享文件存储，数据 IO 会通过以太网络，较本地盘网络上的开销从 us 级到 ms 级，但根本能够满足大部分业务需要。但用户须要在日常应用中同时保护硬件和软件的稳定性，并且单台物理机有存储下限，如果部署多节点或分布式文件系统也会导致更大运维精力投入。

咱们把后面论断放到一起看：

1. 隐形论断：Tensorflow、Mxnet 框架无问题。
2. 调整 Pytorch 框架参数对性能简直没有影响。

3、Pytorch+UFS 的场景下，UFS 文件存储较本地 SSD 盘、单机 SSHFS 性能差距大。

联合以上几点信息并与用户确认后的明确论断：

UFS 联合非 Pytorch 框架应用没有性能瓶颈，Pytorch 框架下用本地 SSD 盘没有性能瓶颈，用 SSHFS 性能可承受。那起因就很显著了，就是 Pytorch+UFS 文件存储这个组合存在 IO 性能问题。

深刻排查优化

看到这里，大家可能会有个疑难：是不是不必 UFS，用本地盘就解决了？

答案是不行，起因是训练所需的数据总量很大，很容易超过了单机的物理介质容量，另外也出于数据安全思考，寄存单机有失落危险，而 UFS 是三正本的分布式存储系统，并且 UFS 能够提供更弹性的 IO 性能。

依据以上的信息疾速排查 3 个论断，基本上能够判断出：Pytorch 在读 UFS 数据过程中，文件读取逻辑或者 UFS 存储 IO 耗时导致。于是咱们通过 strace 察看 Pytorch 读取数据整体流程：

通过 strace 发现，CV2 形式读取 UFS 里的文件（NFSV4 协定）有很屡次 SEEK 动作，即使是单个小文件的读取也会“分片”读取，从而导致了屡次不必要的 IO 读取动作，而最耗时的则是网络，从而导致整体耗时成倍增长。这也是合乎咱们的猜想。

简略介绍一下 NFS 协定特点：

NAS 所有的 IO 都须要通过以太网，个别局域网内延时在 1ms 以内。以 NFS 数据交互为例，通过图中能够看出，针对一次残缺的小文件 IO 操作将波及元数据查问、数据传输等至多 5 次网络交互，每次交互都会波及到 client 与 server 集群的一个 TTL，其实这样的交互逻辑会存在一个问题，当单文件越小、数量越大时则延时问题将越显著，IO 过程中有过多的工夫耗费在网络交互，这也是 NAS 类存储在小文件场景下面临的经典问题。

对于 UFS 的架构而言，为了达到更高扩展性、更便当的维护性、更高的容灾能力，采纳接入层、索引层和数据层的分层架构模式，一次 IO 申请会先通过接入层做负载平衡，client 端再拜访后端 UFS 索引层获取到具体文件信息，最初拜访数据层获取理论文件，对于 KB 级别的小文件，理论在网络上的耗时比单机版 NFS/SSHFS 会更高。

从 Pytorch 框架下两种读图接口来看：CV2 读取文件会“分片”进行，而 PIL 尽管不会“分片”读取，然而基于 UFS 分布式架构，一次 IO 会通过接入、索引、数据层，网络耗时也占比很高。咱们存储共事也理论测试过这 2 种办法的性能差别：通过 strace 发现，相比 OpenCV 的形式，PIL 的数据读取逻辑效率绝对高一些。

优化方向一：如何升高与 UFS 交互频次，从而升高整体存储网络延时

CV2：对单个文件而言，“分片读取”变“一次读取”

通过对 Pytorch 框架接口和模块的调研，如果应用 OpenCV 形式读取文件能够用 2 个办法，cv2.imread 和 cv2.imdecode。

默认个别会用 cv2.imread 形式，读取一个文件时会产生 9 次 lseek 和 11 次 read，而对于图片小文件来说屡次 lseek 和 read 是没有必要的。cv2.imdecode 能够解决这个问题，它通过一次性将数据加载进内存，后续的图片操作须要的 IO 转化为内存拜访即可。

两者的在零碎调用上的对比方下图：

咱们通过应用 cv2.imdecode 形式替换客户默认应用的 cv2.imread 形式，单个文件的总操作耗时从 12ms 降落到 6ms。然而内存无奈 cache 住过大的数据集，不具备任意规模数据集下的训练，然而整体读取性能还是晋升显著。应用 cv2 版本的 benchmark 对一个小数据集进行加载测试后的各场景耗时如下(提早的非线性降落是因为其中蕴含 GPU 计算工夫):

PIL：优化 dataloader 元数据性能，缓存文件句柄

通过 PIL 形式读取单张图片的形式，Pytorch 解决的均匀提早为 7ms(不含 IO 工夫)，单张图片读取 (含 IO 和元数据耗时) 均匀提早为 5 -6ms，此性能程度还有优化空间。

因为训练过程会进行很多个 epoch 的迭代，而每次迭代都会进行数据的读取，这部分操作从屡次训练任务上来看是反复的，如果在训练时由本地内存做一些缓存策略，对性能应该有晋升。但间接缓存数据在集群规模回升之后必定是不事实的，咱们初步只缓存各个训练文件的句柄信息，以升高元数据拜访开销。

咱们批改了 Pytorch 的 dataloader 实现，通过本地内存 cache 住训练须要应用的文件句柄，能够防止每次都尝试做 open 操作。测试后发现 1w 张图片通过 100 次迭代训练后发现，单次迭代的耗时曾经根本和本地 SSD 持平。然而当数据集过大，内存同样无奈 cache 住所有元数据，所以应用场景绝对无限，仍然不具备在大规模数据集下的训练伸缩性。

UFS server 端元数据预加载

以上 client 端的优化成果比拟显著，然而客户业务侧须要更改大量训练代码，最次要是 client 端无奈满足较大数据量的缓存，利用场景无限，咱们持续从 server 端优化，尽量升高整个链路上的交互频次。

失常 IO 申请通过负载平衡达到索引层时，会先通过索引接入 server，而后到索引数据 server。思考到训练场景具备目录拜访的空间局部性，咱们决定加强元数据预取的性能。通过客户申请的文件，引入该文件及相应目录下所有文件的元数据，并预取到索引接入 server，后续的申请将命中缓存，从而缩小与索引数据 server 的交互，在 IO 申请达到索引层的第一步即可获取到对应元数据，从而升高从索引数据 server 进行查问的开销。

通过这次优化之后，元数据操作的提早较最后可能降落一倍以上，在客户端不做更改的状况下，读取小文件性能已达到本地 SSD 盘的 50%。看来单单优化 server 端还是无奈满足预期，通过执行 Pytorch 的 benchmark 程序，咱们失去 UFS 和本地 SSD 盘在整个数据读取耗时。

此时很容易想到一个问题：非 Pytorch 框架在应用 UFS 做训练集存储时，为什么应用中没有遇到 IO 性能瓶颈？

通过调研其余框架的逻辑发现：无论是 MXNet 的 rec 文件，Caffe 的 LMDB，还是 TensorFlow 的 npy 文件，都是在训练前将大量图片小文件转化为特定的数据集格局，所以应用 UFS 在存储网络交互更少，绝对 Pytorch 间接读取目录小文件的形式，防止了大部分网络上的耗时。这个区别在优化时给了咱们很大的启发，将目录树级别小文件转化成一个特定的数据集存储，在读取数据做训练时将 IO 施展出最大性能劣势。

优化方向二：目录级内的小文件转换为数据集，最大水平降到 IO 网络耗时

基于其余训练框架数据集的共性功能，咱们 UFS 存储团队连忙动工，几天开发了针对 Pytorch 框架下的数据集转换工具，将小文件数据集转化为 UFS 大文件数据集并对各个小文件信息建设索引记录到 index 文件，通过 index 文件中索引偏移量可随机读取文件，而整个 index 文件在训练任务启动时一次性加载到本地内存，这样就将大量小文件场景下的频繁拜访元数据的开销齐全去除了，只剩下数据 IO 的开销。该工具后续也可间接利用于其余 AI 类客户的训练业务。

工具的应用很简略，只波及到两步：

• 应用 UFS 自研工具将 Pytorch 数据集以目录模式存储的小文件转化为一个大文件存储到 UFS 上，生成 date.ufs 和 index.ufs。
• 应用我方提供 Folder 类替换 pytorch 原有代码中的 torchvision.datasets.ImageFolder 数据加载模块（即替换数据集读取办法），从而应用 UFS 上的大文件进行文件的随机读取。只需更改 3 行代码即可。

20 行：新增 from my_dataloader import *

205 行：train_dataset = datasets.ImageFolder 改为 train_dataset = MyImageFolder

224 行：datasets.ImageFolder 改为 MyImageFolder

通过 github 上 Pytorch 测试 demo 对 imagenet 数据集进行 5、10、20 小时模拟训练，别离读取不同存储中的数据，具体看下 IO 对整体训练速度的影响。（数据单位：实现的 epoch 的个数）

测试条件：

GPU 服务器：P40 4 物理机，48 核 256G，数据盘 800G6 SATA SSD RAID10

SSHFS：X86 物理机 32 核 /64G，数据盘 480G*6 SATA SSD RAID10

Demo：https://github.com/pytorch/examples/tree/master/imagenet

数据集：总大小 148GB、图片文件数量 120w 以上

通过理论后果能够看出：UFS 数据集形式效率曾经达到甚至超过本地 SSD 磁盘的成果。而 UFS 数据集转化形式，客户端内存中只有大量目录构造元数据缓存，在 100TB 数据的体量下，元数据小于 10MB，能够满足任意数据规模，对于客户业务上的硬件应用无影响。

UFS 产品

针对 Pytorch 小文件训练场景，UFS 通过屡次优化，吞吐性能已失去极大晋升，并且在后续产品布局中，咱们也会联合现有 RDMA 网络、SPDK 等存储相干技术进行继续优化。具体请拜访：https://docs.ucloud.cn/storage_cdn/ufs/overview

本文作者：UCloud 解决方案架构师 马杰

欢送各位与咱们交换无关云计算的所有~~~