GShard：大模型训练的自动分片与条件计算革新

从 6000 亿参数多语言翻译模型出发，揭秘 GShard 如何破解大模型规模化训练的计算、并行与编程难题

一、引言

1.1 研究背景

神经网络规模化已成为提升模型效果的核心路径，在计算机视觉、自然语言处理等领域，模型质量与训练数据、计算量、参数量呈显著的幂律关系。但随着模型参数量突破千亿甚至万亿级别，三大核心挑战成为规模化的拦路虎：

• 计算成本超线性增长：模型深度 / 宽度的直接缩放导致训练步长时间激增，模型并行带来的设备通信开销进一步加剧问题；
• 并行部署难度大：主流深度学习框架缺乏通用高效的模型并行支持，朴素的图划分策略造成设备利用率极低；
• 编程门槛高：手动实现模型分片需协调多设备通信，架构微小变更会引发通信逻辑的连锁修改，对开发者工程能力要求极高。

与此同时，传统密集型模型的规模化已触碰到效率天花板 —— 更大的参数量意味着更高的计算与内存成本，且无法适配多任务、低资源的实际业务场景。在此背景下，Google 提出 GShard 框架，通过 条件计算 与 自动分片 的融合创新，为大模型训练提供了一套高效、通用、低门槛的解决方案。

1.2 论文核心贡献

GShard 是由轻量级标注 API 和 XLA 编译器扩展组成的大模型训练框架，其核心论文《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》的核心贡献可概括为四大突破：

• 架构创新：将稀疏门控混合专家（MoE）与 Transformer 结合，提出 MoE Transformer 架构，实现计算成本与模型容量的亚线性缩放，让千亿级模型训练的计算开销可控；
• 编程抽象：设计轻量级的张量分片标注 API，实现模型描述与并行实现的解耦，开发者仅需少量标注即可实现大模型的并行化，大幅降低编程门槛；
• 编译优化：提出基于 SPMD（单程序多数据）的自动分片编译策略，使编译时间与设备数量无关（O (1)），支持数千台加速器的高效并行；
• 实践验证：基于 GShard 训练出6000 亿参数的多语言机器翻译模型，在 2048 个 TPU v3 加速器上仅用 4 天完成训练，在 100 种语言到英语的翻译任务中实现远超基线模型的效果，且计算成本远低于传统密集型模型。

---

二、核心概念铺垫

术语名词总结

术语 / 概念	核心定义
条件计算（Conditional Computation）	让模型根据输入激活不同的子网络，而非全量网络参与计算，实现计算资源的按需分配，核心是“稀疏激活”
混合专家（MoE, Mixture-of-Experts）	由多个独立的前馈网络（专家）和一个门控函数组成，门控函数为每个输入选择少量专家参与计算，是条件计算的经典实现
稀疏门控MoE	门控函数仅为每个输入选择1~2个专家，大幅降低计算量，同时通过负载均衡机制避免少数专家过载
自动分片（Automatic Sharding）	由编译器根据开发者的张量标注，自动将模型的计算和参数分配到多个设备上，无需手动实现设备通信逻辑
SPMD（单程序多数据）	所有设备运行同一个程序，处理不同的数据分片，相比MPMD（多程序多数据）大幅降低编译复杂度，编译时间与设备数无关
XLA	针对TensorFlow的优化编译器，可将张量计算转换为高效的机器码，GShard基于XLA实现编译层的自动分片扩展
通信原语	多设备并行的基础通信操作，GShard核心包含AllReduce、AllGather、AllToAll、CollectivePermute，支撑设备间高效数据交互

---

三、GShard 核心优化方案

GShard 的核心设计围绕亚线性缩放、抽象解耦、编译器可扩展三大原则展开，由MoE Transformer 架构、轻量级分片标注 API、XLA SPMD 自动分片编译器三部分组成，三者相互配合形成完整的大模型训练体系。

3.1 MoE Transformer：亚线性缩放的架构基础

GShard 的核心架构创新是将稀疏门控 MoE 层与经典 Transformer 结合，替换 Transformer 编码器 / 解码器中的隔层前馈层，实现计算成本与模型容量的亚线性缩放，这是千亿级模型高效训练的算法基础。
MoE 层核心结构：由 E 个独立的前馈网络（专家）和一个门控函数组成，每个输入 token 经门控函数选择最多 2 个专家参与计算，最终输出为所选专家输出的加权和，未被选择的专家不参与计算，大幅降低单步计算量；
负载均衡关键机制：为解决稀疏门控中 “少数专家过载、多数专家闲置” 的问题，GShard 设计了四大核心机制：

