达摩院3个月打造“低碳版”万亿参数模型M6，怎么做到的？

阿里巴巴达摩院今天（6月25日）发布“低碳版”巨模型M6，在全球范围内首次大幅降低万亿参数超大模型训练能耗，更加符合业界对低碳、高效训练 AI 大模型的迫切需求。voYednc

通过一系列突破性的技术创新，达摩院团队仅使用 480 卡 V100 32G GPU，即训练出了规模达人类神经元 10 倍的万亿参数多模态大模型 M6。voYednc

据悉，相比此前英伟达使用 3072 A100 GPU 实现万亿参数、谷歌使用 2048 TPU 实现 1.6 万亿参数大模型，此次达摩院仅使用 480 卡 V100 32G GPU 就实现了万亿模型 M6，节省算力资源超 80%，且训练效率提升近 11 倍。voYednc

不仅如此，达到这种惊人效果，从千亿到万亿参数规模，阿里只花了3个月时间。voYednc

万亿参数的 M6-MoE 模型

自从大模型变得流行起来之后，它所具备的创意能力，一直是被世人所津津乐道。voYednc

那么与国内外已经诞生了几个万亿“体量”的大模型相比，阿里此次提出的大模型，又有什么独到的特点？voYednc

据悉，M6不仅公开了实现的详尽细节、模型的收敛情况（详见文末论文链接），而且还是国内首个实现商业落地的万亿参数多模态大模型。voYednc

M6拥有超越传统AI的认知和创造能力，擅长绘画、写作、问答，在电商、制造业、文学艺术等诸多领域拥有广泛应用前景。voYednc

据了解，经过一段时间的试用，M6将作为AI助理设计师正式上岗阿里新制造平台犀牛智造，通过结合潮流趋势进行快速设计、试穿效果模拟，有望大幅缩短快时尚新款服饰设计周期。M6还已应用于支付宝、淘宝等平台，参与跨模态搜索、文案撰写、图片设计等工作。voYednc

目前，M6主要参与一些基础款的设计。但可预期的是，随着实践经验的丰富，M6的水平将不断进化。voYednc

据了解，M6计划在一年内生成上万款高清服装设计图。voYednc

3个月打造万亿参数模型，怎么做到的？

早在今年1月份，阿里便推出了百亿参数模型，而当时谷歌就已借助MoE （Mixture of Experts）的架构，提出来了1.6万亿参数的Switch Transformer。voYednc

MoE架构能够做到的在扩展模型容量并提升模型效果的基础上，不显著增加运算FLOPs，这样就可以实现高效训练大规模模型的目的。voYednc

于是，阿里从百亿开始的“规模升级”过程中，便借鉴了这个架构。仅耗时2个月，便发布千亿参数大模型，而且只用了32个 V100 GPU。voYednc

普通 Transformer 与 MoE 的对比如下图所示。在经典的数据并行 Transformer 分布式训练中，各 GPU 上同一 FFN 层使用同一份参数。当使用图中最右侧所示的 MoE 策略时，则不再将这部分参数在 GPU 之间共享，一份 FFN 参数被称为 1 个 expert，每个 GPU 上将存放若干份参数不同的 experts。voYednc

在前向过程中，对于输入样本的每个 token，达摩院团队使用 gate 机制为其选择分数最高的 k 个 experts，并将其 hidden states 通过 all-to-all 通信发送到这些 experts 对应的 GPU 上进行 FFN 层计算，待计算完毕后发送回原 GPU，k 个 experts 的输出结果根据 gate 分数加权求和，再进行后续运算。为了避免部分 experts 在训练中接收过多 tokens 从而影响效率，MoE 往往设定一个 capacity 超参指定每个 expert 处理 token 的最大数量，超出 capacity 的 token 将在 FFN 层被丢弃。voYednc

不同的 GPU 输入不同的训练数据分片。通过这种 expert 并行的策略，模型的总参数和容量大大扩增。由于单个样本经过 gate 稀疏激活后只使用有限个 experts，每个样本所需要的计算量并没有显著增加，这带来了突破千亿乃至万亿规模的可能性。voYednc

但阿里在却在研究过程中发现了一个问题：MoE负载不均衡。voYednc

简单来说，原理是这样的。voYednc

大模型常用到的Transformer分布式训练中，通常是各个GPU同一FFN层中，使用同一份参数。voYednc

而MoE就不同了，上述的这部分参数会在GPU之间共享，一份FFN参数被称为1个“专家”（expert），每个GPU上将存放若干份参数不同的“专家”。voYednc

voYednc

（如上图中标红框部分所示）voYednc

但阿里却发现，在原来MoE的训练过程中，非常容易只选择top的几位“专家”，这就使得头部效应非常严重。voYednc

