清华大学舒继武:大模型时代下的存储系统挑战与技术发展

【编者按:由DOIT传媒主办,中国计算机学会信息存储专委会、武汉光电国家研究中心、百易存储研究院支持的2024中国数据与存储峰会,2024年11月8日,在北京新云南皇冠假日圆满举行。

大模型离不开数据,在大模型技术快速发展的时代,存储系统如何应对海量数据的处理需求,成为技术发展的重要课题。峰会主论坛上,清华大学长聘教授、闽江学院院长舒继武带来了关于《大模型时代下的存储系统挑战与技术发展》的主题分享。

以下内容根据速记整理,未经本人审定】

清华大学长聘教授、闽江学院院长舒继武

舒继武:各位好,很感谢DOIT给我这个机会跟大家一起来交流.我给大家交流的题目是”大模型时代下的存储系统的挑战与技术”.

刚才宋主编介绍了AI大模型时代存储系统技术的重要性,我主要讲存储系统到底对这些有哪些挑战。

背景与趋势

大模型现在发展很快很热,不管是从工业界还是学术界,大家都在讨论。

从工业界角度可以看到,大模型的市场规模很大。虽然目前的规模不算很大,但是升的很快,2023年达到国内147亿人民币,预计到2028年过千亿。学术界发表的论文,也是很大量的、跳跃式的往上涨,非常受关注。

为什么大模型如此受关注?这与政策和需求相关,政治、经济、文化、社会的各个领域可以产生很大的定力,产生很大的影响,成为国家战略的支撑。

大模型发展,得益于参数规模很大、上涨速度很快。因为规模很大,所以对存储的需求就不一样了。最左边的图,在训练期间,它的参数已经级别很高,一般都是TB级的数据,在训练阶段它要保持这些参数之外,它的很多中心接口都要保存,这里面也是TB级的数据,以GPT-3为例,如果说大小为8,一个系列长度是128KB,至少每产生一次中间接口就需要7个TB的数据等。

大模型的训练计算要求很高,消耗大量的计算量,这么高的数据量只有GPU可以胜任,GPU成为大模型训练的范式。左边的几个图可以看看,在大模型计算过程中计算量是翻457倍,通过一些测试可以看到,每一次大模型的颁布,GPU的销量马上就呈一个跳跃。

GPU计算量很丰富,但存储形成的容量有限,而又是强耦合的,一块GPU的卡计算量、存储量都是绑定的、数据是易失的,所以大模型的应用带来一系列的问题。计算量要求很高,但是存储的容量每两年才增加2倍,所以这个是严重不匹配。

挑战与思路以GPU为中心的模型数据存储服务

以GPU为中心,或者GPU成为一种模式,这种受众模型的存储带来了哪些挑战?

第一,大模型推理利用 KV Cache 保存上下文信息,其随序列生成不断增长。这意味着有大量的显存碎片并被保存下来,而且碎片还有内部碎片,必须在模型中开辟预留空间,但预留多少容量,每个请求的长度与最终生成的序列长度等都很难预测,另外,GPU卡有很多的请求,但它们之间很难完整分配。

第二,异构存储资源难利用。GPU服务器还离不开大量丰富的存储资源,DRAM、SSD,但GPU没有很强的耦合性,所以扩展起来也存在挑战,如访问时链路、延迟的问题、带宽的问题等等。

以训练阶段来说,计算过程可能只需要10秒左右,但是要把TB级的数据传输过来,耗费的时间可能需要几十秒,大量的时间用在等待,推理阶段也一样。

第三是容错。大模型训练时间长、GPU数量多,这会增加故障次数,Meta团队在训练OPT-175B时就出现了百余次故障。另外,显存是易失存储介质,GPU故障会导致训练所得的参数丢失。大模型的分布式训练模式会导致故障扩散到整个集群中,而且大模型的容错开销高,可达训练时长的70%。

如何破局?

在碎片管理方面,有很多的方法。可依据大模型训推模式,细粒度管理GPU显存的分配,提高GPU显存利用率,以支持大规模训推模型数据的存储。如在推理过程中,将用于存放KV cache的显存划分成定长的块(block),并随着生成序列的增长,按需分配显存块。