• 专家容量限制：为每个专家设置 token 处理阈值，超出阈值的 token 通过残差连接传递，避免单专家负载过高；
• 局部分组调度：将训练批次的 token 均匀分组并并行处理，为每组分配专家的部分容量，保证整体负载均衡；
• 辅助损失：引入辅助损失项加入总损失函数，约束门控函数的选择行为，防止其始终选择少数专家；
• 随机路由：对次优专家的选择加入概率机制，仅当其次优权重满足一定条件时才分发 token，节省专家容量；
• 亚线性缩放特性：MoE Transformer 的参数量随专家数线性增长，但每个 token 仅激活少量专家，单步计算量几乎与专家数无关，实现了 “参数量线性增长、计算量亚线性增长” 的核心目标。

3.2 轻量级分片标注 API：模型与并行的解耦

GShard 设计了一套极简的张量分片标注 API，让开发者无需关注底层的设备通信和分片逻辑，仅需对模型中的关键张量进行少量标注，即可实现自动并行化，核心是让模型开发者聚焦架构设计，将并行实现交给编译器。框架提供 3 类核心 API，基于 TensorFlow/Lingvo 实现，操作简单且不改变张量的逻辑形状：

• replicate(tensor)：将张量在所有设备上复制，适用于非 MoE 层的权重、注意力层参数等，保证所有设备持有相同的参数副本；
• split(tensor, split_dimension, num_partitions)：沿指定维度将张量分片到指定数量的设备，是最常用的分片 API，适用于 MoE 层的专家参数、批次数据等；
• shard(tensor, device_assignment)：通用分片 API，支持多维度分片和自定义设备映射，可根据设备拓扑优化分片策略，减少设备间通信开销。
  此外，GShard 的编译器具备自动分片推断能力—— 开发者仅需标注少量关键张量（如输入、输出、核心权重），编译器通过迭代数据流分析，自动推断其余张量的分片策略，最小化重分片操作，进一步降低编程门槛。

3.3 XLA SPMD 自动分片编译器：高效并行的编译保障

GShard 的核心编译层创新是基于 XLA 实现SPMD 自动分片编译器，摒弃传统的 MPMD 编译策略，解决了大模型并行中 “编译时间随设备数线性增长” 的问题，同时实现了核心算子的高效分片和设备间的低开销通信。
SPMD 编译核心优势：所有设备运行同一个程序，处理不同的张量分片，编译器仅需编译一次，编译时间与设备数量无关（O (1)），即使在 2048 个 TPU 的场景下，编译时间仍保持恒定；
核心通信原语：编译器内置四类高效的分布式通信原语，适配大模型训练的各类通信场景，且针对 TPU 集群做了深度优化：

• AllReduce：实现设备间张量的元素级归约，用于梯度聚合、部分计算结果融合，成本与设备数无关；
• AllGather：将各设备的张量分片拼接为完整张量，用于将分片张量转换为复制张量；
• AllToAll：实现设备间张量的双向分片交换，用于 MoE 层的 token 分发和结果融合，成本呈 O (√D) 亚线性增长；
• CollectivePermute：实现设备间的一对一数据传输，用于光晕交换、设备拓扑优化等场景；
• 通用算子分片支持：编译器实现了对 Einsum、卷积、矩阵乘法、归约等核心算子的自动分片，支持非均匀分片、静态算子配置、光晕交换等复杂场景，保证分片后计算的正确性和高效性，不仅适用于 NLP 的 Transformer 模型，还可扩展到计算机视觉的卷积模型。

3.4 MoE 层的线性代数高效实现

为让 MoE 层适配 TPU 的硬件优化特性，GShard 将 MoE 层的核心逻辑通过爱因斯坦求和（tf.einsum） 实现线性代数表达，把门控计算、token 分发、专家前馈、结果融合等步骤转化为高效的张量操作，大幅提升硬件利用率。同时，GShard 对 MoE 层的计算复杂度做了严格分析，在合理的假设下（单设备 token 数 O (1)、专家数 E=O (D)、分组数 G=O (D)），单设备计算复杂度为 O (1)，跨设备通信成本仅呈 O (√D) 增长，确保了框架在数千台设备上的可扩展性。

---

四、实验结果与性能分析

GShard 的论文以大规模多语言机器翻译（M4） 为核心实验任务，验证框架的有效性 —— 目标是训练一个单模型支持100 种语言到英语的翻译，使用 web 规模的平行语料（250 亿训练示例），对比传统密集型模型和双语基线模型，从翻译质量、训练效率、内存 / 运行时性能三个维度做了全面验证。

