话说人工智能江湖门派众多且繁杂,好似那大脑中的神经网络一般,但繁杂中却仍然有着一统的迹象......许久之前,ML派坐落美利坚合众山中,百年来武学奇才辈出,隐然成江湖第一大名门正派,门内有三套入门武功,曰:图模型加圈,神经网加层,优化目标加正则。有童谣为证:熟练ML入门功,不会作文也会诌。不期数年,北方DL神教异军突起,内修表示学习,外练神经网络,心法众多,曰门,曰注意,曰记忆,曰对抗,曰增强。经ImageNet一役威震武林,豢Alpha犬一匹无人可近。一时家家筑丹炉,人人炼丹忙,门徒云集,依附者众,有一统江湖之势。有童谣为证:左手大数据,右手英伟达,每逢顶会炼丹忙。转至今时,竟知DL神教仍不是源头,它背靠「大力神教」,Transformer、BERT、GPT-3 、DALL.E等等神丹相继纷纷被祭出......丹炉越来越大、炉火越烧越旺,「大力出奇迹」已被众多教徒们信奉为至高无上的教义。............................就这样时间来到了2021年,训练大模型已几乎成为了AI界的共识,近期MLP的涅槃重生与Transformer的一场争斗大戏,让大家不禁悲叹「Attention is all you need」是假的,别的神马也都是不确定的,唯有「Money is all you need 」是真的。然而,因为虚拟货币挖矿等众多原因,显卡的价格一直居高不下,不容乐观,且一度缺货。随手打开某电商平台,可以看到一张2080Ti 涨到一万五起步 ,一张 V100 更是得四五万。可就在这样的背景下,如果现在有人告诉你2080Ti可以当作 V100来用,那得节省多少Money?你会说这怎么可能呢?V100的显存可是有16GB ,而2080Ti的显存只有11GB,相比于计算能力,显存容量可是一票否决制的,虽然计算能力决定了模型训练的快慢,但是显存容量决定了模型能否训练。即使 2080Ti 的计算能力高于 V100,但是由于显存容量不够,batchsize设置为128 的 ResNet50 根本无法进行训练,单就这点儿而言对于吃大显存的 AI 模型就不可能绕的过去。然而人定胜天,世间上最令人兴奋的事莫过于把看似不可能变成可能,好似魔法一般,其秘诀就在于下面这两行代码:没错,就是上面这两行看似简简单单的代码,它就像一个“点石成金”的按钮,蕴含着巨大的能量,只要一键开启之后,2080Ti 就可以当作 V100来用!这么神奇的事情背后究竟是怎么一回事呢?在一番探寻之后发现,原来是旷视在开源深度学习框架MegEngine 的最新版本V1.4中,新增了动态图的显存优化技术(DTR)。“炼丹”时使用这项新功能后可以实现显著降低显存占用的魔法般效果。在使用DTR功能后,11G显存的2080Ti 可以训练出原本32G显存的V100才能训练的模型,可以节省很多Money!另外很多学生党只有 1060 之类的4~6GB的小显存卡,而靠 DTR就能在便宜的民用显卡上训练原本需要吃10GB显存以上的大模型了。 1
对这个训练 AI 大模型神器的感觉还是不够直观?那就不如看一下更直观的实验数据对比。首先是训练耗时对比。下图是MegEngine的DTR实现与原论文在PyTorch中的实现在ResNet1202上的训练情况对比,注意到实验用的显卡不同,所以从数据上看MegEngine稍快一些。不难看出,在显存管理上MegEngine要更好一些,因为在11G的显卡上它仍然能跑batchsize=100 的模型训练。除了论文中尝试的最大batchsize=140之外,MegEngine团队还尝试了更大的batchsize,也都是可以运行的。下面是在MegEngine框架上开启不同显存优化的训练耗时对比,baseline是在动态图模式下不加任何显存优化运行的结果。首先是两个常见的模型——ResNet50和ShuffleNet的对比,可以发现开启DTR优化后极限batchsize超过了静态图Sublinear和baseline,且在batchsize相同时耗时和Sublinear持平。需要说明的是这里拿视觉模型只是为了方便做实验和举例,实际上这项优化功能适用于CV、NLP等不同领域的任意 AI 模型训练。上面的两个模型都是比较偏静态的,所以可以用静态图的Sublinear显存优化来做对比,而下面这个SPOS网络就比较特殊,它是一个从输入到输出有多条路径可以更新的大网络。在训练过程中,每一轮会随机采样去更新某一条路径,这就导致每轮执行的语句可能不相同。对于这种网络,在动态图里实现会比较自然。因此,这里只取了动态图DTR优化的结果与Baseline比较。不论是单卡还是八卡,动态图的极限batchsize都在100,如果打开DTR可以跑到250甚至更大。硬币都是两面的,在MegEngine框架上实现的这项DTR 技术非常之香——2080Ti 可以当作V100来用,但固然也有代价——牺牲了计算算力和计算时长: