摘要:导读从 Presto 专家到 AI 基建破局者罗震霄博士现任Pinterest高级首席工程师(Sr. Staff Software Engineer),负责大数据实时处理引擎、监控平台以及大模型数据预处理。在加入 Pinterest 之前,他曾在 Cloude
导读从 Presto 专家到 AI 基建破局者罗震霄博士现任Pinterest高级首席工程师(Sr. Staff Software Engineer),负责大数据实时处理引擎、监控平台以及大模型数据预处理。在加入 Pinterest 之前,他曾在 Cloudera、Uber、Twitter、Facebook 负责大数据引擎和机器学习平台的研发与运维。罗博士是开源项目 Presto 的Committer及技术指导委员会(Technical Steering Committee)成员。深耕大数据领域多年,罗博士对技术的痛点与瓶颈有着敏锐的洞察。当 Spark、Hadoop 等传统数据处理框架趋于成熟时,大语言模型(LLM)浪潮暴露了“最后一公里”的难题。GPU 利用率长期低于 40%,而 90% 的模型训练卡在数据预处理环节。传统批处理架构难以满足百亿参数模型对数据吞吐量的苛刻需求,亟需新技术来解决效率问题。开源的分布式计算框架 Ray 凭借其出色的资源管理优化能力,成为破局的关键。罗博士带领团队成功将 Ray 落地 Pinterest。本次分享将介绍这一数据预处理平台的技术特性、实践案例、优化效果以及未来的演进方向。主要内容包括以下几个部分:1. Pinterest 的 LLM 处理痛点2. Ray 的破局之道3.落地实践:从概念验证到公司级基建4. 效果显著:性能提升驱动全面落地
分享嘉宾|罗震霄博士 Pinterest Sr. Staff Software Engineer
编辑整理|牛昱琦
内容校对|郭慧敏
出品社区|DataFun
01
Pinterest 是全球领先的图片与视频分享社交平台,月活跃用户超5亿。为了兼顾数据隐私、安全性和服务可用性,Pinterest 从使用 OpenAI 接口转向自研私有化大模型服务。在模型推理效率优化方面,公司应用了包括 Flash Attention(一种优化 Transformer 模型中自注意力机制的算法,能在不损失精度下提升计算效率并减少显存占用)和 vLLM(一种针对大语言模型的高效推理与显存优化框架)在内的先进技术。
尽管自研模型能满足大部分业务需求,但在资源调度层面仍存在显著痛点。整个数据处理流水线(涉及 Spark、TensorFlow、PyTorch 等多种技术栈)被复杂的依赖关系嵌入到 Airflow 中,数据量达 PB 级。这导致新增数据集或特征难以及时应用到模型训练中。团队发现,即使增加 GPU 数量,对提升训练效率也收效甚微。深入排查后,锁定问题根源在于集群采用的固定 CPU:GPU 配比策略,导致数据加载环节成为瓶颈。
Ray 是一款用于扩展 AI 应用程序的统一计算框架,其核心优势包括:
分布式处理:高效地在多节点间并行处理海量数据。智能资源管理:可精细调度和管理 CPU、GPU 等异构资源,确保任务在最优硬件上运行。生态兼容性:原生支持 TensorFlow、PyTorch 等主流 AI 框架,用户无需在多个系统间切换。Ray 的核心架构(Head Node 与 Worker Node):
Head Node(头节点):Ray 集群的控制中心,核心功能包括:通过 GCS(全局控制服务,Global Control Service)维护集群元数据(任务、对象、Actor 位置);协调 Worker 资源分配;提供 Dashboard 可视化界面。它是集群初始化和管理的核心枢纽。Worker Node(工作节点):计算任务的具体执行单元,核心功能包括:通过 Raylet 进程执行远程任务(Task);管理 Actor(有状态计算单元)的生命周期;与分布式对象存储进行数据交互。Worker Node 是无状态计算资源的提供者(CPU/GPU 等),支撑 Ray 分布式计算的执行。协同机制:Head Node 通过 GCS 将任务调度分配至 Worker Node 执行。Worker 将处理结果存入分布式对象存储,并通过 Raylet 同步状态信息,确保分布式任务的高效、一致性执行。Ray 的特性使其成为解决大模型团队瓶颈的理想方案,但引入过程充满挑战。彼时(几年前),还没有像如今 Spark 3 那样的性能提升方案。罗博士经同行推荐了解到 Ray,对它在模型训练优化方面的潜力产生浓厚兴趣。然而,当时 Ray 在大型生产环境中的落地案例尚少。
罗博士团队通过反复沟通和论证,说服管理层同意投入资源进行验证和落地。前期调试阶段投入了大量精力。如今,Pinterest 已成功实现通过 SDK 无缝调用 Ray 服务,流程简单清晰、适配性强,获得大规模应用。其中关键亮点技术包括:
1. map_batches: Ray Data 的核心批处理 API
API 将数据集划分为批次(batch),并针对每个批次并行地应用用户自定义函数(或可调用类)。与逐行处理的 map不同,批处理**是提升 Ray Data 大规模数据处理吞吐量的关键手段,尤其适用于批量数据预处理、模型推理和聚合计算等场景。
核心价值:
显著提升吞吐量:减少函数调用开销,批量处理(如模型推理)可比单条处理快数倍。支持有状态计算:通过可调用类(callable class)维护批次间的状态(如模型权重、窗口数据)。资源感知调度:支持指定 num_cpus、num_gpus 等参数,确保任务调度至合适节点(如将模型推理任务调度到 GPU 节点)。格式灵活性:支持 numpy、pandas、pyarrow 等多种批次数据格式。容错机制:基于 Ray 的分布式引擎,自动处理节点故障(任务重试、数据恢复)。2. 多模型推理(Multimodel Inference):规模化处理的利器
Ray 通过其分布式能力,高效协调部署在同一集群中的多个机器学习模型(如大语言模型、计算机视觉模型、推荐模型)进行并行推理与结果聚合。作为 Ray AI 生态(Ray Core, Ray Train, Ray Serve)的核心组件,它解决了大规模多模型推理的部署和管理挑战。
性能优势:
智能资源分配:自动识别模型资源需求(num_cpus, num_gpus),精准调度到合适节点,避免争抢。弹性伸缩性:支持从单机到千节点规模的无缝扩展,处理 TB 级数据和高并发请求。统一编程模型:提供 Actor、Ray Serve 等一致的 API,大幅简化多模型服务的开发流程。3. Carry Over Columns:跨批次上下文传递优化器
Carry Over Columns 是 Ray 数据预处理中解决跨批次上下文传递开销的关键机制。在批处理或流式处理中,某些关键上下文信息列(如用户 ID、会话 ID、时间窗口标识)不参与当前批次的模型计算(如特征转换),但对后续批次的分析或状态更新(如会话跟踪)至关重要。传统方式需重复加载这些列,造成冗余传输与存储
优化效果:
大幅降低网络传输:避免不参与计算的列跨节点传输,尤其在 TB 级数据集和多节点集群中效果显著。减少内存占用:无需在计算节点本地保存非计算列,释放宝贵内存资源。简化开发流程:开发者无需手动管理上下文传递,降低代码复杂度。4. 累加器(Accumulators):分布式状态共享利器
Ray 的累加器用于在分布式任务中高效安全地进行跨节点状态累加,典型应用场景包括:
实时流处理:在 Pinterest 的流式平台中,用于实现跨批次的滑动窗口聚合(如计算用户 1 小时点击量),支撑实时监控与报警。分布式机器学习:在分布式训练中,用于累加梯度总和或全局指标(如样本数),辅助参数更新。多 Actor 协同任务:在需要多个 Actor 共享和更新同一状态(如全局计数器、累计奖励)时,提供简单高效的解决方案。内部多个团队试用 Ray 后,其带来的性能变革极其显著:
关键业务指标飙升:负责 Related Pins(反映用户内容分享增长的核心营收指标)的团队,在将 AI 处理框架切换为 Ray 后,该指标实现了 4.7 倍的惊人提升。处理速度飞跃:具体案例显示,对比旧版基于 Torch DataLoader 的框架(注:修正 TorchDataLoader 为常见写法),通过 Ray 优化 GPU/CPU 资源调度策略(在减少 GPU 数量的情况下),数据处理作业运行时间从 90 分钟缩短至 19 分钟。成本大幅优化:负责搜索质量的团队,将核心 Spark 任务迁移至 Ray 后,利用其多模型推理并行能力,实现年度成本缩减 30 倍。卓越的实践效果带来了强大的口碑效应。自 2024 年初起,Pinterest 内部各团队纷纷开始采用 Ray 平台。管理层已决定,公司所有机器学习相关任务将全面转向 Ray 执行。
回顾 Ray 在 Pinterest 的落地之路,正是罗博士团队对痛点(数据管道瓶颈)的精准洞察、在技术空白期(当时社区对 Transformer 模型本身的优化关注度远超数据管道)果断引入开源技术并持续建设的决心,以及将传统数据库优化经验(Ray 部分优化借鉴了此思路)迁移到新框架的能力,共同驱动了 Ray 的成功应用,为公司带来了巨大的效率与成本优化。
展望未来,罗博士团队将持续优化 Ray 平台,重点关注:
弹性伸缩性(Auto-scaling):更智能地动态调整集群规模。微调能力(Fine-tuning):增强大模型微调任务的体验和效率。数据加载器性能(DataLoader):进一步提升数据加载吞吐量。容错性(Fault Tolerance):强化大规模集群下的健壮性与可靠性。总结回顾:
罗博士向我们介绍了 Pinterest 团队如何 运用 Ray 框架成功解决了 LLM 数据管道瓶颈问题。该方案通过优化资源调度、提升批处理效率(map_batches)、简化多模型推理管理、减少跨批次数据冗余(Carry Over Columns)以及提供分布式状态工具(Accumulators),显著提升了数据处理效率。这验证了团队早期对数据瓶颈的精准定位和技术选型的有效性。Ray 现已成为 Pinterest 机器学习任务的核心平台,团队将持续优化其扩展性和稳定性。期待 Ray 在 Pinterest 的实践中取得更多突破性成果!
以上就是本次分享的内容,谢谢大家。
来源:DataFunTalk