Aarki 竞价大脑的进化，第 2 部分：利用 GPU 实现十亿行规模的扩展

在 本系列第 1 部分我们分享了 Aarki 如何从经典机器学习演进到深度神经网络，并教会我们的 DSP 利用深度学习进行更智能的思考。然而，如果没有相应的基础设施支持，更智能的模型也只能走这么远。

这篇文章涵盖了该故事的后半部分：我们如何决定投资 GPU，为什么公共云无法为我们扩展，以及如何构建一个每天训练数十亿条记录的系统。

在 CPU 上进行训练已经不够了

Aarki 的机器学习平台每天训练数十个模型，其中一些最大的模型每天训练超过 10 亿条记录。为了高效地训练模型，我们使用硬件加速，利用 GPU 来训练深度学习模型。

我们利用 TensorFlow 生态系统和 NVIDIA 的 CUDA 平台来处理海量数据，以确定哪些广告展示能够为我们的广告主带来最佳价值。当我们开始将需求方平台 (DSP) 迁移到深度学习模型时，我们面临一个抉择：是租用 GPU 还是购买 GPU？

在这篇文章中，我们将重点介绍我们探索过的选项、我们为展示 GPU 的价值所采取的途径以及我们支持大规模深度学习的未来之路。

为什么我们的基础设施不是从 GPU 开始的

我们的数据中心已经针对速度进行了优化

我们的DSP部署在四个私人数据中心，这些数据中心与主要程序化广告交易平台的交易地点位于同一地点。我们投资了硬件，以支持我们拥有和运营的机器每秒处理超过500万个出价请求。

虽然这种方法比使用公共云选项（如 AWS 或 GCP）灵活性差得多，因为您可以根据需要启动新机器，但我们避免了为网络出口支付大笔费用，并且每个请求的摊销成本低于使用其中一个平台。

我们早期的基础设施专注于核心 DSP 吞吐量，而不是训练

我们最初专注于支持 DSP 的硬件，并没有在 GPU 上进行前期投资。数据中心没有 GPU 对我们上一代深度学习模型来说并不成问题，因为我们用于训练模型的库针对 Rust 进行了优化，并且不支持硬件加速。

DeepFM 在没有加速的情况下也能工作——直到它失效

我们之前的模型使用了深度分解机 (Deep Factorization Machines)，它结合了学习所有输入特征的嵌入的深度阶段和学习特征之间交互的分解机阶段。这些模型在访客找回营销活动中效果显著，这类营销活动侧重于安装后的转化事件，例如识别哪些安装最有可能转化为付费用户。

然而，当我们开始研究预安装模型（例如点击预测）时，我们的模型训练平台开始遇到瓶颈。我们能够每天使用数亿条记录来训练模型，但模型训练速度太慢，无法扩展到每天超过十亿条记录。

向深度神经网络 (DNN) 的转变使得 GPU 成为必需

我们的新模型需要硬件加速

我们的下一代定价模型采用支持硬件加速的深度学习框架，用于模型训练和模型服务。我们最初专注于使用 GPU 进行模型训练，因为训练数据集的大小是主要瓶颈。

TensorFlow 和 PyTorch 为我们的 DNN 训练堆栈提供支持

我们目前在生产环境中的定价模型使用 TensorFlow 框架，同时使用 PyTorch 框架构建流水线，用于创建嵌入数据并将其输入到模型中。我们使用硬件加速来训练模型，并将这些模型转换为 ONNX 格式，这样就可以使用 Rust 语言，在我们现有的 DSP 基础设施上使用 CPU 进行处理。

证明 GPU 的价值

数据中心外的测试培训

在为我们的私有数据中心投资 GPU 之前，我们希望确保我们的下一代深度学习模型能够提高我们的 DSP 的性能，并确认我们的新模型可以在 GPU 上比在 CPU 上更高效地进行训练。

为了确保训练效率，我们探索了几种在数据中心外使用 GPU 进行训练的方案。我们最初尝试在办公室的机器上使用采样数据集进行模型训练，然后尝试在 AWS 上启动 GPU 实例，最终决定使用 Google Colab 来验证使用 GPU 是否能显著缩短模型训练时间。

新旧模型的 A/B 测试

