海量小文件场景下训练加速优化之路

作者:
辰算力平

1. 背景

随着大数据、人工智能技术的蓬勃发展,人类对于算力资源的需求也迎来大幅度的增长。在腾讯内部,星辰算力平台以降本增效为目标,整合了公司的GPU训练卡资源,为算法工程师们提供统一的底层GPU算力服务。借助于虚拟化、算力挖掘等技术,平台服务公司内各BG的AI训练场景,GPU利用率业界领先。同时,通过云原生任务化的方式,对接了内部各大业务,促进了AI技术研究效率的提升和创新研究。
当下,由于AI训练时的高性能计算设备(如NVIDIA GPU)成本高昂,如果任务在训练过程中不能保证数据IO的速度,将会导致计算设备低载甚至空载,这无疑在时间和资源上都是一种极大的浪费。

在星辰算力平台内部,用户的训练数据大多存放在平台提供的 CephFS
中,训练时将对应的CephFS目录挂载至容器内部,从而使用户在训练时能够像使用本地文件系统一样使用CephFS。但在平台运营过程中我们发现,在 训练数据集文件数较多时,训练任务使用CephFS会使训练速度变得异常缓慢
。基于这个普遍存在的问题,本文剖析其产生的原理,然后介绍相应的优化方案。最后,通过延伸思考来发散思维,简要介绍了不同场景下AI训练加速的技术。

2. 基本概念

2.1. CephFS IO流程

CephFS IO流程如下图所示。

CephFS IO路径

当客户端进行文件系统调用时(如 open
read
readdir
等),需要先从元数据服务器( Metadata Server
, MDS
)中获取请求文件的元数据信息,元数据信息主要包括文件的 Inode
号、权限、 uid
gid
和访问更改时间等。为了加快元数据的访问效率, MDS将大部分热点元数据都缓存在自己的内存中
,从而避免低效地通过访问 RADOS
Reliable, Autonomic Distributed Object Store
)层来获取元数据。客户端在从MDS中获取元数据后,通过计算的方式( CRUSH
算法)得到数据在 RADOS
中的位置,最后与远程的存储设备进行交互。

从这个架构来看,CephFS是一个 元数据和用户数据分离的文件系统
。文件的元数据和数据存储在 RADOS
中的不同 Pool
中,客户端需要先与 MDS
进行元数据交互,再与 RADOS
进行数据交互。

2.2. Ceph-FUSE

