AI21 Labs 使用 Amazon EC2 P4d 实例 PyTorch 训练拥有 1780 亿个参数的语言模型

AI21 Labs 使用机器学习来开发专注于理解含义的语言模型，并在 2021 年将 Jurassic-1 Jumbo 设定为训练目标，这是近期发布的一种具有 1780 亿个参数的自回归语言模型。注册参加测试版测试的开发人员将获得 Jurassic-1 Jumbo 的访问权限，还可以立即开始根据使用场景自定义该模型。为了高效地训练模型，这家软件初创公司选择了 Amazon Web Services（AWS），并使用 Amazon Elastic Compute Cloud（Amazon EC2）构建了解决方案，Amazon EC2 是一项 Web 服务，可在云中提供安全、可调整大小的计算容量。由于选择了 Amazon EC2，该公司得以控制训练过程，包括节点分配。

为了获得强大的计算和联网功能，该公司选择了 Amazon EC2 P4d 实例，该实例在云中为机器学习训练和高性能计算应用程序提供高吞吐量和低延迟网络。AI21 Labs 使用 Amazon EC2 P4d 实例，通过 Jurassic-1 Jumbo 模型在数百个 GPU 上分发模型训练，提供自然语言处理服务，从而获得了所需的性能和内存。该公司现在训练并控制着自主开发的大型模型，因此公司能够开发同样规模的新模型、更轻松地进行创新。

AI21 Labs 成立于 2017 年，以多项使命为己任：进行自然语言处理研究，同时开发由人工智能提供支持的阅读和写作产品。该公司于 2020 年 10 月推出旗舰产品 Wordtune，这是一款智能写作和编辑助手，现已拥有近一百万用户。公司的另一主要产品 AI21 Studio 为该公司的 Jurassic-1 语言模型提供 API 访问以及自定义模型开发。AI21 Labs 联合创始人兼联合首席执行官 Yoav Shoham 说：“现在像我们这样提供语言模型即服务的公司并不多，不管是独立开发人员还是跨国企业，这些公司让所有人都能利用先进的自然语言处理技术构建应用程序和服务。”“此外，我们也在力求科学创新，以应对大型、复杂的模型在软件工程方面带来的挑战。”

AI21 Labs 需要强大的计算容量和高效的联网速度，还需要得到技术支持和指导，以便高效训练公司的首个深度学习超级模型，并支持该模型的高扩展和性能需求。出于这些原因，该公司于 2021 年初开始在 AWS 上实施解决方案，选择使用 Amazon EC2 P4d 实例来训练该模型。这些实例部署在称为 Amazon EC2 UltraClusters 的超大规模集群中，提供了四千多个 NVIDIA A100 GPU、PB 级非阻塞联网基础设施、以及高吞吐能力、低延迟存储。 

该公司的方法在低延迟、高带宽 GPUDirectRDMA 以及 Elastic Fabric Adapter（EFA）上得到了进一步优化，EFA 是 Amazon EC2 实例的网络接口，让客户可以在 AWS 上大规模运行需要高级别节点间通信的应用程序。由于该模型规模巨大，该团队需要使用并行处理来实现快速高效的训练，因此他们希望借助 AWS 上的联网功能支持其分布式训练和模型并行。AI21 Labs 的 Jurassic 技术主管 Opher Lieber 说：“Amazon EC2 P4d 实例在 EFA 上提供 400Gbps 的高性能联网。”“当 GPU 扩展到数百个时，GPU 到 GPU 的联网速度直接影响高效扩展同时保持成本高效的能力。”

首先，AI21 Labs 在为 EFA 激活的 Amazon EC2 P4d 实例上启动了代码库。然后，公司团队对多节点训练方法的性能和有效扩展进行了测试和验证。接下来，该团队启动了对最大模型（使用了数百个 GPU）的快速训练，目的是验证该训练方法的功能和性能。之后该团队就可以开始在 AWS 上训练 Jurassic-1 Jumbo 模型。在编排方面，该公司选择了一种内部解决方案，该解决方案使用 AWS 软件开发工具包——适用于 Python 的 AWS SDK（Boto3）来分配实例，可以轻松地将客户的 Python 应用程序、库或脚本与各种 AWS 服务集成。

在存储方面，AI21 Labs 选择了 Amazon Simple Storage Service（Amazon S3），该服务提供行业领先的可扩展性、数据可用性、安全性和性能。Lieber 说：“在 AWS 团队的帮助下，我们得以在 Amazon S3 上达到非常好的性能，因此不论是从性能还是价格方面考虑，AWS 都是当之无愧的选择。”该团队使用 Amazon S3 存储桶以分布式方式高效地存储和加载检查点。为了记录训练进度和事件，该团队使用了一项监控和可观测性服务 Amazon CloudWatch。 

在实施解决方案时，AI21 Labs 也从 AWS 获得了支持。公司团队咨询了 AWS 专家，这些专家就服务级别、和架构与硬件相关的问题和疑问提供了指导。此外，该公司使用 PyTorch on AWS 提高了 Jurassic-1 Jumbo 的性能。PyTorch 是一个开源深度学习框架，可以轻松开发机器学习模型并将模型部署到生产中。 

AI21 Labs 只用了几个月的时间就完成了训练，训练于 2021 年 6 月结束。这个新的超级模型是一种自回归语言模型，具有 1780 亿个参数，与该公司竞争对手的产品相当。与同类产品不同的是，该模型有 25.6 万个词汇量，能提供扩展的文本陈述功能，同时支持命名实体。该公司现在通过 AI21 Studio 产品提供 Jurassic-1 Jumbo（以及具有 70 亿个参数的同类产品 Jurassic-1 Large）的公开测试版。许多开发人员可以使用该服务在 Jurassic-1 Jumbo 模型上构建产品，许多行业也已经采用 AI21 Labs，包括营销、内容创作、游戏、医学研究、汽车以及电信和金融行业。

AI21 Labs 有自己的模型，而且可以直接访问模型，因此可以自主进行调整和创新，而不必依赖第三方，也可以不断探索创新目标，这也是该公司的重要使命。AI21 Labs 目前正在构建其他模型的原型，还计划对这些模型进行大规模训练。Shoham 说：“训练和拥有自主开发的超级模型仍然是 Wordtune 和 AI21 Studio 产品的关键差异化因素。”