100 万亿参数 AI 训练模型

推荐人工智能系统是当今互联网服务的重要组成部分：亚马逊和 Netflix 等数十亿美元的收入业务直接由推荐服务驱动。

AI 推荐器会随着规模的扩大而变得更好。 最近已经发布了几个模型，其参数高达十亿甚至万亿。 模型容量的每一次跳跃都带来了质量的显着提升。 100万亿参数的时代指日可待。

复杂、密集的静止神经网络的计算量越来越大，每次训练迭代都有超过 100 个 TFLOP。 因此，对于此类训练任务，拥有一些复杂的机制来管理具有异构资源的集群非常重要。

最近，来自苏黎世联邦理工学院的 Kwai Seattle AI Lab 和 DS3 Lab 合作提出了一个名为“Persia”的新系统，通过对训练算法和训练系统的仔细协同设计来解决这个问题。 在算法层面，Persia 采用混合训练算法，对嵌入层和密集神经网络模块进行不同的处理。 嵌入层异步训练以提高训练样本的吞吐量，而其余神经网络同步训练以保持统计效率。 在系统层面，已经实施了针对内存管理和通信减少的广泛系统优化​​，以释放混合算法的全部潜力。

100 万亿参数 AI 模型的云资源

Persia 100 万亿参数 AI 工作负载运行在以下异构资源上：

3,000 个计算密集型虚拟机核心
8 个 A2 虚拟机，共添加 64 个 A100 Nvidia GPU
30 台高内存虚拟机，每台具有 12 TB RAM，总计 360 TB
使用 Kubernetes 进行编排
所有资源必须在同一区域中同时启动，以最大限度地减少网络延迟。 Google Cloud 能够在很少通知的情况下提供所需的容量。

人工智能培训需要爆发式的资源。

Google Kubernetes Engine (GKE) 被用来协调 138 个虚拟机和软件容器的部署。 将工作负载容器化还允许培训的移植和可重复性。

结果和结论
在 Google Cloud 基础架构的支持下，该团队展示了 Persia 高达 100 万亿个参数的可扩展性。 混合分布式训练算法引入了复杂的系统松弛，以有效利用异构集群，同时收敛速度与普通 SGD 一样快。 谷歌云对于克服本地硬件的限制至关重要，并被证明是大规模分布式机器学习训练的最佳计算环境。

Persia 已在 github 上作为一个开源项目发布，其中包含 Google Cloud 的设置说明——学术界和工业界的每个人都会发现训练 100 万亿参数规模的深度学习推荐模型很容易。