【无广告】模型的发展“又长又宽”，集群的配置却不是“又大又圆”

工作上一直在攻坚大模型的并行化，遇到不少困难，也解决了不少问题。正值最近搬家，有了更多的思考时间，因此写下一些个人观点。

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巨头们故意这样设计模型，引导模型的发展，来推销自己的硬件产品，我们无法臆测，但能够肯定的是，只有做好软硬件层面的“碗”，才能吃好模型算法的“面”。