[论文解读] PanGu-Σ: Towards Trillion Parameter Language Model with Sparse Heterogeneous Computing

Xiaozhe Ren, Pingyi Zhou+15 more|arXiv (Cornell University)|Mar 20, 2023

Topic Modeling被引用 7

一句话总结

PanGu-Σ 是一个在 3290 亿个 token 上使用 512 个 Ascend 910 AI 加速器训练的 1.085 万亿参数稀疏专家混合模型（MoE）。它引入了随机路由专家（RRE）以实现高效的专家路由，并采用专家计算与存储分离（ECSS）技术，使训练吞吐量提升 6.3 倍，优于先前模型在零样本中文 NLP 任务和微调应用中的表现。

ABSTRACT

The scaling of large language models has greatly improved natural language understanding, generation, and reasoning. In this work, we develop a system that trained a trillion-parameter language model on a cluster of Ascend 910 AI processors and MindSpore framework, and present the language model with 1.085T parameters named PanGu-Σ. With parameter inherent from PanGu-α, we extend the dense Transformer model to sparse one with Random Routed Experts (RRE), and efficiently train the model over 329B tokens by using Expert Computation and Storage Separation(ECSS). This resulted in a 6.3x increase in training throughput through heterogeneous computing. Our experimental findings show that PanGu-Σ provides state-of-the-art performance in zero-shot learning of various Chinese NLP downstream tasks. Moreover, it demonstrates strong abilities when fine-tuned in application data of open-domain dialogue, question answering, machine translation and code generation.

研究动机与目标

• 在保持高训练效率和系统吞吐量的前提下，将大语言模型扩展至超过一万亿参数。
• 通过一种新型路由机制，解决稀疏 MoE 模型中工作负载不平衡和高通信延迟的挑战。
• 通过系统级分离专家计算与存储，减少主机-设备通信和优化器开销，实现分布式训练的高效性。
• 在零样本和微调下游 NLP 任务中实现最先进性能，尤其在中文和多模态应用中表现突出。
• 在硬件资源有限的条件下，利用异构计算实现万亿参数模型的高效、可扩展且成本可控的训练。

提出的方法

• 提出随机路由专家（RRE），一种两级路由机制：首先按领域/任务对专家分组，然后在每组内随机且均匀地将 token 分配给专家，无需可学习的门控函数。
• 引入专家计算与存储分离（ECSS），将专家计算与存储解耦，以减少主机到设备及设备到主机的通信开销以及优化器计算开销。
• 在 MindSpore 深度学习框架中结合张量并行、流水线并行和数据并行的混合并行策略，实现跨 512 个 Ascend 910 加速器的训练扩展。
• 通过复用其参数并使用 3290 亿个 token 在 40 多种语言上进行微调，将密集的 PanGu-α 模型扩展为稀疏 MoE 架构。
• 结合模型并行与系统级优化，在模型规模巨大的情况下仍保持高训练吞吐量。
• 在 SuperGLUE 和中文 NLP 基准测试中采用零样本评估协议，以评估模型的泛化能力和推理能力。

实验结果

研究问题

• RQ1 能否在一个由 512 个 Ascend 910 加速器组成的有限集群上，以高吞吐量和低通信开销高效训练一个万亿参数的稀疏 MoE 模型？
• RQ2 与具有可学习门控的常规 MoE 相比，随机路由专家（RRE）机制在训练效率和模型性能方面表现如何？
• RQ3 专家计算与存储分离（ECSS）在多大程度上减少了主机-设备通信和优化器更新等系统级瓶颈？
• RQ4 PanGu-Σ 在多样化的 NLP 任务中，特别是在低资源或少样本设置下的零样本和少样本性能如何？
• RQ5 通过从密集基础模型共享参数构建的稀疏 MoE 模型，是否能在包括对话、翻译和代码生成在内的多个下游任务中实现最先进性能？

主要发现

• PanGu-Σ 在 512 个 Ascend 910 加速器上实现了每秒 69,905 个 token 的训练吞吐量，相较于具有相同超参数的 MoE 模型提升了 6.3 倍。
• 该模型在中文 NLP 基准测试中展现出最先进水平的零样本性能，在多项任务（如文本分类和问答）中优于先前的 SOTA 模型。
• PanGu-Σ 的英文子模型（380 亿参数）在零样本设置下的 SuperGLUE 平均得分为 64.62，尽管仅使用了 1120 亿个英文 token，已接近 GPT-3 13B 模型的 57.2 平均分。
• PanGu-Σ 的微调应用在开放域对话、机器翻译、问答和代码生成任务中表现强劲，显示出优异的迁移能力。
• RRE 机制支持对子模型进行有效提取，适用于对话或代码生成等特定任务，提升模型适应性而无需重新训练。
• ECSS 显著减少了主机-设备通信和优化器计算开销，为 6.3 倍吞吐量提升做出贡献，并使在受限硬件集群上的高效训练成为可能。