工作上一直在攻坚大模型的并行化,遇到不少困难,也解决了不少问题。正值最近搬家,有了更多的思考时间,因此写下一些个人观点。
1
回顾模型变“深”这件事
自从ResNet出现之后,深度学习模型在“深”的维度上真正发展起来。能够简单上手搭建层数更深的模型,必然引入大量的重复结构。显然,这对工业界是个极大利好的消息。给定基本结构,再写一些for循环,就可以在各种业务上尝试模型Deeper的效果。
于是,大兴堆层工程,在工业界盛行起来。先是在CV领域,暴堆Residual Block;随后Attention结构在NLP领域流行起来,通过暴堆Attention Layer,模型质量也提高了。此时,若有同僚来相问,堆叠何时止?答曰:“单卡放不下”。
直到现在为止,单个GPU显存的增长速度,依然远远不能满足模型参数的增长需求。拿32GB规格的V100来说,不做任何优化的情况下,跑个Bert-large都是勉勉强强的事。而Bert-large已经是2018年的模型了。
2
应对深度?堆卡做Pipeline
应对深度的增长,堆多张卡采用模型并行其实是很有效的方式。如果仅仅考虑将单卡放不下的很深的模型训练起来,则将模型横向切成前后依赖的Pipeline,分配到多个GPU上依次做串行的计算,在工程上是是非常好实现的。TensorFlow和PyTorch都有控制device placement的语句,工程师只需要在构建模型的语言上加一些用于控制placement的修饰即可。
当然,这仅仅是将模型跑起来,如果追求更高的训练效率和GPU利用率,那么针对于Pipeline效率研究就必不可少。Pipeline的研究在计算机体系结构领域中早已存在,回想一下计算机指令的流水线,再结合Deep Learning训练的特点,才有了GPipe,Pipe Dream,DAPPLE等工作的出现。这里最棘手的问题就是每个Device上计算量的均匀性,以及Device之间数据传输是否足够小。
总之,当模型不断变深,将模型训练起来并追求更高的GPU效率并不是一件十分困难的事。但我们得总结一下,迫使我们将模型做复杂并行的根本原因有二:
- 单个GPU的显存上限实在有限
- 跨GPU的数据传输带宽也实在有限
我们必须承认,工业界大部分的集群大多具有“单卡算力强,跨机通信弱”的特征。否则,我们可能并不需要这么复杂的模型并行方式,甚至根本不需要模型并行,也不需要那么tricky的工程。不管怎样,在工业界流行的GPU cluster上,往“深”发展的模型确实好解一些。
3
模型变“宽”了,难玩了
可偏偏深度学习也在往另一条路发展,模型开始变“宽”了。典型地,搜广推领域带有大规模Embedding的模型,就是一种典型的例子,这是电商业务需求产生的必然结果,并且大量的商品增长还在不断迫使Embedding的宽度爆炸。训练这种模型依然有一定的门槛,但好在Embedding具有Sparse的特质,并且one-hot与Embedding的巨大计算量的乘法,本质上可以用sparse id做embedding_lookup这种纯访存操作等价替代,所以只需要将Embedding层做sharding,然后均匀的放在多个Host memory上,就可以在保持Data parallelism的语义基础上直接训练起来。实现也非常简单,TensorFlow有自己相对精简的Parameter Server功能,专门用于处理这一种case。
当CV和推荐领域的工程师们在各自熟悉的工程领域上享受着舒适区时,NLP领域的模型率先将dense层变宽了,比如T5模型。这很麻烦,因为这模型它“又长又宽”,而GPU集群却不是“又大又圆”。宽度增长来自于Attention Block中的Attention head数量,以及FFN的维度。面对Dense宽度的增长,将Dense Layer做Sharding放在CPU上不再可行,因为它依然是做Dense计算的,依然是要在GPU上存放的。除此之外,Sharding破坏了单个Op的原子性,必须引入一些额外通信维持原有语义的等价性。然而引入的通信是否会很大,通信的链路走网卡会不会是bottleneck,Sharding到各个GPU上均摊的计算量是否是原来shape计算速度的N倍(卡数决定),这都是新的问题。况且,解决深度那一侧的方法引入的Pipeline,也必须包含到模型的并行策略中。因此,为了适配工业界的GPU集群,使出各种策略混合的超必杀,才是可能让这种模型训练起来的唯一方法。
4
TPU集群这“碗”,确实适合吃“宽面”
此时应该注意,Release这种宽度增长模型的Google,并不存在这些问题。因为他们的TPU集群是“又大又圆”。如果类比GPU集群,单个服务器上挂多块GPU的单机多卡服务器,在TPU面前弱爆了。首先,2048个TPU被组织成一个pod,而跨pod的带宽也做成了和pod内通信速度相同的链路,彻彻底底拓扑同构,“又大又圆”。
另一方面,Google又在Attention Block之后引入了MoE,通过疯狂增加Experts数量,使模型达到史无前例的宽,迫使工程师必须选择将MoE层做Sharding的并行方案,同时还声称拓展“宽”比“深”更有效。MoE并行引入的All-to-All通信更加过分,每轮的数据交换甚至高达几十GB。这样一来,低带宽连接的GPU集群就基本丧失了“盛面”能力,只有持有TPU集群这个“碗”,才能吃好这面。
到底是不是AI巨头们故意这样设计模型,引导模型的发展,来推销自己的硬件产品,我们无法臆测,但能够肯定的是,只有做好软硬件层面的“碗”,才能吃好模型算法的“面”。