在 Kubernetes 云平台上训练和部署生成式 AI 大模型

近年来，生成式 AI 凭借其卓越的技术在各个领域掀起了巨大的轰动。通过大模型，AI 能够创造出逼真的文本、图像、音频等作品，甚至涉足音乐、绘画等艺术领域。这项技术不仅在娱乐、创意产业中大放异彩，还在医疗、教育等领域展现出广阔的应用前景。生成式 AI 的蓬勃发展引发了人们对智能创造力、语言理解等问题的深入思考，成为人工智能领域的热门话题之一。 生成式 AI 备受开发者关注，他们不断探索大模型的性能和应用范围。通过开发创新的应用和解决方案，他们积极将生成式 AI 技术运用到各个领域。对于大多数开发者而言，他们期望能够利用大模型来迅速支持应用开发，而不是深入研究其训练和部署的细节。然而，也有一部分开发者可能对大模型的底层机制感兴趣，或者出于整个团队的 AI 开发需求，他们会深入学习以便进行架构设计和部署工作。 提及大模型，就不得不提到 Anthropic 的 Claude 3。一经推出，这款生成式 AI 大模型立即引起了轰动效应。在各项基准测试中，Claude 3 表现出色，超越了市面上许多同类产品。它在语言、数学、编程、知识问答、文本推理等多个领域的基准测试中都表现优异，已经达到了行业领先水平。 本文不深入讨论 Claude 3 的细节，但我想简要介绍一下它背后成功的底层支撑平台。除了其独特的算法之外，Claude 3 之所以能够获得如此出色的性能，还得益于其底层平台的支持。Anthropic 利用亚马逊云科技构建、训练和部署其拥有数十亿参数的大模型 Claude，这一过程高效可靠。Anthropic 采用了带有 Karpenter 自动伸缩功能的 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail)，并使用 [Amazon S3](https://aws.amazon.com/cn/s3/?trk=cndc-detail) 进行模型检查点的构建，同时还包括使用 Apache Spark 进行数据处理、模型训练和推理等应用。这一架构的效果非常明显，使得训练作业可扩展至数万个 Pod，并且通过 Karpenter 和 Spot 实例的使用，成本降低了 40%。这是 Kubernetes 支撑生成式 AI 的一个典型案例。更多的技术细节请参考下方二维码里的链接。  众所周知，[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)面临着诸多挑战，不仅限于当前生成式 AI 阶段，过去的深度学习和[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)也同样如此。其中一个重要挑战是依赖管理。由于[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)训练或推理环境通常需要大量的软件包，并且这些包之间存在着复杂的依赖关系，因此良好的依赖管理至关重要，以确保[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)环境的顺畅运行。此外，[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)通常需要进行分布式训练，无论是模型还是数据的分布式，都需要对数据或模型进行分块，并在不同的 GPU 上进行分布式训练，因此需要考虑到计算能力的调配和扩展问题。此外，[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)还涉及一系列安全性和合规性问题。 我们一直在努力解决这些问题。因此，早在多年前，我们就开始将容器应用于[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)环境。通过这一做法，我们能够更好地确保环境的一致性，包括训练和推理环境中的软件包、库以及依赖关系。将它们打包到相同的独立单元中，确保具有相同的配置，这样就能够保证整个容器化的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)环境有更好的一致性。与此同时，容器也是一种方便的部署方式，可以轻松部署在不同的云平台、本地环境、混合环境甚至边缘环境中。容器还能够更好地帮助进行资源隔离，从而使得[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)任务能够更好地执行。此外，容器还带来了诸如可扩展性、版本控制等优势，能够确保[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型和环境更加简单和透明。因此，近年来，我们一直将容器应用于[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)中。 Kubernetes 在[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)领域的应用日益广泛。作为一个容器的自动化编排工具，Kubernetes 能够根据负载动态扩展和缩减资源需求，帮助我们灵活地调整计算资源，以适应不同的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)任务。它提供了更灵活、更便捷的资源管理和调度功能，确保底层计算资源和存储资源能够充分应用于训练或推理环境中。Kubernetes 内置了多种容错机制，保障[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)的高可用性，并且支持多种环境，包括多云、混合云、本地甚至边缘环境。此外，Kubernetes 中的负载均衡等功能能够更方便地管理多个[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)任务，提升通信效率。另外，Kubernetes 作为一个庞大的生态系统，提供了许多工具、服务和软件的支持，包括存储、日志、可观测性和监控等，为[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)提供全方位的支持。这些是 Kubernetes 被广泛应用于[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)的重要原因之一。  因此，我们不禁思考：既然我们已经进入了生成式 AI 的时代，为什么不将 Kubernetes 应用于大模型的训练和部署呢？显然，答案是肯定的。我们可以看到越来越多的开发者和企业正在使用 Kubernetes 的基础设施来进行大模型的训练和部署。 在训练传统的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型和大模型之间存在许多不同之处。首先是数据规模，训练大模型需要大量数据，而传统[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型的训练集通常规模较小。当然，这并非绝对，但针对大模型的算法实现通常更加复杂，需要更多底层设施，如更多 GPU、存储和高速通信能力。因此，自动化配置和管理这些基础设施，使其顺利地应用于训练和推理工作负载，是一个非常值得考虑和研究的过程。其次，这涉及到分布式的问题：如何在节点之间进行数据和模型的分布式处理？在构建好分布式架构之后，如何让节点更好地协同，以实现模型的快速训练？在训练过程中，我们如何减少故障？由于训练大模型可能需要大量的 GPU 时间，因此如何快速发现和解决故障，或通过其他方式减少故障，也是需要考虑的重要问题。 对于亚马逊云科技而言，我们可以利用 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 来处理大规模的分布式训练工作负载，全面提升计算能力、存储网络和实例性能。首先，我们着眼于加速计算能力。亚马逊云科技与英伟达长期合作，提供最新的 EC2 GPU 实例。此外，我们的自研芯片也能有效加速训练和推理过程。我们还借助自动伸缩能力，实现对计算节点的快速调整，为计算提供更强大的支持。除了硬件层面，软件优化也至关重要。我们提供了专为 Kubernetes 和[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)优化的操作系统镜像和容器镜像，以最大程度地发挥软硬件协同优势。在存储方面，通过一系列驱动程序的优化，我们实现了存储和计算节点之间更高效、更高性能的互联，使得计算存储能够协同高速运行。在网络方面，我们不断提升性能，降低延迟，提供更高效的网络互联，以确保分布式训练的性能达到最佳水平。  接下来，我想分享一个优秀实践，这是我们官方博客上的一个案例，展示了如何利用 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 上的 Trainium 实例来训练 Llama2 大模型。Trainium 是亚马逊云科技自研的一款面向训练工作负载的芯片。在这个架构中，我们利用 Terraform 构建了 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 集群和托管 EKS 节点组，其中包含 Trainium 实例，为大模型的训练提供了可扩展的计算资源。通过构建容器镜像，我们将 Neuron SDK 集成到 Docker 容器中，并在 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 集群中部署这些容器，作为模型训练作业的运行环境。此外，我们使用 [Amazon FSx for Lustre](https://aws.amazon.com/cn/fsx/lustre/?trk=cndc-detail) 作为共享文件系统，用于存储数据集、tokenizer 文件、模型训练脚本、训练日志和模型检查点，确保训练作业中的数据共享和一致性。在开始训练之前，我们进行了前期准备和配置工作，包括下载 tokenizer 文件、准备数据集和修改训练脚本。通过 Kubernetes MPI Operator 协调分布式训练，每个工作 Pod 在单个 Trainium 实例上运行，并通过高速、低延迟的网络通信进行协作。我们利用 Tensorboard 和 neuron-top 等工具监控训练作业的进度和 Trainium 芯片的利用率，提供了对训练性能和资源利用的可视化。最后，我们编写了清理脚本，在完成训练后清理所有部署的资源，以节省成本并避免资源浪费。  架构的示意图如上，您可以通过扫描下方的二维码查看完整的博客。  对于大模型而言，推理阶段也面临着多种挑战，其中规模是首要挑战。与训练不同，我们更关注推理集群的自动扩展和缩减能力。确保推理集群能够根据用户需求实时扩展或缩小规模至关重要。例如，某领先的大模型提供商，用户的使用量可能会在不同时间出现差异，在去年会频繁遇到推理容量不足，无法为用户提供服务的情况。因此，如何根据后台推理的规模实现自动扩展和缩减，以提供更多有效的推理资源，是一个需要解决的核心问题。将大模型部署在云端会更便捷、更快速地解决这一问题，提供更好的响应能力。另外，性能也是至关重要的考量因素。当然，性能和响应时间本身可能存在一种妥协，需要根据实际应用场景确定。另一个同样重要的因素是成本。有研究表明，[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)中大约 90% 的成本都用于推理阶段。因此，如何优化这一成本是非常需要考虑的问题。 我想提到的是亚马逊云科技推出的开源项目 [Karpenter](https://karpenter.sh/?trk=cndc-detail)，这是一个针对 Kubernetes 的自动化伸缩工具。Karpenter 旨在帮助开发者和企业根据不同的工作负载实现弹性扩展和伸缩。它提供了一整套自动伸缩组和用户自定义的策略，通过这些策略的组合来自动调整集群中节点的数量，从而更有效地利用资源，确保 Kubernetes 集群具有足够的容量来运行工作负载。虽然 Karpenter 最初并不是专门为生成式 AI 或[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)工作负载而构建的，但在实践中，我们发现它能够更好地支持生成式 AI 工作负载，特别是推理工作负载的弹性伸缩。这不仅提升了性能，降低了成本，也使我们逐渐将 Karpenter 应用于[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)环境中。它能够自动化节点管理，相较于传统的 auto scaler 更为智能，能够进行更准确的预测和伸缩。同时，它提供了更强大的弹性展开和成本优化能力，帮助开发者更好地在容器上部署[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)推理。  亚马逊云科技提出了名为 JARK 的综合解决方案，用于端到端的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)编排。JARK 堆栈包含四个核心组件：JupyterHub、Argo Workflows、Ray 和 Kubernetes。这个堆栈是在 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 上运行 Gen AI 模型的关键组成部分。 - JupyterHub：JupyterHub 提供了一个共享平台，可在业务、教育和研究中运行流行的笔记本。在 Gen AI 领域，JupyterHub 加速了实验过程，特别是在反馈循环中。数据工程师可以在此平台上协作创建用于 Prompt Engineering 的模型。 - Argo Workflows：Argo Workflows 是一个开源的容器本地工作流引擎，用于在 Kubernetes 上编排并行作业。它提供了一个结构化和自动化的流水线，专门用于模型的微调。 - Ray：Ray 是一个开源的分布式计算框架，可以轻松扩展应用程序并使用最先进的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)库。Ray 用于在多个节点上分布训练生成模型，加速训练过程并处理更大的数据集。 - Kubernetes：Kubernetes 是一个强大的容器编排平台，可自动化容器化应用程序的部署、扩展和管理。在 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 上运行 Gen AI 模型的基础设施建立在 Kubernetes 之上。  我再分享一个来自亚马逊云科技官方博客的案例，展示了如何在 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 上部署 Stable Diffusion 模型。使用 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 部署 Stable Diffusion 模型的步骤大致如下：首先，创建 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 集群，可使用 Terraform 或亚马逊云科技管理控制台完成，并确保网络和安全组设置正确。接着，配置 Kubernetes 命令行工具（kubectl），与 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 集群进行交互，包括应用程序部署和工作负载管理等。然后，部署必要的基础组件到 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 集群中，如 JupyterHub、Argo Workflows 和 Ray，这些组件构成了 Stable Diffusion 模型部署的基础设施。随后，配置 Ingress 控制器和负载均衡器，将外部流量引导到 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 中的服务。接下来，在 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 中创建 GPU 节点组，用于托管需要 GPU 加速的模型训练和推理任务。然后，在 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 集群中部署 Stable Diffusion 模型，可利用 Hugging Face 的 Accelerate 和 Diffusers 库来简化模型部署和优化训练过程。最后，使用工具如 Ray Serve 配置模型服务，实现模型的在线推理功能，并对模型服务进行验证和优化。  架构的示意图如上，您可通过扫描下方的二维码获取完整博客，进一步了解细节。  总结来说，在 Kubernetes 平台上，尤其是在亚马逊云科技的 [Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail) 上，我们已经建立了完善的数据和生成式 AI 生态系统，快速支持生成式 AI 的训练和推理。我们提供了示例代码、优秀实践和基础设施即代码（IAC）模板等资源，欢迎大家学习和探索。同时，我们也欢迎大家就我们的代码提出更多修改和改进意见：https://github.com/awslabs/data-on-eks?trk=cndc-detail **前述特定亚马逊云科技生成式人工智能相关的服务仅在亚马逊云科技海外区域可用，亚马逊云科技中国仅为帮助您了解行业前沿技术和发展海外业务选择推介该服务。*