雷的批量推理在Pinterest上的第三部分系列

Alex Wang; 软件工程师 I | Lei Pan; 软件工程师 II | Raymond Lee; 高级软件工程师 | Saurabh Vishwas Joshi; 高级软件工程师 | Chia-Wei Chen; 高级软件工程师 |

基于专家建议的翻译： introduction

介绍

在我们系列的第一篇博客文章中，我们讨论了为什么选择使用Ray作为最后阶段的数据处理框架以及它如何帮助我们解决重要的商业问题。在我们系列的第二篇博客文章中，我们描述了如何将Ray集成到我们现有的机器学习基础设施中。在这篇博客文章中，我们将讨论Pinterest上Ray的第二种流行应用：离线批量推理机器学习模型。我们还将分享我们的实现如何实现了4.5倍的吞吐量提升和节省了30倍的成本。

背景:

离线批量推理涉及在一个大规模数据集上操作，并将数据分批次传递给一个机器学习模型，模型会为每个批次生成结果。离线批量推理任务通常包括一系列步骤，比如数据加载、预处理、推理、后处理和结果写入。这些离线批量推理任务对I/O有高需求，也可能对计算资源要求很高。

在 Pinterest，以前的批量推理解决方案是基于 Apache Spark(™) 或 Torch Dataloader. 构建的。但这些解决方案存在一些缺点。

• 缺少不同类型的异构节点实例，用于预处理和推理
• 实现数据加载、推理和结果写入等步骤的流水线和重叠步骤困难
• 由于缺乏流处理支持，Spark 必须进行批处理数据处理
• 我们的 Spark 集群对 GPU 支持不足
• 离线 Spark 批处理任务通过 RPC 调用在线 GPU 集群，效率低下

为什么选择Ray(™)进行批处理呢？

雷成功解决了Pinterest之前批次推理方案中存在的问题。我们基于Ray的批次推理解决方案，即Ray 批次推理，具有几个关键能力。

• Streaming execution：Ray Data 使用流式执行来高效处理非常大的数据集。数据集执行是流水线化的，因此多个执行阶段可以并行运行。这对于批处理推理特别有用，在批处理推理中，我们希望将数据加载、推理和结果写入重叠起来。
• Heterogeneous clusters：Ray 能够管理由不同资源组成的异构集群。这使我们能够独立扩展 CPU 密集型数据加载/写入任务和 GPU 推理任务。我们之前在数据加载阶段遇到瓶颈的批处理推理任务，现在在吞吐量方面有了显著提升。
• Training framework agnostic：我们使用 Ray Data，它自称为可扩展的数据处理库，用于机器学习工作负载。Ray Data 不依赖于任何特定的机器学习库。数据转换可以通过 Pyarrow Table 进行表示，从而保持了与具体库的无关性。Ray 批量推理任务可以使用任何机器学习框架，例如 PyTorch、TensorFlow、HuggingFace 等，进行推理。
• Job composition：批量推理任务通常作为更大规模的机器学习工作负载的一部分。例子包括训练后的离线评估和知识蒸馏。使用 Ray Data，批量推理可以与其他数据集操作一起，在同一个 ray 任务中执行。

深度解析

图1：我们的Ray批处理推理SDK高层架构图

Pinterest 的 Ray Batch Inference 是一个 SDK，它运行在我们现有的 Ray Infrastructure 之上。内部，它利用 Ray Data 来构建流式执行管道，以完成批量推理所需的所有步骤：数据加载（dataloading）、预处理、推理、后处理和写入结果。下面是一个简短的例子，展示了我们的 SDK 如何使用。

    # 创建一个ServableInfo对象  
    # 这代表一个ML模型和它相关的元数据  
    servable_info =   
    TorchScriptServableInfo(  
        full_path="s3://path_to_trained_model.pt",  
        framework=Framework.TORCHSCRIPT,  
        device=DeviceType.CUDA,  
    )  

    # 使用servable_info创建批处理预测器  
    batch_predictor = BatchPredictor.from_servable_info(servable_info)  

    # 加载数据集  
    dataset = read_parquet("s3://path_to_dataset")  

    # 进行预测  
    inference_results: ray.data.Dataset = batch_predictor.predict(dataset)

幕后揭秘

本质上，批量推理是对数据集进行转换的过程。 Ray Data 提供了一个名为 map_batches 的函数，允许你在一个数据批次上定义映射操作。该操作将通过流式处理方式在整个数据集上执行。对于我们来说，映射操作是模型推理。

图2：我们如何使用Ray Data的高层图示。四个Ray actor被调度到一个AWS g5.24xlarge实例，每个Ray actor都分配了一块GPU。

在幕后，Ray Data 调度 ray 任务 或 ray actor 来执行在 map_batches 中定义的映射函数。由于模型推理需要存储模型权重和元数据，我们使用 ray actor 的 map_batches 实现。

当我们的Ray批推理SDK在Ray主节点上执行时，它将调用map_batches。Ray Data将创建Ray actor实例并将其调度到Ray工作节点上。在进行GPU推理时，这些Ray actors们必须被调度到GPU节点上。在实例化过程中，Ray actor实例将下载模型权重并加载到设备上。在实例化完成后，Ray actors们准备好接受数据批次来进行推理。

延续列

批处理推理的一个特性是保留列。保留列是指输入数据集中的一些列，这些列不会在批处理推理管道中处理，但为了满足下游任务的需要，这些列仍需包含在最终输出中。一个常见的保留列的例子是图像签名（图像唯一ID）。虽然图像签名可能不会作为模型的输入特征，但在输出中保留它对于数据分析和索引很有用。

图3：Ray批推理中延续列的工作方式

如果你记得的话，我们使用Ray Data的map_batches来进行批量推理。在Ray Data里，一批数据可以表示为一个pyarrow表。将批量表示为pyarrow表的一个好处是我们可以执行高效且零复制的内存操作，这非常节省内存。

在我们的批量推理管道中，输入批量包含 特征列 ，这些列将输入模型，以及 非传递列 ，这些列不会输入模型。非传递列也将跳过预处理和后处理阶段。使用 pyarrow，将非传递列分离成单独的 pyarrow 表是一个零拷贝过程。包含非传递列的 pyarrow 表也可以通过另一个零拷贝过程附加到输出批量中。

根据不同的使用场景，我们的批量推理作业可能涉及大量延续字段。有效地处理延续字段对于减少数据开销至关重要。

多模型推断

在许多情况下，同时在一次任务中运行多个模型的推理通常是有好处的。我们的Ray批量推理解决方案允许在一个Ray作业中同时运行多达N个模型的推理。通过将N个并行任务合并为一个任务，这样可以减少N倍的数据重复读取和转换的开销。

图4：多模型推理的高层示意图。每个ray actor（Ray actor）可以在同一个设备上加载多个模型。

在Ray批量推理中，我们通过一个称为Servable（服务包）的概念来共同表示一个机器学习模型及其元数据。在我们的多模型推理的实现中，每个Ray actor可以持有多个Servables。

    # 第一个模型
    servable_info_1 = TorchScriptServableInfo()   

    # 第二个模型
    servable_info_2 = TorchScriptServableInfo()   

    # 从这些 servable_infos 中实例化batch predictor
    batch_predictor = BatchPredictor.from_servable_info(
    servable_info=[servable_info_1, servable_info_2],  
    )  

    # 加载数据集
    dataset = read_parquet("s3://path_to_dataset")  

    # 调用预测函数
    inference_results: ray.data.Dataset = batch_predictor.predict(dataset)  

    # 输出数据集的各列将包含所有N个模型的预测结果，如下所示:

    # -----------------------------------------------------------------------------------
    # 保留列 | 模型1结果列 | 模型2结果列 | ... | 模型N结果列
    # -----------------------------------------------------------------------------------

多模态推断在很多领域都有应用，例如：

• 生成不同的嵌入向量。每个模型将从相同的输入生成唯一的嵌入。
• 进行多个模型实验的离线评估多个模型的实验结果。

供能装置

离线批量推理通常用于计算评估指标，例如[AUC-ROC]和[交叉熵损失]。这些评估指标可以使用累加器高效计算，累加器将来自不同推理任务的多个值汇总到一个值，例如在Spark中。我们在Ray上基于现有分布式累加器实现了累加器。在Ray作业完成后，这些指标会被保存到实验跟踪器，如MLFlow或Weights & Biases中。

图5：利用Ray Actors的分布式累加器。内存开销减少了X倍（变为原来的1/X）。

我们对分布式累加器的天真实现会为每个计算出的评估指标执行一个ray任务来发送给一个Ray Actor进行累加这一过程。然而，由于大量的ray任务，导致了显著的开销，导致节点内存不足或ray对象存储溢出。

我们的优化方法采用了一个两阶段累积策略，在该策略中，部分值先被聚合，然后再进行合并。这有助于减少射线远程任务并减少内存占用。

如下图所示的两步积累示例中，所计算的度量为所有积累值的合计。

语言模型推断

鉴于我们的Ray批推理解决方案是框架无关的，它成为构建大型语言模型批推理系统的理想基础。

在 Pinterest，我们目前使用vLLM™ 作为我们内部 LLM 的推理优化引擎。vLLM 通过PagedAttention和KV Cache优化 LLM 的注意力机制，帮助我们以更高的效率处理极大规模的数据，使我们能够以更高的吞吐量处理数据。

LLaMA-7B 模型使用 vLLM™ 的处理速度与 Hugging Face 平台（HF）和 Text Generation Interface 工具（TGI）相比。实验是在 Nvidia A10G GPU 上完成的。结果取自 vLLM 团队的 博客。