于是乎，阿里便对MoE的这个问题进行了改良。voYednc

考虑到负载均衡的问题，需要采用启发式的方法解决该问题，如上述的 expert capacity 和对应的 residual connection 的方法。Google 的 Gshard 和 Switch Transformer 沿用了 MoE 原文经典的做法加入了 auxiliary load balancing loss。目前还没有相关工作观察负载均衡的情况究竟有多严重，以及它是不是真的会影响模型的效果。达摩院团队在小规模的 M6 模型上进行了对 auxiliary loss 的消融实验，观察到该 loss 对最终模型效果影响甚微，甚至没有带来正向效果，然而它确实对 load balance 这个问题非常有效。如下图所示：voYednc

voYednc

上图彩色曲线线表示各个层的 expert 接收有效 token 的变异系数随着训练进行的变化，灰色曲线表明训练阶段的 log PPL。图中变异系数 CV 表明每一层 expert 负载均衡情况，各曲线表明其随着训练步数的变化。不难发现，训练初期所有模型均有较严重的负载不均衡问题，刚开始少数的 expert 接收了绝大部分的 token，导致很多 token 直接被丢弃，但它们均能实现快速下降，尤其具备 auxiliary loss 的模型 CV 能降低到 0.3 左右，也可观察到在该水平下均衡程度很高，每个 expert 都能接收大量有效 token。然而与之相反，不加 auxiliary loss 的模型表现非常不同，有的层甚至在训练后期出现 CV 的飙升。但不管对比训练阶段的 log PPL，还是对比下游语言模型任务的 PPL，不带 auxiliary loss 的模型都表现更优。这一定程度上反映其实负载均衡对最终效果的影响并不大。voYednc

达摩院 M6 团队进一步探索了关键的 top-k gating 策略 k 值和 capacity（C）的选择。首先，他们简单地将 k 值扩大，发现 k 值越大其实效果越好。但考虑到选用不同的 k 值，C 则对应根据下图公式进行调整。通过对 C 调整到 k=1 的水平，观察不同 k 值的 MoE 模型的表现，达摩院团队观察到 k 值更大模型依然表现越好，尽管 k 值增加带来的优势逐渐不太明显。voYednc

voYednc

但 k 值的增加根据 Gshard top-2 gating 的实现，除了存在实现层面上一定的冗余和困难外，循环 argmax 的操作也会导致速度变慢。此外，第二个 expert 的行为会受到第一个 expert 的影响，让训练和测试存在差异。达摩院团队用 expert prototyping 的简单方式替代，相较 baseline 实现了效果提升，且未显著增加计算成本。expert prototyping，即将 expert 分成 k 组，在每组中再进行 top-k 的操作（通常采用 top-1，便于理解），然后将 k 组的结果进行组合，也称之为 k top-1。这种方式实现上更直接简便，并且允许组和组之间并行做 top-k 操作，更加高效。voYednc

达摩院团队观察到，在不同规模的模型上，expert prototyping 都能取得比 baseline 更好的效果，同时速度和计算上也相比 top-k 更有优势。且其在更大规模的模型上优势变得更大，在百亿模型下游 image captioning 任务上甚至能观察到优于 top-k 的表现:voYednc

voYednc

因此达摩院团队将该方法推广到万亿参数 M6 超大模型，并对应和上述的万亿 baseline 做了对比。目前，万亿参数模型训练了大约 3 万步，已经显著优于同等规模的基线模型，呈现约 5 倍的收敛加速。voYednc

沿着这个方向，值得做的工作还有很多：考虑到分组的特性，应当让组和组之间产生足够的差异，让每个组选出来的 experts 尽可能实现组合的效果。达摩院团队对此也在探索对应的有效方案。voYednc

除此之外，算子精度也是阿里此次改良的工作之一。voYednc

谷歌在做Switch Transformer时，为了将模型体积压下来，选择了BF16。voYednc

但精度的降低会带来非常大的技术挑战，就是如何保证模型收敛的问题。voYednc

而且阿里还要做到“低碳版”，不能烧太多的GPU，因此相比谷歌在算子精度方面的工作，阿里可谓走了一条更加“极端”的路线。voYednc

具体而言，XLA优化、混合精度训练、半精度通信等训练效率优化技术，并采用了Adafactor优化器，成功在480张NVIDIA V100-32GB上完成万亿模型的训练。voYednc

并且在训练中，他们采用绝对值更小的初始化，适当减小学习率，保证了训练的稳定性，实现正常的模型收敛，而训练速度也达到了约480samples/s。voYednc

以上便是阿里“低碳版”万亿参数大模型的核心奥秘了。voYednc

论文地址：

https://arxiv.org/abs/2105.15082voYednc

责编：DemivoYednc