大模型与AI底层技术揭秘 (11) 变形记

2023-11-27 16:43:35 浏览数 (1)

有一天,小H在吃完午饭回到办公室,旁边几位同学在打《王者荣耀》,并且在挑拨匹配到的一对情侣队友分手。

“你看你死了他才来救你,他心里一点都不在乎你!”

“你被人追着打的时候,他说不定是在和别的小姑娘打情骂俏呢!”

“他肯定是因为经常和别的女孩子一起玩,才和你配合不好的~”

“你好好问问他,他还有多少事情瞒着你?”

小H觉得这些人太不瑞雪,把他们赶到了会议室去玩,自己钻进睡袋。很快,小H闭上了眼睛。

小H睁开眼睛的时候,发现自己变成了一只甲虫……

当然,小H立即从噩梦中吓醒了。当小H揉着眼睛爬起来,才想起来这是奥地利作家Franz Kafka成名作《变形记》里面的情节。

Franz Kafka是奥地利著名作家,以高产而闻名。在《变形记》发表的96年后,他的迷弟Jay Kreps在Linkedin任职期间,和另外几个朋友开发了一种高性能的流式消息中间件,并以偶像的名字Kafka为之起名。

Kafka并不是今天的主角,但kafka采用的一种加速机制,却成为了高性能与AI计算软硬件体系中的基本操作——

这种机制叫做零拷贝。

在传统模式下,数据从网络和文件之间的传输需要经过4次拷贝,4次上下文切换和4次系统调用:

为了减少上下文切换以及数据拷贝带来的性能开销,Kafka使用了Linux系统调用mmap来处理其索引文件,将其映射到用户态进程可见的内存空间,从而实现快速访问和搜索内存中的索引,加速在日志文件中定位消息的过程。

mmap的具体机制是,将内核中读缓冲区(read buffer)的地址与用户空间的缓冲区(user buffer)进行映射,从而实现内核缓冲区与应用程序内存的共享,省去了将数据从内核读缓冲区(read buffer)拷贝到用户缓冲区(user buffer)的过程。也就是说,CPU只需要获取到读缓冲区的地址,就可以直接访问索引文件的数据了。

而Kafka在发送消息的时候,采用的是Linux的sendfile。如果我们不使用sendfile,传统方式也需要先读取磁盘,再用socket发送,那么也需要4次内存拷贝,其中2次是CPU在内核缓冲器和用户态缓冲区之间拷贝:

而使用sendfile只需要进出一次内核,实现零拷贝:

Jay Kreps 利用了内核提供的零拷贝机制,实现了这样一个像自己偶像Kafka一样,具备高产高性能的消息中间件,并且以自己偶像的名字起名致敬。

我们发现,Kafka是运行在CPU上的,CPU本身硬件具备零拷贝能力,如DMA、MMU的内存重映射等,因此能够实现零拷贝。而AI程序运行在GPU上,如何让GPU也利用内存零拷贝的技术来快速直接存取磁盘上的内容呢?

NVidia给出的答案是:GPU Direct Storage。

在GPU Direct Storage出现之前,如果需要将训练数据从磁盘加载到GPU内存,需要经过以下步骤:

1. GPU向CPU发起中断;(上下文切换,进入内核)

2. CPU读取NVMe SSD,获得存有NVMe SSD数据的一块内存缓冲区地址;(可以使用DMA,但也有上下文切换和进入内核的环节)

3. CPU将这段地址给GPU,并让GPU发起DMA,从这块内存缓冲区中拉取数据;

4. GPU通过PCI-E向RC发起DMA,拉取数据;

5. GPU将数据存到DRAM;

在这个过程中,CPU至少会进出2次系统内核,这是因为,读取磁盘和向GPU传输发起DMA的指令都需要在内核态中实现。

有没有好的方法绕过CPU,让GPU直接读取磁盘呢?

这涉及到磁盘数据的组织问题。我们知道,Linux或其他操作系统中,读取磁盘文件,首先要在文件系统中找到文件所在的磁盘LBA(Logical Block Address),也就是依赖于文件系统提供的数据索引。而GPU上并没有操作系统,也无从解析磁盘前部的文件系统数据索引,是没有办法找到LBA的。

NVidia的工程师们如何解决这一问题呢?

请看下期分解。

0 人点赞