对于 TorchScript 模型，我们实现了一个名为 TorchScriptServable 的可服务化模型。这一理念可以扩展到所有类型的机器学习模型，包括大型语言模型 (LLM)。对于 vLLM，我们实现了一个名为 VLLMServable 的可服务化模型。通过采用这种设计模式，Ray 批量推理可以支持任何类型的机器学习模型，不管其内部实现如何。

    # 创建一个VLLMServableInfo对象  
    # 这代表了一个大型语言模型及其相关元数据  
    servable_info = VLLMServableInfo(  
        framework=Framework.VLLM,  
        device=DeviceType.CUDA,  
        input_features=["prompt"],  
        output_features=["generated_text"],  
        vllm_config = VLLMConfig(  
             engine_args=EngineArgs(  
                  model="my_fine_tuned_llm"  
             ),  
             sampling_params=SamplingParams(  
               ...  
         )  
      )  
    )  

    # 根据servable_info实例化批量预测器  
    batch_predictor = BatchPredictor.from_servable_info(servable_info)  

    # 加载提示数据集  
    dataset = ray.data.from_items(  
            ["弗吉尼亚州的首都是哪里？", "今天的天气如何？", "完成这个句子：你好， ", "如何在Python中对列表进行排序？"]  
    )  

    # 运行预测  
    generated_text: ray.data.Dataset = batch_predictor.predict(dataset)

难题:

在将vLLM™集成到Ray批处理推理的过程中，我们在异构Ray集群环境中使用vLLM时遇到了一个问题。升级到vLLM 0.4.0后，发现该库无法在非GPU节点上导入。尝试在仅CPU的Ray头节点上使用vLLM时，却遇到了异常。

解决方案是将vLLM从我们的面向用户的API解耦出来，这些API仅在ray头部使用，即VLLMServableInfo 和 VLLMConfig。我们只将vLLM与VLLMServable对象耦合，该对象只包含模型工件，并且只由调度到GPU节点上的ray演员使用。

结果部分

Pinterest 的机器学习工程师使用 Ray 批量推理，实现了显著的吞吐量、GPU 利用率和成本节省的提升。

火炬变Ray(™)

Pinterest 的相关 Pins (p2p) 团队在将一个推理作业从 TorchDataloader 迁移到 Ray 平台后，吞吐量提高了 4.5 倍（即提高了 450%）。他们还减少了用于该任务的 GPU 数量，从 8 个（即 2 台 g5.24xlarge 机器）减少到 4 个（即 1 台 g5.24xlarge 机器），这也使 GPU 的利用率得到了提升。迁移还使作业成本大约降低了 25%。

这一改进主要是由于能够利用一个异构光线集群（heterogeneous ray cluster）的能力。相关pins团队的推理任务在管道的数据加载阶段遇到了瓶颈。在向光线集群添加额外的CPU资源后，吞吐量显著提升，因为数据加载阶段能够在CPU上调度更多光线任务。

Spark(™)到Ray(™)

Pinterest 的搜索质量小组在将工作从 Spark(™) 迁移到 Ray(™) 的之后，其中一个推理任务的年度成本减少了近30倍。这一改进是由两个因素促成的，

• 迁移到 Ray(™) 后，作业能够利用 GPU 资源的优势。这极大地提高了管道中的推理阶段吞吐量。
• 该作业利用了我们的多模型推理功能。之前，搜索质量领域的推理作业由 N 个 Spark 作业组成。每个 Spark 作业从离线存储加载相同的模型，并将其结果分别写入独立的表。这 N 个 Spark 作业被合并成了一个 Ray 作业。读取数据集的开销减少了 N 分之一，并且来自多个模型的推理结果能够写入同一个表中。

收养

自2023年第四季度Ray Batch Inference成为正式可用（GA）以来，Pinterest的ML团队已经迅速地采用了它。我们目前在13个不同团队中运行Ray Batch Inference作业，平均每天启动超过60个作业。

接下来会怎样

我们计划将KubeRay整合到我们的Ray™基础设施中。这将使我们能够利用[自动缩放]功能处理批量推理和训练作业，从而大大简化性能调优过程。

即将讨论的其他话题有：

• Pinterest 的 Ray Tune
• 改进 Ray 数据加载器性能
• Ray 的批量推理和训练容错性
• Ray 的特征重要性

感谢

相关推文：Shivin Thukral, Liyao Lu, Travis Ebesu

ATG: 阿玛尔·阿尔卡塔里，维什瓦卡尔马·和辛格

搜索质量：巴万娜·朱奈贾，甘德哈夫·卡波尔。

M10n: 何浩宇，卡蒂克·卡帕尔

ML平台：赖庆贤和维什努·阿伦

想了解更多关于 Pinterest 工程的信息，请查看我们其余的_工程博客 ，并访问我们的Pinterest Labs 网站。想看看并申请开放职位，请访问我们的职位页面 _。