4.1 实验基础配置

硬件平台：Google TPU v3 加速器，实验采用 128/512/2048 个核心的集群配置；
模型配置：训练了多款不同深度（12/36/60 层）、不同专家数（128/512/2048）的 MoE Transformer 模型，核心模型为6000 亿参数（36 层 + 2048 专家），同时训练了 23 亿参数的 96 层密集型 Transformer（GPipe 流水线并行）作为基线；
评价指标：翻译质量采用 ∆BLEU（相对双语基线模型的 BLEU 值提升），训练效率采用TPU 核心年（核心数 × 训练时间）、墙钟时间，性能分析采用单设备内存消耗、训练步长时间、硬件利用率。

4.2 核心指标对比

4.2.1 翻译质量：规模化带来显著提升

模型质量随参数量、深度、专家数的提升呈显著增长趋势，核心结论如下：
6000 亿参数模型表现最优：平均∆BLEU 达 13.5，远超密集型基线模型的 6.1，在高资源和低资源语言上均实现大幅提升；
深度提升的正向影响：相同专家数下，36 层模型的平均∆BLEU 比 12 层模型高 2~3 个点，且低资源语言的提升更为显著，体现了深度对跨语言正迁移的促进作用；
容量瓶颈的缓解：专家数从 128 增加到 512 时，模型∆BLEU 跃升 3.3 个点，512 到 2048 时提升放缓（1.3 个点），验证了 128~512 专家是该任务的容量瓶颈区间，专家数的增加有效缓解了多任务的容量竞争。

4.2.2 训练效率：亚线性缩放实现成本可控

GShard 的 MoE Transformer 在训练效率上实现了对传统密集型模型的碾压，核心数据如下：
6000 亿参数模型训练仅 4 天：在 2048 个 TPU v3 核心上，处理 1 万亿 token 仅需 4 天，总计算成本为22 个 TPU 核心年；
远低于基线模型成本：100 个双语基线模型的训练成本为 29 个 TPU 核心年，同硬件下的密集型基线模型成本高达 235.5 个 TPU 核心年，是 GShard 模型的 10 倍以上；
样本效率更高：36 层模型达到相同训练损失所需的 token 数，仅为 12 层模型的 1/2~1/3，深度提升显著提高了样本利用率，减少了训练数据的消耗。

4.2.3 内存与运行时性能：高效可扩展

GShard 在硬件利用率和可扩展性上表现优异，核心性能指标如下：

• 内存消耗 O (1)：单设备的内存消耗（权重 + 激活）随设备数 / 专家数的增加保持恒定，仅随模型层数线性增长；当内存超出设备容量时，编译器自动重计算激活值，重计算开销仅 28%~34%；
• 步长时间亚线性增长：模型规模从 128 专家（375 亿参数）提升 16 倍至 2048 专家（6000 亿参数），训练步长时间仅增加 1.7 倍，仍能达到理论性能的 48%；
• 硬件利用率高：Transformer 层和 MoE 前馈层的硬件利用率超 85%，即使是受内存带宽限制的注意力层，利用率也达到 30% 以上；
• 通信开销可控：核心通信原语 AllReduce 成本与设备数无关，AllToAll 成本呈 O (√D) 亚线性增长，2048 设备相比 16 设备，AllToAll 时间仅增加 9 倍。

4.3 关键结论

从实验结果可得出三大核心结论，验证了 GShard 框架的有效性和优越性：
MoE + 条件计算是大模型规模化的高效路径：通过稀疏激活实现计算成本的亚线性缩放，让千亿级模型的训练成本远低于传统密集型模型；
自动分片大幅提升并行效率：SPMD 编译策略让编译时间与设备数无关，轻量级 API 降低了编程门槛，使数千台加速器的高效并行成为可能；
模型深度和专家数的提升具有互补效应：深度提升促进跨语言正迁移，利好低资源语言；专家数增加缓解容量瓶颈，利好高资源语言，二者结合实现全场景的翻译质量提升。

---

五、产业落地与实际价值

GShard 作为 Google 推出的大模型训练框架，不仅在学术上实现了技术突破，更在产业落地中展现出巨大的实际价值，为大模型的工业化训练提供了可复制的方案，其价值主要体现在技术、成本、业务三个维度：

5.1 技术维度：推动大模型训练的通用化与低门槛化

GShard 打破了 “大模型训练仅能由少数大厂实现” 的技术壁垒，通过模型与并行的解耦，让普通开发者无需深入掌握分布式并行原理，仅需少量标注即可实现千亿级模型的训练。同时，框架的通用算子分片支持，使其不仅适用于 NLP 的 Transformer 模型，还可扩展到计算机视觉、语音识别等领域的深度学习模型，成为大模型训练的通用化框架。

