AI训练新趋势：基于Kubernetes的云上深度学习
作者简介
- 车漾，阿里云高级技术专家，从事Kubernetes和容器相关产品的开发。尤其关注利用云原生技术构建机器学习平台系统，是GPU共享调度的主要作者和维护者。
- 顾荣，南京大学副研究员，Alluxio项目核心开发者，研究方向大数据处理，2016年获南京大学博士学位，曾在微软亚洲研究院、英特尔、百度从事大数据系统实习研发。
背景介绍
近些年，以深度学习为代表的人工智能技术取得了飞速的发展，正落地应用于各行各业。随着深度学习的广泛应用，众多领域产生了大量强烈的高效便捷训练人工智能模型方面的需求。另外，在云计算时代，以Docker、Kubernetes以主的容器及其编排技术在应用服务自动化部署的软件开发运维浪潮中取得了长足的发展。Kubernetes社区对于GPU等加速计算设备资源的支持方兴未艾。鉴于云环境在计算成本和规模扩展方面的优势，以及容器化在高效部署和敏捷迭代方面的长处，基于“容器化弹性基础架构+云平台GPU实例”进行分布式深度学习模型训练成为了业界生成AI模型的主要趋势。
为了兼顾资源扩展的灵活性，云应用大多采用计算和存储分离的基本架构。其中，对象存储因为能够有效地降低存储成本、提升扩展弹性，经常用来存储管理海量训练数据。除了采用单一云上存储之外，很多云平台的用户因为安全合规、数据主权或者遗产架构方面的因素，大量数据还存储在私有数据中心。这些用户希望基于混合云的方式构建人工智能训练平台，利用云平台的弹性计算能力满足高速增长的AI业务模型训练方面的需求，然而这种“本地存储+云上训练”的训练模式加剧了计算存储分离架构带来的远程数据访问的性能影响。计算存储分离的基本架构虽然可以为计算资源和存储资源的配置和扩展带来更高的灵活性，但是从数据访问效率的角度来看，由于受限于网络传输带宽，用户不经调优简单使用这种架构通常会遇到模型训练性能下降的问题。
常规方案面临的数据访问挑战
目前云上深度学习模型训练的常规方案主要采用手动方式进行数据准备，具体是将数据复制并分发到云上单机高效存储（例如，NVMe SSD）或分布式高性能存储（例如，GlusterFS并行文件系统）上。这种由用户手工或者脚本完成的数据准备过程通常面临如下三个问题：
- 数据同步管理成本高： 数据的不断更新需要从底层存储定期进行数据同步，这个过程管理成本较高。
- 云存储成本开销更多： 需要为云上单机存储或高性能分布式存储支付额外费用。
- 大规模扩展更加复杂： 随着数据量增长，难以将全部数据复制到云上单机存储；即使复制到GlusterFS这样的海量并行文件系统也会花费大量的时间。
基于容器和数据编排的模型训练架构方案
针对云上深度学习训练常规方案存在的上述问题，我们设计并实现了一种基于容器和数据编排技术的模型训练架构方案。具体系统架构如图1所示：
系统架构核心组件
- Kubernetes 是一种流行的深度神经网络训练容器集群管理平台，它提供了通过容器使用不同机器学习框架的灵活性以及按需扩展的敏捷性。阿里云容器服务ACK（Alibaba Cloud Kubernetes）是阿里云提供的Kubernetes服务，可以在阿里云平台的CPU、GPU、NPU(含光800芯片)、神龙裸金属实例上运行Kubernetes工作负载。
- Kubeflow 是开源的基于Kubernetes云原生AI平台，用于开发、编排、部署和运行可扩展的便携式机器学习工作负载。Kubeflow支持两种TensorFlow框架分布式训练，分别是参数服务器模式和AllReduce模式。基于阿里云容器服务团队开发的Arena,用户可以提交这两种类型的分布式训练框架。
- Alluxio 是面向混合云环境的开源数据编排与存储系统。通过在存储系统和计算框架之间增加一层数据抽象层，提供统一的挂载命名空间、层次化缓存和多种数据访问接口，可以支持大规模数据在各种复杂环境（私有云集群、混合云、公有云）中的数据高效访问。
Alluxio 发轫于大数据时代，流觞自诞生了Apache Spark的UC Berkeley AMP实验室。Alluxio系统设计的初衷是为了解决大数据处理流水线中不同计算框架在通过磁盘文件系统（如HDFS）互换数据，造成整个分析性能瓶颈耗时在I/O操作方面的问题。Alluxio项目开源于2013年，经过7年的不断开发迭代，在大数据处理场景下的应用日趋成熟。另外，近些年随着深度学习的崛起，Alluxio分布式缓存技术正逐步成为业界解决云上I/O性能问题的主流解决方案。进一步地，Alluxio推出基于FUSE的POSIX文件系统接口，为云上AI模型训练提供了高效的数据访问手段。
为了能够更好的将Alluxio融入Kubernetes生态系统发挥两者结合的优势，Alluxio团队和阿里云容器服务团队协作开发提供了Alluxio的Helm Chart方案, 极大地简化了在Kubernetes内的部署和使用。
云上训练——Alluxio分布式缓存初探
深度学习实验环境
- 我们使用ResNet-50模型与ImageNet数据集，数据集大小144GB，数据以TFRecord格式存储，每个TFRecord大小约130MB。每个GPU的batch_size设置为256
- 模型训练硬件选择的是4台V100（高配GPU机型），一共32块GPU卡。
- 数据存储在阿里云对象存储服务中，模型训练程序通过Alluxio读取数据，并在读取过程中将数据自动缓存到Alluxio系统。Alluxio缓存层级配置为内存，每台机器提供40GB内存作为内存存储，总的分布式缓存量为160GB，没有使用预先加载策略。
初遇性能瓶颈
在性能评估中，我们发现当GPU硬件从NVidia P100升级到NVidia V100之后，单卡的计算训练速度得到了不止3倍的提升。计算性能的极大提升给数据存储访问的性能带来了压力。这也给Alluxio的I/O提出了新的挑战。
下图是在分别在合成数据(Synthetic Data)和使用Alluxio缓存的性能对比，横轴表示GPU的数量，纵轴表示每秒钟处理的图片数。合成数据指训练程序读取的数据有程序自身产生，没有I/O开销，代表模型训练性能的理论上限; 使用Alluxio缓存指训练程序读取的数据来自于Alluxio系统。在GPU数量为1和2时，使用Alluxio和合成数据对比，性能差距在可以接受的范围。但是当GPU的数量增大到4时，二者差距就比较明显了，Alluxio的处理速度已经从4981 images/second降到了3762 images/second。 而当GPU的数量达到8的时候，Alluxio上进行模型训练的性能不足合成数据的30%。而此时通过系统监控，我们观察到整个系统的计算、内存和网络都远远没有达到瓶颈。这间接说明了简单使用Alluxio难以高效支持V100单机8卡的训练场景。
为了能够深入了解是什么因素影响了性能并进行调优，需要首先研究分析Alluxio在Kubernetes下支持FUSE的整个技术栈。如下图所示
原因剖析
通过深度分析整个技术栈和Alluxio内核，我们将造成相关性能影响的原因总结如下：
1. Alluxio文件操作引入多次RPC交互，在训练场景下引入性能开销。
Alluxio不只是一个单纯的缓存服务。它首先是一个分布式虚拟文件系统，包含完整的元数据管理、块数据管理、UFS管理（UFS是底层文件系统的简称）以及健康检查机制，尤其是它的元数据管理实现比很多底层文件系统更加强大。这些功能是Alluxio的优点和特色，但也意味着使用分布式系统带来的开销。例如，在默认设置下使用Alluxio客户端来读一个文件，即便数据已经缓存在本地的Alluxio Worker中，客户端也会和Master节点有多次RPC交互来获取文件元信息以保证数据的一致性。完成整个读操作的链路额外开销在传统大数据场景下并不明显，但是深度面对学习场景下高吞吐和低延时的需求就显得捉襟见肘了。
2. Alluxio的数据缓存和驱逐策略会频繁触发节点数据缓存震荡。
深度学习场景数据冷热经常不明显，因此每个Alluxio Worker都会完整读取数据。而Alluxio默认模式会优先数据本地读取，即使数据已经保存在Alluxio集群中，也会从其他缓存节点拉取到本地存一份副本。这个特性在我们的场景下会带来两个额外开销： 1.异步数据缓存的额外开销 2.本地空间不足会触发自动数据驱逐的开销，特别当节点缓存数据接近饱和的情况下性能开销巨大。
3. 基于FUSE进行文件系统的开发、部署、使用都很简单，但是默认性能并不理想，原因如下：
- FUSE读操作效率不高，每次read最多只能读128KB，读一个128MB的文件需要1000次调用read。
- FUSE读操作属于非阻塞行为，由libfuse非阻塞线程池处理，一旦并发请求数量远超过线程池(
max_idle_threads)的大小，就会触发频繁的大量线程创建和删除，从而影响读性能。而在FUSE中，这个默认配置是10. - 元数据的频繁访问，因为FUSE内核模块是个桥梁角色，连接了应用程序和Alluxio的文件系统，而每一次读获取文件/目录的inode以及dentry，FUSE内核模块都会到Alluxio系统运行一趟，增加了系统压力。
4. Alluxio和FUSE的集成（下文简称为AlluxioFUSE）在深度学习中常见的多线程高并发场景下性能有待优化，甚至需要深度定制：
- Alluxio目前仅支持在FUSE中使用
direct_io模式，而不能使用kernel_cache模式来借助page cache进一步提高I/O效率。这是因为Alluxio当前设计要求在多线程场景下，每个线程都必须使用自己的文件输入句柄(FileInputStream)。而如果打开page cache，当前的AlluxioFUSE会有些并发预先读到cache的操作，从而产生报错。 - 数据从被Alluxio客户端读入后，到进入FUSE要经历多次拷贝。这些额外的拷贝通常是由于AlluxioFUSE使用到的第三方Java库API限制。
- AlluxioFUSE实现中使用到的第三方库JNRFuse只能适配较低版本的FUSE，并且在高并发场景下有较大的性能负担。
5. Kubernetes对于Alluxio的线程池影响。
Alluxio基于Java 1.8版本实现，其中的一些线程池的计算会依赖于Runtime.getRuntime().availableProcessors()，但是在Kubernetes环境下，默认配置中cpu_shares的值为2，而JVM对于cpu的核心数的计算公式 cpu_shares()/1024，导致结果是1。这会影响java进程在容器内的并发能力。
云上模型训练的性能优化
在分析了上述性能问题和因素之后，我们将设计了一系列性能优化策略以提升云上模型训练的性能。首先，需要明白数据访问的“多快好省”是无法全部兼顾，我们针对的主要是模型训练下只读数据集的数据访问加速。优化的基本思路是关注高性能和数据一致性，而牺牲一部分灵活的自适应性（比如读写同时发生，数据内容不断更新等场景）。
基于上述思路，我们设计了具体的性能优化策略，这些策略遵循以下核心原则：
- 寻找资源限制，包括线程池以及JVM在容器中的配置
- 借助各级缓存，包括FUSE层和Alluxio元数据缓存
- 避免额外开销，减少非必须的调用链路。比如避免不必要的元数据交互，引入上下文切换的GC线程和compiler进程；以及Alluxio内部的一些可以简化的操作
下面将从各层的组件优化角度，对这些优化策略逐一介绍：
对FUSE的优化
升级Linux Kernel版本
FUSE实现分为两层：运行在用户态的libfuse和运行在内核态的FUSE Kernel。高版本的Linux Kernel针对FUSE做了大量的优化。我们对比了Kernel 3.10和4.19的性能，发现读性能可以达到20%的提升。
优化FUSE参数
- 延长FUSE元数据有效时间
Linux中每个打开文件在内核中拥有两种元数据信息：struct dentry和struct inode，它们是文件在内核的基础。所有对文件的操作，都需要先获取文件这两个结构。所以，每次获取文件/目录的inode以及dentry时，FUSE内核模块都会从libfuse以及Alluxio文件系统进行完整操作，这样会带来数据访问的高延时和高并发下对于Alluxio Master的巨大压力。可以通过配置–o entry_timeout=T –o attr_timeout=T进行优化。
2. 配置max_idle_threads避免频繁线程创建销毁引入CPU开销。
这是由于FUSE在多线程模式下，以一个线程开始运行。当有两个以上的可用请求，则 FUSE 会自动生成其他线程。每个线程一次处理一个请求。处理完请求后，每个线程检查目前是否有超过max_idle_threads (默认10)个线程；如果有，则该线程回收。而这个配置实际上要和用户进程生成的 I/O 活跃数相关，可以配置成用户读线程的数量。而不幸的是 max_idle_threads本身只在libfuse3才支持，而AlluxioFUSE只支持libfuse2， 因此我们修改了libfuse2的代码支持了max_idle_threads的配置。
对Alluxio的优化
Alluxio和FUSE的集成通过一个名为AlluxioFuse的进程实现。该进程在运行期会通过调用内嵌的Alluxio客户端和运行的Alluxio Master以及Worker交互。我们针对深度学习的场景，定制AlluxioFuse所使用的Alluxio属性来优化性能。
避免频繁逐出（Cache Eviction）造成缓存抖动
由于深度学习训练场景下，每次训练迭代都是全量数据集的迭代，缓存几个TB的数据集对于任何一个节点的存储空间来说都是捉襟见肘。而Alluxio的默认缓存策略是为大数据处理场景（例如，查询）下的冷热数据分明的需求设计的，数据缓存会保存在Alluxio客户端所在的本地节点，用来保证下次读取的性能最优。具体来说
3. alluxio.user.ufs.block.read.location.policy默认值为alluxio.client.block.policy.LocalFirstPolicy， 这表示Alluxio会不断将数据保存到Alluxio客户端所在的本地节点，就会引发其缓存数据接近饱和时，该节点的缓存一直处于抖动状态，引发吞吐和延时极大的下降，同时对于Master节点的压力也非常大。因此需要 location.policy 设置为 alluxio.client.block.policy.LocalFirstAvoidEvictionPolicy 的同时，指定 alluxio.user.block.avoid.eviction.policy.reserved.size.bytes 参数，这个参数决定了当本地节点的缓存数据量达到一定的程度后，预留一些数据量来保证本地缓存不会被驱逐。通常这个参数应该要大于 节点缓存上限 X （100% - 节点驱逐上限的百分比） 。
4. alluxio.user.file.passive.cache.enabled设置是否在Alluxi的本地节点中缓存额外的数据副本。这个属性是默认开启的。因此，在Alluxio客户端请求数据时，它所在的节点会缓存已经在其他Worker节点上存在的数据。可以将该属性设为false，避免不必要的本地缓存。
5. alluxio.user.file.readtype.default默认值为CACHE_PROMOTE。这个配置会有两个潜在问题，首先是可能引发数据在同一个节点不同缓存层次之间的不断移动，其次是对数据块的大多数操作都需要加锁，而Alluxio源代码中加锁操作的实现不少地方还比较重量级，大量的加锁和解锁操作在并发较高时会带来不小的开销，即便数据没有迁移还是会引入额外开销。因此可以将其设置为CACHE以避免moveBlock操作带来的加锁开销，替换默认的CACHE_PROMOTE
缓存元数据和节点列表
在深度学习训练场景下，每次训练任务开始前会列出所有训练数据文件并读取其元数据，然后运行训练任务的进程会进一步读取训练数据文件。通过Alluxio读取文件访问时默认会完成如下操作：首先从Master获取文件元数据，从中获取block元数据，再从Worker获取block的具体位置，最后真正从获取的位置读取block数据。完成完整的操作链路包括多次RPC开销，引入明显的文件访问延时。如果能将该数据文件的block信息缓存到客户端内存中，会非常明显的提升文件的访问性能。
6. 将alluxio.user.metadata.cache.enabled设置为true, 可以在Alluxio客户端开启文件以及目录的元数据缓存，避免二次访问时仍需要通过RPC访问元数据的问题。结合分配给AlluxioFUSE的堆大小，用户可以配置alluxio.user.metadata.cache.max.size来设置最多缓存文件和目录的元数据数量，也可以配置alluxio.user.metadata.cache.expiration.time调整元数据缓存的有效时间。
同时在每次选择读取数据的Worker节点时，Alluxio Master节点也会不断去查询所有Worker节点的状态，这也会在高并发场景下引入额外开销。
7. 将alluxio.user.worker.list.refresh.interval设置为2min或者更长。
8. 读取文件也会不断更新last accesstime，实际上在高并发的场景下，这会对Alluxio Master造成很大压力。我们通过修改Alluxio代码增加了开关，可以关闭掉last accesstime的更新。
充分利用数据本地性
9. 数据本地性就是尽量将计算移到数据所在的节点上进行，避免数据在网络上的传输。分布式并行计算环境下，数据的本地性非常重要。在容器环境下支持两种短路读写方式：Unix socket方式和直接文件访问方式。
- Unix Socket的方式好处在于隔离性好，不需要Alluxio Client和Alluxio Worker容器运行在同样的Network，UTS，Mount的Namespace。但是它的性能比直接文件访问要差一些，同时会引发netty的OutOfDirectMemoryError
- 而直接访问文件的方式则所以需要确保同一台机器上运行的Alluxio Worker和AlluxioFUSE的主机名和IP地址一致，同时要保证Alluxio Client和Worker共享同样缓存目录，这种方式性能更好同时更加稳定。但是它实际上牺牲了隔离性，需要二者共享Network，UTS，Mount的Namespace
我们目前选择的方案是优先采用后者。
对Java & Kubernetes的优化
配置ActiveProcessorCount
10. Runtime.getRuntime().availableProcessors()控制的；而如果通过Kubernetes部署容器而不指定cpu资源的request数量，容器内Java进程读到proc文件系统下的cpushare数量为2， 而此时的availableProcessors()来自于cpu_shares()/1024，会被算成1。实际上限制了容器内Alluxio的并发线程数。考虑到Alluxio Client属于I/O密集型的应用，因此可以通过-XX:ActiveProcessorCount设置处理器数目。这里的基本原则是 ActiveProcessorCount尽量设置得高些。
调整GC，JIT线程
11. JVM的缺省GC， JIT编译线程数量取决于-XX:ActiveProcessorCount的数量，但实际上也可以通过-XX:ParallelGCThreads -XX:ConcGCThreads -XX:CICompilerCount等参数配置，可以将其设置的小些，避免这些进程频繁的抢占切换，导致性能下降。
性能优化效果
在优化Alluxio之后，ResNet50的训练性能单机八卡性能提升了236.1%，并且扩展性问题得到了解决，训练速度在不但可以扩展到了四机八卡，而且在此场景下和合成数据相比性能损失为3.29%(31068.8 image/s vs 30044.8 image/s)。而实际训练时间方面，四机八卡在合成数据场景下需要63分钟，而使用Alluxio需要65分钟。
总结与进一步工作
在本文中，我们总结了Alluxio在高性能分布式深度学习模型训练场景中落地的挑战点，以及我们在优化Alluxio的实践。进一步地，我们介绍了如何从多个层面提升AlluxioFUSE在高并发读场景下性能优化的经验。最后，我们实现的基于Alluxio优化的分布式模型训练方案，并在4机8卡的ResNet50场景下进行了性能验证，取得了很好的效果。
在进一步工作方面，对于高吞吐海量规模的小文件和高并发读场景，Alluxio还有一些在page cache的支持和FUSE层的稳定性方面的工作，我们阿里云容器服务团队也会和Alluxio开源社区以及南京大学戴海鹏、顾荣等老师一起继续合作努力改进。我们相信通过工业界、开源社区和学术界和联合的创新力量，能够逐步降低计算存储分离场景下深度学习训练的数据访问高成本和复杂度，进一步助力云上普惠AI模型训练。
致谢
感谢Alluxio团队的范斌，邱璐，Calvin Jia，常铖在整个方案的设计和优化过程中的巨大帮助，从Alluxio自身能力上对于元数据缓存系统做了显著的提升，为Alluxio落地AI场景开启了可能性。