大模型将成为下一代人工智能基础设施是 AI 圈内已达成的共识。与生物体神经元越多往往越聪明类似,参数规模越大的 AI 模型,往往拥有更高的智慧上限,训练大模型或将让人类在探索通用人工智能上更进一步。然而,大模型极为高昂的算力成本在很大程度上阻碍了学界、工业界对大模型潜力的深入研究。
针对这一难题,达摩院联合阿里云机器学习 PAI 平台、EFLOPS 计算集群等团队改进了 MOE(Mixture-of-Experts)框架,创造性地通过专家并行策略,大大扩增了单个模型的承载容量。同时,通过加速线性代数、混合精度训练、半精度通信等优化技术,达摩院团队大幅提升了万亿模型的训练速度,且在效果接近无损的前提下有效降低了所需计算资源。
6 月 25 日,阿里巴巴达摩院发布「低碳版」巨模型 M6,在全球范围内首次大幅降低了万亿参数超大模型训练能耗,更加符合业界对低碳、高效训练 AI 大模型的迫切需求。
通过一系列突破性的技术创新,达摩院团队仅使用 480 卡 V100 32G GPU,即训练出了规模达人类神经元 10 倍的万亿参数多模态大模型 M6。相比此前英伟达使用 3072 A100 GPU 实现万亿参数、谷歌使用 2048 TPU 实现 1.6 万亿参数大模型,此次达摩院仅使用 480 卡 V100 32G GPU 就实现了万亿模型 M6,节省算力资源超 80%,且训练效率提升近 11 倍。
万亿参数的 M6-MoE 模型
今年 3 月,达摩院发布了国内首个千亿参数多模态大模型 M6。彼时,OpenAI 前政策主管 Jack Clark 曾公开点评道:「这个模型的规模和设计都非常惊人。这看起来像是众多中国的 AI 研究组织逐渐发展壮大的一种表现。」
此次又成功将 M6 升级至万亿规模,且在训练能耗、训练效率、文到图生成等多项下游任务上为业界带来了值得借鉴的关键技术突破。
大模型研究的一大技术挑战是,模型扩展到千亿及以上参数的超大规模时,将很难放在一台机器上。如果使用模型 + 流水并行的分布式策略,一方面在代码实现上比较复杂,另一方面由于前向和反向传播 FLOPs 过高,模型的训练效率将非常低,在有限的时间内难以训练足够的样本。
为了解决这一难题,达摩院智能计算实验室团队采用了 Mixture-of-Experts(MoE)技术方案,该技术能够在扩展模型容量并提升模型效果的基础上,不显著增加运算 FLOPs,从而实现高效训练大规模模型的目的。
普通 Transformer 与 MoE 的对比如下图所示。在经典的数据并行 Transformer 分布式训练中,各 GPU 上同一 FFN 层使用同一份参数。当使用图中最右侧所示的 MoE 策略时,则不再将这部分参数在 GPU 之间共享,一份 FFN 参数被称为 1 个 expert,每个 GPU 上将存放若干份参数不同的 experts。在前向过程中,对于输入样本的每个 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 层被丢弃。不同的 GPU 输入不同的训练数据分片。通过这种 expert 并行的策略,模型的总参数和容量大大扩增。由于单个样本经过 gate 稀疏激活后只使用有限个 experts,每个样本所需要的计算量并没有显著增加,这带来了突破千亿乃至万亿规模的可能性。
在 MoE 模型的具体实现上,谷歌的工作依赖 mesh tensorflow 和 TPU。达摩院则使用阿里云自研框架 Whale 开发万亿 M6-MoE 模型。将 FFN 层改造成 expert 并行,达摩院主要使用了 Whale 的算子拆分功能。在实现基本 MoE 策略的基础上,达摩院团队进一步整合 Gradient checkpointing、XLA 优化、混合精度训练、半精度通信等训练效率优化技术,并采用了 Adafactor 优化器,成功在 480 张 NVIDIA V100-32GB 上完成万亿模型的训练。在训练中,他们采用绝对值更小的初始化,并且适当减小学习率,保证了训练的稳定性,实现正常的模型收敛,而训练速度也达到了约 480 samples/s。通过对比 1000 亿、2500 亿和 10000 亿参数规模的模型收敛曲线(如下图所示),达摩院团队发现参数规模越大确实能带来效果上的进一步提升。然而,值得注意的是,目前扩参数的方式还是横向扩展(即增加 expert 数和 intermediate size),而非纵向扩展(即扩层数),未来他们将进一步探索纵向扩展,寻求模型深度与宽度的最优平衡。
探索 MoE,进一步提升模型效果
除了规模扩展外,达摩院对 MoE 架构开展了更进一步的探索研究,观察哪些因素对 MoE 模型的效果和效率影响较大。利用 MoE 架构扩大模型规模的一大关键是 expert 并行。而在 expert 并行中,几大因素决定着模型的计算和通信,包括负载均衡策略,topk 策略及对应的 capacity 设计等。
下面先简述达摩院 M6 团队对负载均衡在 MoE 实验的观察,然后介绍 topk 策略和 capacity 如何影响模型效果,并根据观察推出简单的 k top1 策略来达到更优的效果以及它在万亿模型上的效果。
考虑到负载均衡的问题,需要采用启发式的方法解决该问题,如上述的 expert capacity 和对应的 residual connection 的方法。Google 的 Gshard 和 Switch Transformer 沿用了 MoE 原文经典的做法加入了 auxiliary load balancing loss。目前还没有相关工作观察负载均衡的情况究竟有多严重,以及它是不是真的会影响模型的效果。达摩院团队在小规模的 M6 模型上进行了对 auxiliary loss 的消融实验,观察到该 loss 对最终模型效果影响甚微,甚至没有带来正向效果,然而它确实对 load balance 这个问题非常有效。如下图所示:
上图彩色曲线线表示各个层的 expert 接收有效 token 的变异系数随着训练进行的变化,灰色曲线表明训练阶段的 log PPL。图中变异系数 CV 表明每一层 expert 负载均衡情况,各曲线表明其随着训练步数的变化。不难发现,训练初期所有模型均有较严重的负载不均衡问题,刚开始少数的 expert 接收了绝大部分的 token,导致很多 token 直接被丢弃,但它们均能实现快速下降,尤其具备 auxiliary loss 的模型 CV 能降低到 0.3 左右,也可观察到在该水平下均衡程度很高,每个 expert 都能接收大量有效 token。然而与之相反,不加 auxiliary loss 的模型表现非常不同,有的层甚至在训练后期出现 CV 的飙升。但不管对比训练阶段的 log PPL,还是对比下游语言模型任务的 PPL,不带 auxiliary loss 的模型都表现更优。这一定程度上反映其实负载均衡对最终效果的影响并不大。
达摩院 M6 团队进一步探索了关键的 top-k gating 策略 k 值和 capacity(C)的选择。首先,他们简单地将 k 值扩大,发现 k 值越大其实效果越好。但考虑到选用不同的 k 值,C 则对应根据公式进行调整。通过对 C 调整到 k=1 的水平,观察不同 k 值的 MoE 模型的表现,达摩院团队观察到 k 值更大模型依然表现越好,尽管 k 值增加带来的优势逐渐不太明显。
但 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 操作,更加高效。
达摩院团队观察到,在不同规模的模型上,expert prototyping 都能取得比 baseline 更好的效果,同时速度和计算上也相比 top-k 更有优势。且其在更大规模的模型上优势变得更大,在百亿模型下游 image captioning 任务上甚至能观察到优于 top-k 的表现:
因此达摩院团队将该方法推广到万亿参数 M6 超大模型,并对应和上述的万亿 baseline 做了对比。目前,万亿参数模型训练了大约 3 万步,已经显著优于同等规模的基线模型,呈现约 5 倍的收敛加速。
沿着这个方向,值得做的工作还有很多:考虑到分组的特性,应当让组和组之间产生足够的差异,让每个组选出来的 experts 尽可能实现组合的效果。达摩院团队对此也在探索对应的有效方案。
M6 的商业化应用
AI 设计师与智能新制造
经过试用期,M6 将作为 AI 助理设计师正式上岗阿里新制造平台犀牛智造,通过结合潮流趋势进行快速设计、试穿效果模拟,有望大幅缩短快时尚新款服饰设计周期。随着实践经验的增长,M6 的设计能力还将不断进化。
结合阿里的电商背景,M6 团队希望通过 M6 大模型优异的文到图生成能力,和电商领域产业链深度融合,挖掘潜在的应用价值。具体来说,他们已深入到从服饰设计 & 生成、线上展示 & 测款的完整链路,期望利用 M6 的高清图像生成能力,缩短服饰企业的存货周转率,帮助商家对潮流趋势有更好的掌控力和更快速的反应力。
此前 OpenAI DALL·E 生成图片清晰度达 256×256,M6 则将图片生成清晰度提升至 1024×1024。
以下是 M6 生成高清服装设计图的示例,设计和图案均为 AI 创作:
M6 还可以生成现实中不存在的衣服类型,风格可盐可甜,可搞怪:
以下为 M6 参与新款服装设计的流程图:
工业级文案生成
除文生图外,M6 也已具备可在工业界直接落地的图生文能力,能够快速为商品等图片提供描述文案。该能力目前已在淘宝、支付宝部分业务上试应用。
在参数规模不断升级的过程中,达摩院团队发现,M6 的认知和表达能力也在不断提升:它能够观察到图片中更丰富的细节,并使用更精准的语言进行表达。
比如,在对下述风衣图片的描述中,更大参数规模的 M6 相比基础版,注意到了 “经典翻领设计”“腰间系带装饰”“两侧大口袋点缀” 等细节,生成文案信息量更大、措词更精准。
跨模态搜索
M6 对图片、文本的精准理解及匹配能力,已在支付宝、手机淘宝中初步试应用,有望帮助提升用户跨模态搜索的效果。
M6 团队观察到,淘宝上有很多长尾词,主要因为很多 95 后、00 后用户有非常特别的商品需求,这些需求带来了很多长尾的搜索词。比如,有用户可能想要一个表面凹凸的咖啡杯,也就是日式风格凹凸咖啡杯,因为商家一般不会把这样的细节写在商品名和描述中,单纯基于文本的搜索很难搜出对应商品。
多模态大模型为精准的跨模态搜索带来可能。目前 M6 已建立从文本到图片的匹配能力,未来,或将建立从文字到视频内容的认知能力,为搜索形态带来变革。
此外,达摩院此次发布的 M6 巨模型,成为国内首个实现商业化落地的多模态大模型。M6 拥有超越传统 AI 的认知和创造能力,擅长绘画、写作、问答,在电商、制造业、文学艺术等诸多领域拥有广阔的应用前景。
未来规划
今年以来,阿里在超大规模预训练模型领域屡出成果。除发布多模态巨模型 M6 外,阿里巴巴达摩院近期还发布了中文社区领先的语言大模型 PLUG,实现了在 AI 大模型底层技术及应用上的深入布局。
达摩院资深算法专家杨红霞表示:「接下来,M6 团队将继续把低碳 AI 做到极致,推进应用进一步落地,并探索对通用大模型的理论研究。」
未来 M6 团队的主要关注方向包括:
- 继续将 GreenAI 做到极致,让更多学者和企业能参与对下一代 AI 的研究和应用;
- 继续推进大模型应用落地,让下一代 AI 进入包括社会公益在内的更多领域;
- 优化通用模型下游任务训练,让大模型在更多任务上拥有更好表现;
- 探索对通用大模型的理论研究,期望揭开「How it works」;
- 探索大模型训练的软硬件结合,启发下一代人工智能硬件设计。