Ceph-FUSE
是CephFS客户端的一种形式,通过用户空间文件系统( Filesystem in Userspace
, FUSE
)的方式来实现CephFS客户端的功能。 FUSE
是一个面向类Unix计算机操作系统的软件接口,它使无特权的用户能够无需编辑内核代码而创建自己的文件系统。目前Linux通过内核模块对此进行支持。通过这种方式,我们可以编写用户态的应用程序,只需要实现Linux定义的一组文件系统接口,即可在用户态实现一个完整的文件系统。
当用户需要与CephFS进行交互时,客户端的整个IO流程如下:

  1. 用户程序通过 syscall
    glibc
    库进行系统调用
  2. 进程陷入内核态,文件系统操作请求到达Linux虚拟文件系统( Virtual Filesystem
    , VFS
  3. VFS
    Dentry Cache
    Inode Cache
    Page Cache
    dentry
    inode
    
  4. 若缓存不命中,则将请求转发至 FUSE Driver
  5. Ceph-FUSE进程通过 libfuse
    监听到来自于 /dev/fuse
    的请求,与Ceph集群进行交互并返回结果。
Ceph-FUSE IO路径

当用户态程序发起 FUSE
请求时, Ceph-FUSE
在经过处理后会将元数据信息缓存在内存中,提升后续访问的性能。同时,Linux的 Dentry Cache
Inode Cache
Page Cache
也会分别缓存该文件的 dentry
inode
和页,提升热点数据的读取性能。

3. 问题

3.1. 问题源起

星辰算力平台服务了公司内部各个BG和部门的AI算法工程师,因此平台上运行的训练任务场景也各不相同。在运营过程中我们发现,有用户反映某些任务中CephFS的读取速度较慢,使整个训练的时间拉长,其中属CV类的任务较为明显。
平台上CV类的任务数据集,一般都是海量的图片文件。这类数据集的特点是:

  • 文件个数多,小数据集达到十万级别,大数据集达到百万、千万甚至上亿级别。
  • 单个文件占用空间不大,大多是小文件。

3.2. 理论分析

AI训练场景与许多复杂的文件操作场景不同,其数据读写的逻辑较为简单。一般来说,用户会在每个epoch训练 相同的数据
,然后训练多个epoch直至模型达到收敛条件。因此,AI训练场景下, 训练文件在训练过程中保持不变,且被读取的频率相对固定,同时写文件的频率较低

针对这种特点,由于 Ceph-FUSE
会对访问过的元数据进行缓存,同时Linux的 Dentry Cache
Inode Cache
Page Cache
也会充分缓存读取过的文件元数据和文件数据。通常来说,在第二个epoch开始时,由于数据集文件在第一个epoch已被访问过,训练时的IO速度应当有非常明显的提升。然而,事与愿违,对于较多数量的文件,我们发现训练速度没有明显提升,且每个epoch的训练速度都很慢。

为了查出其中的原因,接下来我们复制一个一模一样的任务,打开 Ceph-FUSE
日志进行分析。

3.3. 原因排查

3.3.1. Ceph-FUSE日志分析

在训练任务开始时,打开母机上的 Ceph-FUSE
日志进行查看。
疑点现象:

  1. 在第一个epoch接近末尾时,发现出现了日志 trim_caps mds.x max xxx caps xxx
  2. 每次 trim_caps
    执行,清除的dentry个数为5000。
  3. 该日志每隔5s会打印一次,往后的训练过程中会一直持续。

注: CAPS
是指 capabilities
MDS
CAPS
授予客户端对不同文件进行操作的许可,因此 MDS
需要实时维护每个客户端文件操作的 CAPS
。这就意味着,如果客户端端持有了某个文件的 CAPS
并进行了缓存, MDS
需要知道每个客户端缓存了哪些文件。

3.3.2. 提出猜想

根据疑点现象大概能够提出以下的猜想:

  1. 在第一个epoch结束时发生了 trim_caps
    现象,且多次测试结果均是如此,猜测可能是缓存数量到达了某个阈值。
  2. 日志每隔5s会打印一次,可能是定时器触发了 trim_caps
  3. MDS
    需要维护每个客户端的 CAPS
    ,当客户端读取文件数较多时, MDS
    的cache总会达到 oversize
    的状态,必定会触发 trim_caps

3.3.3. 代码验证

根据上述猜想,可以在茫茫的Ceph源码中直奔主题,分别找出 MDS
Ceph-FUSE
的关键代码。

3.3.3.1. MDS端

根据现象2,在 MDS
中的 tick
函数内找到如下代码:

void MDSRankDispatcher::tick()
{
......
if (is_active() || is_stopping()) {
server->recall_client_state(nullptr, Server::RecallFlags::ENFORCE_MAX); // 选中该MDS下持有较多caps数量的客户端,执行caps回收
mdcache->trim();
mdcache->trim_client_leases();
mdcache->check_memory_usage(); // 当内存使用量过大时,选中该MDS下所有客户端,执行caps回收(recall_client_state)
mdlog->trim();
}
......
}

从中可以看出, MDS
端定时对客户端的 CAPS
进行回收,如果回收后内存使用量仍然过高,就对所有客户端再执行一次 CAPS
回收。在 check_memory_usage
函数中会根据cache试用情况决定是否再执行 recall_client_state

void MDCache::check_memory_usage()
{
......
if (cache_toofull()) {
mds->server->recall_client_state(nullptr);
}
......
}

进入关键函数 recall_client_state
进行查看。

/**
* Call this when the MDCache is oversized, to send requests to the clients
* to trim some caps, and consequently unpin some inodes in the MDCache so
* that it can trim too.
*/

std::pair<bool, uint64_t> Server::recall_client_state(MDSGatherBuilder* gather, RecallFlags flags)
{
......
const bool enforce_max = flags&RecallFlags::ENFORCE_MAX;
const auto max_caps_per_client = g_conf->get_val<uint64_t>("mds_max_caps_per_client"); // 默认为1_M
const auto min_caps_per_client = g_conf->get_val<uint64_t>("mds_min_caps_per_client"); // 默认为100
const auto recall_max_caps = g_conf->get_val<uint64_t>("mds_recall_max_caps"); // 默认为5000
......
/* trim caps of sessions with the most caps first */
std::multimap<uint64_t, Session*> caps_session;
auto f = [&caps_session, enforce_max, max_caps_per_client](Session* s) {
auto num_caps = s->caps.size(); // 当前caps总量
// 当flags为RecallFlags::ENFORCE_MAX时,只把caps数量超过max_caps_per_client的客户端找出来,否则找出所有客户端
if (!enforce_max || num_caps > max_caps_per_client) {
caps_session.emplace(std::piecewise_construct, std::forward_as_tuple(num_caps), std::forward_as_tuple(s));
}
};
mds->sessionmap.get_client_sessions(std::move(f));
......
for (const auto p : boost::adaptors::reverse(caps_session)) {
......
// 计算每个客户端的最大caps数量
uint64_t newlim;
if (num_caps < recall_max_caps || (num_caps-recall_max_caps) < min_caps_per_client) {
newlim = min_caps_per_client;
} else {
newlim = num_caps-recall_max_caps;
}
if (num_caps > newlim) {
/* now limit the number of caps we recall at a time to prevent overloading ourselves */
uint64_t recall = std::min<uint64_t>(recall_max_caps, num_caps-newlim); // 这里可以看出,每次最多回收mds_recall_max_caps个
newlim = num_caps-recall;
......
auto m = new MClientSession(CEPH_SESSION_RECALL_STATE); // 新建一个类型为CEPH_SESSION_RECALL_STATE的请求
m->head.max_caps = newlim; // 设置客户端的最大caps数量
mds->send_message_client(m, session); // 向客户端发送请求
......
}
......
}
......
}

从上述代码基本可以确定 CAPS
被清除的原因, MDS
每隔5s执行了一次 recall_client_state
。由于 mds_max_caps_per_client
默认被设置为 1_M
(也就是 1048576
),当训练程序读取文件个数达到 1_M
后该客户端就会被加入 caps_session
队列发起 CAPS
回收请求。由于 recall_max_caps
默认被设置为 5000
,所以每次 CAPS
回收的个数为 5000

3.3.3.2. Ceph-FUSE端

首先,根据 MDS
端发起的类型为 CEPH_SESSION_RECALL_STATE
的请求,找到客户端接受请求的代码。

void Client::handle_client_session(MClientSession *m) 
{
......
switch (m->get_op()) {
......
case CEPH_SESSION_RECALL_STATE:
trim_caps(session, m->get_max_caps()); // max_caps,值为上述的newlim
break;
......
}
......
}

Ceph-FUSE
接收到 MDS
的请求后,进入 trim_caps
函数。

void Client::trim_caps(MetaSession *s, uint64_t max)
{
mds_rank_t mds = s->mds_num;
size_t caps_size = s->caps.size(); // 客户端caps总量
......
uint64_t trimmed = 0;
auto p = s->caps.begin();
std::set to_trim; // 将需要执行caps回收的Dentry放入其中等待回收

// 以下内容通过迭代器p将caps清理至max以下,将需要清理的Dentry放入to_trim中
while ((caps_size - trimmed) > max && !p.end()) {
......
}

for (const auto &dn : to_trim) {
trim_dentry(dn); // 执行Ceph-FUSE内的dentry缓存
}
to_trim.clear();

caps_size = s->caps.size();
if (caps_size > max)
_invalidate_kernel_dcache(); // 这是关键函数,调用了Linux的remount操作来清理所有的dentries

Ceph-FUSE
接收到 MDS
的请求后,会将 CAPS
总量清理至 max
以下(本例中就是清理 5000
CAPS
)。同时,将这些 CAPS
对应的 dentry
缓存全部清除,并调用操作系统命令来清除 Dentry Cache
Inode Cache
Page Cache
,执行命令为:

static int remount_cb(void *handle)
{
// used for trimming kernel dcache. when remounting a file system, linux kernel
// trims all unused dentries in the file system
char cmd[1024];
CephFuse::Handle *cfuse = (CephFuse::Handle *)handle;
snprintf(cmd, sizeof(cmd), "mount -i -o remount %s", cfuse->opts.mountpoint); // 调用remount,清理文件系统的缓存
int r = system(cmd);
......
}

3.4. 小结

至此,基本真相大白。整体流程如下图所示:

  1. 训练程序启动,开始读取文件。
  2. Ceph-FUSE
    CAPS
    1_M
    
  3. CAPS
    1_M
    CAPS
    5000
    
  4. Ceph-FUSE
    CEPH_SESSION_RECALL_STATE
    5000
    CAPS
    CAPS
    dentry
    
  5. Ceph-FUSE
    调用Linux的 remount
    命令来清除Linux文件系统的cache。
  6. MDS
    检查自身内存使用情况,若超过阈值则重复上述回收操作。
  7. 训练程序第二个epoch后,由于文件系统的cache被清除,导致缓存失效。
CAPS回收流程

4. 解决方案

从上述分析来看,最直观的改进方法就是将 MDS
端的参数 mds_max_caps_per_client
增大,可以使得 MDS
能够维护更多的 CAPS
。然而,这是一种治标不治本的方法。接下来提出一种 Ceph-FUSE
客户端缓存的方案,避免客户端 CAPS
清除导致训练速度变慢。

4.1. 元数据缓存方案

4.1.1. 元数据缓存

Ceph针对的是通用场景,设计复杂的 CAPS
机制来保证多客户端对同一文件读写时的一致性。但在我们的场景中,读写方式却较为固定。主要表现为:

  1. 训练过程中读取的数据集在训练过程中不会发生改变,且读取频率很高。
  2. 写文件的频率较低,主要是 ckpt
    log
    文件,且不会读。

在这个特殊的场景下,可以部分牺牲一致性来获取性能上的提升。具体表现为, Ceph-FUSE
侧可以将以只读方式打开的文件进行元数据缓存,减少与 MDS
的交互,同时在 trim_caps
发生时不去真正删除这部分元数据对应的缓存。核心改造如下所示:

  1. Ceph-FUSE
    open
    I_CACHED
    MDS
    
  2. 如果一个文件被只读打开后,将无法被读写打开,这是为了保证写数据的一致性。
  3. trim_caps
    Ceph-FUSE
    CAPS
    Inode
    I_ORPHAN
    MDS
    CAPS
    MDS
    Inode
    Ceph-FUSE
    CAPS
    
元数据缓存方案

以上优化建立的前提是:只读方式打开的文件不会进行修改。在我们的AI训练场景下,训练任务完美契合了这个条件。

4.1.2. 缓存淘汰算法

Ceph-FUSE会将元数据缓存在本地,但其缓存淘汰算法是一种带高低优先级的 LRU
算法。 LRU
算法核心思想是 如果数据最近被访问过,那么将来被访问的几率也更高
,但这种思想不符合AI训练的场景。在大多任务训练过程中,训练数据文件会被均匀地访问,每一个epoch中被访问过的文件反而是这个epoch中不会再被读取的文件。采用 LRU
算法会使 缓存队列中即将被用到的文件元数据被删除
,如下图所示。

LRU淘汰方式

下图模拟了 LRU
淘汰策略下训练数据集命中率分布曲线。

LRU淘汰策略下训练数据集命中率分布

从该图中可以看出, LRU
淘汰策略下缓存队列长度越接近数据集大小,命中率提升才越明显。 当队列长度只有数据集大小的一半时,命中率只有15%左右

在AI训练的场景下,采用不替换策略( Not Replacement
, NR
)将是命中率最高的算法。在训练的第一个epoch时, Ceph-FUSE
将元数据放到缓存中。当缓存队列已满时, Ceph-FUSE
将不替换现有缓存的数据,保持缓存不变。在第二个epoch时, Ceph-FUSE
从缓存队列中读取文件元数据,若未命中则请求 MDS
获取。

NR算法

4.1.3. 优化结果

结合两点针对 Ceph-FUSE
的优化改动,我们对示例任务进行了测试,得到如下的性能测试数据。

训练任务测试结果

从图中可以看出,经过优化后针对海量小文件训练场景,训练速度的提升非常明显。在第二个epoch后,元数据缓存优化版本的 训练速度提升为原来的3~4倍
,且训练速度较为稳定。相比于之前的版本,经过优化后的 Ceph-FUSE
能够充分利用Linux文件系统的cache,且避免了每个epoch与 MDS
之间的交互。经过优化后的版本训练速度能与本地SSD较为贴近。

4.2. 文件缓存方案

文件缓存方案实际上是一种在元数据缓存优化的基础上,利用本地SSD对文件进行缓存的方案。针对文件数量特别多,利用Linux文件系统cache但是内存不充足的情况,该方法会有一定效果。

训练程序在第一个epoch训练时, Ceph-FUSE
在处理完 read
请求后将文件写入本地SSD中。为了避免海量小文件直接写入本地造成较多的 lookup
操作,同时也为了避免任务完成后文件缓存难以进行清理的问题,考虑将所有读取后的文件进行聚合缓存至一个本地Cache大文件中,由 Ceph-FUSE
来记录每个文件在本地Cache文件中的偏移。
文件缓存方案的详细步骤如下所示:

  • 文件缓存命中:
    • Metadata Cache
      中找出文件在本地Cache文件中的偏移。
    • 通过 pread
      从本地SSD缓存文件中读取指定范围的字节。
  • 文件缓存不命中:
    • 按照正常流程,与Ceph集群进行交互,得到读取的字节流。
    • 写本地Cache文件,并记录该文件在其中的偏移。
    • 更新 Metadata Cache
      ,将文件元数据和偏移量加入其中。
文件缓存方案

该方案虽然能够充分利用本地SSD,但也有一些缺点,具体表现为:

  1. 由于第一个epoch读取文件时, Ceph-FUSE
    会写本地Cache文件,可能会使得第一个epoch训练速度变慢。但当epoch数较多时这部分时间牺牲是值得的。
  2. IO路径变得更长, Ceph-FUSE
    需要读本地文件。

4.3. 方案对比

方案 适用场景
原版 在训练过程中需要修改数据集
元数据缓存 在训练过程中不修改只读打开的文件
元数据缓存 + 文件缓存 内存紧张,无法充分使用文件系统缓存

5. 延伸方案

上述分析和方案主要针对的是海量小文件的IO密集型计算场景,接下来发散思维,简要介绍一下多种AI加速的解决方案。
我们将AI训练任务分为IO密集型、GPU计算密集型和CPU计算密集型三类任务。

延伸思考

5.1. IO密集型任务

IO密集型任务指的是训练瓶颈在数据IO上的任务。这类任务一般会读取较多的数据集文件,数据量较大,GPU由于数据IO的瓶颈一直处于 饥饿
状态,因此GPU利用率较低。总结以下几种解决方案:

  1. 元数据缓存

元数据缓存方案能够将读取过的文件元数据缓存在内存中。在元数据和用户数据分离的文件系统中,高效的元数据性能对整个系统性能至关重要。在数据集只读场景下,元数据缓存可以在FUSE侧完成,也可以在用户侧完成。该方案一方面能够大大较少与元数据服务器之间的交互,缓存热点元数据,同时也能降级元数据服务器的压力。

  1. 文件缓存

文件缓存方案充分利用了本地SSD进行文件缓存。在数据集只读场景下,文件缓存仍然是可以在FUSE侧完成,也可以在用户侧完成。通过缓存文件元数据并聚合小文件进行本地存储,能使训练任务的IO方式从网络IO逐渐演变为本地IO。

  1. 聚合数据集文件

聚合数据集文件方案主要指的是 lmdb
TFRecord
等技术。在这种方案下,文件数目大大减少,可以有效地缓解深度学习场景下数据存取的问题,进而提高集群资源利用率。但文件聚合存储的方式对场景有一些限制,比如:数据更新修改会相对麻烦;数据集全局shuffle比较困难,只能做部分的shuffle。

  1. GPUDirect Storage

GPUDirect Storage
是NVIDIA公司在2019年推出的有关GPU显存和存储设备之间直接进行交互的技术。传统方式下磁盘中的数据需要先加载至内存中,再拷贝到GPU显存进行训练。在这项技术下,可以绕过CPU让GPU直接与存储设备进行交互,在本地或远程存储(NVMe磁盘)与GPU显存之间建立直接的数据IO路径。该方案一方面可以避免主存内数据冗余副本的产生,另一方面也缓解了CPU和内存的压力。

5.2. GPU计算密集型任务

GPU计算密集型任务指的是训练瓶颈在GPU计算上的任务,通常需要保证数据IO和梯度同步的低延时,使得GPU时刻处于忙碌状态。简要介绍以下几种解决方案:

  1. 数据预取

数据预取是最容易实现的方案。在每一个 iteration
计算过程中,事先对下一个或几个 iteration
所需的数据进行预取并预处理,保证下一个 iteration
开始时特征已处于就绪状态。

  1. GPUDirect RDMA(Remote direct memory access)

GPUDirect RDMA
Kepler GPU
CUDA 5.0
期间被提出,现在已得到较为广泛的支持。在多机训练过程中,这项技术能让多个GPU之间直接进行通信,同样也是避免了主存内数据冗余副本的产生,减少数据拷贝环节。配合Mellanox RDMA设备,数据可以从GPU显存经RDMA网卡发送出去,经另一台设备的RDMA网卡后传输至GPU,大大较少了IO路径。目前 Horovod
等分布式训练工具均以提供对 GPUDirect RDMA
的支持。

5.3. CPU计算密集型任务

CPU计算密集型任务指的是训练瓶颈在CPU计算上的任务,这类任务通常的计算瓶颈在于数据的预处理。此类任务CPU处于高负载状态,但GPU利用率和磁盘IO可能并不高。有以下几种解决方案:

  1. NVIDIA DALI(Data Loading Library)

NVIDIA DALI
是一个经过优化的数据加载的开源库,提供数据从磁盘加载到训练的完整pipeline。同时该库中还提供了音频处理、图像处理、视频处理等预处理方法,能够将在CPU上执行等预处理步骤放到GPU上快速执行,从而加速AI训练输入数据的预处理。

  1. 特征存储

特征存储方式是一种直观有效的方案,本质是进行CPU-GPU算力分离。对于某些大规模数据集,事先利用CPU算力对原始数据进行预处理,将样本特征打包后写入云存储设备中,然后多个GPU任务均可共享这些样本特征数据。但这类方法缺点在于当特征选取发生变化时,需要重新进行预处理。

6. 总结与展望

本文从实际训练场景出发,首先简单介绍了CephFS相关的基本概念,接着通过现象和源码分析训练过程中读取文件缓存失效的原因,然后给出了相应的解决方案。经过优化后,测试任务的训练速度能提升至原来的3~4倍。最后,通过延伸思考来发散思维,简要介绍了不同场景下AI训练加速的技术。

未来,针对IO密集型任务,利用 GPUDirect Storage
和Ceph的 RADOS API
等技术,结合本地SSD的高速缓存,可以在用户侧探索更极致的加速方案。这种方式理论上能够拥有更快的文件读取速度,能在用户侧对文件的元数据和数据进行充分缓存,减少用户态和内核态转换,是未来可以继续研究的方向。

欢迎关注腾讯程序员视频号