我们在 Google Colab 上训练了挑战者模型，并在现有和下一代深度学习模型之间进行了 A/B 测试，以增强我们迁移到 GPU 进行模型训练的信心。在选择是否购买 GPU 时，我们必须做出的关键决策之一是确定哪种类型的 GPU 最适合我们的模型训练用例。

L4 GPU 比 CPU 加速 10 倍

我们利用 Google Colab 收集了使用 NVIDIA T4、L4 和 A100 实例类型的模型训练统计数据，并发现 L4 卡非常适合我们的模型架构。与使用 CPU 进行训练相比，使用 NVIDIA L4 卡，我们能够获得 10 倍的速度提升，将最大模型的批处理时间从 200 毫秒缩短至 20 毫秒。

我们发现 L4 卡是我们投资 GPU 的良好起点。随着我们探索新的架构、增加模型参数并向模型添加更多功能，我们将继续评估哪些实例类型能够为我们的模型提供最佳性能。

原型证明了这一点，但我们需要生产规模

我们使用 Google Colab 进行的 A/B 测试证明了迁移到 GPU 的价值，我们能够在更大的数据集上进行训练，并超越我们当前这一代的深度学习模型。我们的下一步是实现这些流程的自动化，以支持每天对数十亿条记录进行模型再训练。

Google Colab 非常适合原型设计和实验，但它并不适用于生产工作负载，Google 建议使用 Vertex AI 来处理这些流程。我们需要决定未来的发展方向，是利用公有云方案（例如 AWS SageMaker 或 GCP Vertex），还是购买自己的硬件来加速模型训练。

租用 GPU 不适合我们的规模或设置

数据传输和 TensorFlow I/O 使得云训练变得不切实际

我们在租用显卡进行模型训练时面临的最大问题是将数据从我们的私有数据中心传输到 GPU 所在位置。我们在私有数据中心的机器集群上使用 Spark 准备训练数据集，并以 TensorFlow 的 TFRecord 格式输出编码后的数据集。

即使使用 gzip 压缩记录，我们最大的数据集也有数百 GB 的大小，我们每天使用不同的训练数据集训练数十个这种大小的模型。

另一个问题是，直接使用云存储（例如 AWS S3 和 GCP 云存储）上的文件训练 TensorFlow 模型的功能已从核心 TensorFlow 库中移除，并迁移至 TensorFlow I/O 库，该库现已不再维护。在日常模型训练中移动所有这些数据会给我们的模型流水线带来编排挑战，而且我们每天需要为数 TB 的数据支付云存储和传输费用。

我们将 GPU 引入内部并围绕它们构建管道

鉴于我们之前使用 Google Colab 进行过测试，我们已经知道 NVIDIA L4 卡将为我们的模型训练提供一个良好的起点。我们与一家供应商合作，确定了在我们的私有数据中心添加服务器（每台机器配备多张 L4 卡）的方案价格，并从一小批 GPU 驱动的机器开始。

一旦机器安装完毕，我们就能将所有资源都存放在自己的私有数据中心，从而避免了数据传输这一最大的痛点。这种方法使我们能够在开源工具上投入更多，包括用于编排的 Prefect、用于性能分析的 TensorBoard 以及用于模型流水线 ML Ops 的 MLflow。

对于 GPU 来说，购买选项并不总是最好的选择，但由于我们的私有数据中心部署和模型训练的可预测工作负载，它对我们来说效果很好。

我们现在每天在内部进行数十亿行数据的训练

需要一些工程工作来协调用于数据准备的 Spark 工作流程和用于模型训练的 GPU 作业之间的交接，但这是值得的，现在我们每天能够基于数十亿个示例来训练我们最大的模型。

购买选项并不总是正确的途径，但对于我们的私有数据中心和可预测的模型工作负载，它为我们提供了所需的控制、效率和规模。

毕竟，如果不扩展基础设施，就无法扩展智能

第一部分是关于大脑的进化。第二部分是关于为大脑提供其成长所需的基础设施。通过投资 GPU 并将模型训练引入内部，我们实现了性能的飞跃，而云成本却没有大幅增加。

现在，深度学习已成为我们竞价逻辑的核心，而 GPU 则为每日十亿行训练提供支持，我们构建了一个旨在更快地学习、更智能地优化和毫不妥协地扩展的系统。

错过了故事的开头？阅读第一部分 这里.

如果您喜欢这个博客，请给 Ben 留言 bweber@aarki.com。他很想听到您的消息！