在异构资源的利用方面。对策之一是,基于流水线并行的显存扩容,在不支持NVLink的GPU服务器上(如消费级GPU 服务器),上述方案中频繁的GPU间通信会带来巨大的开销(IO时间占训练总时长的70%),流水线并行模式下,单卡仅需向相邻GPU传输少量激活数据。清华大学团队在探索如何在流水线并行模式下卸载模型参数至DRAM,将数据传输开销引入流水线并行中的模型切分算法,即通信开销最低的流水线切分方案,将数据传输与模型计算并行化,以掩盖数据传输带来的开销。对策之二是减少GPU与异构存储介质之间的KV Cache传输量,例如,KV Cache语义感知的数据压缩,或者存算协同卸载,但由此也带来缺点:前者损失了精度,后者对存储侧的计算资源需求高。

大量的推理需要用到终极的结果KVCache。如何减少KVCache传输过程中的传输量,清华大学基于隐藏状态的KV Cache恢复方面展开研究,核心思想是在推理时保存每层输入隐藏状态至DRAM或SSD中,之后的推理过程中仅传输隐藏状态至GPU,并在GPU侧通过矩阵乘恢复KV Cache,其优势是不影响推理精度,对存储侧的计算无要求,而且传输体积相较卸载KV Cache减少一半,计算开销仅为重计算的1/6。

在容错层面,采取冗余计算与故障后重定向的方式,以多张GPU中冗余计算多份相同的模型数据,如果发生故障,冗余的模型参数已是最新版本的参数,不需要重复对数据集训练,无需重新对数据集进行开始,或者将故障GPU上计算任务直接路由到带有冗余数据的GPU上。

以上讲的就是以GPU为中心的模型存储服务、碎片的管理、资源的利用和开销的挑战等内容,有很多的破局思路值得参考和借鉴。

总结和展望

时下大模型火热,需要重点关注以下几个问题。

一是大模型全周期绿色云存储问题。市场上用云做了很多大模型的训练和推理,导致大量的存储和缓存产生,碳排放的问题也暴露出来,因此应在训练阶段和推理阶段来解决存储的量,另外如何采用低价值的存储,既满足训练和推理阶段新的需求,还能够降低成本,降低排放量。这个说起来容易,但是有很多实际问题需要解决。

二是云端协同的存储。训练阶段存在很多个性化的微调,如何降低成本;在推理阶段如何把单侧端发挥出来,并有效保护用户的隐私。有很多问题要讨论,如任务的切分,哪些在单侧做,哪些在云端做,传输开销等的问题。

三是系统与算法的协调设计问题。很多的训练是在特定的系统上进行,如果将训练、算法和系统的设计结合起来,效率会更高更好。由于现有系统里有大量的新兴硬件,可编程硬件里面有很多的计算、存储功能,涉及到感存算或存算的问题,需要在现实中结合起来把新兴硬件特点发挥出来。同时,针对大模型的算法特性重新设计领域专用的训推硬件和系统,从而提高大模型训推性能。

在这些方面,清华大学开展了一些代表性的探索工作,例如样本的存储,海量文件系统包括存储IO的加速,训练和推理阶段的扩容管理等,也发表了相关的论文,有些工作还在继续推进。

总之,大模型训推对数据存储需求高,以 GPU 为中心的大模型训推已成为范式。但是GPU是显存计算和存储强耦合的,有易失、增长远低于大模型数据存储需求的增长等特点,而且大模型推理中显存资源利用不充分,难以利用GPU服务器中所包含的大量异构存储资源,大模型训练数据容错需求高、容错开销大,满足不了当下存储需求,即使万卡的集成,这些问题依旧存在。

如何顺利集成和推理集成存储资源,降低成本、容错等一系列挑战,需要重点关注面向大模型推理的显存资源管理、基于异构存储介质的显存扩容以及低IO开销的数据容错等挑战。

谢谢大家!