《GShard：用条件计算与自动分片扩展巨型模型》

要解决的问题

做多语言翻译，一种办法是每种语言对训练一个模型。GShard 论文正是拿这种安排做对照：如果要把100种语言翻译成英语，按老办法就要训练100个双语基线模型^[1]。用一个多语言大模型取代这一群小模型，意味着模型容量要大得多，参数量随之被推到几千亿的量级。到了这个规模，模型无法舒舒服服地放在一块加速器上，把它手工切分到大量设备上的工程工作量就成了真正的瓶颈——模型代码改一处，切分代码就得跟着改一处。

这篇论文想拆掉的正是这个瓶颈：让模型代码保持原样可读，同时又能描述分布在成千上万块芯片上的计算。

怎么做的

论文把 GShard 描述为一个模块，由一组轻量级注释 API（annotation API）和对 XLA 编译器的扩展两部分组成。二者合起来，提供了一种只对现有模型代码做极小改动、就能表达多种并行计算模式的办法^[1]。思路是：模型作者标注出网络中需要跨设备切分的部分，编译器负责算出如何划分张量、并补上切分所需的通信——论文把这件事称为自动分片（automatic sharding）。

第二个要素是条件计算（conditional computation），具体形式是稀疏门控的专家混合（sparsely-gated mixture-of-experts）。在专家混合层里，若干专家子网络并排放置，门控决定每个输入被送给哪些专家，于是对任何一个词元（token）来说，只有一部分网络被用到。参数量和每个词元的计算量由此解耦：模型可以非常大，而处理单个词元仍然只用其中一小部分。作者用这一机制，把多语言神经机器翻译 Transformer 通过自动分片扩展到了6000亿参数以上^[1]。

两半合起来，分工很清楚：专家混合提供容量，而不用让计算量按比例上涨；注释加编译器这一层，则让庞大的模型无需手写切分代码就能被表达和训练。

结果

这个巨型模型在2048块 TPU v3 加速器上高效训练了4天^[1]。在从100种语言翻译成英语的任务上，它的质量远优于此前的方法^[1]。

成本核算和质量的成绩同样重要。作者给出的数字是:6000亿参数模型的训练成本为22个 TPU v3 核心年，而训练全部100个双语基线模型需要29个 TPU v3 核心年^[1]。也就是说，一个超大稀疏模型不仅质量更好，训练起来还比它所取代的那一批较小的稠密模型更便宜。

为什么重要

大模型研究常以参数量来叙述，但参数量本身说明不了模型究竟能不能训练出来。GShard 的贡献正对着这个缺口。一方面，它给出了从稀疏条件计算走到几千亿参数可用模型的具体路径，并公开报告了计算开销，使规模与训练预算之间的取舍可以被判断，而不必靠猜。另一方面，它把模型在加速器上的分布当作编译器问题，而不是改写代码的问题：分片通过注释表达，因此可以描述多种并行模式，而模型代码仍接近原来的样子^[1]。

仍然存在的取舍

稀疏路由换来了容量，也对系统提出了自己的要求。词元被送往不同的专家，而这些专家可能位于不同设备上，训练因此需要稠密模型所没有的芯片间通信；专家之间的负载需要管理，否则会出现一部分设备空转、另一部分满负荷的情况；而「每个词元只用部分网络」能带来多少好处，取决于底下的硬件与互连。论文给出的是在大规模下表达并训练这类模型的办法，它确立的这套模式——条件计算加自动分片——把这些工程问题留给了后续搭建者。