用于加速基础模型创新的云原生开源软件栈

作者：

TaliaGershon,IBM研究院混合云架构研究总监

PriyaNagpurkar,IBM研究院混合云平台研发副总裁

CarlosCosta,IBM研究院基础模型和多云业务首席研究科学家

DarrellReimer,IBM研究院杰出工程师

本博文将介绍在RedHatOpenShiftContainerPlatform上运行的高性能AI训练堆栈，该平台是新推出的的构建基础。

作为训练堆栈的补充，我们的技术栈可降低成本与优化性能，对基础模型进行调整并为其提供服务。以下多项技术已贡献给开源社区在开源社区中发布，如PyTorch、Ray、Kserve和 OpenDataHub (ODH)（一个可在Kubernetes上构建、部署和管理数据密集型应用的开源平台）。在ODH上发展成熟的技术随后被引入RedHat的OpenShiftAI，IBM目前的新一代AI平台watsonx.ai使用RedHatOpenShiftAI。借助这种方法，IBM和RedHat为客户提供最先进的开源基础模型栈，以支持在他们选择的任何环境（本地、IBMCloud或其他公有云）中运行。

我们的基础模型训练方法

我们在开始重构AI训练堆栈时有两个目标。首先，一是希望保持传统HPC系统的实用性：最大程限度提高硬件利用率，并高效利用高性能基础架构。其次，二是希望提供混合云开发体验的灵活性和生产力优势：提高开发敏捷性和代码重用率，简化基础架构和软件的管理与扩展。为了实现第二个目标，我们使用Kubernetes构建解决方案，通过容器复用代码以及扩展软件。但这一决定意味着我们需要将Kubernetes转化为可处理高性能工作负载的平台。

我们还需要一个可以应对每一步AI训练工作流程的解决方案：进行数据预处理、分布式训练和模型验证。我们确定了合作的关键开源社区，共同处理端到端工作流程，确定了需要克服的关键用户体验障碍，以便用户启动、运行和扩展工作。

下图1的左侧介绍了我们的训练软件堆栈的整体情况，该堆栈自2022年底以来一直在IBMCloud的Vela上运行，由IBM研究院使用。图1右侧描述了我们用于调整和提供基础模型的堆栈，稍后将在博客中详细讨论。

高效、可扩展的数据预处理，模型训练与验证

为了使AI管道中的所有运行步骤简单高效，我们主要运用和贡献PyTorch和Ray这两项关键的开源技术，并做出自己的贡献。借助Ray，我们使用数据科学家友好的PythonAPI实现了可扩展的数据预处理（例如使用仇恨、滥用和亵渎过滤器过滤数据）和后处理步骤（例如模型微调和验证）。通过运行RaywithMCAD，我们支持通过并发运行的异构Ray作业高效共享资源池。

我们正与PyTorch合作推进对分布式训练的支持，包括通过引入rate_limiter来改进对完全分片数据并行(FSDP)训练API的支持。我们最近展演示了在基于以太网的环境（如IBMCloud中的Vela）上为参数为10B+的模型高效扩展分布式训练作业。此外，通过集成MCAD与TorchX（TorchX是PyTorch应用的通用作业启动器），我们能够使用不同API和框架以透明方式支持多种基于PyTorch的作业。这些工作都受底层工作管理系统的支持，无需AI从业者修改代码。

简化用户体验

工作流程本身的训练分为以下三个步骤：模型探索（通常是使用一些GPU运行的缩小试验）、扩展分布式训练作业（消耗数百个GPU）以及模型验证。对于许多AI从业者来说，编排这些步骤可能很复杂，而且配置和管理的时间长。我们通过CodeFlare项目解决了这一挑战，它能够提供引导式、简化的用户体验，以有效地提高训练、测试和监控模型训练生命周期的效率。

CodeFlareCLI（基于控制台和UI）指导用户应对复杂的远程OpenShift集群运行，同时自动执行作业配置、存储设置、日志记录以及管理用于监控和分析的端点。CodeFlareSDK（基于Jupyter）为用户提供直观的Python界面，进行批处理资源请求处理、作业提交和观察。借助这些功能，我们大大降低了AI研究同事进入云原生堆栈的门槛。

在Vela上运行我们的软件栈

到2022年底，IBM的所有基础模型训练工作都过渡到在IBMCloud的Vela上运行此软件栈。如今，MCAD管理这些AI作业的队列，从单GPU作业到利用超过512个GPU的作业，并处理作业优先级和配额管理。我们在过程中发现了其他方法，可以让团队在以GPU为中心的环境中（如Vela）中管理OpenShift集群更轻松，例如，通过增强OpenShiftInstaller配置基础架构(IPI)，以便更轻松地在高性能基础架构上部署和管理OpenShift。

我们的基础模型调优和服务方法

训练和验证最先进的基础模型是AI价值链的关键早期阶段，但当模型在AI工作流程的调优和推理步骤中投入生产使用时，最终会捕获真正的价值。我们用于推理和模型调优的软件栈主要是在底层硬件上高效执行模型，以最佳方式批处理传入的请求，简化AI与应用的集成，并提供最先进的模型适应技术。上图1的右侧介绍了我们的基础模型调优和服务栈，以下是更为详细的描述。

推理性能

用于优化基础模型在给定硬件平台上运行方式的软件库可以将吞吐量和延迟提高10-100倍。我们的服务软件栈中有一组精心设计的成熟优化路径（包括ONNX和HuggingFaceOptimum），用于推理常见模型架构，并且可扩展，以能够适应新的推理服务器或优化。考虑到AI和开源社区的快速创新步伐，可扩展性是我们软件栈的关键设计点。此外，真正的AI服务会同时收到来自多个用户的、针对多个模型的大量推理请求。我们的服务堆栈动态批处理传入的请求，并通过构建和反馈HuggingFace、Kserve和ModelMesh社区，有效地在模型之间进行多路复用。

简化应用程序集成

目前可用于运行AI模型的推理服务器需要用户具备大量的AI专业知识。模型的输入和输出都是张量。对于希望利用这些模型完成任务的应用开发人员来说，这种格式并不易于理解。为了使此过程对开发人员更加友好，必须将模型输出转换为更易于使用的内容。我们创建了Caikit抽象层，可为应用开发人员提供直观的API和数据模型，并提供一个稳定的接口，支持模型和应用独立发展。这一抽象层在IBM的Watson模型中使用，用于提供基础架构，并将很快贡献给开源很快将作为开源项目提供。

基础模型调优整

基础模型的关键价值主张之一是能够利用预训练的基础模型，并使用专用数据“调优整”或“适应”，以提高其下游任务的性能。我们的目标是打包最先进的技术，以实现计算效率高的模型自适应，让不甚了解其工作原理的人也能够使用。我们的可扩展堆栈目前支持多任务提示调优(MPT) 和微调，通过开源项目参数高效微调(PEFT)集成。在接下来的几个月里，我们将开源一些提示调优算法和实现。

使用和RedHat合作

IBM研究院正在与RedHat合作，我们将开发的功能贡献给主要开源社区，直接发布在OpenDataHub(ODH)，让其他人也从这项工作中受益。ODH是一个全面的开源工具集合，旨在利用OpenShift的优势来推动整个AI开发生命周期。OpenDataHub中引入的许多技术已发展成为RedHatOpenShiftAI的一部分，并作为watsonx.ai的中间件基础提供服务。图2显示了本博客中描述的对各种开源贡献将如何汇集到ODH中以支持基础模型用例。

后续工作

面向基础模型时代重构我们的端到端软件栈对我们的AI社区具有相当大的价值。AI研究人员不再需要非常深入的基础架构知识，就可以让作业高性能运行。他们不再需要弄清楚如何将作业从几个GPU扩展到数百个，或者如何精确地分配作业以实现高工作负载性能，现在这些任务都可以交给软件栈处理。代码可以跨团队得到重新使用，试验也易于重现复制。我们还大幅简化了AI开发人员如何以高计算效率和开发人员友好的方式提供和调整基础模型的方式。

或许最重要的是最重要的一点可以说是，在OpenShift上构建此堆栈提供了对其他环境的可移植性，让合作伙伴可以在本地和任何公有云中使用这些功能。我们很高兴与红帽一起通过开放数据中心将这些创新引入开源社区我们很高兴与RedHat开展合作，我们通过OpenDataHub将这些创新内容引入开源社区，在Kubernetes上推进AI工作流程方面的最新技术，并为在RedHatOpenShiftAI和watsonx.ai中使用这些创新内容奠定了基础。通过这种方法，我们为基础模型的端到端生命周期启用了一个企业级平台。我们期待与您在上游社区进行合协作。