—— 用于加速基础模型创新的云原生开源软件栈
作者:
Talia Gershon, IBM 研究院混合云架构研究总监
Priya Nagpurkar, IBM 研究院混合云平台研发副总裁
Carlos Costa, IBM 研究院基础模型和多云首席研究科学家
Darrell Reimer, IBM 研究院杰出工程师
北京2023年5月15日 /美通社/ -- 基础模型和生成式 AI 激发人类的集体想象力,推动我们发现改善生活和工作方式的新方法。 从通过自然语言实现更加无缝的技术交互,到自动生成代码或其他数据,再到跨科学各个领域的用例,基础模型的应用场景与日俱增。 IBM 的目标是将该技术融入我们的产品组合,帮助我们的客户快速高效且安全地将基础模型应用到他们自己的产品中。
在这一过程中,我们分享了为什么在 IBM Cloud 中构建 AI 超级计算机 Vela。 这项工作大力推动我们重构整个技术堆栈,从而加快我们训练、微调和部署尖端 AI 模型的步伐。 通过这个过程,我们成功构建了一个面向基础模型时代进行优化的现代化、灵活的 AI 软件栈。
本博文将介绍在 Red Hat OpenShift Container Platform 上运行的高性能云原生 AI 训练堆栈,该平台是新推出的 watsonx 平台的构建基础。
作为训练堆栈的补充,我们的技术栈可降低成本与优化性能,对基础模型进行调整并为其提供服务。 以下多项技术已贡献给开源社区,如 PyTorch、Ray、Kserve 和 Open Data Hub (ODH) (一个可在 Kubernetes 上构建、部署和管理数据密集型应用的开源平台)。 在 ODH 上发展成熟的技术随后被引入 Red Hat 的 OpenShift AI,IBM 目前的新一代 AI 平台 watsonx.ai 使用 Red Hat OpenShift AI。 借助这种方法,IBM 和 Red Hat 为客户提供最先进的开源基础模型栈,以支持在他们选择的任何环境(本地、IBM Cloud 或其他公有云)中运行。
我们的基础模型训练方法
我们在开始重构 AI 训练堆栈时有两个目标。 首先,希望保持传统 HPC 系统的实用性:最大程度提高硬件利用率,并高效利用高性能基础架构。 其次,希望提供混合云开发体验的灵活性和生产力优势:提高开发敏捷性和代码重用率,简化基础架构和软件的管理与扩展。 为了实现第二个目标,我们使用 Kubernetes 构建解决方案,通过容器复用代码以及扩展软件。 但这一决定意味着我们需要将 Kubernetes 转化为可处理高性能工作负载的平台。
我们还需要一个可以应对每一步 AI 训练工作流程的解决方案:进行数据预处理、分布式训练和模型验证。 我们确定了合作的关键开源社区,共同处理端到端工作流程,确定了需要克服的关键用户体验障碍,以便用户启动、运行和扩展工作。
下图 1 的左侧介绍了我们的训练软件堆栈的整体情况,该堆栈自 2022 年底以来一直在 IBM Cloud 的 Vela 上运行,由 IBM 研究院使用。 图 1 右侧描述了我们用于调整和提供基础模型的堆栈,稍后将在博客中详细讨论。
高级 Kubernetes 原生资源利用与管理
这项工作伊始,Kubernetes 生态系统在大规模和高性能 AI 工作负载方面仍然存在重大差距。 早期关注的领域之一是将基础架构功能(如网络资源)应用于工作负载,而不会产生额外的开销。 为此,我们创建了一个多 NIC CNI 操作程序,消除了封装,可配置底层网络接口并将网络延迟减少了一半,同时与开箱即用的容器网络解决方案相比带宽增加了七倍。 这些改进对最终用户完全透明。
我们尝试填补的第二个空白领域是使用正确的云原生作业调度程序。 由于大量 AI 开发人员希望提交作业以在 Vela 上运行,我们需要一个调度程序来分配资源并确定作业的优先级,以最大限度地提升资源利用率。 为了解决这个问题,IBM 研究人员创建了多集群应用进程调度程序 (MCAD),可提供作业队列、作业优先级与抢占、超时以及系统用户之间的资源共享编排。 此外,我们启用了工作负载打包和分组调度,以消除资源碎片,这些都通过 OpenShift 实现。 我们进一步开发了 InstaScale,它与 MCAD 共同动态扩展云托管的 OpenShift 集群。 通过自动从云提供商获取 GPU 并按需发布,InstaScale 让从业者不必担心基础架构管理和成本问题。
高效、可扩展的数据预处理,模型训练与验证
为了使 AI 管道中的所有运行步骤简单高效,我们主要运用 和贡献PyTorch 和 Ray 这两项关键的开源技术。 借助 Ray,我们使用数据科学家友好的 Python API 实现了可扩展的数据预处理(例如使用仇恨、滥用和亵渎过滤器过滤数据)和后处理步骤(例如模型微调和验证)。 通过运行 Ray with MCAD,我们支持并发运行的异构 Ray 作业高效共享资源池。
我们正与 PyTorch 合作推进对分布式训练的支持,包括通过引入 rate_limiter 来改进对完全分片数据并行 (FSDP) 训练 API 的支持。 我们最近展示了在基于以太网的环境(如 IBM Cloud 中的 Vela)上为参数为 10B+ 的模型高效扩展分布式训练作业。 此外,通过集成 MCAD 与 TorchX(TorchX 是 PyTorch 应用的通用作业启动器),我们能够使用不同 API 和框架以透明方式支持多种基于 PyTorch 的作业。 这些工作都受底层工作管理系统的支持,无需 AI 从业者修改代码。
简化用户体验
工作流程本身的训练分为以下三个步骤:模型探索(通常是使用一些 GPU 运行的缩小试验)、扩展分布式训练作业(消耗数百个 GPU)以及模型验证。 对于许多 AI 从业者来说,编排这些步骤可能很复杂,而且配置和管理的时间长。 我们通过 CodeFlare 项目解决了这一挑战,它能够提供引导式、简化的用户体验,以有效地提高训练、测试和监控模型训练生命周期的效率。
CodeFlare CLI (基于控制台和 UI)指导用户应对复杂的远程 OpenShift 集群运行,同时自动执行作业配置、存储设置、日志记录以及管理用于监控和分析的端点。CodeFlare SDK(基于 Jupyter)为用户提供直观的 Python 界面,进行批处理资源请求处理、作业提交和观察。 借助这些功能,我们大大降低了 AI 研究同事进入云原生堆栈的门槛。
在 Vela 上运行我们的软件栈
到 2022 年底,IBM 的所有基础模型训练工作都过渡到在 IBM Cloud 的 Vela 上运行此软件栈。 如今,MCAD 管理这些 AI 作业的队列,从单 GPU 作业到利用超过 512 个 GPU 的作业,并处理作业优先级和配额管理。 我们在过程中发现了其他方法,可以让团队在以 GPU 为中心的环境中(如 Vela)中管理 OpenShift 集群更轻松,例如,通过增强 OpenShift Installer 配置基础架构 (IPI),以便更轻松地在高性能基础架构上部署和管理 OpenShift。
我们的基础模型调优和服务方法
训练和验证最先进的基础模型是 AI 价值链的关键早期阶段,但当模型在 AI 工作流程的调优和推理步骤中投入生产使用时,最终会捕获真正的价值。 我们用于推理和模型调优的软件栈主要是在底层硬件上高效执行模型,以最佳方式批处理传入的请求,简化 AI 与应用的集成,并提供最先进的模型适应技术。 上图 1 的右侧介绍了我们的基础模型调优和服务栈,以下是更为详细的描述。
推理性能
用于优化基础模型在给定硬件平台上运行方式的软件库可以将吞吐量和延迟提高 10-100 倍。 我们的服务软件栈中有一组精心设计的成熟优化路径(包括 ONNX 和 Hugging Face Optimum),用于推理常见模型架构,并且可扩展,以适应新的推理服务器或优化。 考虑到 AI 和开源社区的快速创新步伐,可扩展性是我们软件栈的关键设计点。 此外,真正的 AI 服务会同时收到来自多个用户的、针对多个模型的大量推理请求。 我们的服务堆栈动态批处理传入的请求,并通过构建和反馈 Hugging Face、Kserve 和 Model Mesh 社区,有效地在模型之间进行多路复用。
简化应用程序集成
目前可用于运行 AI 模型的推理服务器需要用户具备大量的 AI 专业知识。 模型的输入和输出都是张量。 对于希望利用这些模型完成任务的应用开发人员来说,这种格式并不易于理解。 为了使此过程对开发人员更加友好,必须将模型输出转换为更易于使用的内容。 我们创建了 Caikit 抽象层,可为应用开发人员提供直观的 API 和数据模型,并提供一个稳定的接口,支持模型和应用独立发展。 这一抽象层在 IBM 的 Watson 模型中使用,用于提供基础架构,并将很快贡献给开源。
基础模型调优
基础模型的关键价值主张之一是能够利用预训练的基础模型,并使用专用数据"调优"或"适应",以提高其下游任务的性能。 我们的目标是打包最先进的技术,以实现计算效率高的模型自适应,让不甚了解其工作原理的人也能够使用。 我们的可扩展堆栈目前支持多任务提示调优 (MPT) 和微调,通过开源项目参数高效微调 (PEFT) 集成。 在接下来的几个月里,我们将开源一些提示调优算法和实现。
和 Red Hat合作
IBM 研究院正在与 Red Hat 合作,我们将开发的功能贡献给主要开源社区,直接发布在 Open Data Hub (ODH),让其他人也从这项工作中受益。 ODH 是一个全面的开源工具集合,旨在利用 OpenShift 的优势来推动整个 AI 开发生命周期。 Open Data Hub 中引入的许多技术已发展成为 Red Hat OpenShift AI 的一部分,并作为 watsonx.ai 的中间件基础提供服务。 图 2 显示了本博客中描述的对各种开源贡献将如何汇集到 ODH 中以支持基础模型用例。
后续工作
面向基础模型时代重构我们的端到端软件栈对我们的 AI 社区具有相当大的价值。 AI 研究人员不再需要非常深入的基础架构知识,就可以让作业高性能运行。 他们不再需要弄清楚如何将作业从几个 GPU 扩展到数百个,或者如何精确地分配作业以实现高工作负载性能,现在这些任务都可以交给软件栈处理。 代码可跨团队重用,试验也易于重现。 我们还大幅简化了 AI 开发人员如何以高计算效率和开发人员友好的方式提供和调整基础模型的方式。
或许最重要的是,在 OpenShift 上构建此堆栈提供了对其他环境的可移植性,让合作伙伴可以在本地和任何公有云中使用这些功能。 我们很高兴与红帽一起通过开放数据中心将这些创新引入开源社区,在 Kubernetes 上推进 AI 工作流程方面的最新技术,并为在 Red Hat OpenShift AI 和 watsonx.ai 中使用这些创新内容奠定了基础。 通过这种方法,我们为基础模型的端到端生命周期启用了一个企业级平台。 我们期待与您在上游社区进行合作。