12月20日,由中国互联网协会、微博、新浪新闻主办的“数字力量,探索无穹”2023探索大会在北京拉开帷幕。潞晨科技联合创始人兼CTO卞正达分享了题为《Colossal-AI:AI大模型的挑战与系统优化》的演讲。
潞晨科技联合创始人兼CTO卞正达-主题演讲
以下是卞正达演讲实录,内容经编辑略有删减:
感谢各位,我是来自潞晨科技的卞正达。今天主要代表董事长跟大家做一个交流,今天分享这个主题主要是关于AI大模型的分布式系统优化。首先我先介绍一下大模型时代下的发展背景和一些挑战。
近年来AI飞速发展。之前2016的时候,当时比较火的ResNet模型,可能几个小时就可以训练完一个CIFAR10应用,到了后来训练BERT模型可能花一两天可以做完。但是到了如今大模型时代,我们模型的参数量和计算量都是膨胀到一个非常可观数量级。
比如说谷歌PaLM模型,如果用一张显卡训练,据说是要训练时间长达超过300年,同时训练成本高达900万美元。这么高成本因为主要它的参数量和训练所需的数据量的规模都呈现一个逐年飞速增长的趋势。所以它带来的计算量已经达到了一个非常巨大的数量级。
同时,我们为了成功部署大模型的训练和推理,我们需要去维护一个成百上千张显卡的巨大的集群。这个成本也是非常可观的。
因此我们推出Colossal-AI框架,想要在分布式部署这一层去利用高效的分布式算法,解决AI大模型部署一个超高成本问题。这个框架是作为一个软件的基础设施,它是向上可以承接PyTorch、Huggingface、Lightning等不同的AI应用,向下可以兼容包括GPU、TPU、NPU等等不同的硬件加速方案。
我们是想要利用这个高性能分布式的算法,去帮助企业落地分布式大模型,也帮助他们去实现降本增效。我们这个框架一个主要的设计思路,包括以下三层:
1、高效内存管理系统,它的主要的思路是可以利用更低廉的CPU内存和硬盘空间去缓存模型计算中冗余的存储开销,极大地降低GPU的存储压力,相当于降低硬件门槛。
2、N维并行系统,这一层我们利用一些分布式的算法,不仅能够均匀去切分大模型的存储开销,同时也能够实现一个比较高效的计算和通信。
3、低延迟推理系统,我们利用量化、并行计算等等策略最大化的推理部署一个性价比。
尽管我们框架的推出时间也不长,到现在差不多两年这么一个时间,我们还是在开源社区,学术界、工业界获得一定的认可和影响力,我们看到横向对比一些其他厂商分布式框架,增长速度比较快,目前积累了3.5万多颗Star。这些开源社区用户也是来自全世界各地。同时我们的一些核心工作也是被一些顶级的会议所接收。
下面具体介绍一下我们框架的核心设计细节。第一个就是N维并行系统,在我们开发这个Colossal-AI框架之前,已经有非常多的并行的一些技术方案,像是Gpipe的流水线并行方案,Deepspeed的Zero数据并行方案、以及Megatron的3D并行方案等等。但是我们发现用户拿到一个实际需求之后,他很难去正确选择一个合适的加速方案,把这个需求转化成真正一个落地的大模型解决方案。
我们一个思路,把目前业界最SOTA最高效这些并行方案整合进一套系统,借助我们长期做系统优化的经验,去帮助用户他们来选择合适的并行方案,同时提供一个最高效的系统实现。
我们的这套系统里主要有三个维度:
1、数据并行。
2、模型并行。
3、序列并行。
像数据并行,我们解决了一个问题,常见数据并行方案下,如果我们训练批大小超过8000以上,导致最终模型有一个泛化问题。在我们框架里通过LARS、LAMB这样一些针对做优化这些优化器解决这个问题,同时把批大小提高到像64K一个程度。
第二个是模型并行,这个包括我们框架里提供的方案,包括流水线并行,包括张量并行,同时针对不同的物理拓扑,会提供一些不同的模型并行优化算法,达到最佳的计算和通信同步的效率。
第三个是序列并行,因为现在长序列一个模型也是目前比较热门一个发展趋势。其实其他一些框架或多或少有一些类似去解决这个长序列一个方案,比如说Deepspeed等都有序列方案,但是如果仔细研究过他们的实现,会发现他们在计算attention的时候,并没有切分序列长度这个维度。
我们框架里是在全程计算中都实现了序列维度的切分,也就是说只要有足够多的设备,我们就能够实现一个无限长的序列一个计算。
我们这套序列并行方案其实不仅对训练,而且对大模型推理也是有效的。因为根据我们的经验,我们常用这些推理卡,A10部署一个7B左右模型推理,单卡能够承载这个拓扑数2500-3000级别。我们知道看到有很多,像是64K,甚至128K这样一个超长序列模型,它去做这样的模型超长序列推理时候,会需要我们这样的序列并行方案去实现一个超长序列的分布式处理。
第二点,内存管理异构系统。因为我们知道深度学习,特别是在用混合精度训练时候,它的大部分存储其实被用在优化器参数更新上,计算比较重的前向和反向计算过程当中,其实它的存储开销相对较少。因此我们可以通过较为廉价的CPU存储去缓存一部分的模型计算中一些冗余存储,比如可以把这个优化器的部分都迁移到CPU存储上,降低这个GPU的存储压力。但是如果把所有的冗余存储都搬到CPU上,会导致CPU和GPU之间有一个频繁的通信数据交换。我们知道目前一些物理设备上,CPU和GPU之间带宽比较瓶颈,因此我们还是想要尽可能减少异构存储带来额外的通信开销。我们一个思路是,尽可能还是把大部分存储都放在GPU上,只有把超过上限这个部分缓存到CPU上,这样就能减少大部分的数据交换,然后同时我们也尽可能把计算都保留在GPU上去做。
同时,我们还开发了基于chunk的tensor存储机制。这个思路来自于DDP里通过bucket去提升实际梯度同步通信效率。类似地在zero并行等方案中,也可以通过chunk把小的tensor混合到一起通信,这样可以地更好利用带宽来提升通信效率。
基于上述这些系统优化,对于常见的开源模型,例如GPT、Llama等等,我们都能够实现明显的降本和加速。
同时我们在年初2-3月份时候,也是第一个推出ChatGPT的RLHF方案的一个复现和开源。我们还推出中文ColossalChat这么一个多轮对话的应用,同时也包括ColossalEval这个对话模型的验证工具。
另外我们还开源了一个中文LLaMA-2这么一个迁移训练方案。这个方案是我们基于LLaMA-2开源基座,只使用了8.5个Billion数据量,不超过1000美刀的成本,把LLaMA-2语言能力迁移到中文场景里面。我们可以看到在常见的公开中文benchmark上,它的能力和之前7B、6B规模的一些开源大模型,例如百川、ChatGLM的分数也是不相上下的。
刚才有很多专家提到,目前一个主流的训练方案是预训练加上微调这两步。通过这个方案,我们把它拆成预训练+迁移训练+微调这三步。因为第一步预训练是成本最高的环节,我们可以利用目前开源一些大模型,我们也可以看到目前开源模型这个发展也是十分迅速,LLaMA等这些开源模型都非常不错。我们可以借助它们的能力,在垂直领域上只要通过迁移训练加上微调两步把它大模型能力落地到具体场景中,可以实现一个较低的成本下,完成比较定制化一个专业模型。因此我们开源模型之后也是受到了一定的热度。
我们今年还推出一些进一步提升用户体验,降低用户门槛的一些产品,像是一个云平台跟一体机。云平台就是我们集成Colossal-AI的分布式加速能力,以及包括训练、微调、推理、部署等场景,提供一个低代码、低门槛、低成本的云上AI开发体验。
一体机大模型服务器这个产品,是我们会针对软硬件做一些极致优化,同时提供丰富模型库和使用范例,让用户尽量开箱即用,获得一个低成本、低代码的AI开发体验。
最后也欢迎大家加入这个社区,一起贡献ColossalAI这个框架。
谢谢大家!
(1)_20231227155943_738.jpg)