5.2 成本维度：大幅降低大模型的训练与部署成本

GShard 的亚线性缩放特性和高效并行能力，让大模型的训练成本大幅降低：一方面，稀疏激活减少了单步计算量，降低了硬件的计算开销；另一方面，自动分片优化了设备资源利用率，减少了硬件的投入规模。此外，GShard 训练的单模型可支持多任务、多语言，避免了为每个任务单独训练模型，大幅降低了模型的研发和维护成本。

5.3 业务维度：赋能多语言、多任务的实际业务场景

GShard 的核心验证任务是100 种语言到英语的多语言翻译，这一场景直接对接了跨境电商、国际社交、机器翻译等实际业务需求。框架训练的单模型不仅能支持多语言翻译，还能通过跨语言正迁移，提升低资源语言的翻译质量，解决了传统双语模型在低资源语言上效果差的行业痛点。此外，MoE Transformer 的多任务适配能力，可轻松扩展到文本摘要、情感分析、对话生成等 NLP 业务场景，为大模型的产业落地提供了高效的模型基础。

5.4 产业影响：引领大模型训练的技术趋势

GShard 提出的条件计算 + 自动分片组合，成为后续大模型训练框架的重要参考方向，推动了大模型训练技术的两大趋势：
稀疏化成为大模型的核心设计原则：从密集型模型向稀疏型模型转变，通过条件计算、稀疏注意力等方式，实现计算成本的亚线性缩放；
编译层优化成为并行的核心发力点：从手动并行、框架层并行向编译层自动并行转变，利用编译器的能力实现张量的智能分片和设备的高效通信。

---

六、总结与思考

6.1 核心知识点回顾

GShard 作为大模型训练的经典框架，其核心价值在于通过条件计算和自动分片的融合创新，破解了大模型规模化训练的三大核心挑战，本文的核心知识点可总结为一个核心、三大组件、四大优势：

• 一个核心：以亚线性缩放为核心设计原则，让模型参数量线性增长的同时，计算成本仅亚线性增长，这是千亿级模型高效训练的关键；
• 三大组件：MoE Transformer（架构层，实现稀疏激活）、轻量级分片标注 API（编程层，实现模型与并行解耦）、XLA SPMD 编译器（编译层，实现自动分片与高效并行）；
• 四大优势：翻译质量显著提升、训练成本大幅降低、编程门槛低、框架通用性强。
  此外，GShard 的实验验证了模型深度和专家数对多语言任务的互补效应，以及条件计算在提升样本效率、缓解容量瓶颈上的核心作用，为后续大模型的架构设计提供了重要参考。

6.2 延伸思考：GShard 的局限性与未来发展方向

GShard 作为 2020 年的经典工作，虽实现了技术突破，但在大模型技术快速发展的今天，仍存在一些局限性，同时也为后续研究指明了发展方向：
局限性

• MoE 层的训练不稳定性：稀疏门控 MoE 层存在训练震荡、专家利用率不均等问题，需要精细的超参数调优，且在小批量训练场景下效果不佳；
• 硬件适配的局限性：GShard 基于 Google TPU 和 XLA 实现，对 GPU 的适配性较差，而工业界主流的大模型训练硬件以 GPU 为主；
• 推理效率问题：MoE 模型的推理需要进行门控计算和专家选择，相比密集型模型存在额外的开销，且稀疏激活的特性难以利用硬件的稠密计算优化。
  未来发展方向
• MoE 架构的优化：研究更高效的门控函数设计、专家路由策略，提升 MoE 层的训练稳定性和专家利用率，如动态专家数、自适应路由等；
• 跨硬件的通用适配：基于 PyTorch、CUDA 等工业界主流框架，实现 GShard 的核心思想，提升框架对 GPU 集群的适配能力；
• 训推一体化优化：针对 MoE 模型的推理特性，做专门的编译优化和硬件加速，如专家重排序、稀疏计算加速等，提升 MoE 模型的推理效率；
• 与其他技术的融合：将 GShard 的自动分片与 ZeRO 的内存优化、3D 并行的混合并行策略结合，实现更高效的大模型训练，突破万亿级参数的训练瓶颈。

总的来说，GShard 不仅为大模型训练提供了一套高效的技术方案，更重要的是提出了 “稀疏化 + 自动化” 的大模型规模化思路，这一思路至今仍在引领大模型训练技术的发展，为后续的大模型框架研发奠定了重要的理论和实践基础。

---

参考资料

1. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding